Fix namedsharding for replicating params (AI-Hypercomputer#188)

Signed-off-by: Kunjan <kunjanp@google.com>

Author: coolkp

src/maxdiffusion/generate_flux.py

@@ -343,7 +343,7 @@ def run(config):
       config.t5xxl_model_name_or_path, max_length=config.max_sequence_length, use_fast=True
   )
 
-  encoders_sharding = NamedSharding(devices_array, P())
+  encoders_sharding = NamedSharding(mesh, P())
   partial_device_put_replicated = functools.partial(device_put_replicated, sharding=encoders_sharding)
   clip_text_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), clip_text_encoder.params)
   clip_text_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, clip_text_encoder.params)

src/maxdiffusion/generate_flux_multi_res.py

@@ -381,7 +381,7 @@ def run(config):
       config.t5xxl_model_name_or_path, max_length=config.max_sequence_length, use_fast=True
   )
 
-  encoders_sharding = NamedSharding(devices_array, P())
+  encoders_sharding = NamedSharding(mesh, P())
   partial_device_put_replicated = functools.partial(device_put_replicated, sharding=encoders_sharding)
   clip_text_encoder.params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), clip_text_encoder.params)
   clip_text_encoder.params = jax.tree_util.tree_map(partial_device_put_replicated, clip_text_encoder.params)

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -195,7 +195,7 @@ def load_vae(cls, devices_array: np.array, mesh: Mesh, rngs: nnx.Rngs, config: H
     # This replaces random params with the model.
     params = load_wan_vae(config.pretrained_model_name_or_path, params, "cpu")
     params = jax.tree_util.tree_map(lambda x: x.astype(config.weights_dtype), params)
-    params = jax.device_put(params, NamedSharding(devices_array, P()))
+    params = jax.device_put(params, NamedSharding(mesh, P()))
     wan_vae = nnx.merge(graphdef, params)
     p_create_sharded_logical_model = partial(create_sharded_logical_model, logical_axis_rules=config.logical_axis_rules)
     # Shard

src/maxdiffusion/trainers/wan_trainer.py

@@ -26,7 +26,7 @@
 from ..schedulers import FlaxEulerDiscreteScheduler
 from .. import max_utils, max_logging, train_utils, maxdiffusion_utils
 from ..checkpointing.wan_checkpointer import (WanCheckpointer, WAN_CHECKPOINT)
-from multihost_dataloading import _form_global_array
+from maxdiffusion.multihost_dataloading import _form_global_array
 
 
 class WanTrainer(WanCheckpointer):