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
tests/diffusion/model_loader/test_runtime_weight_cache.py → tests/diffusion/model_loader/test_host_weight_cache.py
+6 -6
View File
@@ -10,9 +10,9 @@ import pytest
import torch
from torch import nn
from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
build_runtime_weight_cache_plan,
default_runtime_cache_root,
from vllm_omni.diffusion.model_loader.host_weight_cache import (
build_host_weight_cache_plan,
default_host_weight_cache_root,
)
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -42,7 +42,7 @@ class _Pipeline(nn.Module):
def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
assert default_host_weight_cache_root().endswith("/.cache/vllm-omni/dlo-host-weights")
def _build(pipeline: nn.Module, cache_root: Path, **overrides):
@@ -67,7 +67,7 @@ def _build(pipeline: nn.Module, cache_root: Path, **overrides):
"pipeline_parallel_size": 1,
}
kwargs.update(overrides)
return build_runtime_weight_cache_plan(pipeline, **kwargs)
return build_host_weight_cache_plan(pipeline, **kwargs)
def _multiprocess_builder(cache_root: str, queue) -> None:
@@ -121,7 +121,7 @@ def test_unstable_loader_identity_fails_closed(tmp_path):
assert result.plan is None
assert result.fallback_code == "unstable_identity"
assert result.fallback_reason == "runtime-cache identity does not support values of type builtins.object"
assert result.fallback_reason == "host weight cache identity does not support values of type builtins.object"
assert list(tmp_path.rglob("*.safetensors")) == []
tests/diffusion/offloader/test_distributed_layerwise_backend.py
+16 -16
View File
@@ -808,7 +808,7 @@ class _MmapPostLoadPipeline(nn.Module):
class TestMmapWeightLoading:
def test_runtime_cache_maps_final_weights_without_rerunning_post_load(
def test_host_weight_cache_maps_final_weights_without_rerunning_post_load(
self,
tmp_path,
patched_offload_runtime,
@@ -833,7 +833,7 @@ class TestMmapWeightLoading:
weight_file = tmp_path / "runtime.safetensors"
save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
plan = HostWeightPlan(
backing_kind="runtime_cache",
backing_kind="host_weight_cache",
bindings={
"transformer.proj.weight": TensorBinding(
storage_key="transformer.proj.weight",
@@ -854,7 +854,7 @@ class TestMmapWeightLoading:
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
backend._load_weights_from_runtime_cache(pipeline, modules, plan)
backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
assert pipeline.transformer.post_load_calls == 0
assert pipeline.transformer.proj.weight is old_parameter
@@ -862,7 +862,7 @@ class TestMmapWeightLoading:
assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
backend._release_mmap_handles()
def test_runtime_cache_mapping_failure_leaves_ordinary_tensors_intact(
def test_host_weight_cache_mapping_failure_leaves_ordinary_tensors_intact(
self,
tmp_path,
patched_offload_runtime,
@@ -871,7 +871,7 @@ class TestMmapWeightLoading:
pipeline.transformer = nn.Linear(2, 2, bias=False)
old_parameter = pipeline.transformer.weight
plan = HostWeightPlan(
backing_kind="runtime_cache",
backing_kind="host_weight_cache",
bindings={
"transformer.weight": TensorBinding(
storage_key="transformer.weight",
@@ -893,7 +893,7 @@ class TestMmapWeightLoading:
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
with pytest.raises(Exception, match="missing.safetensors"):
backend._load_weights_from_runtime_cache(pipeline, modules, plan)
backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
assert pipeline.transformer.weight is old_parameter
@@ -991,7 +991,7 @@ class TestMmapWeightLoading:
assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
backend.disable()
def test_runtime_cache_enable_reclaims_private_weights_before_hook_setup(
def test_host_weight_cache_enable_reclaims_private_weights_before_hook_setup(
self,
tmp_path,
patched_offload_runtime,
@@ -1013,7 +1013,7 @@ class TestMmapWeightLoading:
weight_file = tmp_path / "runtime.safetensors"
save_file(weights, str(weight_file))
plan = HostWeightPlan(
backing_kind="runtime_cache",
backing_kind="host_weight_cache",
bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
runtime_layout_key="runtime-layout",
post_load_complete=True,
@@ -1038,7 +1038,7 @@ class TestMmapWeightLoading:
assert len(cleanup_calls) == 2
backend.disable()
def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
def test_host_weight_cache_registration_uses_budget_and_releases_before_mmap(
self,
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
@@ -1067,14 +1067,14 @@ class TestMmapWeightLoading:
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dlo_use_allgather=False,
dlo_runtime_cache_pin_limit_gib=1.5,
dlo_host_weight_cache_pin_limit_gib=1.5,
),
torch.device("cuda"),
)
sources = {"runtime.safetensors": [torch.ones(1)]}
backend._runtime_cache_mapped_sources = sources
backend._host_weight_cache_mapped_sources = sources
assert backend._try_register_runtime_cache_mmap()
assert backend._try_register_host_weight_cache_mmap()
assert create_calls == [(sources, int(1.5 * 1024**3))]
backend.enabled = True
@@ -1085,7 +1085,7 @@ class TestMmapWeightLoading:
assert events == ["unregister", "mmap"]
def test_runtime_cache_registration_falls_back_on_unsupported_platform(
def test_host_weight_cache_registration_falls_back_on_unsupported_platform(
self,
patched_offload_runtime,
):
@@ -1094,13 +1094,13 @@ class TestMmapWeightLoading:
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dlo_use_allgather=False,
dlo_runtime_cache_pin_limit_gib=1.0,
dlo_host_weight_cache_pin_limit_gib=1.0,
),
torch.device("cpu"),
)
backend._runtime_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
assert not backend._try_register_runtime_cache_mmap()
assert not backend._try_register_host_weight_cache_mmap()
assert backend._host_registration is None