[Diffusion] Tighten runtime cache compatibility guards

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-16 23:26:25 +08:00
1 parent db86a6a078
commit 99347e1ba7
6 files changed
+43 -9

No files matched your search

@@ -12,6 +12,7 @@ import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
from vllm.config.load import LoadConfig
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -425,6 +426,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = UnquantizedLinearMethod()
events: list[str] = []
runtime_plan = HostWeightPlan(
backing_kind="runtime_cache",
@@ -459,6 +461,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
assert kwargs["tensor_parallel_rank"] == 0
assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
assert kwargs["quantization_config"] is None
return HostWeightPlanResult(runtime_plan)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
@@ -502,6 +505,7 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = object()
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = lambda _model: None # type: ignore[method-assign]
loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign]
@@ -511,11 +515,12 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
)
monkeypatch.setattr(
cache_mod,
"build_runtime_weight_cache_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "writer timed out", "lock_timeout"),
)
def build_cache(_pipeline, **kwargs):
assert kwargs["quantization_config"] == "module_quant_method"
return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
original_weight = model.transformer.weight.detach().clone()
assert loader.load_model(load_device="cpu") is model
@@ -61,6 +61,7 @@ def _build(pipeline: nn.Module, cache_root: Path, **overrides):
"tensor_parallel_rank": 0,
"sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
"use_hsdp": False,
"enable_expert_parallel": False,
"quantization_config": None,
"cfg_parallel_size": 1,
"pipeline_parallel_size": 1,
@@ -125,6 +126,14 @@ def test_final_runtime_content_is_authoritative(tmp_path):
assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
def test_expert_parallel_ownership_fails_closed(tmp_path):
result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
assert result.plan is None
assert result.fallback_code == "unsupported_expert_parallel"
assert list(tmp_path.rglob("*.safetensors")) == []
def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
first = _build(_Pipeline(), tmp_path)
assert first.plan is not None