229 lines
9.2 KiB
Python
229 lines
9.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Numerical contracts for the Wan decoder streaming fast paths.
|
|
|
|
Two decoders are driven frame by frame through ``WanStreamingDecoder`` over two chunks of one session and
|
|
a fresh session (the persistent buffers' causal front frames go zeros -> cache -> zeros), on CPU: plain
|
|
diffusers modules in one process, and the spatially sharded wrappers across two gloo ranks with real halo
|
|
exchange, each in fp32 and under bf16 autocast (where the parameter cast applies). The exact level must be
|
|
bit-identical to the untouched decoder on every rank. The fused level deliberately rounds each fused
|
|
norm-SiLU result to the activation dtype before the next convolution, so its bf16 contract bounds both the
|
|
worst output error and the aggregate RMS error instead.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import os
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
from diffusers.models.autoencoders.autoencoder_kl_wan import (
|
|
WanCausalConv3d,
|
|
WanDecoder3d,
|
|
)
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders.wan_decoder_fast_path import (
|
|
install_wan_decoder_fast_path,
|
|
)
|
|
from vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard import (
|
|
WanDistCausalConv3d,
|
|
WanDistConv2d,
|
|
install_wan_spatial_shard_decode,
|
|
)
|
|
from vllm_omni.experimental.ar_diffusion.streaming_decode import WanStreamingDecoder
|
|
|
|
|
|
def _make_vae(seed: int = 0) -> SimpleNamespace:
|
|
torch.manual_seed(seed)
|
|
# Three levels: the default temporal-upsample list gives one upsample2d and one upsample3d resample.
|
|
decoder = WanDecoder3d(dim=8, z_dim=4, dim_mult=[1, 2, 4], num_res_blocks=1, attn_scales=[])
|
|
post_quant_conv = WanCausalConv3d(4, 4, 1)
|
|
for p in list(decoder.parameters()) + list(post_quant_conv.parameters()):
|
|
p.data.uniform_(-0.5, 0.5)
|
|
decoder.eval()
|
|
count = sum(1 for m in decoder.modules() if m.__class__.__name__ == "WanCausalConv3d")
|
|
return SimpleNamespace(
|
|
decoder=decoder,
|
|
post_quant_conv=post_quant_conv,
|
|
_cached_conv_counts={"decoder": count},
|
|
)
|
|
|
|
|
|
def _run(vae, latents: list[torch.Tensor], autocast: bool) -> list[torch.Tensor]:
|
|
decoder = WanStreamingDecoder(vae)
|
|
outs = []
|
|
ctx = torch.autocast("cpu", dtype=torch.bfloat16) if autocast else torch.no_grad()
|
|
with torch.no_grad(), ctx:
|
|
state = decoder.new_decode_state("a")
|
|
outs.append(decoder.decode_chunk(latents[0], state))
|
|
outs.append(decoder.decode_chunk(latents[1], state))
|
|
fresh = decoder.new_decode_state("b")
|
|
outs.append(decoder.decode_chunk(latents[2], fresh))
|
|
return [o.float() for o in outs]
|
|
|
|
|
|
def _latents() -> list[torch.Tensor]:
|
|
torch.manual_seed(1)
|
|
return [
|
|
torch.randn(1, 4, 2, 6, 10),
|
|
torch.randn(1, 4, 1, 6, 10),
|
|
torch.randn(1, 4, 2, 6, 10),
|
|
]
|
|
|
|
|
|
def _assert_fused_output_close(expected: torch.Tensor, actual: torch.Tensor, autocast: bool) -> None:
|
|
if not autocast:
|
|
torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-4)
|
|
return
|
|
|
|
# FusedWanRMSNormSiLU returns the activation dtype immediately, whereas the eager Wan norm's fp32
|
|
# affine result stays fp32 through SiLU and is rounded by the following convolution. Those intentional
|
|
# bf16 boundaries accumulate across the decoder, so a per-element relative tolerance is misleading near
|
|
# zero. Keep a strict quality budget on both the worst pixel and the output as a whole. Decoder outputs
|
|
# are tanh-bounded to [-1, 1], making these absolute limits directly meaningful.
|
|
error = (actual - expected).abs()
|
|
eps = torch.finfo(torch.bfloat16).eps
|
|
max_error = error.amax().item()
|
|
rms_error = error.square().mean().sqrt().item()
|
|
assert max_error <= 4 * eps, f"fused bf16 max error {max_error} exceeds {4 * eps}"
|
|
assert rms_error <= eps, f"fused bf16 RMS error {rms_error} exceeds {eps}"
|
|
|
|
|
|
def _check_pair(reference, candidate, autocast: bool, sharded: bool, level: str = "exact") -> None:
|
|
dtype = torch.bfloat16 if autocast else None
|
|
counts = install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level)
|
|
assert counts["level"] == level
|
|
assert counts["upsamples"] == 2
|
|
convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, torch.nn.Conv3d | torch.nn.Conv2d)) + 1
|
|
assert counts["conv_params_cast"] == (convs if autocast else 0)
|
|
dist_convs = sum(1 for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d))
|
|
assert counts["persistent_input_buffers"] == dist_convs
|
|
assert (dist_convs > 0) == sharded
|
|
if level == "fused":
|
|
assert (
|
|
counts["fused_norm_silu"]
|
|
== 2 * sum(1 for m in candidate.decoder.modules() if m.__class__.__name__ == "WanResidualBlock") + 1
|
|
)
|
|
assert counts["channels_last_convs"] == convs
|
|
assert counts["resample_forwards"] == 2
|
|
with pytest.raises(ValueError, match="already installed"):
|
|
install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level="exact")
|
|
# Idempotent: a second install is a no-op that reports the same counts.
|
|
assert install_wan_decoder_fast_path(candidate, conv_dtype=dtype, level=level) is counts
|
|
latents = _latents()
|
|
expected = _run(reference, latents, autocast)
|
|
actual = _run(candidate, latents, autocast)
|
|
assert len(expected) == len(actual) == 3
|
|
for e, a in zip(expected, actual):
|
|
assert e.shape == a.shape
|
|
if level == "exact":
|
|
assert torch.equal(e, a)
|
|
else:
|
|
_assert_fused_output_close(e, a, autocast)
|
|
if sharded:
|
|
bufs = [m._input_buf for m in candidate.decoder.modules() if isinstance(m, WanDistCausalConv3d | WanDistConv2d)]
|
|
assert all(b is not None for b in bufs)
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("autocast", [False, True])
|
|
def test_fast_path_is_bit_identical(autocast: bool) -> None:
|
|
reference = _make_vae()
|
|
candidate = copy.deepcopy(reference)
|
|
_check_pair(reference, candidate, autocast, sharded=False)
|
|
|
|
|
|
def _swap_norms_to_rmsnorm_vae(vae) -> int:
|
|
"""Ensure the decoder uses RMSNormVAE and return the number of matching norms.
|
|
|
|
Collecting pipeline tests can already apply patch_wan_rms_norm process-wide before this test runs.
|
|
"""
|
|
from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
|
|
|
|
count = 0
|
|
for module in list(vae.decoder.modules()):
|
|
for name, child in list(module.named_children()):
|
|
if isinstance(child, RMSNormVAE):
|
|
count += 1
|
|
continue
|
|
if child.__class__.__name__ == "WanRMS_norm":
|
|
images = child.gamma.dim() == 3 # (dim, 1, 1) for the attention block's 4D norm
|
|
norm = RMSNormVAE(
|
|
child.gamma.shape[0],
|
|
channel_first=child.channel_first,
|
|
images=images,
|
|
bias=False,
|
|
)
|
|
norm.gamma.data.copy_(child.gamma.data.reshape(norm.gamma.shape))
|
|
setattr(module, name, norm)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("norm", ["default", "RMSNormVAE"])
|
|
@pytest.mark.parametrize("autocast", [False, True])
|
|
def test_fused_level_plumbing_on_cpu(autocast: bool, norm: str) -> None:
|
|
reference = _make_vae()
|
|
if norm == "RMSNormVAE":
|
|
assert _swap_norms_to_rmsnorm_vae(reference) > 0
|
|
candidate = copy.deepcopy(reference)
|
|
_check_pair(reference, candidate, autocast, sharded=False, level="fused")
|
|
|
|
|
|
def _sharded_worker(
|
|
rank: int,
|
|
world_size: int,
|
|
port: str,
|
|
autocast: bool,
|
|
return_dict,
|
|
level: str = "exact",
|
|
) -> None:
|
|
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
|
os.environ["MASTER_PORT"] = port
|
|
dist.init_process_group("gloo", rank=rank, world_size=world_size)
|
|
try:
|
|
torch.manual_seed(0)
|
|
reference = _make_vae()
|
|
candidate = copy.deepcopy(reference)
|
|
for vae in (reference, candidate):
|
|
install_wan_spatial_shard_decode(vae, dist.group.WORLD, split_dim="width", dst=None)
|
|
_check_pair(reference, candidate, autocast, sharded=True, level=level)
|
|
return_dict[rank] = "ok"
|
|
except BaseException as exc: # report, the parent asserts
|
|
return_dict[rank] = f"{type(exc).__name__}: {exc}"
|
|
raise
|
|
finally:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("level", ["exact", "fused"])
|
|
@pytest.mark.parametrize("autocast", [False, True])
|
|
def test_sharded_fast_path_on_every_rank(autocast: bool, level: str) -> None:
|
|
manager = mp.get_context("spawn").Manager()
|
|
return_dict = manager.dict()
|
|
port = str(29610 + int(autocast) + 2 * (level == "fused"))
|
|
mp.spawn(
|
|
_sharded_worker,
|
|
args=(2, port, autocast, return_dict, level),
|
|
nprocs=2,
|
|
join=True,
|
|
)
|
|
for rank in range(2):
|
|
assert return_dict.get(rank) == "ok", f"rank {rank}: {return_dict.get(rank)}"
|
|
|
|
|
|
@pytest.mark.core_model
|
|
@pytest.mark.cpu
|
|
def test_conv_dtype_requires_half() -> None:
|
|
with pytest.raises(ValueError, match="conv_dtype"):
|
|
install_wan_decoder_fast_path(_make_vae(), conv_dtype=torch.float32)
|