Skip to content

Fall back to sharding the height in the WAN VAE. #521

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ninatu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fall back to sharding the height in the WAN VAE. #521

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 18 additions & 12 deletions src/maxdiffusion/models/wan/autoencoder_kl_wan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,21 @@ def _update_cache(cache, idx, value):
return cache[:idx] + (value,) + cache[idx + 1 :]


def _spatial_sharding(mesh: Optional[jax.sharding.Mesh], shape: Sequence[int]) -> Optional[NamedSharding]:
"""Returns the sharding of a (N, D, H, W, C) array over the vae_spatial axis.

Shards the width, or the height if only the height is divisible by
vae_spatial. If neither is, the width is still sharded and XLA pads it.
"""
if mesh is None or "vae_spatial" not in mesh.shape:
return None
_, _, height, width, _ = shape
Comment on lines +57 to +59

@gemini-code-assist gemini-code-assist Bot Nov 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To prevent potential ValueError during unpacking, we should defensively check that the input shape is indeed 5D (i.e., has a length of 5) before attempting to unpack it. If it is not 5D, we can safely return None to skip applying the sharding constraint.

Suggested change
if mesh is None or "vae_spatial" not in mesh.shape:
return None
_, _, height, width, _ = shape
if mesh is None or "vae_spatial" not in mesh.shape or len(shape) != 5:
return None
_, _, height, width, _ = shape

num_shards = mesh.shape["vae_spatial"]
if width % num_shards != 0 and height % num_shards == 0:
return NamedSharding(mesh, P("redundant", None, "vae_spatial", None, None))
return NamedSharding(mesh, P("redundant", None, None, "vae_spatial", None))


# Helper to ensure kernel_size, stride, padding are tuples of 3 integers
def _canonicalize_tuple(x: Union[int, Sequence[int]], rank: int, name: str) -> Tuple[int, ...]:
"""Canonicalizes a value to a tuple of integers."""
Expand Down Expand Up @@ -121,7 +136,7 @@ def __init__(

def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) -> jax.Array:
if self.mesh is not None and "vae_spatial" in self.mesh.shape:
spatial_sharding = NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
spatial_sharding = _spatial_sharding(self.mesh, x.shape)
if spatial_sharding is not None:
x = jax.lax.with_sharding_constraint(x, spatial_sharding)

Expand Down Expand Up @@ -1169,12 +1184,6 @@ def _encode(self, x: jax.Array, feat_cache: AutoencoderKLWanCache):
iter_ = 1 + ((t - 1 + CHUNK_SIZE - 1) // CHUNK_SIZE) if t > 1 else 1
enc_feat_map = feat_cache._enc_feat_map

spatial_sharding = (
NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
if self.mesh is not None and "vae_spatial" in self.mesh.shape
else None
)

def finalize(out, enc_feat_map):
feat_cache._enc_feat_map = enc_feat_map
enc = self.quant_conv(out)
Expand All @@ -1184,6 +1193,7 @@ def finalize(out, enc_feat_map):
with jax.named_scope("AutoencoderKLWan_encode_chunk_0"):
chunk_0 = x[:, :1, ...]
out_0, enc_feat_map, _ = self.encoder(chunk_0, feat_cache=enc_feat_map, feat_idx=0)
spatial_sharding = _spatial_sharding(self.mesh, out_0.shape)
if spatial_sharding is not None:
out_0 = jax.lax.with_sharding_constraint(out_0, spatial_sharding)

Expand Down Expand Up @@ -1280,11 +1290,7 @@ def _decode(
x = self.post_quant_conv(z)

dec_feat_map = feat_cache._feat_map
spatial_sharding = (
NamedSharding(self.mesh, P("redundant", None, None, "vae_spatial", None))
if self.mesh is not None and "vae_spatial" in self.mesh.shape
else None
)
spatial_sharding = _spatial_sharding(self.mesh, x.shape)

@gemini-code-assist gemini-code-assist Bot Nov 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There is a critical bug in _decode where if spatial_sharding is None (e.g., when running on a single device without a mesh), chunk_in_0 (line 1298) and chunk_in_1 (line 1307) are never defined. This will raise a NameError when they are passed to self.decoder on lines 1299 and 1308.

To fix this, you should define chunk_in_0 and chunk_in_1 unconditionally and apply the sharding constraint conditionally. For example:

# First chunk (i=0)
with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):
  chunk_in_0 = x[:, 0:1, ...]
  if spatial_sharding is not None:
    chunk_in_0 = jax.lax.with_sharding_constraint(chunk_in_0, spatial_sharding)
  out_0, dec_feat_map, _ = self.decoder(chunk_in_0, feat_cache=dec_feat_map, feat_idx=0)

And similarly for chunk_in_1.


# First chunk (i=0)
with jax.named_scope("AutoencoderKLWan_decode_chunk_0"):
Expand Down
Loading