diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index b9f5f85f9..61cd54e2e 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -310,6 +310,7 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch assert loader.load_model(load_device="cpu") is model assert not loaded_ordinary_weights + assert plan.requires_backend assert loader.take_host_weight_plan() is plan assert loader.take_host_weight_plan() is None @@ -449,7 +450,6 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat backing_kind="host_weight_cache", bindings={}, runtime_layout_key="layout-key", - post_load_complete=True, ) def load_weights(_model): @@ -530,7 +530,6 @@ def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeyp backing_kind="host_weight_cache", bindings={}, runtime_layout_key="layout-key", - post_load_complete=True, ) loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] diff --git a/tests/diffusion/model_loader/test_host_weight_cache.py b/tests/diffusion/model_loader/test_host_weight_cache.py index 74f34e1fb..11e787d3e 100644 --- a/tests/diffusion/model_loader/test_host_weight_cache.py +++ b/tests/diffusion/model_loader/test_host_weight_cache.py @@ -90,7 +90,7 @@ def test_equivalent_process_layouts_reuse_one_entry(tmp_path): assert first.plan is not None assert second.plan is not None assert first.plan.runtime_layout_key == second.plan.runtime_layout_key - assert first.plan.post_load_complete + assert not first.plan.requires_backend assert {binding.file_path for binding in first.plan.bindings.values()} == { binding.file_path for binding in second.plan.bindings.values() } diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index c03e58c1d..40234064d 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -842,7 +842,6 @@ class TestMmapWeightLoading: ) }, runtime_layout_key="runtime-layout", - post_load_complete=True, expected_file_digests={str(weight_file): file_digest(weight_file)}, ) backend = DistributedLayerwiseOffloadBackend( @@ -881,7 +880,6 @@ class TestMmapWeightLoading: ) }, runtime_layout_key="runtime-layout", - post_load_complete=True, expected_file_digests={str(tmp_path / "missing.safetensors"): "unused"}, ) backend = DistributedLayerwiseOffloadBackend( @@ -926,7 +924,6 @@ class TestMmapWeightLoading: ) }, runtime_layout_key="runtime-layout", - post_load_complete=True, expected_file_digests={str(weight_file): expected_digest}, ) backend = DistributedLayerwiseOffloadBackend( @@ -1065,7 +1062,6 @@ class TestMmapWeightLoading: 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, expected_file_digests={str(weight_file): file_digest(weight_file)}, ) backend = DistributedLayerwiseOffloadBackend( @@ -1113,7 +1109,6 @@ class TestMmapWeightLoading: 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, expected_file_digests={str(weight_file): "stale-digest"}, ) backend = DistributedLayerwiseOffloadBackend( @@ -1269,6 +1264,47 @@ class TestMmapWeightLoading: assert not backend._try_register_host_weight_cache_mmap() assert backend._host_registration is None + def test_host_weight_cache_transfer_is_configured_once_for_all_consumers( + self, + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, + ): + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dlo_use_allgather=False, + ), + torch.device("cuda"), + host_weight_plan=HostWeightPlan( + backing_kind="host_weight_cache", + bindings={}, + ), + ) + backend._using_rank_local_mmap = True + backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]} + streamed_hook = SimpleNamespace(registered_mmap=False) + resident_group = SimpleNamespace( + registered_mmap=False, + _cpu_staging_buffers=[{"buffer": torch.ones(1)}], + ) + backend._resident_layer_group = resident_group + register_calls: list[None] = [] + + def register() -> bool: + register_calls.append(None) + return True + + monkeypatch.setattr(backend, "_try_register_host_weight_cache_mmap", register) + + backend._configure_host_weight_cache_transfer([streamed_hook]) # type: ignore[list-item] + + assert register_calls == [None] + assert streamed_hook.registered_mmap + assert resident_group.registered_mmap + assert resident_group._cpu_staging_buffers == [] + assert backend._host_weight_cache_mapped_sources == {} + def test_host_weight_cache_registration_falls_back_on_unsupported_platform( self, patched_offload_runtime, diff --git a/vllm_omni/diffusion/model_loader/host_weight_cache.py b/vllm_omni/diffusion/model_loader/host_weight_cache.py index e0e4ec2ab..c1fe9da4c 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_cache.py +++ b/vllm_omni/diffusion/model_loader/host_weight_cache.py @@ -401,7 +401,6 @@ def _validate_entry( backing_kind="host_weight_cache", bindings=bindings, runtime_layout_key=cache_key, - post_load_complete=True, expected_file_digests={ str(entry_dir / filename): metadata["sha256"] for filename, metadata in file_manifest.items() }, diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index ae724dcf0..a776971aa 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -9,6 +9,7 @@ import json import os from collections.abc import Callable, Sequence from dataclasses import dataclass +from typing import Literal import torch from safetensors import safe_open @@ -37,13 +38,17 @@ class TensorBinding: class HostWeightPlan: """Complete, prevalidated host backing consumed by an offload backend.""" - backing_kind: str + backing_kind: Literal["checkpoint_mmap", "host_weight_cache"] bindings: dict[str, TensorBinding] planned_source_prefixes: frozenset[str] = frozenset() runtime_layout_key: str | None = None - post_load_complete: bool = False expected_file_digests: dict[str, str] | None = None + @property + def requires_backend(self) -> bool: + """Whether ordinary loading was skipped and this plan must be consumed.""" + return self.backing_kind == "checkpoint_mmap" + @dataclass(frozen=True) class HostWeightPlanResult: diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index bdeded55d..c6820f060 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -86,7 +86,7 @@ def get_offload_backend( # Extract and validate configuration config = OffloadConfig.from_od_config(od_config) - plan_requires_backend = host_weight_plan is not None and not host_weight_plan.post_load_complete + plan_requires_backend = host_weight_plan is not None and host_weight_plan.requires_backend if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE: if plan_requires_backend: raise RuntimeError( diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 3105496d9..742964197 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1179,8 +1179,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): """Replace final ordinary-loader tensors with validated cache views.""" from safetensors import safe_open - if not plan.post_load_complete: - raise RuntimeError("A host weight cache plan must describe fully post-processed weights") if any(binding.transform is not None for binding in plan.bindings.values()): raise RuntimeError("Host weight cache bindings cannot carry deferred checkpoint transforms") @@ -1307,6 +1305,29 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) return True + def _configure_host_weight_cache_transfer( + self, + hooks: list[DistributedLayerwiseOffloadHook], + ) -> None: + """Select registered mmap or staging once for every cache consumer.""" + plan = self.host_weight_plan + if ( + (not hooks and self._resident_layer_group is None) + or not self._using_rank_local_mmap + or plan is None + or plan.backing_kind != "host_weight_cache" + ): + return + + self._using_registered_mmap = self._try_register_host_weight_cache_mmap() + self._host_weight_cache_mapped_sources.clear() + for hook in hooks: + hook.registered_mmap = self._using_registered_mmap + if self._resident_layer_group is not None: + self._resident_layer_group.registered_mmap = self._using_registered_mmap + if self._using_registered_mmap: + self._resident_layer_group._cpu_staging_buffers.clear() + def _release_registered_mmap(self) -> bool: registration = self._host_registration if registration is None: @@ -1866,19 +1887,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) pipeline._dlo_residency_controller = self + all_hooks = [hook for group in self._all_hook_groups for hook in group] + self._configure_host_weight_cache_transfer(all_hooks) + if not self._all_hook_groups: self.enabled = bool(self._resident_blocks) - if ( - self.enabled - and self._using_rank_local_mmap - and self.host_weight_plan.backing_kind == "host_weight_cache" - ): - self._using_registered_mmap = self._try_register_host_weight_cache_mmap() - self._host_weight_cache_mapped_sources.clear() - assert self._resident_layer_group is not None - self._resident_layer_group.registered_mmap = self._using_registered_mmap - if self._using_registered_mmap: - self._resident_layer_group._cpu_staging_buffers.clear() if self._using_mmap and not self.enabled: self._release_mmap_handles() return @@ -1886,23 +1899,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Unified allocation: 2 shared output buffers + 2 shared shard buffers # sized to the max block across ALL module groups (gen_layers + # language_model). Groups execute sequentially, so 2 buffers suffice. - all_hooks: list[DistributedLayerwiseOffloadHook] = [] - for group in self._all_hook_groups: - all_hooks.extend(group) - unified_buffers = self._allocate_shared_buffers(all_hooks) unified_shard_buffers = None if self.dp_size > 1: unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) - if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "host_weight_cache": - self._using_registered_mmap = self._try_register_host_weight_cache_mmap() - self._host_weight_cache_mapped_sources.clear() - for hook in all_hooks: - hook.registered_mmap = self._using_registered_mmap - if self._resident_layer_group is not None: - self._resident_layer_group.registered_mmap = self._using_registered_mmap - unified_cpu_staging = None cpu_staging_events = None if self._using_rank_local_mmap and not self._using_registered_mmap: