Merge pull request #407 from AI-Hypercomputer:ninatu/wan_vae_sharding_bug

PiperOrigin-RevId: 918635461

Author: Google-ML-Automation

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -674,11 +674,8 @@ def _create_common_components(cls, config, vae_only=False, i2v=False):
     vae_spatial = getattr(config, "vae_spatial", -1)
     total_devices = math.prod(devices_array.shape)
 
-    if vae_spatial <= 0:
-      dp_size = mesh.shape.get("data", 1)
-      if dp_size == -1 or dp_size == 0:
-        dp_size = 1
-      vae_spatial = (2 * total_devices) // dp_size
+    if vae_spatial == -1:
+      vae_spatial = total_devices
 
     assert (
         total_devices % vae_spatial == 0

src/maxdiffusion/pyconfig.py

@@ -281,8 +281,8 @@ def user_init(raw_keys):
         raw_keys["global_batch_size_to_train_on"],
     ) = _HyperParameters.calculate_global_batch_sizes(raw_keys["per_device_batch_size"])
 
-    if raw_keys.get("vae_spatial", -1) == -1:
-      raw_keys["vae_spatial"] = 1
+    if "vae_spatial" not in raw_keys:
+      raw_keys["vae_spatial"] = -1
 
 
 def get_num_slices(raw_keys):