Renaming VAE sharding axis to vae_spatial

Author: eltsai

src/maxdiffusion/configs/base_wan_14b.yml

@@ -43,6 +43,7 @@ activations_dtype: 'bfloat16'
 
 # Replicates vae across devices instead of using the model's sharding annotations for sharding.
 replicate_vae: False
+vae_spatial: -1 # default to total_device * 2 // (dp)
 
 # matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision
 # Options are "DEFAULT", "HIGH", "HIGHEST"

src/maxdiffusion/configs/base_wan_27b.yml

@@ -43,6 +43,7 @@ activations_dtype: 'bfloat16'
 
 # Replicates vae across devices instead of using the model's sharding annotations for sharding.
 replicate_vae: False
+vae_spatial: -1
 
 # matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision
 # Options are "DEFAULT", "HIGH", "HIGHEST"

src/maxdiffusion/generate_wan.py

@@ -166,6 +166,7 @@ def run(config, pipeline=None, filename_prefix=""):
   max_logging.log(f"hardware: {jax.devices()[0].platform}")
   max_logging.log(f"number of devices: {jax.device_count()}")
   max_logging.log(f"per_device_batch_size: {config.per_device_batch_size}")
+  max_logging.log(f"vae_spatial: {config.vae_spatial}")
   max_logging.log("============================================================")
 
   compile_time = time.perf_counter() - s0

src/maxdiffusion/models/attention_flax.py

@@ -255,7 +255,7 @@ def _tpu_flash_attention(
     kv_max_block_size = key.shape[1]
   else:
     kv_max_block_size = q_max_block_size
-    
+
   # ensure that for cross attention we override the block sizes.
   if flash_block_sizes and key.shape[1] == query.shape[1]:
     block_sizes = flash_block_sizes

src/maxdiffusion/models/vae_flax.py

@@ -28,7 +28,7 @@
 from ..configuration_utils import ConfigMixin, flax_register_to_config
 from ..utils import BaseOutput
 from .modeling_flax_utils import FlaxModelMixin
- 
+
 
 
 @flax.struct.dataclass

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -99,10 +99,10 @@ def __init__(
         self.mesh = mesh
 
         # Weight sharding (Kernel is sharded along output channels)
-        num_fsdp_devices = mesh.shape["fsdp"]
+        num_fsdp_devices = mesh.shape["vae_spatial"]
         kernel_sharding = (None, None, None, None, None)
         if out_channels % num_fsdp_devices == 0:
-            kernel_sharding = (None, None, None, None, "fsdp")
+            kernel_sharding = (None, None, None, None, "vae_spatial")
 
         self.conv = nnx.Conv(
             in_features=in_channels,
@@ -121,7 +121,7 @@ def __init__(
     def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) -> jax.Array:
         # Sharding Width (index 3)
         # Spec: (Batch, Time, Height, Width, Channels)
-        spatial_sharding = NamedSharding(self.mesh, P(None, None, None, "fsdp", None))
+        spatial_sharding = NamedSharding(self.mesh, P(None, None, None, "vae_spatial", None))
         x = jax.lax.with_sharding_constraint(x, spatial_sharding)
 
         current_padding = list(self._causal_padding)
@@ -1098,7 +1098,7 @@ def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache):
     iter_ = 1 + (t - 1) // 4
     enc_feat_map = feat_cache._enc_feat_map
 
-    spatial_sharding = NamedSharding(self.mesh, P(None, None, None, "fsdp", None))
+    spatial_sharding = NamedSharding(self.mesh, P(None, None, None, "vae_spatial", None))
 
     # First iteration (i=0): size 1
     chunk_0 = x[:, :1, ...]
@@ -1180,7 +1180,7 @@ def _decode(
 
     dec_feat_map = feat_cache._feat_map
     # NamedSharding for the Width axis (axis 3)
-    spatial_sharding = NamedSharding(self.mesh, P(None, None, None, "fsdp", None))
+    spatial_sharding = NamedSharding(self.mesh, P(None, None, None, "vae_spatial", None))
 
     # First chunk (i=0)
     chunk_in_0 = jax.lax.with_sharding_constraint(x[:, 0:1, ...], spatial_sharding)
@@ -1264,4 +1264,4 @@ def decode(
     decoded = self._decode(z, feat_cache).sample
     if not return_dict:
       return (decoded,)
-    return FlaxDecoderOutput(sample=decoded)
+    return FlaxDecoderOutput(sample=decoded)

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -54,6 +54,7 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
           scheduler_state=common_components["scheduler_state"],
           devices_array=common_components["devices_array"],
           mesh=common_components["mesh"],
+          vae_mesh=common_components["vae_mesh"],
           config=config,
         )
 

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -63,6 +63,7 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
           scheduler_state=common_components["scheduler_state"],
           devices_array=common_components["devices_array"],
           mesh=common_components["mesh"],
+          vae_mesh=common_components["vae_mesh"],
           config=config,
         )
     return pipeline, low_noise_transformer, high_noise_transformer

src/maxdiffusion/pyconfig.py

@@ -248,6 +248,13 @@ def user_init(raw_keys):
         _HyperParameters.calculate_global_batch_sizes(raw_keys["per_device_batch_size"])
     )
 
+    if getattr(raw_keys, "vae_spatial", -1) == -1 or "vae_spatial" in raw_keys and raw_keys["vae_spatial"] == -1:
+      total_device = len(jax.devices())
+      dp = raw_keys.get("ici_data_parallelism", 1) * raw_keys.get("dcn_data_parallelism", 1)
+      if dp == -1 or dp == 0:
+        dp = 1
+      raw_keys["vae_spatial"] = (total_device * 2) // dp
+
 
 def get_num_slices(raw_keys):
   if int(raw_keys["compile_topology_num_slices"]) > 0: