Merge commit '6fe1ff7709c72c08bbfc33a7e14b73ec94135549' into codex-issue-7035-h100-determinism
Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
david6666666 committed
commit
f9253ae1ad
96 files changed
+8480
-822
No files matched your search
@@ -0,0 +1,679 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint import (
|
||||
FASTH3_BASE_MODEL,
|
||||
FASTH3_V2_BASE_SCHEDULE,
|
||||
FASTH3_V2_MODEL_ID,
|
||||
FastH3CheckpointSpec,
|
||||
)
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3DiTArchConfig,
|
||||
MiniMaxH3DiTModel,
|
||||
)
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
resolve_minimax_h3_diffusion_model_path,
|
||||
)
|
||||
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
|
||||
is_minimax_h3_modular,
|
||||
resolve_minimax_h3_model_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _fastvideo_metadata() -> dict[str, object]:
|
||||
"""Return the release-owned metadata read from fastvideo_inference.json."""
|
||||
return {
|
||||
"model_id": FASTH3_V2_MODEL_ID,
|
||||
"schema_version": "fasth3-inference-contract-v1",
|
||||
"dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125],
|
||||
"transformer_forwards": 8,
|
||||
"video_scheduler_shift": 10.0,
|
||||
"audio_scheduler_shift": 3.0,
|
||||
"guidance_scale": 1.0,
|
||||
"attention_backend": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa_sparsity": 0.8,
|
||||
"vsa_tile_size": 64,
|
||||
"task": "t2av",
|
||||
}
|
||||
|
||||
|
||||
def _sampling(
|
||||
*,
|
||||
num_inference_steps: int | str | None = 8,
|
||||
guidance_scale: float | None = 1.0,
|
||||
extra_args: dict[str, object] | None = None,
|
||||
lora_request: LoRARequest | None = None,
|
||||
timesteps: torch.Tensor | None = None,
|
||||
sigmas: list[float] | None = None,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
return OmniDiffusionSamplingParams(
|
||||
# Keep the string form to exercise malformed client input validation.
|
||||
num_inference_steps=num_inference_steps, # type: ignore[arg-type]
|
||||
guidance_scale=guidance_scale,
|
||||
extra_args=extra_args or {},
|
||||
lora_request=lora_request,
|
||||
timesteps=timesteps,
|
||||
sigmas=sigmas,
|
||||
)
|
||||
|
||||
|
||||
def _od_config(
|
||||
*,
|
||||
backend: str = "FASTVIDEO_VSA",
|
||||
topk: int | None = None,
|
||||
per_role: dict[str, AttentionSpec] | None = None,
|
||||
lora_path: str | None = None,
|
||||
ring_degree: int = 1,
|
||||
allgather_degree: int = 1,
|
||||
) -> OmniDiffusionConfig:
|
||||
attention_config = AttentionConfig(
|
||||
default=AttentionSpec(backend=backend, fastvideo_vsa_topk=topk),
|
||||
per_role=per_role or {},
|
||||
)
|
||||
return OmniDiffusionConfig(
|
||||
diffusion_attention_config=attention_config,
|
||||
lora_path=lora_path,
|
||||
parallel_config=DiffusionParallelConfig(ring_degree=ring_degree, allgather_degree=allgather_degree),
|
||||
)
|
||||
|
||||
|
||||
def test_v2_release_has_nine_sigma_nodes_and_eight_intervals():
|
||||
assert FASTH3_V2_BASE_SCHEDULE.base_schedule == (
|
||||
0.999,
|
||||
0.874,
|
||||
0.749,
|
||||
0.624,
|
||||
0.5,
|
||||
0.375,
|
||||
0.25,
|
||||
0.125,
|
||||
0.0,
|
||||
)
|
||||
assert len(FASTH3_V2_BASE_SCHEDULE.base_schedule) == 9
|
||||
assert FASTH3_V2_BASE_SCHEDULE.num_inference_steps == 8
|
||||
assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(10.0) == pytest.approx(
|
||||
[
|
||||
0.9998999099189271,
|
||||
0.985788405143244,
|
||||
0.9675752486758817,
|
||||
0.94316807738815,
|
||||
0.9090909090909091,
|
||||
0.8571428571428571,
|
||||
0.7692307692307693,
|
||||
0.5882352941176471,
|
||||
0.0,
|
||||
],
|
||||
abs=1e-12,
|
||||
)
|
||||
assert FASTH3_V2_BASE_SCHEDULE.shifted_sigmas(3.0) == pytest.approx(
|
||||
[
|
||||
0.9996664442961973,
|
||||
0.9541484716157204,
|
||||
0.8995196156925539,
|
||||
0.8327402135231315,
|
||||
0.75,
|
||||
0.6428571428571429,
|
||||
0.5,
|
||||
0.3,
|
||||
0.0,
|
||||
],
|
||||
abs=1e-12,
|
||||
)
|
||||
|
||||
|
||||
def test_fastvideo_metadata_builds_release_sampling_metadata():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
|
||||
assert spec.vsa_sparsity == 0.8
|
||||
assert spec.release_metadata() == {
|
||||
"partition": "fl2va",
|
||||
"tasks": ["t2va"],
|
||||
"sigma_shift_scales": {"video": 10.0, "audio": 3.0},
|
||||
"base_schedule": list(FASTH3_V2_BASE_SCHEDULE.base_schedule),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value"),
|
||||
[
|
||||
("model_id", "FastVideo/FastH3-other"),
|
||||
("schema_version", "fasth3-inference-contract-v0"),
|
||||
("dmd_denoising_steps", [999, 749, 500, 250]),
|
||||
("transformer_forwards", 4),
|
||||
("video_scheduler_shift", 12.0),
|
||||
("audio_scheduler_shift", 4.0),
|
||||
("guidance_scale", 2.0),
|
||||
("attention_backend", "FLASH_ATTN"),
|
||||
("vsa_sparsity", 0.9),
|
||||
("vsa_tile_size", 32),
|
||||
("task", "t2va"),
|
||||
],
|
||||
)
|
||||
def test_fastvideo_metadata_rejects_contract_drift(field, value):
|
||||
metadata = _fastvideo_metadata()
|
||||
metadata[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=field):
|
||||
FastH3CheckpointSpec.from_metadata(metadata)
|
||||
|
||||
|
||||
def test_fastvideo_metadata_rejects_missing_contract_field():
|
||||
metadata = _fastvideo_metadata()
|
||||
metadata.pop("task")
|
||||
|
||||
with pytest.raises(ValueError, match="task"):
|
||||
FastH3CheckpointSpec.from_metadata(metadata)
|
||||
|
||||
|
||||
def test_legacy_h3_metadata_is_not_treated_as_fastvideo_release():
|
||||
with pytest.raises(ValueError, match="model_id"):
|
||||
FastH3CheckpointSpec.from_metadata(
|
||||
{
|
||||
"schema_version": 1,
|
||||
"partition": "fl2va",
|
||||
"tasks": ["t2va", "fl2va"],
|
||||
"sigma_shift_scales": {"video": 12.0, "audio": 3.0},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_v2_serving_contract_accepts_backend_selected_for_h3_self_role():
|
||||
config = _od_config(
|
||||
backend="TORCH_SDPA",
|
||||
per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA")},
|
||||
)
|
||||
FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).check_serving_contract(
|
||||
partition="fl2va", od_config=config
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "partition", "match"),
|
||||
[
|
||||
({}, "ref2va", "task-type fl2va"),
|
||||
({"lora_path": "/tmp/adapter"}, "fl2va", "LoRA"),
|
||||
({"backend": "TORCH_SDPA"}, "fl2va", "FASTVIDEO_VSA"),
|
||||
({"topk": 2}, "fl2va", "fixed fastvideo_vsa_topk"),
|
||||
({"ring_degree": 2}, "fl2va", "local attention"),
|
||||
({"allgather_degree": 2}, "fl2va", "local attention"),
|
||||
],
|
||||
)
|
||||
def test_v2_serving_contract_rejects_incompatible_runtime(kwargs, partition, match):
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
with pytest.raises(ValueError, match=match):
|
||||
spec.check_serving_contract(partition=partition, od_config=_od_config(**kwargs))
|
||||
|
||||
|
||||
def test_v2_serving_contract_reads_a_per_role_fixed_topk_override():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
config = _od_config(
|
||||
per_role={"self": AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=1)},
|
||||
)
|
||||
with pytest.raises(ValueError, match="fixed fastvideo_vsa_topk"):
|
||||
spec.check_serving_contract(partition="fl2va", od_config=config)
|
||||
|
||||
|
||||
def test_v2_request_contract_accepts_omitted_steps_and_default_shifts():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
spec.check_request(_sampling(num_inference_steps=None))
|
||||
spec.check_request(_sampling(extra_args={"flow_shift": 10, "audio_flow_shift": 3}, guidance_scale=1))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sampling",
|
||||
[
|
||||
_sampling(num_inference_steps=None),
|
||||
_sampling(timesteps=torch.tensor([0.9, 0.5, 0.0])),
|
||||
_sampling(sigmas=[1.0, 0.5, 0.0]),
|
||||
],
|
||||
)
|
||||
def test_v2_step_execution_requires_pinned_count_and_schedule(sampling):
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
with pytest.raises(OmniClientError, match="step execution"):
|
||||
spec.check_request(sampling, step_execution=True)
|
||||
|
||||
|
||||
def test_v2_step_execution_accepts_only_explicit_pinned_count():
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
spec.check_request(_sampling(num_inference_steps=8), step_execution=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "match"),
|
||||
[
|
||||
({"num_inference_steps": 7}, "num_inference_steps=8"),
|
||||
({"num_inference_steps": 9}, "num_inference_steps=8"),
|
||||
({"num_inference_steps": "8"}, "num_inference_steps=8"),
|
||||
({"extra_args": {"flow_shift": 9}}, "flow_shift=10"),
|
||||
({"extra_args": {"audio_flow_shift": 4}}, "audio_flow_shift=3"),
|
||||
({"extra_args": {"flow_shift": "bad"}}, "flow_shift=10"),
|
||||
({"guidance_scale": 2.0}, "guidance_scale=1"),
|
||||
(
|
||||
{"lora_request": LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/test.safetensors")},
|
||||
"per-request LoRA",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_v2_request_contract_rejects_sampling_drift(kwargs, match):
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
with pytest.raises(OmniClientError, match=match):
|
||||
spec.check_request(_sampling(**kwargs))
|
||||
|
||||
|
||||
def test_resolve_native_vaes_uses_pinned_revision_and_only_vae_patterns(tmp_path, monkeypatch):
|
||||
model_root = tmp_path / "FastH3"
|
||||
model_root.mkdir()
|
||||
revision = "base-revision-123"
|
||||
(model_root / "provenance.json").write_text(
|
||||
json.dumps({"base_model": f"hf://{FASTH3_BASE_MODEL}@{revision}"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
snapshot = tmp_path / "hf-snapshot"
|
||||
calls: dict[str, object] = {}
|
||||
|
||||
def fake_download(**kwargs: object) -> str:
|
||||
calls.update(kwargs)
|
||||
snapshot.mkdir()
|
||||
return str(snapshot)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.models.minimax_h3.fasth3_checkpoint.download_weights_from_hf_specific",
|
||||
fake_download,
|
||||
)
|
||||
|
||||
result = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
|
||||
|
||||
assert result == snapshot / "FL2VA"
|
||||
assert calls == {
|
||||
"model_name_or_path": FASTH3_BASE_MODEL,
|
||||
"cache_dir": None,
|
||||
"allow_patterns": ["FL2VA/video_vae/**", "FL2VA/audio_vae/**"],
|
||||
"revision": revision,
|
||||
"require_all": True,
|
||||
}
|
||||
assert sorted(path.name for path in model_root.iterdir()) == ["provenance.json"]
|
||||
assert not (model_root / "model_index.json").exists()
|
||||
assert not (model_root / "conversion.json").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_model",
|
||||
[
|
||||
"MiniMaxAI/MiniMax-H3@main",
|
||||
"hf://MiniMaxAI/MiniMax-H3@",
|
||||
"hf://Other/MiniMax-H3@revision",
|
||||
"",
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_resolve_native_vaes_rejects_unpinned_provenance(tmp_path, base_model):
|
||||
model_root = tmp_path / "FastH3"
|
||||
model_root.mkdir()
|
||||
(model_root / "provenance.json").write_text(json.dumps({"base_model": base_model}), encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="provenance"):
|
||||
FastH3CheckpointSpec.from_metadata(_fastvideo_metadata()).resolve_native_vaes(model_root)
|
||||
|
||||
|
||||
def test_modular_local_model_root_uses_text_encoder_without_partition_files(tmp_path):
|
||||
root = tmp_path / "FastH3"
|
||||
(root / "text_encoder").mkdir(parents=True)
|
||||
(root / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
|
||||
(root / "modular_model_index.json").write_text(
|
||||
json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert is_minimax_h3_modular(str(root)) is True
|
||||
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "text_encoder")
|
||||
assert not (root / "FL2VA").exists()
|
||||
|
||||
|
||||
def test_modular_hub_model_root_downloads_only_text_encoder(monkeypatch, tmp_path):
|
||||
from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
|
||||
|
||||
snapshot = tmp_path / "modular-snapshot"
|
||||
calls: dict[str, object] = {}
|
||||
|
||||
def fake_index(model: str, *, revision: str | None = None) -> dict[str, str]:
|
||||
assert model == "FastVideo/FastVideo-FastH3-8-Step-V2"
|
||||
assert revision == "release-revision"
|
||||
return {"_class_name": "MiniMaxH3ModularPipeline"}
|
||||
|
||||
def fake_download(**kwargs: object) -> str:
|
||||
calls.update(kwargs)
|
||||
snapshot.mkdir()
|
||||
return str(snapshot)
|
||||
|
||||
monkeypatch.setattr(checkpoint_module, "get_diffusion_model_index", fake_index)
|
||||
monkeypatch.setattr(checkpoint_module, "download_weights_from_hf_specific", fake_download)
|
||||
|
||||
result = resolve_minimax_h3_model_root(
|
||||
"FastVideo/FastVideo-FastH3-8-Step-V2",
|
||||
"release-revision",
|
||||
"fl2va",
|
||||
)
|
||||
|
||||
assert result == str(snapshot / "text_encoder")
|
||||
assert calls == {
|
||||
"model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
|
||||
"cache_dir": None,
|
||||
"allow_patterns": ["text_encoder/**"],
|
||||
"revision": "release-revision",
|
||||
"require_all": True,
|
||||
}
|
||||
|
||||
|
||||
def test_native_hub_model_with_modular_class_keeps_partition_layout(monkeypatch):
|
||||
from vllm_omni.model_executor.models.minimax_h3 import checkpoint as checkpoint_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
checkpoint_module,
|
||||
"get_diffusion_model_index",
|
||||
lambda _model, *, revision=None: {"_class_name": "MiniMaxH3ModularPipeline"},
|
||||
)
|
||||
|
||||
assert checkpoint_module.is_minimax_h3_modular("MiniMaxAI/MiniMax-H3") is False
|
||||
|
||||
|
||||
def test_native_local_model_root_keeps_partition_layout(tmp_path):
|
||||
root = tmp_path / "MiniMax-H3"
|
||||
(root / "FL2VA" / "text_encoder").mkdir(parents=True)
|
||||
(root / "FL2VA" / "text_encoder" / "config.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
assert is_minimax_h3_modular(str(root)) is False
|
||||
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
|
||||
|
||||
|
||||
def test_diffusion_resolver_returns_local_modular_root_without_partition_files(tmp_path):
|
||||
root = tmp_path / "FastH3"
|
||||
root.mkdir()
|
||||
(root / "modular_model_index.json").write_text(
|
||||
json.dumps({"_class_name": "MiniMaxH3ModularPipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(root)
|
||||
assert not (root / "FL2VA").exists()
|
||||
assert not (root / "Ref2VA").exists()
|
||||
|
||||
|
||||
def test_diffusion_resolver_hub_modular_downloads_transformer_contract_only(monkeypatch, tmp_path):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
|
||||
|
||||
snapshot = tmp_path / "modular-snapshot"
|
||||
calls: dict[str, object] = {}
|
||||
|
||||
def fake_modular(_model: str, _revision: str | None = None) -> bool:
|
||||
return True
|
||||
|
||||
def fake_download(**kwargs: object) -> str:
|
||||
calls.update(kwargs)
|
||||
snapshot.mkdir()
|
||||
return str(snapshot)
|
||||
|
||||
monkeypatch.setattr(pipeline_module, "is_minimax_h3_modular", fake_modular)
|
||||
monkeypatch.setattr(pipeline_module, "download_weights_from_hf_specific", fake_download)
|
||||
|
||||
result = resolve_minimax_h3_diffusion_model_path(
|
||||
"FastVideo/FastVideo-FastH3-8-Step-V2", "release-revision", "fl2va"
|
||||
)
|
||||
|
||||
assert result == str(snapshot)
|
||||
assert calls == {
|
||||
"model_name_or_path": "FastVideo/FastVideo-FastH3-8-Step-V2",
|
||||
"cache_dir": None,
|
||||
"allow_patterns": [
|
||||
"modular_model_index.json",
|
||||
"fastvideo_inference.json",
|
||||
"provenance.json",
|
||||
"transformer/**",
|
||||
],
|
||||
"revision": "release-revision",
|
||||
"require_all": True,
|
||||
}
|
||||
patterns = calls["allow_patterns"]
|
||||
assert isinstance(patterns, list)
|
||||
assert all("video_vae" not in pattern and "audio_vae" not in pattern for pattern in patterns)
|
||||
|
||||
|
||||
def test_diffusers_arch_config_normalizes_all_aliases():
|
||||
arch = MiniMaxH3DiTArchConfig.from_mapping(
|
||||
{
|
||||
"num_refiner_layers": 3,
|
||||
"ffn_dim": 12,
|
||||
"in_channels": 4,
|
||||
"audio_in_channels": 6,
|
||||
"freq_dim": 8,
|
||||
"time_embed_hidden_dim": 10,
|
||||
"rope_freq_dim": 5,
|
||||
"patch_size": [1, 2, 2],
|
||||
}
|
||||
)
|
||||
|
||||
assert arch.token_refiner_num_layers == 3
|
||||
assert arch.ffn_hidden_size == 12
|
||||
assert arch.latents_dim == 4
|
||||
assert arch.audio_latents_dim == 6
|
||||
assert arch.timestep_input_dim == 8
|
||||
assert arch.time_embed_hidden_size == 10
|
||||
assert arch.rope_inv_freq_len == 5
|
||||
assert arch.patch_size == (1, 2, 2)
|
||||
|
||||
|
||||
class _WeightTarget(nn.Module):
|
||||
def __init__(self, shape: tuple[int, ...], loader: Callable[..., None]) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.zeros(shape))
|
||||
self.weight.weight_loader = loader
|
||||
|
||||
|
||||
class _PipelineLoadTransformer(nn.Module):
|
||||
def __init__(self, loaded: set[str]) -> None:
|
||||
super().__init__()
|
||||
self.arch = MiniMaxH3DiTArchConfig(num_layers=1)
|
||||
self._loaded = loaded
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
list(weights)
|
||||
return self._loaded
|
||||
|
||||
def post_load_weights(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def test_pipeline_load_weights_rejects_missing_fast_h3_compression_gate():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.transformer = _PipelineLoadTransformer({"blocks.0.attn.qkv_proj.weight"})
|
||||
pipeline.video_vae = None
|
||||
pipeline.audio_vae = None
|
||||
pipeline._fasth3 = None
|
||||
pipeline._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
|
||||
with pytest.raises(ValueError, match="missing compression gates"):
|
||||
pipeline.load_weights([("transformer.blocks.0.attn.qkv_proj.weight", torch.zeros(1))])
|
||||
|
||||
|
||||
def _load_fixture(
|
||||
*, diffusers_weights: bool
|
||||
) -> tuple[MiniMaxH3DiTModel, list[tuple[str | None, torch.Tensor]], list[tuple[int, torch.Tensor]]]:
|
||||
model = object.__new__(MiniMaxH3DiTModel)
|
||||
nn.Module.__init__(model)
|
||||
model.arch = MiniMaxH3DiTArchConfig(
|
||||
num_layers=1,
|
||||
hidden_size=3,
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=2,
|
||||
ffn_hidden_size=6,
|
||||
rope_inv_freq_len=2,
|
||||
)
|
||||
model._diffusers_weights = diffusers_weights
|
||||
model._rope_theta = 10000.0
|
||||
qkv_calls: list[tuple[str | None, torch.Tensor]] = []
|
||||
fc1_calls: list[tuple[int, torch.Tensor]] = []
|
||||
|
||||
def qkv_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str | None = None) -> None:
|
||||
if shard_id is None:
|
||||
assert loaded_weight.shape == (12, 3)
|
||||
param.data.copy_(loaded_weight)
|
||||
qkv_calls.append((None, loaded_weight.clone()))
|
||||
return
|
||||
assert loaded_weight.shape == (4, 3)
|
||||
assert shard_id in {"q", "k", "v"}
|
||||
offset = {"q": 0, "k": 4, "v": 8}[shard_id]
|
||||
param.data[offset : offset + 4].copy_(loaded_weight)
|
||||
qkv_calls.append((shard_id, loaded_weight.clone()))
|
||||
|
||||
def fc1_loader(param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: int) -> None:
|
||||
assert loaded_weight.shape == (3, 3)
|
||||
param.data[shard_id * 3 : (shard_id + 1) * 3].copy_(loaded_weight)
|
||||
fc1_calls.append((shard_id, loaded_weight.clone()))
|
||||
|
||||
model.blocks = nn.ModuleList([nn.Module()])
|
||||
model.blocks[0].attn = nn.Module()
|
||||
model.blocks[0].attn.qkv_proj = _WeightTarget((12, 3), qkv_loader)
|
||||
model.blocks[0].mlp = nn.Module()
|
||||
model.blocks[0].mlp.fc1 = _WeightTarget((6, 3), fc1_loader)
|
||||
model.rope = nn.Module()
|
||||
model.rope.register_buffer("inv_freq", torch.zeros(2))
|
||||
return model, qkv_calls, fc1_calls
|
||||
|
||||
|
||||
def test_diffusers_load_weights_uses_shape_aware_qkv_and_mlp_slices():
|
||||
model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=True)
|
||||
q = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 10
|
||||
k = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100
|
||||
v = torch.arange(12, dtype=torch.float32).reshape(4, 3) + 200
|
||||
mlp = torch.arange(18, dtype=torch.float32).reshape(6, 3) + 300
|
||||
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", q),
|
||||
("transformer_blocks.0.attn.to_k.weight", k),
|
||||
("transformer_blocks.0.attn.to_v.weight", v),
|
||||
("transformer_blocks.0.ff.net.0.proj.weight", mlp),
|
||||
]
|
||||
)
|
||||
|
||||
assert loaded == {
|
||||
"blocks.0.attn.qkv_proj.weight",
|
||||
"blocks.0.mlp.fc1.weight",
|
||||
"rope.inv_freq",
|
||||
}
|
||||
assert [shard_id for shard_id, _ in qkv_calls] == ["q", "k", "v"]
|
||||
assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
|
||||
torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, torch.cat((q, k, v)))
|
||||
torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, torch.cat((mlp[3:], mlp[:3])))
|
||||
torch.testing.assert_close(
|
||||
model.rope.inv_freq,
|
||||
1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
|
||||
def test_native_load_weights_keeps_grouped_qkv_and_gate_first_layouts():
|
||||
model, qkv_calls, fc1_calls = _load_fixture(diffusers_weights=False)
|
||||
grouped_qkv = torch.arange(36, dtype=torch.float32).reshape(12, 3)
|
||||
fc1 = torch.arange(18, dtype=torch.float32).reshape(6, 3)
|
||||
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
("blocks.0.attn.qkv_proj.weight", grouped_qkv),
|
||||
("blocks.0.mlp.fc1.weight", fc1),
|
||||
]
|
||||
)
|
||||
|
||||
expected_qkv = torch.cat(
|
||||
(
|
||||
grouped_qkv[:2],
|
||||
grouped_qkv[6:8],
|
||||
grouped_qkv[2:4],
|
||||
grouped_qkv[8:10],
|
||||
grouped_qkv[4:6],
|
||||
grouped_qkv[10:],
|
||||
)
|
||||
)
|
||||
assert loaded == {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight"}
|
||||
assert [shard_id for shard_id, _ in qkv_calls] == [None]
|
||||
assert [shard_id for shard_id, _ in fc1_calls] == [0, 1]
|
||||
torch.testing.assert_close(model.blocks[0].attn.qkv_proj.weight, expected_qkv)
|
||||
torch.testing.assert_close(model.blocks[0].mlp.fc1.weight, fc1)
|
||||
assert torch.equal(model.rope.inv_freq, torch.zeros(2))
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_unknown_name():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
|
||||
with pytest.raises(ValueError, match="unsupported Diffusers H3 weight"):
|
||||
model.load_weights([("transformer_blocks.0.attn.not_a_projection.weight", torch.zeros(4, 3))])
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_duplicate_source_name():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
weight = torch.zeros(4, 3)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate Diffusers H3 weight"):
|
||||
model.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", weight),
|
||||
("transformer_blocks.0.attn.to_q.weight", weight),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_incomplete_qkv_group():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
|
||||
with pytest.raises(ValueError, match="incomplete Diffusers H3 QKV"):
|
||||
model.load_weights(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.weight", torch.zeros(4, 3)),
|
||||
("transformer_blocks.0.attn.to_k.weight", torch.zeros(4, 3)),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_diffusers_load_weights_rejects_missing_compression_gate_parameter():
|
||||
model, _, _ = _load_fixture(diffusers_weights=True)
|
||||
|
||||
with pytest.raises(ValueError, match="has no model parameter"):
|
||||
model.load_weights([("transformer_blocks.0.attn.to_gate_compress.weight", torch.zeros(4, 3))])
|
||||
|
||||
|
||||
def test_resolve_sigma_positions_uses_release_metadata_schedule():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline._fasth3 = None
|
||||
pipeline._lora_sigma_schedules = {}
|
||||
spec = FastH3CheckpointSpec.from_metadata(_fastvideo_metadata())
|
||||
pipeline._base_schedule_by_partition = {
|
||||
"fl2va": DMD2SigmaSchedule.from_metadata(spec.release_metadata()),
|
||||
}
|
||||
|
||||
positions, steps = pipeline._resolve_sigma_positions("t2va", _sampling())
|
||||
assert positions == FASTH3_V2_BASE_SCHEDULE.base_schedule
|
||||
assert steps == 8 == len(positions) - 1
|
||||
with pytest.raises(OmniClientError, match="must be 8"):
|
||||
pipeline._resolve_sigma_positions("t2va", _sampling(num_inference_steps=9))
|
||||
@@ -423,6 +423,7 @@ def test_pipeline_loads_task_selected_components_with_encoder_ownership(
|
||||
assert pipeline.partition == expected_partition
|
||||
assert pipeline.supported_tasks == expected_tasks
|
||||
assert len(created["dit"]) == expected_dits
|
||||
assert all(kwargs["diffusers_weights"] is False for _, kwargs in created["dit"])
|
||||
component_path = tmp_path / component_partition
|
||||
assert created["video_vae"] == [str(component_path / "video_vae")]
|
||||
assert created["audio_vae"] == [str(component_path / "audio_vae")]
|
||||
|
||||
@@ -262,7 +262,10 @@ def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
|
||||
|
||||
def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
|
||||
path = tmp_path / "fasth3" / "adapter_model.safetensors"
|
||||
_write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
|
||||
_write_adapter(
|
||||
path,
|
||||
tensors={"transformer_blocks.0.attn.not_a_projection.lora_A.weight": torch.ones((_RANK, _HIDDEN))},
|
||||
)
|
||||
# Dropping it silently would load a model that is not the distilled student.
|
||||
with pytest.raises(FastH3AdapterError, match="name no known"):
|
||||
_load(path.parent)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
@@ -41,7 +42,7 @@ class _FakeAttention(nn.Module):
|
||||
super().__init__()
|
||||
|
||||
|
||||
def _small_od_config():
|
||||
def _small_od_config(*, class_name=None):
|
||||
arch = {
|
||||
"num_layers": 1,
|
||||
"token_refiner_num_layers": 1,
|
||||
@@ -60,6 +61,8 @@ def _small_od_config():
|
||||
"final_adaln_out_features": 2 * 8,
|
||||
"rope_inv_freq_len": 2,
|
||||
}
|
||||
if class_name is not None:
|
||||
arch["_class_name"] = class_name
|
||||
return SimpleNamespace(
|
||||
tf_model_config=arch,
|
||||
parallel_config=SimpleNamespace(ulysses_degree=1),
|
||||
@@ -133,6 +136,24 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_explicit_native_weight_format_overrides_diffusers_class_name(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
|
||||
|
||||
monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "Attention", _FakeAttention)
|
||||
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
model = h3.MiniMaxH3DiTModel(
|
||||
_small_od_config(class_name="MiniMaxH3Transformer3DModel"),
|
||||
diffusers_weights=False,
|
||||
)
|
||||
|
||||
assert model._diffusers_weights is False
|
||||
|
||||
|
||||
class _WeightTarget(nn.Module):
|
||||
def __init__(self, loader):
|
||||
super().__init__()
|
||||
|
||||
@@ -187,8 +187,14 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
|
||||
"model": model_ref,
|
||||
"enforce_eager": True,
|
||||
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
|
||||
"stage_0_tensor_parallel_size": 2,
|
||||
# This test loads the direct FL2VA diffusion pipeline (not the
|
||||
# disaggregated two-stage MiniMax-H3 pipeline), so configure both the
|
||||
# DiT world and the Qwen3-VL encoder TP explicitly.
|
||||
"tensor_parallel_size": 2,
|
||||
"text_encoder_tp_size": 2,
|
||||
"vae_use_tiling": True,
|
||||
# MiniMax H3's VAE is implemented by code shipped with the checkpoint.
|
||||
"trust_remote_code": True,
|
||||
}
|
||||
|
||||
omni_bl = Omni(**common_kwargs)
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Regression tests for transformer quant-config auto-detection.
|
||||
|
||||
The loader path at pipeline_wan2_2.py carries two quantization contracts:
|
||||
@@ -27,6 +27,7 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -381,3 +382,213 @@ def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_
|
||||
|
||||
# od_config must remain unchanged — _create_transformer does not modify it.
|
||||
assert od_config.quantization_config is None
|
||||
|
||||
|
||||
def _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch):
|
||||
import json
|
||||
|
||||
import torch
|
||||
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v as i2v_module
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as vace_module
|
||||
|
||||
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform.is_npu", lambda: True)
|
||||
for component in ("transformer", "transformer_2"):
|
||||
directory = tmp_path / component
|
||||
directory.mkdir()
|
||||
config = dict(_MIN_CFG)
|
||||
expert_quant_config = (
|
||||
disk_quant_config[0 if component == "transformer" else 1]
|
||||
if isinstance(disk_quant_config, tuple)
|
||||
else disk_quant_config
|
||||
)
|
||||
if expert_quant_config is not None:
|
||||
config["quantization_config"] = expert_quant_config
|
||||
(directory / "config.json").write_text(json.dumps(config))
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"transformer_2": ["diffusers", "WanTransformer3DModel"]}))
|
||||
|
||||
class FakeComponent(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
|
||||
|
||||
class FakeTransformer(FakeComponent):
|
||||
def __init__(self, quant_config=None, **kwargs):
|
||||
super().__init__()
|
||||
self.received_quant_config = quant_config
|
||||
|
||||
for module in (wan22_module, i2v_module):
|
||||
monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu"))
|
||||
monkeypatch.setattr(module, "prefetch_subfolders", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(module, "from_pretrained_with_prefetch", lambda *args, **kwargs: FakeComponent())
|
||||
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
|
||||
monkeypatch.setattr(vace_module, "WanVACETransformer3DModel", FakeTransformer)
|
||||
pipeline_cls = {
|
||||
"t2v": wan22_module.Wan22Pipeline,
|
||||
"i2v": i2v_module.Wan22I2VPipeline,
|
||||
"vace": vace_module.Wan22VACEPipeline,
|
||||
}[pipeline_kind]
|
||||
monkeypatch.setattr(pipeline_cls, "setup_diffusion_pipeline_profiler", lambda *args, **kwargs: None)
|
||||
od_config = OmniDiffusionConfig(
|
||||
model=str(tmp_path),
|
||||
model_class_name=pipeline_cls.__name__,
|
||||
dtype=torch.bfloat16,
|
||||
quantization_config=active,
|
||||
boundary_ratio=0.875,
|
||||
flow_shift=12.0,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
od_config.enrich_config()
|
||||
pipeline = pipeline_cls(od_config=od_config)
|
||||
return pipeline
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
|
||||
@pytest.mark.parametrize("method, flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
|
||||
@pytest.mark.parametrize("offline", [False, True])
|
||||
def test_pipeline_constructor_routes_each_expert_quantization(
|
||||
pipeline_kind, method, flag, offline, tmp_path, monkeypatch, mocker
|
||||
):
|
||||
"""Run the actual pipeline constructors, expert factories and quant-method dispatch."""
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
active = build_quant_config(
|
||||
{
|
||||
"transformer": {
|
||||
"method": method,
|
||||
"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [0, 1],
|
||||
},
|
||||
"transformer_2": {
|
||||
"method": method,
|
||||
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [],
|
||||
},
|
||||
}
|
||||
)
|
||||
disk_quant_config = None
|
||||
if offline:
|
||||
disk_quant_config = {
|
||||
"quant_method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ["blocks.12.ffn.net_2"],
|
||||
"w4a8_fallback_layers": ["blocks.9.attn1.to_qkv"],
|
||||
"w4a8_fallback_steps": [9],
|
||||
}
|
||||
pipeline = _construct_quantized_pipeline(pipeline_kind, active, disk_quant_config, tmp_path, monkeypatch)
|
||||
for index, transformer in enumerate((pipeline.transformer, pipeline.transformer_2)):
|
||||
received = transformer.received_quant_config
|
||||
assert received.get_name() == method
|
||||
assert getattr(received, flag) == offline
|
||||
assert received.w4a8_fallback_steps == ([0, 1] if index == 0 else [])
|
||||
selected = f"blocks.{10 + index}.attn1.to_qkv"
|
||||
other = f"blocks.{11 - index}.attn1.to_qkv"
|
||||
assert received.w4a8_fallback_layers == [selected]
|
||||
assert received.get_quant_method(mocker.Mock(spec=LinearBase), selected).is_w4a8_fallback_layer
|
||||
assert not received.get_quant_method(mocker.Mock(spec=LinearBase), other).is_w4a8_fallback_layer
|
||||
if offline:
|
||||
assert isinstance(
|
||||
received.get_quant_method(mocker.Mock(spec=LinearBase), "blocks.12.ffn.net_2"), UnquantizedLinearMethod
|
||||
)
|
||||
assert pipeline.od_config.quantization_config is active
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit_none", [False, True])
|
||||
def test_pipeline_disabled_component_requires_unquantized_checkpoint(explicit_none, monkeypatch):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
FakeTransformer, captured = _make_fake_transformer()
|
||||
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
|
||||
components = {"transformer": {"method": "mxfp4"}, "transformer_2": None} if explicit_none else {"vae": None}
|
||||
active = build_quant_config(components)
|
||||
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
|
||||
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
|
||||
assert captured[-1].get("quant_config") is None
|
||||
with pytest.raises(ValueError, match="disabled for component 'transformer_2'"):
|
||||
pipeline._create_transformer(
|
||||
{**_MIN_CFG, "quantization_config": {"quant_method": "mxfp4", "is_checkpoint_mxfp4_serialized": True}},
|
||||
component="transformer_2",
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_component_default_reaches_expert_factory(monkeypatch):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
FakeTransformer, captured = _make_fake_transformer()
|
||||
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
|
||||
active = build_quant_config({"default": {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}})
|
||||
pipeline = _FakePipeline(OmniDiffusionConfig(model="", quantization_config=active))
|
||||
pipeline._create_transformer(_MIN_CFG, component="transformer_2")
|
||||
assert captured[-1]["quant_config"] is active.default_config
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
|
||||
@pytest.mark.parametrize(
|
||||
"configured_expert, with_default",
|
||||
[
|
||||
("transformer", False),
|
||||
("transformer", True),
|
||||
("transformer_2", False),
|
||||
("transformer_2", True),
|
||||
],
|
||||
)
|
||||
def test_pipeline_partial_expert_config_does_not_match_sibling(
|
||||
pipeline_kind, configured_expert, with_default, tmp_path, monkeypatch
|
||||
):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
component_spec = {
|
||||
configured_expert: {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
|
||||
}
|
||||
if with_default:
|
||||
component_spec["default"] = {"method": "mxfp4", "w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"]}
|
||||
active = build_quant_config(component_spec)
|
||||
pipeline = _construct_quantized_pipeline(pipeline_kind, active, None, tmp_path, monkeypatch)
|
||||
sibling = "transformer_2" if configured_expert == "transformer" else "transformer"
|
||||
selected = getattr(pipeline, configured_expert).received_quant_config
|
||||
other = getattr(pipeline, sibling).received_quant_config
|
||||
assert selected is active.component_configs[configured_expert]
|
||||
assert other is active.default_config
|
||||
if with_default:
|
||||
assert other.w4a8_fallback_layers == ["blocks.11.attn1.to_qkv"]
|
||||
else:
|
||||
assert other is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_kind", ["t2v", "i2v", "vace"])
|
||||
@pytest.mark.parametrize("different_ignored_layers", [False, True])
|
||||
def test_enriched_auto_quantization_preserves_each_expert_policy(
|
||||
pipeline_kind, different_ignored_layers, tmp_path, monkeypatch
|
||||
):
|
||||
"""Startup auto-detection must not turn the first expert into a user override."""
|
||||
high = {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"ignored_layers": ["blocks.0.ffn.net_2"],
|
||||
"w4a8_fallback_steps": [],
|
||||
"w4a8_fallback_layers": [],
|
||||
"mxfp4_scale_alg": 0,
|
||||
"require_smooth_scale": False,
|
||||
}
|
||||
low = {
|
||||
**high,
|
||||
"ignored_layers": ["blocks.1.ffn.net_2"] if different_ignored_layers else high["ignored_layers"],
|
||||
"w4a8_fallback_steps": [37],
|
||||
"w4a8_fallback_layers": ["blocks.11.attn1.to_qkv"],
|
||||
"mxfp4_scale_alg": 2,
|
||||
"require_smooth_scale": True,
|
||||
}
|
||||
pipeline = _construct_quantized_pipeline(pipeline_kind, None, (high, low), tmp_path, monkeypatch)
|
||||
assert pipeline.od_config.quantization_config is not None # enrich_config really detected the first expert.
|
||||
for transformer, expected in [(pipeline.transformer, high), (pipeline.transformer_2, low)]:
|
||||
config = transformer.received_quant_config
|
||||
for policy in (
|
||||
"ignored_layers",
|
||||
"w4a8_fallback_steps",
|
||||
"w4a8_fallback_layers",
|
||||
"mxfp4_scale_alg",
|
||||
"require_smooth_scale",
|
||||
):
|
||||
assert getattr(config, policy) == expected[policy], policy
|
||||
@@ -20,9 +20,11 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cpu_platform(monkeypatch):
|
||||
def _cpu_pipeline_runtime(monkeypatch):
|
||||
# These lightweight pipelines have no accelerator or distributed groups.
|
||||
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
|
||||
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
|
||||
monkeypatch.setattr("vllm_omni.diffusion.distributed.parallel_state._PP", SimpleNamespace(world_size=1))
|
||||
|
||||
|
||||
class _StubTransformer(nn.Module):
|
||||
@@ -477,6 +479,100 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
|
||||
assert torch.equal(result, torch.full_like(latents, 10.0))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_first_request", [False, True])
|
||||
def test_runner_publishes_wan_step_context_across_requests(fail_first_request, monkeypatch):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
get_forward_context,
|
||||
is_forward_context_available,
|
||||
override_forward_context,
|
||||
)
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
|
||||
|
||||
contexts: list[ForwardContext] = []
|
||||
events = []
|
||||
|
||||
class InjectedPredictionError(Exception):
|
||||
pass
|
||||
|
||||
class RecordingTransformer(_StubTransformer):
|
||||
def __init__(self, expert):
|
||||
super().__init__()
|
||||
self.expert = expert
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states, **kwargs):
|
||||
context = get_forward_context()
|
||||
assert context is contexts[-1]
|
||||
request_index = len(contexts) - 1
|
||||
step = context.denoise_step_idx
|
||||
positive = bool(encoder_hidden_states[0, 0, 0] > 0)
|
||||
events.append((request_index, step, self.expert, positive, _is_w4a8_fallback_step([0, 2])))
|
||||
if fail_first_request and request_index == 0 and step == 1 and not positive:
|
||||
raise InjectedPredictionError
|
||||
return (torch.ones_like(hidden_states) * (1 if positive else -1),)
|
||||
|
||||
def encode_prompt(**kwargs):
|
||||
context = get_forward_context()
|
||||
assert context.denoise_step_idx is None
|
||||
contexts.append(context)
|
||||
embeds = torch.ones(1, 8, 8)
|
||||
return embeds, -embeds
|
||||
|
||||
pipeline = _make_pipeline()
|
||||
pipeline.transformer = RecordingTransformer("high")
|
||||
pipeline.transformer_2 = RecordingTransformer("low")
|
||||
pipeline.scheduler = _StubScheduler([900, 500, 100])
|
||||
monkeypatch.setattr(pipeline, "encode_prompt", encode_prompt)
|
||||
monkeypatch.setattr(pipeline, "scheduler_step_maybe_with_cfg", lambda pred, t, latents, cfg: latents)
|
||||
|
||||
# Use the real runner context manager, Wan forward/diffuse, and CFG
|
||||
# dispatch. Only model computation, scheduler math and unrelated KV I/O
|
||||
# are replaced; no test code publishes a denoise step.
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.pipeline = pipeline
|
||||
runner.od_config = OmniDiffusionConfig(model="", dtype=torch.float32)
|
||||
runner.vllm_config = None
|
||||
runner.cache_backend = None
|
||||
monkeypatch.setattr(runner, "_prepare_request_for_forward", lambda *args, **kwargs: None)
|
||||
with override_forward_context(None):
|
||||
for request_index in range(2):
|
||||
request = OmniDiffusionRequest(
|
||||
prompt="a fox walks",
|
||||
request_id=f"request-{request_index}",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
num_frames=1, num_inference_steps=3, guidance_scale=2.0, output_type="latent"
|
||||
),
|
||||
)
|
||||
|
||||
def execute():
|
||||
return runner._execute_request_list(
|
||||
[request],
|
||||
od_config=runner.od_config,
|
||||
allow_single_output=True,
|
||||
require_request_batch_support=False,
|
||||
record_name="test_wan_step_context",
|
||||
record_output_peak_memory=False,
|
||||
)
|
||||
|
||||
if fail_first_request and request_index == 0:
|
||||
with pytest.raises(InjectedPredictionError):
|
||||
execute()
|
||||
else:
|
||||
execute()
|
||||
assert not is_forward_context_available()
|
||||
|
||||
assert len(contexts) == 2 and contexts[0] is not contexts[1]
|
||||
expected = [
|
||||
(request, step, "high" if step == 0 else "low", positive, step in (0, 2))
|
||||
for request in range(2)
|
||||
for step in range(2 if fail_first_request and request == 0 else 3)
|
||||
for positive in (True, False)
|
||||
]
|
||||
assert events == expected
|
||||
|
||||
|
||||
class _StubDMDScheduler:
|
||||
def __init__(self) -> None:
|
||||
self.predict_clean_calls: list[tuple[float, float, float]] = []
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CPU coverage for MXFP4 checkpoint loading and the W4A8 dtype contract."""
|
||||
|
||||
import sys
|
||||
from types import ModuleType
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.config.load import LoadConfig
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
override_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.quantization.mxfp4_config import (
|
||||
DiffusionMXFP4Config,
|
||||
DiffusionMXFP4DualScaleMixedConfig,
|
||||
NPUMxfp4DualScaleLinearMethod,
|
||||
NPUMxfp4LinearMethod,
|
||||
NPUMxfp4OnlineLinearMethod,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
class _CheckpointLayer(torch.nn.Module):
|
||||
def __init__(self, method, dtype=torch.bfloat16):
|
||||
super().__init__()
|
||||
method.create_weights(
|
||||
self,
|
||||
input_size_per_partition=512,
|
||||
output_partition_sizes=[2],
|
||||
input_size=512,
|
||||
output_size=2,
|
||||
params_dtype=dtype,
|
||||
weight_loader=default_weight_loader,
|
||||
)
|
||||
|
||||
def load_weights(self, weights):
|
||||
params = dict(self.named_parameters())
|
||||
loaded = set()
|
||||
for name, value in weights:
|
||||
if name in params:
|
||||
params[name].weight_loader(params[name], value)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
|
||||
|
||||
def _load_checkpoint(layer, config, weights, monkeypatch):
|
||||
od_config = OmniDiffusionConfig(model="", dtype=torch.bfloat16, quantization_config=config)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
monkeypatch.setattr(loader, "get_all_weights", lambda model: iter(weights.items()))
|
||||
loader.load_weights(layer)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
def test_single_scale_requires_checkpoint_scale_only_when_serialized(serialized, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized)
|
||||
# Exercise the shared implementation directly as well as the online MRO
|
||||
# below: checkpoint requirements must follow the configured storage mode.
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {"weight": torch.ones((2, 512), dtype=torch.bfloat16)}
|
||||
if serialized:
|
||||
with pytest.raises(ValueError, match="Required weights.*weight_scale"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
else:
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_online_mxfp4_loads_bf16_checkpoint_without_scale(scale_alg, monkeypatch):
|
||||
config = DiffusionMXFP4Config(mxfp4_scale_alg=scale_alg)
|
||||
method = NPUMxfp4OnlineLinearMethod(config)
|
||||
layer = _CheckpointLayer(method)
|
||||
assert set(dict(layer.named_parameters())) == {"weight"}
|
||||
assert layer.weight.is_meta
|
||||
|
||||
weight = torch.arange(1024, dtype=torch.bfloat16).reshape(2, 512)
|
||||
packed_weight = torch.zeros((2, 256), dtype=torch.uint8)
|
||||
generated_scale = torch.full((2, 16), 127, dtype=torch.uint8)
|
||||
quantized_inputs = []
|
||||
|
||||
def quantize(value, **kwargs):
|
||||
quantized_inputs.append(value.clone())
|
||||
expected = dict(dst_type="fp4", axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
|
||||
if scale_alg == 2:
|
||||
expected["dst_type_max"] = 7.25
|
||||
assert kwargs == expected
|
||||
return packed_weight, generated_scale
|
||||
|
||||
# Keep real creation, lazy loading, processing, and strict loader checks;
|
||||
# substitute only the device quantization operator for this CPU test.
|
||||
npu = ModuleType("torch_npu")
|
||||
npu.__dict__.update(float4_e2m1fn_x2="fp4", npu_dynamic_mx_quant=quantize)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", npu)
|
||||
_load_checkpoint(layer, config, {"weight": weight}, monkeypatch)
|
||||
|
||||
assert len(quantized_inputs) == 1
|
||||
torch.testing.assert_close(quantized_inputs[0], weight, rtol=0, atol=0)
|
||||
torch.testing.assert_close(layer.weight, packed_weight, rtol=0, atol=0)
|
||||
torch.testing.assert_close(layer.weight_scale, generated_scale.reshape(2, 8, 2), rtol=0, atol=0)
|
||||
assert not getattr(layer.weight_scale, "is_checkpoint_required", False)
|
||||
method.process_weights_after_loading(layer)
|
||||
assert len(quantized_inputs) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("require_smooth", [False, True])
|
||||
@pytest.mark.parametrize("has_smooth", [False, True])
|
||||
def test_single_scale_checkpoint_smooth_default_and_required(require_smooth, has_smooth, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=require_smooth)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
scale = torch.linspace(0.5, 3.0, 512)
|
||||
if has_smooth:
|
||||
weights["mul_scale"] = scale
|
||||
else:
|
||||
weights.pop("mul_scale")
|
||||
|
||||
if require_smooth and not has_smooth:
|
||||
with pytest.raises(ValueError, match="mul_scale"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
else:
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
torch.testing.assert_close(layer.mul_scale, scale if has_smooth else torch.ones(512), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_scale", ["shape", "name"])
|
||||
def test_declared_single_scale_smooth_rejects_malformed_checkpoint(bad_scale, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
if bad_scale == "shape":
|
||||
weights["mul_scale"] = torch.ones(511)
|
||||
with pytest.raises(AssertionError):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
else:
|
||||
weights["multiply_scale"] = weights.pop("mul_scale")
|
||||
with pytest.raises(ValueError, match="mul_scale"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("required", [False, True])
|
||||
@pytest.mark.parametrize("bad_value", [float("nan"), float("inf"), 0.0, -1.0])
|
||||
def test_single_scale_rejects_invalid_smooth_values(required, bad_value, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=required)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
weights["mul_scale"][7] = bad_value
|
||||
with pytest.raises(ValueError, match="finite, strictly positive"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_scale", [torch.ones(512, dtype=torch.int32), torch.ones(1, 512)])
|
||||
def test_single_scale_rejects_invalid_smooth_type_or_rank(bad_scale, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(config))
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters()}
|
||||
weights["mul_scale"] = bad_scale
|
||||
with pytest.raises(ValueError, match="one-dimensional floating-point"):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_value", [1e10, 1e-10])
|
||||
def test_single_scale_rejects_smooth_overflow_or_underflow_after_cast(bad_value, monkeypatch):
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
|
||||
method = NPUMxfp4LinearMethod(config)
|
||||
layer = _CheckpointLayer(method, torch.float16)
|
||||
layer.mul_scale.data.fill_(bad_value)
|
||||
# Validation must fail before any device packing or GEMM.
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu"))
|
||||
with pytest.raises(ValueError, match="finite, strictly positive"):
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "missing"),
|
||||
[
|
||||
("single", "weight"),
|
||||
("single", "weight_scale"),
|
||||
("dual", "weight"),
|
||||
("dual", "weight_scale"),
|
||||
("dual", "weight_dual_scale"),
|
||||
("dual", "mul_scale"),
|
||||
],
|
||||
)
|
||||
def test_offline_mxfp4_required_tensors_do_not_use_legacy_scale_tolerance(mode, missing, monkeypatch):
|
||||
if mode == "single":
|
||||
config = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)
|
||||
method = NPUMxfp4LinearMethod(config)
|
||||
else:
|
||||
config = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=True)
|
||||
method = NPUMxfp4DualScaleLinearMethod(config)
|
||||
layer = _CheckpointLayer(method)
|
||||
weights = {name: torch.ones_like(param) for name, param in layer.named_parameters() if name != missing}
|
||||
with pytest.raises(ValueError, match=missing):
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
def test_single_scale_smooth_uses_real_row_parallel_loader(rank):
|
||||
from vllm.model_executor.layers.linear import RowParallelLinear
|
||||
|
||||
layer = _CheckpointLayer(NPUMxfp4LinearMethod(DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True)))
|
||||
layer.tp_rank = rank
|
||||
checkpoint_scale = torch.arange(1024, dtype=torch.float32)
|
||||
RowParallelLinear.weight_loader(layer, layer.mul_scale, checkpoint_scale)
|
||||
torch.testing.assert_close(layer.mul_scale, checkpoint_scale[rank * 512 : (rank + 1) * 512], rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("fallback", [False, True])
|
||||
@pytest.mark.parametrize("with_bias", [False, True])
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_single_scale_smooth_precedes_quantization_and_preserves_output(
|
||||
dtype, fallback, with_bias, scale_alg, monkeypatch
|
||||
):
|
||||
config = DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=True, w4a8_fallback_steps=[0], mxfp4_scale_alg=scale_alg
|
||||
)
|
||||
method = NPUMxfp4LinearMethod(config)
|
||||
layer = _CheckpointLayer(method, dtype)
|
||||
weights = {
|
||||
"weight": torch.linspace(-0.5, 1.5, 1024, dtype=dtype).reshape(2, 512),
|
||||
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
|
||||
"mul_scale": torch.linspace(0.5, 3.0, 512),
|
||||
}
|
||||
_load_checkpoint(layer, config, weights, monkeypatch)
|
||||
x = torch.linspace(-1.0, 2.0, 3072, dtype=dtype).reshape(2, 3, 512)
|
||||
bias = torch.tensor([0.5, -1.0], dtype=dtype) if with_bias else None
|
||||
quantized_inputs = []
|
||||
|
||||
# Only the two device operators are substituted. Exercise real apply(),
|
||||
# step selection, Smooth multiplication, bias handling, and reshape logic.
|
||||
npu = ModuleType("torch_npu")
|
||||
|
||||
def quantize(value, **kwargs):
|
||||
quantized_inputs.append(value.clone())
|
||||
expected = dict(
|
||||
dst_type=torch.float8_e4m3fn if fallback else "fp4",
|
||||
axis=-1,
|
||||
block_size=32,
|
||||
round_mode="rint",
|
||||
scale_alg=0 if fallback else scale_alg,
|
||||
)
|
||||
if not fallback and scale_alg == 2:
|
||||
expected["dst_type_max"] = 7.25
|
||||
assert kwargs == expected
|
||||
return value, torch.ones((6, 16), dtype=torch.uint8)
|
||||
|
||||
def matmul(x_q, weight, weight_scale, *, bias, output_dtype, **kwargs):
|
||||
if fallback:
|
||||
assert output_dtype == torch.bfloat16
|
||||
if bias is not None:
|
||||
assert bias.dtype == output_dtype
|
||||
assert bias.shape == (1, 2)
|
||||
result = x_q.float() @ weight.float()
|
||||
if bias is not None:
|
||||
result += bias.float()
|
||||
return result.to(output_dtype)
|
||||
|
||||
npu.__dict__.update(
|
||||
float8_e4m3fn=torch.float8_e4m3fn,
|
||||
float8_e8m0fnu=torch.uint8,
|
||||
float4_e2m1fn_x2="fp4",
|
||||
npu_dynamic_mx_quant=quantize,
|
||||
npu_quant_matmul=matmul,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", npu)
|
||||
with override_forward_context(ForwardContext()):
|
||||
set_forward_context_denoise_step_idx(0 if fallback else 1)
|
||||
result = method.apply(layer, x, bias)
|
||||
|
||||
expected_input = x.reshape(-1, 512) * weights["mul_scale"].to(dtype)
|
||||
assert len(quantized_inputs) == 1
|
||||
torch.testing.assert_close(quantized_inputs[0], expected_input, rtol=0, atol=0)
|
||||
expected = expected_input.float() @ weights["weight"].float().T
|
||||
if bias is not None:
|
||||
expected += bias.float()
|
||||
if fallback:
|
||||
expected = expected.bfloat16()
|
||||
torch.testing.assert_close(result, expected.to(dtype).reshape(2, 3, 2), rtol=0, atol=0)
|
||||
assert result.dtype == dtype
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
|
||||
|
||||
import pytest
|
||||
@@ -27,12 +27,42 @@ def test_mxfp4_config_get_name():
|
||||
assert DiffusionMXFP4Config.get_name() == "mxfp4"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [True, False, None, "2", 2.0, 1, -1, 3])
|
||||
def test_mxfp4_scale_alg_rejects_invalid_values(value):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="mxfp4_scale_alg"):
|
||||
build_quant_config({"method": "mxfp4", "mxfp4_scale_alg": value})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("runtime_alg", [0, 2])
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
def test_runtime_algorithm_survives_offline_storage_rebuild(runtime_alg, serialized):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, mxfp4_scale_alg=runtime_alg)
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"mxfp4_scale_alg": 2,
|
||||
"ignored_layers": ["proj_out"],
|
||||
},
|
||||
)
|
||||
assert resolved.is_checkpoint_mxfp4_serialized
|
||||
assert resolved.mxfp4_scale_alg == runtime_alg
|
||||
assert resolved.w4a8_fallback_layers == resolved.w4a8_fallback_steps == []
|
||||
|
||||
|
||||
def test_mxfp4_config_from_config_defaults():
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
cfg = DiffusionMXFP4Config.from_config({})
|
||||
assert cfg.is_checkpoint_mxfp4_serialized is False
|
||||
assert cfg.ignored_layers == []
|
||||
assert cfg.w4a8_fallback_layers == []
|
||||
|
||||
|
||||
def test_mxfp4_config_from_config_serialized():
|
||||
@@ -42,6 +72,98 @@ def test_mxfp4_config_from_config_serialized():
|
||||
assert cfg.is_checkpoint_mxfp4_serialized is True
|
||||
|
||||
|
||||
def test_mxfp4_config_requires_offline_smooth():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
values = {"is_checkpoint_mxfp4_serialized": True, "require_smooth_scale": True}
|
||||
assert DiffusionMXFP4Config.from_config(values).require_smooth_scale
|
||||
assert build_quant_config({"method": "mxfp4", **values}).require_smooth_scale
|
||||
with pytest.raises(ValueError, match="offline"):
|
||||
build_quant_config({"method": "mxfp4", "require_smooth_scale": True})
|
||||
with pytest.raises(ValueError, match="boolean"):
|
||||
DiffusionMXFP4Config.from_config({**values, "require_smooth_scale": "false"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
def test_mxfp4_disk_smooth_requirement_preserves_step_policy(serialized):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=serialized, w4a8_fallback_steps=[1])
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"require_smooth_scale": True,
|
||||
},
|
||||
)
|
||||
assert resolved.require_smooth_scale
|
||||
assert resolved.is_checkpoint_mxfp4_serialized
|
||||
assert resolved.w4a8_fallback_steps == [1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disk_requirement", [None, False, True])
|
||||
def test_mxfp4_expert_rebuild_cannot_weaken_explicit_smooth_requirement(disk_requirement):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True)
|
||||
disk = {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"ignored_layers": ["blocks.0.attn2.to_q"],
|
||||
}
|
||||
if disk_requirement is not None:
|
||||
disk["require_smooth_scale"] = disk_requirement
|
||||
resolved = resolve_quant_config_from_disk(active, disk)
|
||||
assert resolved.require_smooth_scale
|
||||
assert resolved.ignored_layers == disk["ignored_layers"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disk_qc", ["mxfp4", {"quant_method": "mxfp4"}])
|
||||
def test_mxfp4_method_only_disk_config_preserves_caller_layer_policy(disk_qc):
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
active = DiffusionMXFP4Config(
|
||||
ignored_layers=["blocks.0.attn2.to_q"],
|
||||
w4a8_fallback_steps=[1],
|
||||
w4a8_fallback_layers=["blocks.1.attn1.to_qkv"],
|
||||
)
|
||||
|
||||
resolved = resolve_quant_config_from_disk(active, disk_qc)
|
||||
|
||||
assert resolved is active
|
||||
assert resolved.ignored_layers == ["blocks.0.attn2.to_q"]
|
||||
assert resolved.w4a8_fallback_steps == [1]
|
||||
assert resolved.w4a8_fallback_layers == ["blocks.1.attn1.to_qkv"]
|
||||
|
||||
|
||||
def test_quantization_disk_string_rejects_mismatched_active_method():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active = build_quant_config("mxfp4")
|
||||
|
||||
with pytest.raises(ValueError, match="mxfp4_dualscale.*active quantization config is 'mxfp4'"):
|
||||
resolve_quant_config_from_disk(active, "mxfp4_dualscale")
|
||||
|
||||
|
||||
def test_quantization_disk_string_accepts_equivalent_method_alias():
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
class ActiveAliasConfig:
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "inc"
|
||||
|
||||
active = ActiveAliasConfig()
|
||||
|
||||
assert resolve_quant_config_from_disk(active, "auto-round") is active
|
||||
|
||||
|
||||
def test_mxfp4_config_from_config_ignored_layers():
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
@@ -57,6 +179,202 @@ def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
|
||||
assert cfg.ignored_layers == ["proj_out"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"policy",
|
||||
[
|
||||
{"w4a8_fallback_steps": [0]},
|
||||
{"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
|
||||
{"w4a8_fallback_steps": [1], "w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]},
|
||||
],
|
||||
)
|
||||
def test_dualscale_rejects_w4a8_configuration(serialized, policy):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
|
||||
|
||||
config = {"quant_method": "mxfp4_dualscale", "is_checkpoint_serialized": serialized, **policy}
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=serialized, **policy)
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
DiffusionMXFP4DualScaleMixedConfig.from_config(config)
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
build_quant_config({"transformer": config})
|
||||
with pytest.raises(ValueError, match="does not support W4A8"):
|
||||
resolve_quant_config_from_disk(None, config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
|
||||
@pytest.mark.parametrize("steps", [[-1], [True], [1.5], ["2"], "0,2", (0, 2)])
|
||||
def test_mxfp4_rejects_invalid_fallback_steps(method, steps):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="non-negative integer"):
|
||||
build_quant_config({"method": method, "w4a8_fallback_steps": steps})
|
||||
|
||||
|
||||
def test_mxfp4_fallback_steps_follow_request_context():
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
ForwardContext,
|
||||
override_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.mxfp4_config import _is_w4a8_fallback_step
|
||||
|
||||
steps = [2, 0, 2]
|
||||
cfg = build_quant_config({"method": "mxfp4", "w4a8_fallback_steps": steps})
|
||||
steps.append(1)
|
||||
assert cfg.w4a8_fallback_steps == [0, 2]
|
||||
with override_forward_context(None):
|
||||
assert not _is_w4a8_fallback_step([])
|
||||
with pytest.raises(RuntimeError, match="publishes its denoise step"):
|
||||
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
|
||||
for _request in range(2):
|
||||
with override_forward_context(ForwardContext()):
|
||||
for step, expected in [(0, True), (1, False), (2, True), (3, False)]:
|
||||
set_forward_context_denoise_step_idx(step)
|
||||
# Both CFG passes and both experts read the same index.
|
||||
assert [_is_w4a8_fallback_step(cfg.w4a8_fallback_steps) for _cfg_pass in range(2)] == [
|
||||
expected,
|
||||
expected,
|
||||
]
|
||||
with pytest.raises(RuntimeError, match="publishes its denoise step"):
|
||||
_is_w4a8_fallback_step(cfg.w4a8_fallback_steps)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized")])
|
||||
@pytest.mark.parametrize("steps", [[], [0, 2]])
|
||||
@pytest.mark.parametrize("already_serialized", [False, True])
|
||||
def test_mxfp4_disk_rebuild_preserves_runtime_fallback(method, flag, steps, already_serialized):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
layers = ["blocks.10.attn1.to_qkv"] if steps else []
|
||||
active = build_quant_config(
|
||||
{
|
||||
"method": method,
|
||||
flag: already_serialized,
|
||||
"w4a8_fallback_steps": steps,
|
||||
"w4a8_fallback_layers": layers,
|
||||
}
|
||||
)
|
||||
# Separate experts may have distinct high precision layers. Both must
|
||||
# retain the requested step policy during online/offline reconciliation.
|
||||
for ignored in [["blocks.0.attn1.to_q"], ["blocks.1.attn1.to_q"]]:
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ignored,
|
||||
"w4a8_fallback_steps": [9],
|
||||
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
|
||||
},
|
||||
)
|
||||
assert getattr(resolved, flag)
|
||||
assert resolved.ignored_layers == ignored
|
||||
assert resolved.w4a8_fallback_steps == steps
|
||||
assert resolved.w4a8_fallback_layers == layers
|
||||
assert active.ignored_layers == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"runtime_policy",
|
||||
[{}, {"w4a8_fallback_steps": [1]}, {"w4a8_fallback_layers": ["blocks.10.attn1.to_qkv"]}],
|
||||
)
|
||||
def test_mxfp4_explicit_omission_disables_saved_policies_for_both_experts(runtime_policy):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active = build_quant_config({"method": "mxfp4", **runtime_policy})
|
||||
resolved_experts = []
|
||||
for ignored in (["blocks.0.attn2.to_q"], ["blocks.1.ffn.net_2"]):
|
||||
disk = {
|
||||
"quant_method": "mxfp4",
|
||||
"is_checkpoint_mxfp4_serialized": True,
|
||||
"ignored_layers": ignored,
|
||||
"w4a8_fallback_steps": [9],
|
||||
"w4a8_fallback_layers": ["blocks.9.ffn.net_2"],
|
||||
}
|
||||
resolved = resolve_quant_config_from_disk(active, disk)
|
||||
assert resolved.w4a8_fallback_steps == runtime_policy.get("w4a8_fallback_steps", [])
|
||||
assert resolved.w4a8_fallback_layers == runtime_policy.get("w4a8_fallback_layers", [])
|
||||
assert resolved.ignored_layers == ignored
|
||||
detected = resolve_quant_config_from_disk(None, disk)
|
||||
assert detected.w4a8_fallback_steps == [9]
|
||||
assert detected.w4a8_fallback_layers == ["blocks.9.ffn.net_2"]
|
||||
resolved_experts.append(resolved)
|
||||
assert resolved_experts[0] is not resolved_experts[1]
|
||||
assert active.is_checkpoint_mxfp4_serialized is False
|
||||
assert active.ignored_layers == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
|
||||
)
|
||||
def test_serialized_mxfp4_disk_omission_clears_active_ignored_layers(method, flag):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active_spec = {
|
||||
"method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ["blocks.0.attn2.to_q"],
|
||||
"w4a8_fallback_steps": [1] if method == "mxfp4" else [],
|
||||
"w4a8_fallback_layers": ["blocks.1.attn1.to_qkv"] if method == "mxfp4" else [],
|
||||
}
|
||||
if method == "mxfp4":
|
||||
active_spec["require_smooth_scale"] = True
|
||||
active = build_quant_config(active_spec)
|
||||
|
||||
resolved = resolve_quant_config_from_disk(active, {"quant_method": method, flag: True})
|
||||
|
||||
assert resolved is not active
|
||||
assert resolved.ignored_layers == []
|
||||
assert resolved.w4a8_fallback_steps == active.w4a8_fallback_steps
|
||||
assert resolved.w4a8_fallback_layers == active.w4a8_fallback_layers
|
||||
if method == "mxfp4":
|
||||
assert resolved.require_smooth_scale
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,flag", [("mxfp4", "is_checkpoint_mxfp4_serialized"), ("mxfp4_dualscale", "is_checkpoint_serialized")]
|
||||
)
|
||||
@pytest.mark.parametrize("with_active", [False, True])
|
||||
def test_serialized_mxfp4_disk_uses_modules_to_not_convert(method, flag, with_active):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
active = None
|
||||
disk_steps = [9] if method == "mxfp4" else []
|
||||
active_steps = [1] if method == "mxfp4" else []
|
||||
expected_steps = disk_steps
|
||||
if with_active:
|
||||
active = build_quant_config(
|
||||
{
|
||||
"method": method,
|
||||
flag: True,
|
||||
"ignored_layers": ["blocks.0.attn2.to_q"],
|
||||
"w4a8_fallback_steps": active_steps,
|
||||
}
|
||||
)
|
||||
expected_steps = active_steps
|
||||
|
||||
resolved = resolve_quant_config_from_disk(
|
||||
active,
|
||||
{
|
||||
"quant_method": method,
|
||||
flag: True,
|
||||
"modules_to_not_convert": ["blocks.2.ffn.net_2"],
|
||||
"w4a8_fallback_steps": disk_steps,
|
||||
},
|
||||
)
|
||||
|
||||
assert resolved.ignored_layers == ["blocks.2.ffn.net_2"]
|
||||
assert resolved.w4a8_fallback_steps == expected_steps
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_quant_config integration
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -594,18 +912,31 @@ def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
|
||||
assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
|
||||
|
||||
|
||||
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
|
||||
@pytest.mark.parametrize("steps", [[], [0, 2]])
|
||||
def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch, steps):
|
||||
"""A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
|
||||
cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
|
||||
cfg = DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"], w4a8_fallback_steps=steps
|
||||
)
|
||||
|
||||
assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
|
||||
|
||||
|
||||
def test_rocm_rejects_w4a8_step_fallback(_rocm_platform, monkeypatch):
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
|
||||
_patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
|
||||
cfg = DiffusionMXFP4Config(w4a8_fallback_steps=[0, 2])
|
||||
|
||||
with pytest.raises(NotImplementedError, match="only supported on NPU"):
|
||||
cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
|
||||
|
||||
|
||||
def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
|
||||
"""MXFP4 on ROCm requires gfx950; any other arch must raise."""
|
||||
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
|
||||
@@ -652,3 +983,221 @@ def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
|
||||
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
|
||||
assert layer.input_size_per_partition == _TP2_K // _TP2
|
||||
assert layer.output_size_per_partition == _TP2_N
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["mxfp4", "mxfp4_dualscale"])
|
||||
@pytest.mark.parametrize("layers", ["blocks.1.attn1.to_qkv", [1], [True], [""], [" a"], ["blocks.*"], ["a..b"]])
|
||||
def test_mxfp4_rejects_invalid_fallback_layers(method, layers):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="exact runtime Linear paths"):
|
||||
build_quant_config({"method": method, "w4a8_fallback_layers": layers})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["mxfp4"])
|
||||
def test_mxfp4_layer_policy_from_checkpoint_and_component_config(method):
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.component_config import resolve_component_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
names = ["blocks.10.attn1.to_qkv", "blocks.10.attn1.to_qkv"]
|
||||
components = build_quant_config(
|
||||
{"transformer": {"method": method, "w4a8_fallback_layers": names}, "transformer_2": {"method": method}}
|
||||
)
|
||||
names.append("blocks.11.attn1.to_qkv")
|
||||
assert resolve_component_quant_config(components, "transformer").w4a8_fallback_layers == [names[0]]
|
||||
assert resolve_component_quant_config(components, "transformer_2").w4a8_fallback_layers == []
|
||||
detected = resolve_quant_config_from_disk(None, {"quant_method": method, "w4a8_fallback_layers": [names[0]]})
|
||||
assert detected.w4a8_fallback_layers == [names[0]]
|
||||
assert type(detected).from_config({"w4a8_fallback_layers": [names[0]]}).w4a8_fallback_layers == [names[0]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kind, serialized_flag",
|
||||
[("mxfp4", "is_checkpoint_mxfp4_serialized")],
|
||||
)
|
||||
@pytest.mark.parametrize("offline", [False, True])
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_mxfp4_layer_and_step_union_routes_actual_methods(
|
||||
kind, serialized_flag, offline, scale_alg, monkeypatch, mocker
|
||||
):
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
||||
from vllm_omni.quantization import build_quant_config, mxfp4_config
|
||||
|
||||
monkeypatch.setattr(mxfp4_config.current_omni_platform, "is_npu", lambda: True)
|
||||
selected = "blocks.10.attn1.to_qkv"
|
||||
other = "blocks.11.attn2.to_q"
|
||||
preserved = "blocks.12.ffn.net_2"
|
||||
config = build_quant_config(
|
||||
{
|
||||
"method": kind,
|
||||
"mxfp4_scale_alg": scale_alg,
|
||||
serialized_flag: offline,
|
||||
"w4a8_fallback_layers": [selected, preserved],
|
||||
"ignored_layers": [preserved],
|
||||
"w4a8_fallback_steps": [1],
|
||||
}
|
||||
)
|
||||
# Real config dispatch creates separate methods for selected and other layers.
|
||||
linear = mocker.Mock(spec=LinearBase)
|
||||
assert isinstance(config.get_quant_method(linear, preserved), UnquantizedLinearMethod)
|
||||
x = torch.ones(2, 64, dtype=torch.bfloat16)
|
||||
layer = torch.nn.Module()
|
||||
for name in ("weight", "weight_scale", "mul_scale"):
|
||||
layer.register_buffer(name, torch.ones(64, dtype=x.dtype))
|
||||
calls = []
|
||||
|
||||
def record_a8(*args):
|
||||
calls.append("A8")
|
||||
return x
|
||||
|
||||
def record_a4(*args):
|
||||
calls.append("A4")
|
||||
return x
|
||||
|
||||
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", record_a8)
|
||||
for name in (selected, other):
|
||||
method = config.get_quant_method(linear, name)
|
||||
assert method.is_w4a8_fallback_layer == (name == selected)
|
||||
scales = (x, x)
|
||||
monkeypatch.setattr(method, "_quantize_activation", lambda *args: scales)
|
||||
monkeypatch.setattr(method, "_quant_matmul", record_a4)
|
||||
for step in (0, 1, 2):
|
||||
with override_forward_context(ForwardContext(denoise_step_idx=step)):
|
||||
method.apply(layer, x)
|
||||
assert calls == ["A8", "A8", "A8", "A4", "A8", "A4"]
|
||||
# A permanently selected layer works without a request context, even if steps are also configured.
|
||||
with override_forward_context(None):
|
||||
config.get_quant_method(linear, selected).apply(layer, x)
|
||||
assert calls[-1] == "A8"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("online", [False, True])
|
||||
@pytest.mark.parametrize("selected", [False, True])
|
||||
@pytest.mark.parametrize("scale_alg", [0, 2])
|
||||
def test_single_scale_prepares_one_weight_for_both_activation_precisions(online, selected, scale_alg, monkeypatch):
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
||||
from vllm_omni.quantization import mxfp4_config
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
|
||||
prepared = []
|
||||
packed = []
|
||||
|
||||
def quantize(x, **kwargs):
|
||||
prepared.append(kwargs)
|
||||
return x.to(torch.int8), torch.full((x.shape[0], 2), 127, dtype=torch.uint8)
|
||||
|
||||
def pack(x, dtype):
|
||||
packed.append(x)
|
||||
return x.to(dtype)
|
||||
|
||||
fake_npu = SimpleNamespace(
|
||||
float4_e2m1fn_x2=torch.int8,
|
||||
float8_e8m0fnu=torch.uint8,
|
||||
float8_e4m3fn=torch.float8_e4m3fn,
|
||||
npu_dtype_cast=pack,
|
||||
npu_dynamic_mx_quant=quantize,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
|
||||
config = mxfp4_config.DiffusionMXFP4Config(
|
||||
is_checkpoint_mxfp4_serialized=not online,
|
||||
w4a8_fallback_layers=["blocks.0.attn1.to_qkv"] if selected else [],
|
||||
w4a8_fallback_steps=[1],
|
||||
mxfp4_scale_alg=scale_alg,
|
||||
)
|
||||
cls = mxfp4_config.NPUMxfp4OnlineLinearMethod if online else mxfp4_config.NPUMxfp4LinearMethod
|
||||
method = cls(config, prefix="blocks.0.attn1.to_qkv")
|
||||
layer = torch.nn.Linear(64, 2, bias=False, dtype=torch.bfloat16)
|
||||
layer.orig_dtype = torch.bfloat16
|
||||
if not online:
|
||||
layer.register_parameter("weight_scale", torch.nn.Parameter(torch.full((2, 2), 127, dtype=torch.uint8), False))
|
||||
layer.register_parameter("mul_scale", torch.nn.Parameter(torch.linspace(0.5, 2, 64), False))
|
||||
method.process_weights_after_loading(layer)
|
||||
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
|
||||
method.process_weights_after_loading(layer)
|
||||
assert len(prepared) == int(online)
|
||||
expected = dict(dst_type=torch.int8, axis=-1, block_size=32, round_mode="rint", scale_alg=scale_alg)
|
||||
if scale_alg == 2:
|
||||
expected["dst_type_max"] = 7.25
|
||||
assert prepared == ([expected] if online else [])
|
||||
assert len(packed) == int(not online)
|
||||
branches = []
|
||||
|
||||
def a4(x_q, x_scale, target, bias, dtype):
|
||||
assert (target.weight.data_ptr(), target.weight_scale.data_ptr()) == pointers
|
||||
branches.append("A4")
|
||||
return torch.zeros(x_q.shape[0], 2, dtype=dtype)
|
||||
|
||||
def a8(x, weight, scale, bias, dtype):
|
||||
assert (weight.data_ptr(), scale.data_ptr()) == pointers
|
||||
branches.append("A8")
|
||||
return torch.zeros(x.shape[0], 2, dtype=dtype)
|
||||
|
||||
monkeypatch.setattr(method, "_quant_matmul", a4)
|
||||
monkeypatch.setattr(mxfp4_config, "_npu_w4a8_matmul", a8)
|
||||
monkeypatch.setattr(method, "_quantize_activation", lambda x: (x, x))
|
||||
monkeypatch.setattr(fake_npu, "npu_dynamic_mx_quant", lambda *a, **kw: pytest.fail("weight requantization"))
|
||||
monkeypatch.setattr(fake_npu, "npu_dtype_cast", lambda *a, **kw: pytest.fail("weight repacking"))
|
||||
for step in (0, 1, 2):
|
||||
with override_forward_context(ForwardContext(denoise_step_idx=step)):
|
||||
method.apply(layer, torch.ones(2, 64, dtype=torch.bfloat16))
|
||||
assert branches == (["A8"] * 3 if selected else ["A4", "A8", "A4"])
|
||||
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
|
||||
assert set(layer.state_dict()) == (
|
||||
{"weight", "weight_scale"} if online else {"weight", "weight_scale", "mul_scale"}
|
||||
)
|
||||
assert not hasattr(layer, "w4a8_weight")
|
||||
assert not hasattr(layer, "w4a8_weight_scale")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("online", [False, True])
|
||||
def test_dualscale_without_fallback_keeps_original_a4_path(online, monkeypatch):
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.quantization import mxfp4_config
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "npu", lambda self: self, raising=False)
|
||||
original_ones = torch.ones
|
||||
monkeypatch.setattr(torch, "ones", lambda *args, **kw: original_ones(*args, **{**kw, "device": "cpu"}))
|
||||
calls = []
|
||||
|
||||
def quantize_dual(x, **kw):
|
||||
calls.append(kw.get("smooth_scale"))
|
||||
return x.to(torch.int8), torch.ones(x.shape[0], 1), torch.ones(x.shape[0], 16, dtype=torch.uint8)
|
||||
|
||||
fake_npu = SimpleNamespace(
|
||||
float4_e2m1fn_x2=torch.int8,
|
||||
npu_dtype_cast=lambda x, dtype: x.to(dtype),
|
||||
npu_format_cast=lambda x, *args, **kw: x,
|
||||
npu_dynamic_dual_level_mx_quant=quantize_dual,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch_npu", fake_npu)
|
||||
config = mxfp4_config.DiffusionMXFP4DualScaleMixedConfig(
|
||||
is_checkpoint_serialized=not online, w4a8_fallback_steps=[], w4a8_fallback_layers=[]
|
||||
)
|
||||
cls = mxfp4_config.NPUMxfp4DualScaleOnlineLinearMethod if online else mxfp4_config.NPUMxfp4DualScaleLinearMethod
|
||||
method = cls(config)
|
||||
layer = torch.nn.Module()
|
||||
layer.input_size_per_partition = 512
|
||||
for name, value in {
|
||||
"weight": torch.ones(2, 512, dtype=torch.bfloat16),
|
||||
"weight_scale": torch.full((2, 16), 127, dtype=torch.uint8),
|
||||
"weight_dual_scale": torch.ones(2, 1, 1),
|
||||
"mul_scale": torch.ones(512),
|
||||
}.items():
|
||||
layer.register_parameter(name, torch.nn.Parameter(value, requires_grad=False))
|
||||
method.process_weights_after_loading(layer)
|
||||
method.process_weights_after_loading(layer)
|
||||
assert len(calls) == int(online)
|
||||
assert set(layer.state_dict()) == {"weight", "weight_scale", "weight_dual_scale", "mul_scale"}
|
||||
assert not hasattr(layer, "w4a8_weight")
|
||||
monkeypatch.setattr(method, "_quant_matmul", lambda *args: torch.zeros(2, 2, dtype=torch.bfloat16))
|
||||
method.apply(layer, torch.ones(2, 512, dtype=torch.bfloat16))
|
||||
assert len(calls) == 1 + int(online)
|
||||
torch.testing.assert_close(calls[-1], layer.mul_scale)
|
||||
@@ -0,0 +1,237 @@
|
||||
# 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
|
||||
@@ -25,6 +25,12 @@ from vllm_omni.diffusion.model_metadata import (
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _local_model_paths(monkeypatch):
|
||||
# These tests exercise config transport, not model repository resolution.
|
||||
monkeypatch.setattr("vllm_omni.diffusion.data.get_model_path", lambda model, revision: model)
|
||||
|
||||
|
||||
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
|
||||
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
|
||||
|
||||
|
||||
Reference in new issue
Block a user