[Diffusion] Address DLO host cache review feedback

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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]

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)

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,

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
         ):

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()
+        },
     )
 
 

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)

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: