adding vae spatial axis to conv_in

Author: eltsai

src/maxdiffusion/configs/base_wan_14b.yml

@@ -192,6 +192,7 @@ vae_logical_axis_rules: [
                           ['conv_batch', 'redundant'],
                           ['out_channels', 'vae_spatial'],
                           ['conv_out', 'vae_spatial'],
+                          ['conv_in', 'vae_spatial'],
                         ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -168,6 +168,7 @@ vae_logical_axis_rules: [
                           ['conv_batch', 'redundant'],
                           ['out_channels', 'vae_spatial'],
                           ['conv_out', 'vae_spatial'],
+                          ['conv_in', 'vae_spatial'],
                         ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 

src/maxdiffusion/configs/base_wan_27b.yml

@@ -180,6 +180,7 @@ vae_logical_axis_rules: [
                           ['conv_batch', 'redundant'],
                           ['out_channels', 'vae_spatial'],
                           ['conv_out', 'vae_spatial'],
+                          ['conv_in', 'vae_spatial'],
                         ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -174,6 +174,7 @@ vae_logical_axis_rules: [
                           ['conv_batch', 'redundant'],
                           ['out_channels', 'vae_spatial'],
                           ['conv_out', 'vae_spatial'],
+                          ['conv_in', 'vae_spatial'],
                         ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -175,6 +175,7 @@ vae_logical_axis_rules: [
                           ['conv_batch', 'redundant'],
                           ['out_channels', 'vae_spatial'],
                           ['conv_out', 'vae_spatial'],
+                          ['conv_in', 'vae_spatial'],
                         ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 

src/maxdiffusion/tests/wan_vae_test.py

@@ -21,8 +21,8 @@
 import jax
 import jax.numpy as jnp
 from flax import nnx
-from flax import linen as nn
 from flax.linen import partitioning as nn_partitioning
+from flax.linen import logical_to_mesh_sharding
 from jax.sharding import Mesh
 from .. import pyconfig
 from ..max_utils import (
@@ -224,7 +224,7 @@ def test_zero_padded_conv(self):
     output_torch = resample(input)
     assert output_torch.shape == (1, 96, 240, 360)
 
-    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+    with self.mesh, nn_partitioning.axis_rules(self.config.vae_logical_axis_rules):
       model = ZeroPaddedConv2D(dim=dim, rngs=rngs, kernel_size=(1, 3, 3), stride=(1, 2, 2))
       dummy_input = jnp.ones(input_shape)
       dummy_input = jnp.transpose(dummy_input, (0, 2, 3, 1))
@@ -262,7 +262,7 @@ def test_wan_resample(self):
     torch_wan_resample = TorchWanResample(dim=dim, mode=mode)
     torch_output = torch_wan_resample(dummy_input)
     assert torch_output.shape == (batch, dim, t, h // 2, w // 2)
-    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+    with self.mesh, nn_partitioning.axis_rules(self.config.vae_logical_axis_rules):
       wan_resample = WanResample(dim, mode=mode, rngs=rngs)
       # channels is always last here
       input_shape = (batch, t, h, w, dim)
@@ -305,7 +305,7 @@ def test_3d_conv(self):
     dummy_cache = jnp.zeros((batch_size, cache_depth, in_height, in_width, in_channels))
 
     # Instantiate the module
-    with self.mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    with self.mesh, nn_partitioning.axis_rules(config.vae_logical_axis_rules):
       causal_conv_layer = WanCausalConv3d(
           in_channels=in_channels,
           out_channels=out_channels,
@@ -357,7 +357,7 @@ def test_wan_residual(self):
     dim = 96
     input_shape = (batch, t, height, width, dim)
     expected_output_shape = (batch, t, height, width, dim)
-    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    with mesh, nn_partitioning.axis_rules(config.vae_logical_axis_rules):
       wan_residual_block = WanResidualBlock(in_dim=in_dim, out_dim=out_dim, rngs=rngs, mesh=mesh)
       dummy_input = jnp.ones(input_shape)
       dummy_output, _, _ = wan_residual_block(dummy_input)
@@ -381,7 +381,7 @@ def test_wan_attention(self):
     height = 60
     width = 90
     input_shape = (batch, t, height, width, dim)
-    with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+    with self.mesh, nn_partitioning.axis_rules(self.config.vae_logical_axis_rules):
       wan_attention = WanAttentionBlock(dim=dim, rngs=rngs)
       dummy_input = jnp.ones(input_shape)
       output, _, _ = wan_attention(dummy_input)
@@ -412,7 +412,7 @@ def test_wan_midblock(self):
     height = 60
     width = 90
     input_shape = (batch, t, height, width, dim)
-    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    with mesh, nn_partitioning.axis_rules(config.vae_logical_axis_rules):
       wan_midblock = WanMidBlock(dim=dim, rngs=rngs, mesh=mesh)
       dummy_input = jnp.ones(input_shape)
       output, _, _ = wan_midblock(dummy_input)
@@ -443,7 +443,7 @@ def test_wan_decode(self):
     num_res_blocks = 2
     attn_scales = []
     temperal_downsample = [False, True, True]
-    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    with mesh, nn_partitioning.axis_rules(config.vae_logical_axis_rules):
       wan_vae = AutoencoderKLWan(
           rngs=rngs,
           base_dim=dim,
@@ -494,7 +494,7 @@ def test_wan_encode(self):
     num_res_blocks = 2
     attn_scales = []
     temperal_downsample = [False, True, True]
-    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    with mesh, nn_partitioning.axis_rules(config.vae_logical_axis_rules):
       wan_vae = AutoencoderKLWan(
           rngs=rngs,
           base_dim=dim,
@@ -540,7 +540,7 @@ def vae_encode(video, wan_vae, vae_cache, key):
       # 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):
+    with self.mesh, nn_partitioning.axis_rules(self.config.vae_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)
     video_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4"
@@ -559,7 +559,7 @@ def vae_encode(video, wan_vae, vae_cache, key):
     params = jax.tree_util.tree_map(lambda x: x.astype(jnp.bfloat16), params)
 
     logical_state_spec = nnx.get_partition_spec(state)
-    logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
+    logical_state_sharding = logical_to_mesh_sharding(logical_state_spec, mesh, config.vae_logical_axis_rules)
     logical_state_sharding = dict(nnx.to_flat_state(logical_state_sharding))
 
     state_flat = dict(nnx.to_flat_state(state))