diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index f147e4414..31a2d7f90 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -16,6 +16,7 @@ from typing import Tuple, Optional, Dict, Union, Any import contextlib +import functools import math import jax import jax.numpy as jnp @@ -38,6 +39,90 @@ from ...gradient_checkpoint import GradientCheckpointType BlockSizes = common_types.BlockSizes +BATCH = common_types.BATCH + + +def _batch_sharding(mesh, ndim: int, batch_axis: int, batch_size: Optional[int] = None): + """NamedSharding with the activation-batch mesh axes on ``batch_axis``, rest replicated. + + Returns None (caller skips the constraint, array stays replicated) when there + is no mesh, when the ambient ``axis_rules`` context does not map BATCH to any + mesh axis, or when ``batch_size`` is not divisible by the product of the + batch mesh axes -- e.g. ``data_parallelism > 1`` with ``per_device_batch_size + = 1`` and CFG off, which the pipelines already tolerate for the CFG cache. + ``with_sharding_constraint`` would otherwise raise inside the jit. + """ + if mesh is None: + return None + batch_mesh_axes = nn.logical_to_mesh_axes((BATCH,))[0] + if batch_mesh_axes is None: + return None + if batch_size is not None: + axes = (batch_mesh_axes,) if isinstance(batch_mesh_axes, str) else tuple(batch_mesh_axes) + num_shards = math.prod(mesh.shape[a] for a in axes) + if batch_size % num_shards != 0: + return None + spec = [None] * ndim + spec[batch_axis] = batch_mesh_axes + return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) + + +def _shard_kv_cache_on_batch(kv_cache, mesh): + """Pins every K/V cache array to the activation batch sharding. + + Cache arrays are either [B, S, H*D] (one layer) or [L, B, S, H*D] (stacked), + so the batch axis is the first non-layer axis. The cross-attention consumer + re-constrains K/V to (BATCH, CROSS_ATTN_HEAD, None, D_KV), i.e. batch on + ('data', 'fsdp') and everything else replicated, so this is exactly the + layout it needs — no resharding at use time, and each device only holds its + own B_local slice instead of the full CFG-doubled batch. + """ + if mesh is None: + return kv_cache + + def _constrain(x): + batch_axis = 1 if x.ndim == 4 else 0 + sharding = _batch_sharding(mesh, x.ndim, batch_axis, batch_size=x.shape[batch_axis]) + return x if sharding is None else jax.lax.with_sharding_constraint(x, sharding) + + return jax.tree.map(_constrain, kv_cache) + + +@functools.partial(jax.jit, static_argnames=("num_layers", "mesh")) +def _scanned_compute_kv(blocks_graphdef, blocks_state, encoder_hidden_states, encoder_attention_mask, num_layers, mesh): + """Cross-attention K/V cache for all layers as ONE sharded executable. + + Replaces the eager ``nnx.vmap`` over the stacked blocks. Eager vmap + materialises the K/V projections, the FP32 RMSNorm upcast and the bf16 + downcast for all ``num_layers`` at once *and* on every device (the inputs + are replicated), which is what blew the 32 GiB parts at B=8. + + Here the batch axis is sharded on ('data', 'fsdp') before the projections, + and ``nnx.scan`` processes one layer per iteration so the FP32 temporaries + are bounded to a single layer and freed before the next. The stacked + [L, B, S, H*D] outputs stay batch-sharded. + """ + blocks = nnx.merge(blocks_graphdef, blocks_state) + batch_size = encoder_hidden_states.shape[0] + in_sharding = _batch_sharding(mesh, encoder_hidden_states.ndim, 0, batch_size=batch_size) + if in_sharding is not None: + encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, in_sharding) + if encoder_attention_mask is not None: + encoder_attention_mask = jax.lax.with_sharding_constraint( + encoder_attention_mask, _batch_sharding(mesh, encoder_attention_mask.ndim, 0, batch_size=batch_size) + ) + + def scan_fn(carry, block): + kv = block.compute_kv(encoder_hidden_states, encoder_attention_mask) + return carry, _shard_kv_cache_on_batch(kv, mesh) + + _, kv_cache = nnx.scan( + scan_fn, + length=num_layers, + in_axes=(nnx.Carry, 0), + out_axes=(nnx.Carry, 0), + )(None, blocks) + return _shard_kv_cache_on_batch(kv_cache, mesh) def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int): @@ -588,6 +673,7 @@ def __init__( out_channels = out_channels or in_channels self.num_layers = num_layers self.scan_layers = scan_layers + self.mesh = mesh self.enable_jax_named_scopes = enable_jax_named_scopes attention_config = { "use_base2_exp": False, @@ -741,16 +827,17 @@ def compute_kv_cache( encoder_attention_mask = jnp.concatenate([encoder_attention_mask, text_mask], axis=1) if self.scan_layers: - - @nnx.vmap( - in_axes=(0, None, None), - out_axes=0, - transform_metadata={nnx.PARTITION_NAME: "layers_per_stage"}, + # Jitted + scanned + batch-sharded; see _scanned_compute_kv for why the + # previous eager nnx.vmap was a 6.8 GiB-per-video HBM hit on every device. + blocks_graphdef, blocks_state = nnx.split(self.blocks) + kv_cache = _scanned_compute_kv( + blocks_graphdef, + blocks_state, + encoder_hidden_states, + encoder_attention_mask, + num_layers=self.num_layers, + mesh=self.mesh, ) - def _compute_kv(block, enc_states, enc_mask): - return block.compute_kv(enc_states, enc_mask) - - kv_cache = _compute_kv(self.blocks, encoder_hidden_states, encoder_attention_mask) else: kv_cache_list = [] for block in self.blocks: @@ -761,6 +848,7 @@ def _compute_kv(block, enc_states, enc_mask): k_list = [d[k][0] for d in kv_cache_list] v_list = [d[k][1] for d in kv_cache_list] kv_cache[k] = (jnp.stack(k_list, axis=0), jnp.stack(v_list, axis=0)) + kv_cache = _shard_kv_cache_on_batch(kv_cache, self.mesh) return kv_cache, encoder_attention_mask diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index 4006e2fd1..55e616302 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -994,6 +994,45 @@ def _decode_latents_to_video(self, latents: jax.Array, trace: Optional[dict] = N video = np.asarray(video) return video + @staticmethod + def _auto_vae_spatial(config, total_devices: int, batch_parallel_vae: bool = False) -> int: + """Picks the `vae_spatial` mesh extent when the config leaves it at -1 (auto). + + The VAE mesh is ('redundant', 'vae_spatial') with redundant * vae_spatial == + total_devices. Batch is sharded on 'redundant' and width / conv channels on + 'vae_spatial'. Spatial sharding is what lets a single video fit, but with + vae_spatial == total_devices the batch axis is fully replicated, so at + B == total_devices every chip holds all B videos' activations and the + causal conv stack pays a ~10x slowdown (41.6s -> 6.3s at B=8 on tpu7x-8) + and OOMs on 32 GiB parts. + + Measured on v6e-8 / tpu7x-8 (720p, 81 frames): + * B % total_devices == 0: pure data parallelism (vae_spatial == 1) is + strictly better -- 1 video per device, zero cross-chip collectives. + * 1 < B < total_devices: mixed meshes are NOT a win. At B=2 on v6e-8 + the (2, 4) mesh needed a 21.2 GiB decode program vs 15.5 GiB for + (1, 8), even though per-device activation volume is identical. + So this only switches to vae_spatial == 1 when the batch is a multiple of + the device count and keeps full spatial sharding otherwise. Setting + `vae_spatial` explicitly in the config still overrides this. + + `batch_parallel_vae` gates the batch-aware choice. The VAE hard-constrains + every activation to P('redundant', ..., 'vae_spatial', ...), so a + vae_spatial == 1 mesh requires *every* tensor that goes through the VAE to + have a batch divisible by total_devices. That holds for the T2V pipelines + (the only VAE call is the final decode of the full batch), but not for + pipelines that also encode conditioning through the VAE at a different + batch -- VACE encodes reference images one at a time ([1, 1, H, W, C]), + animate / i2v encode conditioning clips. Those pipelines leave the gate + off and keep the full-spatial (1, N) mesh. + """ + if not batch_parallel_vae: + return total_devices + batch_size = int(getattr(config, "global_batch_size_to_train_on", 1) or 1) + if batch_size >= total_devices and batch_size % total_devices == 0: + return 1 + return total_devices + @classmethod def _create_common_components( cls, @@ -1002,6 +1041,7 @@ def _create_common_components( load_text_encoder=True, load_scheduler=True, i2v=False, + batch_parallel_vae=False, ): devices_array = max_utils.create_device_mesh(config) mesh = Mesh(devices_array, config.mesh_axes) @@ -1010,7 +1050,7 @@ def _create_common_components( total_devices = math.prod(devices_array.shape) if vae_spatial == -1: - vae_spatial = total_devices + vae_spatial = cls._auto_vae_spatial(config, total_devices, batch_parallel_vae=batch_parallel_vae) assert ( total_devices % vae_spatial == 0 @@ -1021,7 +1061,8 @@ def _create_common_components( vae_mesh = Mesh(vae_devices_array, ("redundant", "vae_spatial")) max_logging.log( - f"Created VAE specific mesh with axes ('redundant', 'vae_spatial') to support spatial sharding of {vae_spatial}." + f"Created VAE specific mesh with axes ('redundant', 'vae_spatial') = " + f"({total_devices // vae_spatial}, {vae_spatial}) to support spatial sharding of {vae_spatial}." ) # logical axis rules for VAE encoding/decoding diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py index 0cd0d7f7b..f80ca786a 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py @@ -68,6 +68,9 @@ def _load_and_init( load_vae=load_vae, load_text_encoder=load_text_encoder, load_scheduler=load_scheduler, + # T2V only runs the VAE on the final full-batch decode, so the + # batch-aware (N, 1) mesh is safe here; see _auto_vae_spatial. + batch_parallel_vae=True, ) transformer = None try: diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py index f55bd79e5..109549dde 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py @@ -79,6 +79,9 @@ def _load_and_init( load_vae=load_vae, load_text_encoder=load_text_encoder, load_scheduler=load_scheduler, + # T2V only runs the VAE on the final full-batch decode, so the + # batch-aware (N, 1) mesh is safe here; see _auto_vae_spatial. + batch_parallel_vae=True, ) low_noise_transformer, high_noise_transformer = None, None transformer_executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) diff --git a/src/maxdiffusion/tests/wan/wan_vae_mesh_test.py b/src/maxdiffusion/tests/wan/wan_vae_mesh_test.py new file mode 100644 index 000000000..2e2f63ed9 --- /dev/null +++ b/src/maxdiffusion/tests/wan/wan_vae_mesh_test.py @@ -0,0 +1,132 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +CPU checks for the batch-aware VAE mesh selection and the batch-sharded +cross-attention KV cache helpers. +""" + +import os +import types + +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P + +from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline +from maxdiffusion.models.wan.transformers import transformer_wan + + +def _cfg(batch_size): + return types.SimpleNamespace(global_batch_size_to_train_on=batch_size) + + +@pytest.mark.parametrize( + "batch_size, total_devices, expected", + [ + (0, 8, 8), + (1, 8, 8), + (2, 8, 8), + (3, 8, 8), + (4, 8, 8), + (8, 8, 1), + (12, 8, 8), + (16, 8, 1), + (4, 4, 1), + (1, 1, 1), + ], +) +def test_auto_vae_spatial_batch_parallel(batch_size, total_devices, expected): + """With the gate on, B % N == 0 picks pure data parallelism, otherwise full spatial.""" + assert WanPipeline._auto_vae_spatial(_cfg(batch_size), total_devices, batch_parallel_vae=True) == expected + + +@pytest.mark.parametrize("batch_size", [0, 1, 2, 8, 16]) +def test_auto_vae_spatial_default_is_full_spatial(batch_size): + """Without the gate (i2v / vace / animate) auto always resolves to the old value.""" + assert WanPipeline._auto_vae_spatial(_cfg(batch_size), 8) == 8 + assert WanPipeline._auto_vae_spatial(_cfg(batch_size), 8, batch_parallel_vae=False) == 8 + + +def test_only_t2v_pipelines_enable_batch_parallel_vae(): + """Guard against the gate leaking into pipelines that VAE-encode at batch != global batch.""" + pipeline_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "pipelines", "wan") + enabled = set() + for name in sorted(os.listdir(pipeline_dir)): + if not name.endswith(".py"): + continue + with open(os.path.join(pipeline_dir, name)) as f: + if "batch_parallel_vae=True" in f.read(): + enabled.add(name) + assert enabled == {"wan_pipeline_2_1.py", "wan_pipeline_2_2.py"} + + +@pytest.fixture +def mesh_and_rules(): + all_devices = jax.devices() + if len(all_devices) < 2: + pytest.skip("Requires at least 2 JAX devices to test multi-shard batch sharding") + usable = (len(all_devices) // 2) * 2 + devices = np.array(all_devices[:usable]).reshape(2, 1, -1, 1) + mesh = Mesh(devices, ("data", "fsdp", "context", "tensor")) + rules = (("activation_batch", ("data", "fsdp")),) + return mesh, rules + + +def test_batch_sharding_divisible(mesh_and_rules): + mesh, rules = mesh_and_rules + with mesh, nn_partitioning.axis_rules(rules): + sharding = transformer_wan._batch_sharding(mesh, ndim=3, batch_axis=0, batch_size=4) + assert sharding is not None + assert sharding.spec == P(("data", "fsdp"), None, None) + + +def test_batch_sharding_not_divisible_returns_none(mesh_and_rules): + """B=3 on data*fsdp=2 must skip the constraint instead of raising inside the jit.""" + mesh, rules = mesh_and_rules + with mesh, nn_partitioning.axis_rules(rules): + assert transformer_wan._batch_sharding(mesh, ndim=3, batch_axis=0, batch_size=3) is None + + +def test_batch_sharding_without_rules_or_mesh(mesh_and_rules): + mesh, _ = mesh_and_rules + assert transformer_wan._batch_sharding(None, ndim=3, batch_axis=0) is None + with mesh, nn_partitioning.axis_rules(()): + assert transformer_wan._batch_sharding(mesh, ndim=3, batch_axis=0) is None + + +def test_shard_kv_cache_on_batch_layouts(mesh_and_rules): + """[B,S,D] shards on axis 0, stacked [L,B,S,D] on axis 1; odd batches stay replicated.""" + mesh, rules = mesh_and_rules + cache = { + "text": (jnp.zeros((4, 16, 8)), jnp.zeros((2, 4, 16, 8))), + "odd": (jnp.zeros((3, 16, 8)), jnp.zeros((2, 3, 16, 8))), + } + with mesh, nn_partitioning.axis_rules(rules): + out = jax.jit(lambda c: transformer_wan._shard_kv_cache_on_batch(c, mesh))(cache) + k, v = out["text"] + # JAX canonicalises the spec of the returned array (size-1 'fsdp' is dropped), + # so compare semantically rather than by literal PartitionSpec. + assert k.sharding.is_equivalent_to(NamedSharding(mesh, P(("data", "fsdp"), None, None)), k.ndim) + assert v.sharding.is_equivalent_to(NamedSharding(mesh, P(None, ("data", "fsdp"), None, None)), v.ndim) + assert k.addressable_shards[0].data.shape[0] == 2 + assert v.addressable_shards[0].data.shape[1] == 2 + k_odd, v_odd = out["odd"] + assert k_odd.shape == (3, 16, 8) and v_odd.shape == (2, 3, 16, 8) + assert k_odd.sharding.is_equivalent_to(NamedSharding(mesh, P()), k_odd.ndim)