diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1912628b7..45d1852ea 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -921,6 +921,18 @@ def test_structured_diffusion_config_restricts_host_weight_registration_to_no_al ) +@pytest.mark.parametrize( + "field_name", + [ + "dlo_host_weight_cache_lock_timeout", + "dlo_host_weight_cache_pin_limit_gib", + ], +) +def test_structured_diffusion_config_rejects_non_finite_host_weight_cache_values(field_name): + with pytest.raises(ValueError, match=rf"{field_name} must be finite"): + omni_config_module._DiffusionConfigProjection(**{field_name: float("inf")}) + + def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage(): pipeline = _resolve_pipeline_or_skip("qwen3_tts") legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] diff --git a/tests/diffusion/model_loader/test_host_weight_cache.py b/tests/diffusion/model_loader/test_host_weight_cache.py index 98c6c1c06..634299dcf 100644 --- a/tests/diffusion/model_loader/test_host_weight_cache.py +++ b/tests/diffusion/model_loader/test_host_weight_cache.py @@ -17,6 +17,8 @@ from vllm_omni.diffusion.model_loader.host_weight_cache import ( pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] +_SPAWN_TIMEOUT_SECONDS = 120 + class _LoaderMarker: pass @@ -237,9 +239,11 @@ def test_concurrent_processes_publish_one_immutable_entry(tmp_path): processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)] for process in processes: process.start() - results = [queue.get(timeout=30) for _ in processes] + # Spawn workers import the full vLLM stack before entering the target; + # loaded CI hosts can spend more than 30 seconds in that startup path. + results = [queue.get(timeout=_SPAWN_TIMEOUT_SECONDS) for _ in processes] for process in processes: - process.join(timeout=30) + process.join(timeout=_SPAWN_TIMEOUT_SECONDS) assert all(process.exitcode == 0 for process in processes) assert all(code is None for code, _, _ in results) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 5aa13dc3c..3a735a636 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -18,6 +18,7 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module +from vllm_omni.diffusion.model_loader.host_weight_cache_utils import file_digest from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, TensorBinding, @@ -842,6 +843,7 @@ class TestMmapWeightLoading: }, runtime_layout_key="runtime-layout", post_load_complete=True, + expected_file_digests={str(weight_file): file_digest(weight_file)}, ) backend = DistributedLayerwiseOffloadBackend( OffloadConfig( @@ -880,6 +882,7 @@ class TestMmapWeightLoading: }, runtime_layout_key="runtime-layout", post_load_complete=True, + expected_file_digests={str(tmp_path / "missing.safetensors"): "unused"}, ) backend = DistributedLayerwiseOffloadBackend( OffloadConfig( @@ -897,6 +900,52 @@ class TestMmapWeightLoading: assert pipeline.transformer.weight is old_parameter + def test_host_weight_cache_digest_change_leaves_ordinary_tensors_intact( + self, + tmp_path, + patched_offload_runtime, + ): + pipeline = nn.Module() + pipeline.transformer = nn.Linear(2, 2, bias=False) + old_parameter = pipeline.transformer.weight + old_value = old_parameter.detach().clone() + weight_file = tmp_path / "runtime.safetensors" + save_file({"transformer.weight": torch.full((2, 2), 9.0)}, str(weight_file)) + expected_digest = file_digest(weight_file) + with weight_file.open("r+b") as handle: + handle.seek(-1, os.SEEK_END) + original = handle.read(1) + handle.seek(-1, os.SEEK_END) + handle.write(bytes([original[0] ^ 0xFF])) + plan = HostWeightPlan( + backing_kind="host_weight_cache", + bindings={ + "transformer.weight": TensorBinding( + storage_key="transformer.weight", + file_path=str(weight_file), + ) + }, + runtime_layout_key="runtime-layout", + post_load_complete=True, + expected_file_digests={str(weight_file): expected_digest}, + ) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + + with pytest.raises(RuntimeError, match="changed after loader validation"): + backend._load_weights_from_host_weight_cache(pipeline, modules, plan) + + assert pipeline.transformer.weight is old_parameter + assert torch.equal(pipeline.transformer.weight, old_value) + def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime): pipeline = _MmapPostLoadPipeline() weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()} @@ -1017,6 +1066,7 @@ class TestMmapWeightLoading: 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( OffloadConfig( @@ -1038,6 +1088,78 @@ class TestMmapWeightLoading: assert len(cleanup_calls) == 2 backend.disable() + def test_host_weight_cache_enable_falls_back_to_ordinary_tensors( + self, + tmp_path, + patched_offload_runtime, + ): + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)]) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + pipeline = Pipeline() + weight_file = tmp_path / "runtime.safetensors" + weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()} + save_file(weights, str(weight_file)) + plan = HostWeightPlan( + 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( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + + backend.enable(pipeline) + + assert backend.enabled + assert backend.host_weight_plan is None + assert not backend._using_mmap + assert not backend._using_rank_local_mmap + backend.disable() + + def test_disable_synchronizes_allgather_device_work( + self, + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, + ): + synchronize_calls: list[None] = [] + monkeypatch.setattr( + dist_backend_module.current_omni_platform, + "synchronize", + lambda: synchronize_calls.append(None), + ) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=True, + ), + torch.device("cpu"), + ) + backend.enabled = True + backend._using_rank_local_mmap = False + + backend.disable() + + assert synchronize_calls == [None] + def test_host_weight_cache_registration_uses_budget_and_releases_before_mmap( self, monkeypatch: pytest.MonkeyPatch, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 321aa53a3..84111d883 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -11,6 +11,7 @@ later PRs cut consumers over to these classes. from __future__ import annotations import copy +import math from collections.abc import Mapping from dataclasses import InitVar, dataclass, field, fields from functools import wraps @@ -752,6 +753,11 @@ class _DiffusionConfigProjection: from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode from vllm_omni.quantization import build_quant_config + if not math.isfinite(self.dlo_host_weight_cache_lock_timeout): + raise ValueError("dlo_host_weight_cache_lock_timeout must be finite") + if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib): + raise ValueError("dlo_host_weight_cache_pin_limit_gib must be finite") + if self.dlo_host_weight_cache_pin_limit_gib and ( not self.enable_distributed_layerwise_offload or self.dlo_use_allgather ): diff --git a/vllm_omni/diffusion/model_loader/host_weight_cache.py b/vllm_omni/diffusion/model_loader/host_weight_cache.py index cf1b0f2d3..1b6d4be78 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_cache.py +++ b/vllm_omni/diffusion/model_loader/host_weight_cache.py @@ -402,6 +402,9 @@ def _validate_entry( 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 849f05679..ae724dcf0 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -42,6 +42,7 @@ class HostWeightPlan: 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 @dataclass(frozen=True) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 0d4223bf8..de298794b 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -20,6 +20,7 @@ from __future__ import annotations import time from itertools import chain +from pathlib import Path from typing import Any import torch @@ -28,6 +29,7 @@ from torch import nn from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.diffusion.model_loader.host_weight_cache_utils import file_digest from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, ) @@ -1182,6 +1184,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if any(binding.transform is not None for binding in plan.bindings.values()): raise RuntimeError("Host weight cache bindings cannot carry deferred checkpoint transforms") + cache_files = {binding.file_path for binding in plan.bindings.values()} + if plan.expected_file_digests is None or set(plan.expected_file_digests) != cache_files: + raise RuntimeError("Host weight cache plan does not cover every shard digest") + for file_path in sorted(cache_files): + try: + actual_digest = file_digest(Path(file_path)) + except OSError as exc: + raise RuntimeError(f"Cannot revalidate host weight cache shard {file_path}: {exc}") from exc + if actual_digest != plan.expected_file_digests[file_path]: + raise RuntimeError(f"Host weight cache shard changed after loader validation: {file_path}") + required_names: set[str] = set() for dit_name, dit_module in zip(modules.dit_names, modules.dits): required_names.update(f"{dit_name}.{name}" for name, _ in dit_module.named_parameters()) @@ -1663,15 +1676,25 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self.config.dlo_use_allgather: raise ValueError("Host weight cache backing is no-AllGather only in the first implementation") self._using_rank_local_mmap = True - self._load_weights_from_host_weight_cache( - pipeline, - modules, - self.host_weight_plan, - ) - # Assignments above dropped the model's private tensor - # references. Reclaim those allocations before hook setup so - # steady-state PSS reflects the shared mmap backing. - self._cleanup_after_loading() + try: + self._load_weights_from_host_weight_cache( + pipeline, + modules, + self.host_weight_plan, + ) + except Exception as exc: + logger.warning( + "DLO host weight cache remapping failed; retaining ordinary host tensors: %s", + exc, + ) + self.host_weight_plan = None + self._using_mmap = False + self._using_rank_local_mmap = False + else: + # Assignments above dropped the model's private tensor + # references. Reclaim those allocations before hook setup + # so steady-state PSS reflects the shared mmap backing. + self._cleanup_after_loading() else: raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}") if self._using_rank_local_mmap: @@ -1991,8 +2014,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if not self.enabled and not hasattr(self, "_mmap_file_cache"): return - if self._using_rank_local_mmap: - current_omni_platform.synchronize() + # Copy and communication streams are asynchronous in both AllGather + # and no-AllGather modes. Drain all device work before hooks, process + # groups, registrations, or mapped storage can be torn down. + current_omni_platform.synchronize() registration_released = self._release_registered_mmap() for blocks in self._blocks: