238 lines
11 KiB
Python
238 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Exercise Wan's real name mapping and vLLM's QKV parameter loaders on CPU."""
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.config.load import LoadConfig
|
|
from vllm.model_executor.layers.linear import QKVParallelLinear
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
|
|
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, DiffusionMXFP4DualScaleMixedConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_tp_state(monkeypatch):
|
|
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
|
|
prefix = "vllm_omni.diffusion.models.wan2_2.wan2_2_transformer"
|
|
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(f"{prefix}.get_tensor_model_parallel_world_size", lambda: 1)
|
|
|
|
|
|
class _WanCheckpoint(torch.nn.Module):
|
|
load_weights = WanTransformer3DModel.load_weights
|
|
_log_w4a8_fallback_load_summary = WanTransformer3DModel._log_w4a8_fallback_load_summary
|
|
|
|
def __init__(self, config, *, num_blocks=1):
|
|
super().__init__()
|
|
blocks = []
|
|
for block_idx in range(num_blocks):
|
|
block = torch.nn.Module()
|
|
block.attn1 = torch.nn.Module()
|
|
block.attn1.to_qkv = QKVParallelLinear(
|
|
hidden_size=512,
|
|
head_size=32,
|
|
total_num_heads=4,
|
|
bias=False,
|
|
params_dtype=torch.bfloat16,
|
|
quant_config=config,
|
|
prefix=f"blocks.{block_idx}.attn1.to_qkv",
|
|
disable_tp=True,
|
|
)
|
|
blocks.append(block)
|
|
self.blocks = torch.nn.ModuleList(blocks)
|
|
|
|
|
|
def _checkpoint(mode, fused):
|
|
if mode == "single":
|
|
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
|
|
elif mode == "dual":
|
|
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
|
|
else:
|
|
config = None
|
|
model = _WanCheckpoint(config)
|
|
weights = {}
|
|
expected = {}
|
|
for name, param in model.named_parameters():
|
|
suffix = name.rsplit(".", 1)[1]
|
|
output_dim = getattr(param, "output_dim", None)
|
|
if output_dim is None:
|
|
value = torch.full_like(param, 2.5)
|
|
shards = [value] * 3
|
|
else:
|
|
parts = param.chunk(3, dim=output_dim)
|
|
shards = [
|
|
torch.full_like(part, (121 if suffix == "weight_scale" else 1) + i) for i, part in enumerate(parts)
|
|
]
|
|
value = torch.cat(shards, dim=output_dim)
|
|
expected[name] = value
|
|
if fused:
|
|
weights[name] = value
|
|
else:
|
|
for shard_id, shard in zip(("q", "k", "v"), shards):
|
|
weights[f"blocks.0.attn1.to_{shard_id}.{suffix}"] = shard
|
|
return model, config, weights, expected
|
|
|
|
|
|
def _load(model, config, weights, monkeypatch):
|
|
loader = DiffusersPipelineLoader(
|
|
LoadConfig(), OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
|
|
)
|
|
loader.counter_before_loading_weights = 0.0
|
|
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
|
|
loader.load_weights(model)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
|
|
@pytest.mark.parametrize("fused", [False, True])
|
|
def test_wan_qkv_loads_actual_weights_and_scales(mode, fused, monkeypatch):
|
|
model, config, weights, expected = _checkpoint(mode, fused)
|
|
_load(model, config, weights, monkeypatch)
|
|
for name, param in model.named_parameters():
|
|
torch.testing.assert_close(param, expected[name], rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["single", "dual", "bf16"])
|
|
def test_wan_incomplete_qkv_weight_is_not_marked_fully_loaded(mode, monkeypatch):
|
|
model, config, weights, _ = _checkpoint(mode, fused=False)
|
|
weights.pop("blocks.0.attn1.to_k.weight")
|
|
weights.pop("blocks.0.attn1.to_v.weight")
|
|
with pytest.raises(ValueError, match="to_qkv.weight"):
|
|
_load(model, config, weights, monkeypatch)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["single", "dual"])
|
|
def test_wan_incomplete_qkv_scale_is_not_marked_fully_loaded(mode, monkeypatch):
|
|
model, config, weights, _ = _checkpoint(mode, fused=False)
|
|
weights.pop("blocks.0.attn1.to_k.weight_scale")
|
|
weights.pop("blocks.0.attn1.to_v.weight_scale")
|
|
with pytest.raises(ValueError, match="to_qkv.weight_scale"):
|
|
_load(model, config, weights, monkeypatch)
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["single", "dual"])
|
|
@pytest.mark.parametrize("reverse", [False, True])
|
|
def test_wan_split_qkv_rejects_different_smooth_scales(mode, reverse, monkeypatch):
|
|
model, config, weights, _ = _checkpoint(mode, fused=False)
|
|
weights["blocks.0.attn1.to_k.mul_scale"] = weights["blocks.0.attn1.to_k.mul_scale"] * 2
|
|
if reverse:
|
|
weights = dict(reversed(weights.items()))
|
|
with pytest.raises(ValueError, match="must share the same Smooth tensor"):
|
|
_load(model, config, weights, monkeypatch)
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["to_q.missing", "to_qkv.missing", "to_qkv_extra.weight", "to_query.weight"])
|
|
def test_wan_unknown_qkv_key_is_not_reported_as_loaded(name):
|
|
model, _, _, _ = _checkpoint("single", fused=True)
|
|
loaded = model.load_weights([(f"blocks.0.attn1.{name}", torch.ones(2))])
|
|
assert loaded == set()
|
|
|
|
|
|
@pytest.mark.parametrize("fused", [False, True])
|
|
def test_wan_single_scale_rejects_relabelled_dualscale_checkpoint(fused, monkeypatch):
|
|
model, config, weights, _ = _checkpoint("single", fused=fused)
|
|
projection = "to_qkv" if fused else "to_q"
|
|
weights[f"blocks.0.attn1.{projection}.weight_dual_scale"] = torch.ones(128, 1, 1)
|
|
with pytest.raises(ValueError, match="cannot load DualScale tensor"):
|
|
_load(model, config, weights, monkeypatch)
|
|
|
|
|
|
def test_wan_layer_policy_validates_runtime_qkv_and_bf16_override():
|
|
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv"])
|
|
model = _WanCheckpoint(config)
|
|
# Uses the same validation called at the end of WanTransformer3DModel construction.
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
|
for bad in ("blocks.0.attn1.to_q", "blocks.0.attn1", "blocks.99.attn1.to_qkv", "transformer.blocks.0.attn1.to_qkv"):
|
|
config.w4a8_fallback_layers = [bad]
|
|
with pytest.raises(ValueError, match="exact runtime Linear paths"):
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
|
|
|
|
|
def test_wan_layer_policy_defers_only_other_pp_stage_paths():
|
|
from vllm.model_executor.models.utils import PPMissingLayer
|
|
|
|
config = DiffusionMXFP4Config(w4a8_fallback_layers=["blocks.0.attn1.to_qkv", "blocks.1.attn1.to_qkv"])
|
|
model = _WanCheckpoint(config)
|
|
model.blocks.append(PPMissingLayer())
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
|
config.w4a8_fallback_layers = ["blocks.1"]
|
|
with pytest.raises(ValueError, match="exact runtime Linear paths"):
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
|
|
|
|
|
@pytest.mark.parametrize("is_checkpoint_serialized", [False, True], ids=["online", "offline"])
|
|
def test_loader_logs_wan_expert_summaries_after_weight_processing(is_checkpoint_serialized, mocker):
|
|
selected = "blocks.0.attn1.to_qkv"
|
|
bf16 = "blocks.1.attn1.to_qkv"
|
|
transformer_config = DiffusionMXFP4Config(
|
|
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
|
|
w4a8_fallback_layers=[selected],
|
|
)
|
|
transformer_2_config = DiffusionMXFP4Config(
|
|
is_checkpoint_mxfp4_serialized=is_checkpoint_serialized,
|
|
ignored_layers=[bf16],
|
|
w4a8_fallback_layers=[selected, bf16],
|
|
)
|
|
pipeline = torch.nn.Module()
|
|
pipeline.transformer = _WanCheckpoint(transformer_config)
|
|
pipeline.transformer_2 = _WanCheckpoint(transformer_2_config, num_blocks=2)
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer, transformer_config)
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(pipeline.transformer_2, transformer_2_config)
|
|
|
|
loader = DiffusersPipelineLoader(
|
|
LoadConfig(),
|
|
OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=transformer_config),
|
|
)
|
|
events = []
|
|
summaries = []
|
|
|
|
def process_weights(model, _target_device):
|
|
events.append("process")
|
|
for module in model.modules():
|
|
method = getattr(module, "quant_method", None)
|
|
if getattr(method, "is_w4a8_fallback_layer", False):
|
|
module._already_called_process_weights_after_loading = True
|
|
|
|
def capture_summary(message, *args):
|
|
rendered = message % args
|
|
summaries.append(rendered)
|
|
events.append(f"summary:{args[0]}")
|
|
|
|
loader._init_from_load_format = lambda *_args, **_kwargs: pipeline # type: ignore[method-assign]
|
|
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
|
|
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
|
|
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info",
|
|
side_effect=capture_summary,
|
|
)
|
|
|
|
assert loader.load_model(load_device="cpu") is pipeline
|
|
assert events == ["load", "process", "summary:transformer", "summary:transformer_2"]
|
|
assert "component=transformer; selected=1 ['transformer.blocks.0.attn1.to_qkv']" in summaries[0]
|
|
assert "BF16 overrides=0 []" in summaries[0]
|
|
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[0]
|
|
assert "component=transformer_2; selected=1 ['transformer_2.blocks.0.attn1.to_qkv']" in summaries[1]
|
|
assert "BF16 overrides=1 ['transformer_2.blocks.1.attn1.to_qkv']" in summaries[1]
|
|
assert "weight/A8 processing state=ready (1/1 selected layers processed)" in summaries[1]
|
|
|
|
|
|
def test_wan_load_summary_does_not_claim_unprocessed_layer_is_ready(mocker):
|
|
selected = "blocks.0.attn1.to_qkv"
|
|
config = DiffusionMXFP4Config(w4a8_fallback_layers=[selected])
|
|
model = _WanCheckpoint(config)
|
|
WanTransformer3DModel._validate_w4a8_fallback_layers(model, config)
|
|
logger_info = mocker.patch("vllm_omni.diffusion.models.wan2_2.wan2_2_transformer.logger.info")
|
|
|
|
WanTransformer3DModel._log_w4a8_fallback_load_summary(model, "transformer")
|
|
|
|
message, *args = logger_info.call_args.args
|
|
rendered = message % tuple(args)
|
|
assert "weight/A8 processing state=not-ready (0/1 selected layers processed)" in rendered
|