sharding attempt

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -4,6 +4,7 @@ skip_jax_distributed_system: False
 attention: 'flash'
 attention_sharding_uniform: True 
 precision: 'bf16'
+data_sharding: ['data', 'fsdp', 'context', 'tensor']
 remat_policy: "NONE"
 names_which_can_be_saved: []
 names_which_can_be_offloaded: []

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -530,6 +530,7 @@ def _load_and_init(cls, config: HyperParameters, restored_checkpoint, vae_only=F
           vocoder=components["vocoder"]
       )
       pipeline.mesh = components["mesh"]
+      pipeline.config = config
       if load_transformer:
           pipeline.transformer = cls.quantize_transformer(config, pipeline.transformer, pipeline, pipeline.mesh)
       return pipeline, pipeline.transformer
@@ -1204,6 +1205,16 @@ def __call__(
           prompt_attention_mask_jax = jnp.concatenate([negative_prompt_attention_mask_jax, prompt_attention_mask_jax], axis=0)
           latents_jax = jnp.concatenate([latents_jax] * 2, axis=0)
           audio_latents_jax = jnp.concatenate([audio_latents_jax] * 2, axis=0)
+          
+      if hasattr(self, "mesh") and self.mesh is not None:
+          data_sharding = NamedSharding(self.mesh, P())
+          if hasattr(self, "config") and hasattr(self.config, "data_sharding"):
+              data_sharding = NamedSharding(self.mesh, P(*self.config.data_sharding))
+          if isinstance(prompt_embeds_jax, list):
+              prompt_embeds_jax = [jax.device_put(x, data_sharding) for x in prompt_embeds_jax]
+          else:
+              prompt_embeds_jax = jax.device_put(prompt_embeds_jax, data_sharding)
+          prompt_attention_mask_jax = jax.device_put(prompt_attention_mask_jax, data_sharding)
 
       # GraphDef and State
       graphdef, state = nnx.split(self.transformer)