add shardings to projection and patch embedding.

Author: entrpn

src/maxdiffusion/configs/base_wan_14b.yml

@@ -136,7 +136,7 @@ logical_axis_rules: [
                       ['norm', 'tensor'],
                       ['conv_batch', ['data','fsdp']],
                       ['out_channels', 'tensor'],
-                      ['conv_in', 'fsdp'],
+                      ['conv_out', 'fsdp'],
                     ]
 data_sharding: [['data', 'fsdp', 'tensor']]
 

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -171,6 +171,13 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
+        kernel_init=nnx.with_partitioning(
+            nnx.initializers.xavier_uniform(),
+            (
+                "mlp",
+                "embed",
+            ),
+        ),
     )
 
   def __call__(self, x: jax.Array) -> jax.Array:
@@ -374,6 +381,16 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
+        kernel_init=nnx.with_partitioning(
+            nnx.initializers.xavier_uniform(),
+            (
+                None,
+                None,
+                None,
+                None,
+                "conv_out"
+            ),
+        ),
     )
 
     # 2. Condition embeddings

src/maxdiffusion/trainers/wan_trainer.py

@@ -80,7 +80,7 @@ def __init__(self, config):
       raise ValueError("this script currently doesn't support training text_encoders")
 
     #self.global_batch_size = self.config.per_device_batch_size * jax.device_count()
-    self.global_batch_size = config.global_batch_size if config.global_batch_size > 0 else config.per_device_batch_size * jax.device_count()
+    self.global_batch_size = config.per_device_batch_size * jax.device_count()
 
   def post_training_steps(self, pipeline, params, train_states, msg=""):
     pass
@@ -97,7 +97,7 @@ def calculate_tflops(self, pipeline):
     return 0
 
   def get_data_shardings(self, mesh):
-    data_sharding = jax.sharding.NamedSharding(mesh, P(*self.config.data_sharding[0]))
+    data_sharding = jax.sharding.NamedSharding(mesh, P(*self.config.data_sharding))
     data_sharding = {
       "latents" : data_sharding,
       "encoder_hidden_states" : data_sharding