sharding attempt

Author: prishajain1

src/maxdiffusion/configs/ltx2_video.yml

@@ -66,7 +66,6 @@ logical_axis_rules: [
                       ['conv_out', 'fsdp'],
                       ['conv_in', 'fsdp']
                     ]
-data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
 dcn_fsdp_parallelism: -1
 dcn_context_parallelism: 1

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -342,12 +342,24 @@ def __init__(
     self.dropout_rate = dropout
 
     # 1. Projections
-    self.to_q = nnx.Linear(query_dim, self.inner_dim, use_bias=bias, rngs=rngs, dtype=dtype)
+    self.to_q = nnx.Linear(
+        query_dim, self.inner_dim, use_bias=bias, rngs=rngs, dtype=dtype,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", "heads")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("heads",)),
+    )
 
     # Handle Self vs Cross Attention input dims
     kv_dim = context_dim if context_dim is not None else query_dim
-    self.to_k = nnx.Linear(kv_dim, self.inner_dim, use_bias=bias, rngs=rngs, dtype=dtype)
-    self.to_v = nnx.Linear(kv_dim, self.inner_dim, use_bias=bias, rngs=rngs, dtype=dtype)
+    self.to_k = nnx.Linear(
+        kv_dim, self.inner_dim, use_bias=bias, rngs=rngs, dtype=dtype,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", "heads")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("heads",)),
+    )
+    self.to_v = nnx.Linear(
+        kv_dim, self.inner_dim, use_bias=bias, rngs=rngs, dtype=dtype,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("embed", "heads")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("heads",)),
+    )
 
     # 2. Normalization (Applied to full inner_dim, NOT per-head)
     self.norm_q = nnx.RMSNorm(
@@ -358,7 +370,11 @@ def __init__(
     )
 
     # 3. Output
-    self.to_out = nnx.Linear(self.inner_dim, query_dim, use_bias=out_bias, rngs=rngs, dtype=dtype)
+    self.to_out = nnx.Linear(
+        self.inner_dim, query_dim, use_bias=out_bias, rngs=rngs, dtype=dtype,
+        kernel_init=nnx.with_partitioning(nnx.initializers.xavier_uniform(), ("heads", "embed")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+    )
 
     if self.dropout_rate > 0:
       self.dropout_layer = nnx.Dropout(self.dropout_rate, rngs=rngs)

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -57,8 +57,8 @@ def __init__(
         use_bias=True,
         dtype=dtype,
         param_dtype=weights_dtype,
-        kernel_init=nnx.initializers.zeros,
-        bias_init=nnx.initializers.zeros,
+        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed", "embed")),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
     )
 
   def __call__(
@@ -291,12 +291,22 @@ def __init__(
     key = rngs.params()
     k1, k2, k3, k4 = jax.random.split(key, 4)
 
-    self.scale_shift_table = nnx.Param(jax.random.normal(k1, (6, self.dim), dtype=weights_dtype) / jnp.sqrt(self.dim))
+    self.scale_shift_table = nnx.Param(
+        jax.random.normal(k1, (6, self.dim), dtype=weights_dtype) / jnp.sqrt(self.dim),
+        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed")),
+    )
     self.audio_scale_shift_table = nnx.Param(
-        jax.random.normal(k2, (6, audio_dim), dtype=weights_dtype) / jnp.sqrt(audio_dim)
+        jax.random.normal(k2, (6, audio_dim), dtype=weights_dtype) / jnp.sqrt(audio_dim),
+        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed")),
+    )
+    self.video_a2v_cross_attn_scale_shift_table = nnx.Param(
+        jax.random.normal(k3, (5, self.dim), dtype=weights_dtype),
+        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed")),
+    )
+    self.audio_a2v_cross_attn_scale_shift_table = nnx.Param(
+        jax.random.normal(k4, (5, audio_dim), dtype=weights_dtype),
+        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed")),
     )
-    self.video_a2v_cross_attn_scale_shift_table = nnx.Param(jax.random.normal(k3, (5, self.dim), dtype=weights_dtype))
-    self.audio_a2v_cross_attn_scale_shift_table = nnx.Param(jax.random.normal(k4, (5, audio_dim), dtype=weights_dtype))
 
   def __call__(
       self,