changed vae test logic

Author: eltsai

src/maxdiffusion/pyconfig.py

@@ -257,11 +257,7 @@ def user_init(raw_keys):
     ) = _HyperParameters.calculate_global_batch_sizes(raw_keys["per_device_batch_size"])
 
     if raw_keys.get("vae_spatial", -1) == -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
+      raw_keys["vae_spatial"] = 1
 
 
 def get_num_slices(raw_keys):

src/maxdiffusion/tests/wan_vae_test.py

@@ -340,7 +340,13 @@ def test_wan_residual(self):
     )
     config = pyconfig.config
     devices_array = create_device_mesh(config)
-    mesh = Mesh(devices_array, config.mesh_axes)
+    # Add vae_spatial axis to mesh for VAE operations
+    mesh_axes = list(config.mesh_axes)
+    if "vae_spatial" not in mesh_axes:
+      mesh_axes.append("vae_spatial")
+      # Reshape devices to include vae_spatial (size 1 for test)
+      devices_array = devices_array.reshape(devices_array.shape + (1,))
+    mesh = Mesh(devices_array, mesh_axes)
     # --- Test Case 1: same in/out dim ---
     in_dim = out_dim = 96
     batch = 1
@@ -392,7 +398,13 @@ def test_wan_midblock(self):
     )
     config = pyconfig.config
     devices_array = create_device_mesh(config)
-    mesh = Mesh(devices_array, config.mesh_axes)
+    # Add vae_spatial axis to mesh for VAE operations
+    mesh_axes = list(config.mesh_axes)
+    if "vae_spatial" not in mesh_axes:
+      mesh_axes.append("vae_spatial")
+      # Reshape devices to include vae_spatial (size 1 for test)
+      devices_array = devices_array.reshape(devices_array.shape + (1,))
+    mesh = Mesh(devices_array, mesh_axes)
     batch = 1
     t = 1
     dim = 384
@@ -417,7 +429,13 @@ def test_wan_decode(self):
     )
     config = pyconfig.config
     devices_array = create_device_mesh(config)
-    mesh = Mesh(devices_array, config.mesh_axes)
+    # Add vae_spatial axis to mesh for VAE operations
+    mesh_axes = list(config.mesh_axes)
+    if "vae_spatial" not in mesh_axes:
+      mesh_axes.append("vae_spatial")
+      # Reshape devices to include vae_spatial (size 1 for test)
+      devices_array = devices_array.reshape(devices_array.shape + (1,))
+    mesh = Mesh(devices_array, mesh_axes)
     dim = 96
     z_dim = 16
     dim_mult = [1, 2, 4, 4]
@@ -462,7 +480,13 @@ def test_wan_encode(self):
     )
     config = pyconfig.config
     devices_array = create_device_mesh(config)
-    mesh = Mesh(devices_array, config.mesh_axes)
+    # Add vae_spatial axis to mesh for VAE operations
+    mesh_axes = list(config.mesh_axes)
+    if "vae_spatial" not in mesh_axes:
+      mesh_axes.append("vae_spatial")
+      # Reshape devices to include vae_spatial (size 1 for test)
+      devices_array = devices_array.reshape(devices_array.shape + (1,))
+    mesh = Mesh(devices_array, mesh_axes)
     dim = 96
     z_dim = 16
     dim_mult = [1, 2, 4, 4]
@@ -508,7 +532,13 @@ def vae_encode(video, wan_vae, vae_cache, key):
     )
     config = pyconfig.config
     devices_array = create_device_mesh(config)
-    mesh = Mesh(devices_array, config.mesh_axes)
+    # Add vae_spatial axis to mesh for VAE operations
+    mesh_axes = list(config.mesh_axes)
+    if "vae_spatial" not in mesh_axes:
+      mesh_axes.append("vae_spatial")
+      # Reshape devices to include vae_spatial (size 1 for test)
+      devices_array = devices_array.reshape(devices_array.shape + (1,))
+    mesh = Mesh(devices_array, mesh_axes)
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
       wan_vae = AutoencoderKLWan.from_config(config.pretrained_model_name_or_path, subfolder="vae", rngs=rngs, mesh=mesh)
       vae_cache = AutoencoderKLWanCache(wan_vae)