anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion] Rename runtime cache to host weight cache

Browse Source 
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-19 11:23:36 +08:00
1 parent ee665fcfc3
commit 3bb7547a0f
20 files changed
+317 -302
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/model_loader/test_diffusers_loader.py
+20 -20
View File
@@ -406,9 +406,9 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
assert loader.take_host_weight_plan() is None
def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
@@ -427,9 +427,9 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_enable_runtime_cache=True,
dlo_runtime_cache_dir=str(tmp_path),
dlo_runtime_cache_lock_timeout=1.0,
dlo_enable_host_weight_cache=True,
dlo_host_weight_cache_dir=str(tmp_path),
dlo_host_weight_cache_lock_timeout=1.0,
model="unused",
revision=None,
model_config={},
@@ -440,7 +440,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
model.transformer.quant_method = UnquantizedLinearMethod()
events: list[str] = []
runtime_plan = HostWeightPlan(
backing_kind="runtime_cache",
backing_kind="host_weight_cache",
bindings={},
runtime_layout_key="layout-key",
post_load_complete=True,
@@ -485,16 +485,16 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
)
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
assert loader.load_model(load_device="cpu") is model
assert events == ["load", "process", "calibrate", "eval", "cache"]
assert loader.take_host_weight_plan() is runtime_plan
def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(monkeypatch, tmp_path):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
@@ -509,9 +509,9 @@ def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monke
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_enable_runtime_cache=True,
dlo_runtime_cache_dir=str(tmp_path),
dlo_runtime_cache_pin_limit_gib=1.0,
dlo_enable_host_weight_cache=True,
dlo_host_weight_cache_dir=str(tmp_path),
dlo_host_weight_cache_pin_limit_gib=1.0,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
@@ -519,7 +519,7 @@ def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monke
model.transformer = nn.Linear(2, 2, bias=False)
events: list[str] = []
runtime_plan = HostWeightPlan(
backing_kind="runtime_cache",
backing_kind="host_weight_cache",
bindings={},
runtime_layout_key="layout-key",
post_load_complete=True,
@@ -531,23 +531,23 @@ def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monke
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
def unexpected_checkpoint_plan(*_args, **_kwargs):
pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
pytest.fail("registered host weight cache mode must not select the raw checkpoint mmap plan")
def build_cache(_pipeline, **_kwargs):
events.append("cache")
return HostWeightPlanResult(runtime_plan)
monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
assert loader.load_model(load_device="cpu") is model
assert events == ["load", "process", "cache"]
assert loader.take_host_weight_plan() is runtime_plan
def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
def test_dlo_host_weight_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
@@ -562,8 +562,8 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_enable_runtime_cache=True,
dlo_runtime_cache_dir=str(tmp_path),
dlo_enable_host_weight_cache=True,
dlo_host_weight_cache_dir=str(tmp_path),
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
@@ -584,7 +584,7 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat
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)
monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
original_weight = model.transformer.weight.detach().clone()
assert loader.load_model(load_device="cpu") is model