Optimize batch loading and metrics writing, replace PositionalSharding with NamedSharding (#186)

* fix profiling

* Use torch cpu, async write to tensorboard, script to convert latents to tfrecord, batch iterator for tfrecord cached, namedsharding instead of positional sharding

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

* Replace positional sharding with named sharding

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

* Formatting

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

* Formatting

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

* Fallback to regular tfrecord iterator for datasets without all the processed features

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

* README update

---------

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(mesh, P())
+  encoders_sharding = NamedSharding(devices_array, 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(mesh, P())
+  encoders_sharding = NamedSharding(devices_array, 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

@@ -198,11 +198,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)
-<<<<<<< HEAD
     params = jax.device_put(params, NamedSharding(mesh, P()))
-=======
-    params = jax.device_put(params, NamedSharding(devices_array, P()))
->>>>>>> f344ab0 (Optimize batch loading and metrics writing, replace PositionalSharding with NamedSharding (#186))
     wan_vae = nnx.merge(graphdef, params)
     p_create_sharded_logical_model = partial(create_sharded_logical_model, logical_axis_rules=config.logical_axis_rules)
     # Shard
@@ -403,7 +399,7 @@ def __call__(
             num_channels_latents=num_channel_latents,
         )
 
-      data_sharding = NamedSharding(self.mesh, P())
+      data_sharding = NamedSharding(self.devices_array, P())
       if len(prompt) % jax.device_count() == 0:
         data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self.config.data_sharding))