diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index e254e3a43..1184a70ea 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -20,7 +20,7 @@ implemented but has less end-to-end coverage than the primary path. Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported. | Feature | DLO + AllGather | DLO without AllGather | -|---|---|---| +| --- | --- | --- | | **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. | | **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. | | **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | @@ -169,10 +169,167 @@ This mode means: copy per rank. - The scheduler does not require a synchronized DP request wave for DLO. +### Final-layout Host Weight Runtime consumer + +> **Status:** PR2 implementation. The representation-independent identity, +> producer, and restoration contracts landed in +> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer +> wires the final-layout BF16 path into eligible no-AllGather DLO. + +Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies +only to no-AllGather DLO. Enable it with +`--host-weight-runtime-mode preferred` (or `required`) and +`--host-weight-runtime-root `. Registration of shared mappings +and direct asynchronous H2D remain a separate PR3 transport change. + +The modes express operator fallback policy, not different artifact formats: + +- `preferred` consumes an exact hit, but on a miss it allows canonical loading + followed by post-load publication. It is the normal population path. +- `required` consumes the same exact artifact but fails startup on a miss or + unusable artifact. It never invokes canonical DiT fallback or post-load + publication and therefore cannot populate an empty store. + +PR2 has no separate prewarm command. Operators populate one matching producer +cohort per node-local storage domain in `preferred` mode, then restart the same +model revision and parallel layout with `required`. TP coordinates have distinct +identities, while equivalent DP replicas share them. + +The current producer/consumer boundary is the model-declared final-layout BF16 +contract (MiniMax H3 today). It supports ordinary-loader final layouts for TP1 +and TP2 rank identities plus SP layout identities; online quantization, HSDP, +LoRA/adapted weights, and non-default load formats remain ineligible. HWR mode +`disabled`, DLO-disabled, and DLO AllGather configurations stop before source +identity or store construction and retain the existing checkpoint-mmap or +ordinary-loader path. + +Eligibility is decided before constructing HWR, resolving canonical sources, +hashing identity inputs, probing a filesystem, or emitting an HWR observer +event: + +```text +DLO disabled / HWR disabled / DLO AllGather + -> zero final-layout HWR interaction + -> preserve current checkpoint-mmap or ordinary-loader behavior +``` + +When final-layout HWR is explicitly enabled for no-AllGather DLO, runtime +outcomes are authoritative: + +- `LOCAL_HIT`: plan and commit an exact restoration, then transfer the lease + transactionally to DLO. +- `CANONICAL_FALLBACK`: bypass checkpoint mmap, canonically materialize and + finalize the DiT, then call `publish_after_load()` for a future startup. The + current startup retains its canonical tensors. +- `FAILED`: fail startup. Preferred mode does not reinterpret nonretryable + identity, configuration, or compatibility failures as misses. +- Required mode cannot bootstrap an empty store through a `POST_LOAD_ONLY` + producer. + +Checkpoint mmap remains an unchanged control path whenever final-layout HWR is +not selected. + +#### Pre-service transaction boundary + +The startup transaction does not end at restore commit. A lease-backed model is +disposable until backend setup and initial prefetch complete: + +```text +UNRESOLVED + -> LEASE_OWNED_BY_LOADER + -> RESTORE_PLANNED + -> COMMIT_STARTED + -> CARRIER_OWNED + -> BACKEND_OWNED + -> BACKEND_READY + -> IN_SERVICE +``` + +Failure before `COMMIT_STARTED` closes the loader-owned lease. Preferred mode +may canonically load using the untouched skeleton; required mode fails. + +Failure from `COMMIT_STARTED` through `BACKEND_READY` must: + +1. synchronize or otherwise quiesce partial backend and initial-prefetch work; +2. release hook, staging, restore-plan, and model references; +3. close the lease through its current owner; +4. discard the restored model; and +5. in preferred mode, construct a fresh canonical model while bypassing HWR + lookup, HWR publication, and checkpoint mmap for that recovery attempt. + +Required mode fails instead of constructing the fresh fallback. Once the model +enters service, failures follow normal runtime handling rather than startup +fallback. + +#### Finalization phases + +Cold canonical loading separates byte-changing work from shared runtime-state +finalization: + +```text +Cold-only byte-changing: + casting, reordering, packing, quantization, generated scales, + model-specific weight transforms, and calibration that mutates + parameters or persistent buffers + +Shared cold/warm non-byte: + validation, hook installation, eval state, bookkeeping, + and non-persistent runtime state +``` + +The warm path must snapshot restored tensor bytes and backing pointers before +shared finalization and prove both are unchanged afterward. + +#### Lease ownership + +Lease transfer is single-owner and single-take: + +```text +loader owns + -> restore plan borrows + -> commit succeeds + -> runner carrier takes once + -> backend takes once before asynchronous work + -> backend drains pending H2D work + -> backend releases hooks, staging, and model references + -> backend closes lease +``` + +The runner carrier is process-local, rejects serialization, and rejects a +duplicate `take()`. If the carrier still owns the lease, runner cleanup closes +it. Once the backend takes ownership, only backend abort or teardown may drain +work and close it. + +The implementation keeps the final-layout lease through backend setup and +initial prefetch. The backend uses the existing two bounded rank-local staging +slots for the immutable HWR tensors, and its transactional `enable()` cleanup +removes partial hooks and closes the lease before reporting a setup failure. +Preferred mode then uses the runner's fresh canonical retry; required mode +propagates the failure. + +#### PR2 promotion gates + +- Warm hit performs zero ordinary DiT materialization and zero producer calls. +- Shared warm finalization changes neither restored bytes nor backing pointers. +- Preferred `FAILED` outcomes do not silently fall back. +- Planning failure leaves the skeleton untouched. +- Any pre-service failure after commit begins discards the restored model; + preferred mode uses a fresh canonical model and required mode fails. +- Mixed components load normally and no required tensor remains on `meta`. +- Disabled mode and AllGather emit zero final-layout HWR interaction. +- The lease carrier rejects duplicate take and serialization. +- Backend setup or prefetch failure drains asynchronous work before lease close. +- Compatible checkpoint mmap remains an unchanged benchmark and control path. +- Prewarm uses one matching producer cohort per TP/SP topology and storage + domain. +- Benchmarks compare ordinary loading, cold publication, and warm staging with + output parity, startup latency, aggregate PSS, page-cache state, HBM, and H2D + payload. + ## Parallelism compatibility | Parallelism | DLO + AllGather | DLO without AllGather | -|---|---|---| +| --- | --- | --- | | **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. | | **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. | | **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. | @@ -185,7 +342,8 @@ This mode means: the SP group becomes DLO's sharding group in AllGather mode. - **DP + TP/SP without AllGather:** standard model loading defines the rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight - collective. + collective. When the model declares the final-layout contract, HWR keys the + reusable artifact by TP rank/size and SP layout. - **HSDP + SP:** the general parallel configuration permits HSDP over SP, but DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight materialization and synchronization. @@ -211,10 +369,11 @@ reconstructs those tensors with their recorded layouts. Other online methods must use `--dlo-no-use-allgather` or disable online quantization until their runtime layouts are validated. -A normalized runtime mmap cache, built through the ordinary loader, is the -proposed general mechanism for sharing transformed TP or quantized layouts. -That cache and its publication/lifecycle protocol are intentionally outside -this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). +The Host Weight Runtime representation and publication contracts are merged; +see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and +[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned +no-AllGather consumer above is the PR2 implementation; registration and direct +asynchronous H2D remain a separate transport follow-up. ## Validation coverage @@ -243,7 +402,7 @@ VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap or shared-mmap host-memory savings. | Configuration | Result | Warm E2E | Peak device memory | Host PSS | -|---|---:|---:|---:|---:| +| --- | ---: | ---: | ---: | ---: | | DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB | | DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB | | DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB | @@ -275,7 +434,7 @@ pipeline components, so each worker should be compared with the same worker in the other storage mode. | Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction | -|---|---:|---:|---:|---:|---:| +| --- | ---: | ---: | ---: | ---: | ---: | | DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB | | DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB | | **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** | diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 49139e695..fb13b0391 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -63,6 +63,8 @@ omni = Omni( | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` | | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | +| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` | +| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset | ## Host-weight loading @@ -104,6 +106,88 @@ process. When the effective DLO group size is one, `dlo_use_allgather=True` does not perform a collective and uses the same rank-local transfer behavior. +### Final-layout Host Weight Runtime + +HWR is an opt-in startup optimization for models that declare the final-layout +BF16 restore contract. Use it only with no-AllGather DLO: + +```bash +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --host-weight-runtime-mode preferred \ + --host-weight-runtime-root /var/cache/vllm-omni/hwr +``` + +#### Choosing a mode + +| Mode | Exact local artifact hit | Miss or recoverable artifact/store problem | Intended use | +| --- | --- | --- | --- | +| `disabled` | HWR is not consulted | Use the existing checkpoint-mmap or ordinary-loader path | Default compatibility path | +| `preferred` | Restore the final-layout artifact | Load canonically, serve with those tensors, and attempt to publish an artifact for the next startup | Normal deployment and store population | +| `required` | Restore the final-layout artifact | Fail startup without canonical DiT fallback or publication | Enforce a pre-populated store in controlled rollouts or CI | + +Both enabled modes still fail on non-retryable configuration, identity, or +compatibility errors. `preferred` is a fallback policy for a cache miss or a +recoverable cache problem; it does not hide an invalid deployment. + +#### Populating a store for `required` mode + +`required` is deliberately consume-only. PR2 does not include a separate +prewarm command, so populate each node-local storage domain with one matching +`preferred` producer cohort: + +1. Choose a persistent, writable root visible to every diffusion worker in the + storage domain. Do not use a process-private temporary directory. +2. Start the deployment with `preferred`, using the exact model revision, + dtype, TP size, SP configuration, and other weight-layout settings intended + for serving. No inference request is needed; publication happens during + model startup. Wait for the engine to become healthy, then shut it down + cleanly. +3. Restart with the same arguments and root, changing only the mode to + `required`. A successful startup proves that every worker acquired a valid + artifact; a miss, corrupt artifact, or incompatible identity fails startup. +4. Repeat the population start on every node or storage domain because this + store is node-local. + +For example: + +```bash +# First startup: canonically load and populate the exact artifacts. +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --host-weight-runtime-mode preferred \ + --host-weight-runtime-root /var/cache/vllm-omni/hwr + +# After a healthy startup and clean shutdown, enforce cache hits. +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --host-weight-runtime-mode required \ + --host-weight-runtime-root /var/cache/vllm-omni/hwr +``` + +Include the same TP and SP layout flags in both commands. DP rank and DP size +are excluded from artifact identity, so the population and serving DP sizes may +differ and equivalent DP replicas share artifacts. TP rank is included, so a +TP-N deployment normally needs N rank-specific artifacts in each storage +domain; launching the matching TP cohort creates that set. If the model +revision or layout changes, run `preferred` again for the new identity before +returning to `required`. + +On a cold start, the canonical loader remains authoritative and publishes a +validated final-layout artifact for later workers. A warm start restores the +DiT final tensors without ordinary DiT materialization, then DLO streams them +through the same two bounded host staging slots. The artifact identity includes +the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP +layouts do not alias one another. Publication failure remains separate from a +`preferred` serving startup; the next `required` startup provides the explicit +artifact-availability check. + +When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does +not resolve HWR sources or construct its store. + ## Declarative topology Models may declare an `OffloadPlan` instead of embedding offload logic: @@ -135,7 +219,9 @@ must enter each collective. - Direct checkpoint mmap currently requires TP1. TP greater than one is outside the Phase A shared-mmap support scope and falls back before model mutation to the ordinary TP-aware loader. DLO can stream that runtime layout, - but it provides no shared-mmap host-memory benefit or guarantee. + and eligible no-AllGather configurations may use HWR final-layout artifacts, + but direct checkpoint mmap provides no shared-mmap host-memory guarantee for + TP greater than one. - HSDP plus AllGather is rejected to avoid double sharding. HSDP without AllGather has limited end-to-end validation. - Per-tensor online FP8 linears use the ordinary loader and can run with either @@ -147,10 +233,10 @@ must enter each collective. `OffloadPlan` that declares eligible `resident_dit_paths`. - DP concurrency requires an explicit, identical inference-step count. -Sharing transformed TP or quantized runtime layouts through a normalized mmap -cache is a follow-up design in +Sharing quantized or otherwise unvalidated transformed layouts through a +normalized HWR producer is a follow-up design in [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of -the direct-checkpoint path. +the current BF16 final-layout path. See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for an end-to-end example. diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index ce3c99418..c4d08f3a9 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -1,17 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for the DiffusersPipelineLoader. """ import json +from pathlib import Path from types import SimpleNamespace import pytest import torch import torch.nn as nn from huggingface_hub import snapshot_download +from safetensors.torch import save_file from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none @@ -23,6 +25,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( TensorBinding, ) from vllm_omni.diffusion.models.helios import HeliosPipeline +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.diffusion.registry import initialize_model pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -73,6 +76,73 @@ class _DummyPipelineModel(nn.Module): return loaded +class _HWRTransformer(nn.Module): + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="test-hwr-transformer", + version="1", + ) + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)) + + def validate_restored_host_weights(self): + assert self.weight.dtype is torch.bfloat16 + + +class _HWRPipeline(nn.Module): + def __init__(self, source_root: Path): + super().__init__() + self.transformer = _HWRTransformer() + self.load_count = 0 + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=str(source_root), + subfolder=None, + revision=None, + prefix="transformer.", + fall_back_to_pt=False, + ) + ] + + def load_weights(self, weights): + loaded: set[str] = set() + params = dict(self.named_parameters()) + for name, tensor in weights: + if name in params: + params[name].data.copy_(tensor.to(dtype=params[name].dtype)) + loaded.add(name) + self.load_count += 1 + return loaded + + +def _hwr_config(model: str | Path, root: Path, *, mode: str = "preferred") -> SimpleNamespace: + return SimpleNamespace( + model=str(model), + dtype=torch.bfloat16, + host_weight_runtime_mode=mode, + host_weight_runtime_root=str(root), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + lora_path=None, + quantization_config=None, + diffusion_attention_config=None, + parallel_config=SimpleNamespace( + use_hsdp=False, + data_parallel_size=1, + sequence_parallel_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + cfg_parallel_size=1, + enable_expert_parallel=False, + ulysses_degree=1, + ring_degree=1, + allgather_degree=1, + ulysses_mode="strict", + ), + ) + + def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader: od_config = SimpleNamespace( dtype=torch.float32, @@ -92,6 +162,81 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade return loader +def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( + tmp_path: Path, + monkeypatch, +): + canonical_root = tmp_path / "canonical" + canonical_root.mkdir() + save_file( + {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)}, + str(canonical_root / "model.safetensors"), + ) + store_root = tmp_path / "hwr-store" + + def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]: + loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root)) + pipeline = _HWRPipeline(canonical_root) + monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline) + return loader, pipeline + + cold_loader, cold_model = make_loader() + cold = cold_loader.load_model(load_device="cpu", device=torch.device("cpu")) + assert cold is cold_model + assert cold_model.load_count == 1 + assert cold_loader._hwr_state is not None + + warm_loader, warm_model = make_loader() + monkeypatch.setattr( + warm_loader, + "_process_weights_after_loading", + lambda *args, **kwargs: pytest.fail("warm HWR restore re-entered byte-changing finalization"), + ) + warm = warm_loader.load_model(load_device="cpu", device=torch.device("cpu")) + + assert warm is warm_model + assert warm_model.load_count == 0 + assert torch.equal(warm_model.transformer.weight, cold_model.transformer.weight) + from vllm_omni.diffusion.offloader.startup import take_offload_startup_state + + startup_state = take_offload_startup_state(warm) + assert startup_state is not None + warm_plan = startup_state.host_weight_plan + assert warm_plan is not None + assert warm_plan.lease_carrier is not None + warm_plan.lease_carrier.close() + + +def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch): + from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module + + loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(tmp_path, tmp_path / "store")) + models: list[_DummyPipelineModel] = [] + + def init_model(*args, **kwargs): + del args, kwargs + model = _DummyPipelineModel(source_prefix="transformer.") + models.append(model) + return model + + def commit_error(*args, **kwargs): + del args, kwargs + raise loader_module._HWRCommitError("restore commit failed") + + monkeypatch.setattr(loader, "_init_from_load_format", init_model) + monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: ()) + monkeypatch.setattr(loader, "_resolve_hwr", commit_error) + monkeypatch.setattr(loader, "load_weights", lambda *args, **kwargs: None) + monkeypatch.setattr(loader, "_process_weights_after_loading", lambda *args, **kwargs: None) + monkeypatch.setattr(loader, "_apply_skip_softmax_calibration", lambda *args, **kwargs: None) + + recovered = loader.load_model(load_device="cpu", device=torch.device("cpu")) + + assert len(models) == 2 + assert recovered is models[1] + assert loader.take_host_weight_plan() is None + + def _make_dlo_online_quant_config() -> OmniDiffusionConfig: return OmniDiffusionConfig( model="", @@ -103,6 +248,64 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig: ) +@pytest.mark.parametrize( + ("dist_offload", "use_allgather"), + [ + (False, False), + (True, True), + ], +) +def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction( + monkeypatch, + tmp_path, + dist_offload, + use_allgather, +): + """Disabled and AllGather paths must never construct or probe HWR.""" + from vllm_omni.host_weight_runtime import HostWeightRuntime + + root = tmp_path / "must-not-be-touched" + loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root)) + model = _DummyPipelineModel(source_prefix="transformer.") + modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,)) + + def unexpected_store_construction(*args, **kwargs): + raise AssertionError(f"HWR store interaction was not eligible: {args}, {kwargs}") + + monkeypatch.setattr(HostWeightRuntime, "from_config", unexpected_store_construction) + assert ( + loader._resolve_hwr( + model, + modules, + dist_offload=dist_offload, + use_allgather=use_allgather, + load_format="default", + sources=tuple(model.weights_sources), + ) + is None + ) + assert not root.exists() + + +def test_required_hwr_rejects_a_model_without_a_restore_contract(tmp_path): + loader = DiffusersPipelineLoader( + LoadConfig(), + _hwr_config("dummy-model", tmp_path / "store", mode="required"), + ) + model = _DummyPipelineModel(source_prefix="transformer.") + modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,)) + + with pytest.raises(ValueError, match="restore contract"): + loader._resolve_hwr( + model, + modules, + dist_offload=True, + use_allgather=False, + load_format="default", + sources=tuple(model.weights_sources), + ) + + @pytest.mark.parametrize("offline", [False, True]) def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod diff --git a/tests/diffusion/model_loader/test_final_layout_host_weights.py b/tests/diffusion/model_loader/test_final_layout_host_weights.py index 90cdd4998..b8ebc27f5 100644 --- a/tests/diffusion/model_loader/test_final_layout_host_weights.py +++ b/tests/diffusion/model_loader/test_final_layout_host_weights.py @@ -726,6 +726,36 @@ def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) -> assert identity_v2.key != base_identity.key +@pytest.mark.parametrize( + "parallel", + [ + FinalLayoutParallelIdentity(), + FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0), + FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1), + FinalLayoutParallelIdentity(sequence_parallel_size=2, ulysses_degree=2), + FinalLayoutParallelIdentity(sequence_parallel_size=2, ring_degree=2), + FinalLayoutParallelIdentity( + tensor_parallel_size=2, + tensor_parallel_rank=1, + sequence_parallel_size=2, + ulysses_degree=2, + ), + ], +) +def test_parallel_identity_matrix_is_exact_for_tp_and_sp(tmp_path: Path, parallel: FinalLayoutParallelIdentity) -> None: + """TP rank/size and SP backend are identity coordinates for HWR artifacts.""" + model = _TinyPipeline() + source = _prepared_source(tmp_path) + context = _identity(model, source, request=_request(parallel=parallel)) + layout = context.identity.layout + + assert layout.tensor_parallel_size == parallel.tensor_parallel_size + assert layout.tensor_parallel_rank == parallel.tensor_parallel_rank + assert layout.sequence_parallel_size == parallel.sequence_parallel_size + assert layout.sequence_parallel_backend == parallel.sequence_parallel_backend + assert "data_parallel" not in context.identity.layout.metadata.to_value().get("parallel", {}) + + @pytest.mark.parametrize( ("semantic_request", "message"), [ diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 78dd82b6f..5a3390f89 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -38,6 +38,7 @@ from vllm_omni.diffusion.offloader.offload_plan import ( OffloadPlan, get_offload_plan, ) +from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -695,6 +696,67 @@ class _SingleBlockModel(nn.Module): self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) +class _HWRPipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=3) + + +class _FakeHostWeightLease: + def __init__(self, resolution_id: str): + self.closed = False + self.provenance = SimpleNamespace(resolution_id=resolution_id) + + def close(self): + self.closed = True + + +class _FakeLeaseCarrier: + def __init__(self, lease: _FakeHostWeightLease): + self.lease = lease + self.taken = False + + @property + def closed(self): + return self.taken or self.lease.closed + + def take(self): + if self.taken: + raise RuntimeError("already taken") + self.taken = True + return self.lease + + def close(self): + if not self.taken: + self.lease.close() + + +def _fake_hwr_plan(resolution_id: str = "hwr-test"): + lease = _FakeHostWeightLease(resolution_id) + carrier = _FakeLeaseCarrier(lease) + plan = HostWeightPlan( + backing_kind="host_weight_runtime", + bindings={}, + lease_carrier=carrier, + ) + return plan, carrier, lease + + +def _hwr_backend(): + plan, carrier, lease = _fake_hwr_plan() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + return _HWRPipeline(), backend, carrier, lease + + class _MultiBlockModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] @@ -838,6 +900,47 @@ class TestMmapWeightLoading: assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group) backend.disable() + def test_hwr_carrier_is_taken_and_released_after_bounded_staging(self, patched_offload_runtime): + """A warm HWR plan uses the existing two-slot rank-local transport.""" + pipeline, backend, carrier, lease = _hwr_backend() + model_bytes = sum(param.numel() * param.element_size() for param in pipeline.transformer.parameters()) + + backend.enable(pipeline) + + assert carrier.taken + assert backend._using_rank_local_mmap + assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group) + staging_bytes = sum( + buffer.numel() * buffer.element_size() + for buffer in backend._all_hook_groups[0][0].cpu_staging_buffers[0].values() + ) + assert staging_bytes * 2 < model_bytes + assert not lease.closed + + backend.disable() + assert lease.closed + + def test_hwr_backend_failure_drains_partial_setup_before_lease_close( + self, + monkeypatch, + patched_offload_runtime, + ): + pipeline, backend, carrier, lease = _hwr_backend() + monkeypatch.setattr( + backend, + "_allocate_shared_buffers", + lambda hooks: (_ for _ in ()).throw(RuntimeError("device buffer allocation failed")), + ) + + with pytest.raises(RuntimeError, match="device buffer allocation failed"): + backend.enable(pipeline) + + assert carrier.taken + assert lease.closed + assert not backend.enabled + assert not backend._blocks + assert not backend._all_hook_groups + class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): @@ -1504,6 +1607,48 @@ class TestMmapValidation: class TestConfigValidation: """Tests for configuration validation in OffloadConfig / execute_request.""" + @pytest.mark.parametrize( + ("dp", "sp", "tp", "use_allgather", "expected_dlo_group"), + [ + (1, 1, 1, True, 1), + (2, 1, 1, True, 2), + (1, 4, 1, True, 4), + (2, 2, 2, True, 2), + (2, 2, 2, False, 1), + (1, 4, 2, False, 1), + ], + ) + def test_dlo_group_selection_covers_dp_sp_tp_matrix( + self, + dp, + sp, + tp, + use_allgather, + expected_dlo_group, + ): + """DLO sharding follows DP, falls back to SP, and never uses TP.""" + + config = OffloadConfig.from_od_config( + SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=use_allgather, + dlo_resident_layers=0, + pin_cpu_memory=False, + parallel_config=SimpleNamespace( + data_parallel_size=dp, + use_hsdp=False, + hsdp_shard_size=-1, + hsdp_replicate_size=1, + sequence_parallel_size=sp, + tensor_parallel_size=tp, + ), + ) + ) + assert config.dp_size == expected_dlo_group + assert config.dlo_use_allgather is use_allgather + def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch): import vllm_omni.diffusion.offloader as offloader_module @@ -1538,6 +1683,53 @@ class TestConfigValidation: host_weight_plan=plan, ) + def test_startup_recovery_stays_inside_offloader_boundary(self, monkeypatch): + import vllm_omni.diffusion.offloader as offloader_module + + class FailingBackend: + def __init__(self): + self.disabled = False + + def enable(self, pipeline): + del pipeline + raise RuntimeError("initial prefetch failed") + + def disable(self): + self.disabled = True + + warm_pipeline = nn.Module() + canonical_pipeline = nn.Module() + plan, carrier, _lease = _fake_hwr_plan("hwr-recovery") + failing_backend = FailingBackend() + calls = [] + + def fake_backend(od_config, device, host_weight_plan): + del od_config, device + calls.append(host_weight_plan) + return failing_backend if host_weight_plan is plan else None + + attach_offload_startup_state( + warm_pipeline, + OffloadStartupState( + host_weight_plan=plan, + fresh_model_loader=lambda: canonical_pipeline, + allow_fresh_retry=True, + ), + ) + monkeypatch.setattr(offloader_module, "get_offload_backend", fake_backend) + + recovered, backend = offloader_module.enable_offload_backend( + SimpleNamespace(), + warm_pipeline, + device=torch.device("cpu"), + ) + + assert recovered is canonical_pipeline + assert backend is None + assert failing_backend.disabled + assert carrier.closed + assert calls == [plan, None] + def test_hsdp_with_allgather_rejected(self): """HSDP + DLO + AllGather should raise ValueError (double sharding).""" from vllm_omni.diffusion.offloader.base import OffloadConfig diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index c331f08ee..d2a9f1165 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -825,9 +825,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): del kwargs return SimpleNamespace(transformer=torch.nn.Identity()) - def take_host_weight_plan(self): - return None - class _DummyMemoryProfiler: consumed_memory = 0 @@ -869,8 +866,8 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) monkeypatch.setattr( model_runner_module, - "get_offload_backend", - lambda od_config, device, host_weight_plan: None, + "enable_offload_backend", + lambda od_config, pipeline, device: (pipeline, None), ) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend @@ -883,8 +880,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): assert dummy_cache_backend.enabled is False -@pytest.mark.core_model -@pytest.mark.cpu def test_set_forward_context_enters_vllm_config_contexts(monkeypatch): """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions.""" import vllm.config.vllm as vllm_config_module diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index bbacf7f2c..94c3eb48f 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for step-level diffusion execution across runner / worker / executor / engine.""" import contextlib @@ -879,7 +879,11 @@ class TestRunner: monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler) - monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr( + model_runner_module, + "enable_offload_backend", + lambda _config, pipeline, **_kwargs: (pipeline, None), + ) monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) with pytest.raises(ValueError, match="RequestOnlyPipeline"): diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index d84f49b31..6c5828902 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Integration tests for diffusion pipeline streaming output mode. Covers: @@ -544,7 +544,11 @@ class TestSupportedPipelines: monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler) - monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr( + model_runner_module, + "enable_offload_backend", + lambda _config, pipeline, **_kwargs: (pipeline, None), + ) monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) with pytest.raises(ValueError, match="NoStepPipeline"): diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 64c1ce99a..37d4af8c3 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -362,6 +362,35 @@ def test_serve_cli_forwards_distributed_offload_residency(): assert engine_args["dlo_resident_layers"] == 20 +def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--enable-distributed-layerwise-offload", + "--dlo-no-use-allgather", + "--host-weight-runtime-mode", + "preferred", + "--host-weight-runtime-root", + "/var/cache/vllm-omni/hwr", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + engine_args = stage_cfg["engine_args"] + + assert explicit_kwargs["host_weight_runtime_mode"] == "preferred" + assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" + assert engine_args["host_weight_runtime_mode"] == "preferred" + assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" + + def test_serve_cli_accepts_diffusion_compile_controls(): """Ensure both compile controls reach the diffusion stage.""" parser = TrackingArgumentParser() diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py index 763d2622a..f776f5778 100644 --- a/tests/host_weight_runtime/test_filesystem_store.py +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -389,6 +389,24 @@ def test_build_lookup_share_one_published_backing_and_lease_blocks_cleanup(tmp_p assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS +def test_publication_hardens_artifact_after_atomic_rename(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert built.status is StoreStatus.BUILT + assert built.lease is not None + built.lease.close() + artifact_dir = store.artifacts_dir / identity.key + assert artifact_dir.stat().st_mode & 0o777 == 0o555 + + def test_concurrent_lease_close_waits_for_resource_teardown(tmp_path: Path) -> None: identity = _identity() store = _make_store(tmp_path / "store") diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py index 770a35bfa..21bab37de 100644 --- a/tests/host_weight_runtime/test_runtime_resolution.py +++ b/tests/host_weight_runtime/test_runtime_resolution.py @@ -17,6 +17,7 @@ from vllm_omni.host_weight_runtime import ( FailureCode, HostWeightFailure, HostWeightLease, + HostWeightLeaseCarrier, HostWeightRuntime, HostWeightRuntimeConfig, IntegrityPolicy, @@ -166,6 +167,28 @@ def test_local_production_then_exact_warm_hit_emits_one_terminal_report_each(tmp assert all(report.resolution_id for report in reports) +def test_lease_carrier_is_single_take_and_process_local(tmp_path: Path) -> None: + identity = _identity() + producer = CountingProducer(identity) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")) + ) + resolution = runtime.resolve(identity, producer=producer) + assert resolution.lease is not None + + carrier = HostWeightLeaseCarrier(resolution.lease) + with pytest.raises(TypeError, match="cannot be serialized"): + carrier.__reduce_ex__(4) + + lease = carrier.take() + assert lease is resolution.lease + with pytest.raises(RuntimeError, match="already taken"): + carrier.take() + carrier.close() # ownership has moved; the carrier must not close twice. + assert not lease.closed + lease.close() + + @pytest.mark.parametrize( ("mode", "expected"), [ diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 6f3a2ab65..280b3c86b 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -696,6 +696,8 @@ class _DiffusionConfigProjection: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = Field(default=0, ge=0) + host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled" + host_weight_runtime_root: str | None = None pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) @@ -753,6 +755,7 @@ class _DiffusionConfigProjection: TransformerConfig, build_attention_config, parse_kv_cache_skip_selector, + validate_host_weight_runtime_options, ) from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode from vllm_omni.quantization import build_quant_config @@ -830,6 +833,11 @@ class _DiffusionConfigProjection: elif self.max_cpu_loras < 1: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + validate_host_weight_runtime_options( + mode=self.host_weight_runtime_mode, + root=self.host_weight_runtime_root, + ) + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): raise ValueError( "diffusers_load_kwargs and diffusers_call_kwargs are only " diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index fdfd7c288..1825fbdfe 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -386,6 +386,8 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None dlo_resident_layers: int | None = None + host_weight_runtime_mode: str | None = None + host_weight_runtime_root: str | None = None # Diffusion-specific debug and observability knobs. enable_diffusion_pipeline_profiler: bool | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0f6ae494a..c64330c6e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -100,6 +100,19 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any] return normalized +def validate_host_weight_runtime_options(*, mode: object, root: object) -> None: + """Validate HWR policy without touching the configured storage domain. + + Filesystem locality and store construction belong to the eligible loader + path. Keeping this check purely structural is what lets disabled and + AllGather DLO configurations retain zero HWR interaction. + """ + if mode not in {"disabled", "preferred", "required"}: + raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") + if mode != "disabled" and (not isinstance(root, str) or not root.strip()): + raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") + + def parse_kv_cache_skip_selector( selector: str | list[int] | tuple[int, ...] | set[int] | None, ) -> set[int] | None: @@ -766,6 +779,11 @@ class OmniDiffusionConfig: dlo_use_allgather: bool = True # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 + # Final-layout Host Weight Runtime policy. The loader only activates this + # for eligible no-AllGather DLO; all other configurations preserve their + # existing loader/storage path. + host_weight_runtime_mode: str = "disabled" + host_weight_runtime_root: str | None = None pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -1142,6 +1160,11 @@ class OmniDiffusionConfig: elif self.max_cpu_loras < 1: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + validate_host_weight_runtime_options( + mode=self.host_weight_runtime_mode, + root=self.host_weight_runtime_root, + ) + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): raise ValueError( "diffusers_load_kwargs and diffusers_call_kwargs are only " diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index a7a41088d..d582270cc 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import contextlib import dataclasses import glob @@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) +from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, build_checkpoint_mmap_plan, @@ -125,7 +126,7 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ ) -class DiffusersPipelineLoader: +class DiffusersPipelineLoader(HWRLoaderMixin): """Model loader that can load diffusers pipeline components from disk.""" @dataclasses.dataclass @@ -159,6 +160,9 @@ class DiffusersPipelineLoader: self.quant_config = od_config.quantization_config self.parallel_config = od_config.parallel_config self.host_weight_plan: HostWeightPlan | None = None + self._hwr_state: dict[str, object] | None = None + self._last_load_request: dict[str, object] | None = None + self._force_canonical_load = False def take_host_weight_plan(self) -> HostWeightPlan | None: """Transfer the loader-produced plan to the offload backend.""" @@ -469,6 +473,13 @@ class DiffusersPipelineLoader: ) -> nn.Module: """Load a model with the given configurations.""" self.host_weight_plan = None + self._hwr_state = None + self._last_load_request = { + "load_device": load_device, + "load_format": load_format, + "custom_pipeline_name": custom_pipeline_name, + "device": device, + } if load_format is None: load_format = "default" # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), @@ -511,7 +522,34 @@ class DiffusersPipelineLoader: plan_result = None weight_sources = self._get_weight_sources(model) - if _dist_offload: + hwr_state = None + if not self._force_canonical_load: + try: + hwr_state = self._resolve_hwr( + model, + ModuleDiscovery.discover(model), + dist_offload=_dist_offload, + use_allgather=_use_ag, + load_format=load_format, + sources=weight_sources, + ) + except _HWRCommitError: + from vllm_omni.host_weight_runtime import RuntimeMode + + mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled")) + if mode is not RuntimeMode.PREFERRED: + raise + logger.warning( + "HWR restore commit failed; discarding the model and retrying a fresh canonical load", + exc_info=True, + ) + del model + return self.load_fresh_canonical_model() + self._hwr_state = hwr_state + hwr_active = hwr_state is not None + if hwr_active and hwr_state is not None: + self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan")) + if _dist_offload and not hwr_active and not self._force_canonical_load: modules = ModuleDiscovery.discover(model) plan_result = build_checkpoint_mmap_plan( model, @@ -584,12 +622,43 @@ class DiffusersPipelineLoader: self.load_weights(model) self._process_weights_after_loading(model, target_device) + # A warm final-layout hit has already completed all + # byte-changing work through the restorer. Shared runtime + # finalization happens once at the end for both cold and warm + # paths; the warm path never re-enters the ordinary + # materialization/finalization pipeline. + if offload_after_quant: model.to("cpu") logger.info("Quantization complete, offloaded model back to CPU") - self._apply_skip_softmax_calibration(model) - return model.eval() + try: + self._apply_skip_softmax_calibration(model) + model = model.eval() + if self._hwr_state is not None: + warm_snapshot = self._hwr_state.get("warm_snapshot") + if warm_snapshot is not None: + self._assert_final_layout_tensors_unchanged(model, cast(dict[str, tuple[int, str]], warm_snapshot)) + self._publish_hwr_after_load(model, ModuleDiscovery.discover(model), self._hwr_state) + except Exception: + hwr_plan = self._hwr_state.get("plan") if self._hwr_state is not None else None + if isinstance(hwr_plan, HostWeightPlan): + carrier = hwr_plan.lease_carrier + if carrier is not None: + carrier.close() + from vllm_omni.host_weight_runtime import RuntimeMode + + mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled")) + if mode is RuntimeMode.PREFERRED: + logger.warning( + "HWR warm finalization failed; discarding the model and retrying a fresh canonical load", + exc_info=True, + ) + del model + return self.load_fresh_canonical_model() + raise + self._attach_offload_startup_state(model) + return model @staticmethod def _request_offload_after_quant(model: nn.Module) -> int: diff --git a/vllm_omni/diffusion/model_loader/host_weight_loader.py b/vllm_omni/diffusion/model_loader/host_weight_loader.py new file mode 100644 index 000000000..3094cd899 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weight_loader.py @@ -0,0 +1,479 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Host Weight Runtime integration for diffusion model loaders. + +This mixin keeps HWR identity, restore, publication, and startup-recovery +policy out of the ordinary diffusers loader. The host loader supplies the +canonical source-preparation and weight-loading hooks; this module owns only +the optional final-layout transaction. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import os +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, cast + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan, TensorBinding + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext + from vllm_omni.diffusion.model_loader.host_weights.source_identity import PreparedWeightSource + + +class _HWRCommitError(RuntimeError): + """A committed warm restore made the current model disposable.""" + + +class HWRLoaderMixin: + """Optional final-layout HWR behavior shared by diffusion loaders.""" + + @staticmethod + def _identity_value(value: object) -> object: + """Convert config objects into deterministic identity metadata.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, dict): + return {str(key): HWRLoaderMixin._identity_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [HWRLoaderMixin._identity_value(item) for item in value] + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + try: + return HWRLoaderMixin._identity_value(to_dict()) + except TypeError: + pass + if dataclasses.is_dataclass(value): + return HWRLoaderMixin._identity_value(dataclasses.asdict(value)) + return f"{type(value).__module__}.{type(value).__qualname__}:{value!r}" + + @staticmethod + def _identity_fingerprint(value: object) -> str: + from vllm_omni.host_weight_runtime import CanonicalJson + + return hashlib.sha256(CanonicalJson.from_value(value).encoded).hexdigest() + + @staticmethod + def _snapshot_final_layout_tensors(model: nn.Module, names: Iterable[str]) -> dict[str, tuple[int, str]]: + snapshot: dict[str, tuple[int, str]] = {} + for name in names: + parent_path, _, leaf_name = name.rpartition(".") + parent = model.get_submodule(parent_path) + tensor = parent._parameters.get(leaf_name) + if tensor is None: + tensor = parent._buffers.get(leaf_name) + if tensor is None or tensor.is_meta: + raise RuntimeError(f"cannot snapshot missing or meta final-layout tensor {name!r}") + value = tensor.detach() + if value.device.type != "cpu": + value = value.cpu() + value = value.contiguous() + digest = hashlib.sha256(memoryview(value.view(torch.uint8).numpy())).hexdigest() + snapshot[name] = (tensor.untyped_storage().data_ptr(), digest) + return snapshot + + @classmethod + def _assert_final_layout_tensors_unchanged( + cls, + model: nn.Module, + snapshot: dict[str, tuple[int, str]], + ) -> None: + current = cls._snapshot_final_layout_tensors(model, snapshot) + changed = [name for name in snapshot if current[name] != snapshot[name]] + if changed: + raise RuntimeError( + "shared warm finalization changed restored final-layout tensor bytes or backing pointers: " + f"{changed[:5]}" + ) + + def _hwr_eligibility_mode( + self, + model: nn.Module, + modules: object, + *, + dist_offload: bool, + use_allgather: bool, + load_format: str, + ) -> object | None: + """Return the enabled HWR mode only after all zero-interaction gates.""" + from vllm_omni.host_weight_runtime import RuntimeMode + + raw_mode = getattr(self.od_config, "host_weight_runtime_mode", "disabled") + try: + mode = RuntimeMode(raw_mode) + except ValueError as exc: + raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") from exc + + # These gates intentionally precede HWR imports, source preparation, + # identity construction, and store creation. + if mode is RuntimeMode.DISABLED or not dist_offload or use_allgather: + return None + + parallel = self.parallel_config + reasons: list[str] = [] + if load_format != "default": + reasons.append("load_format must be 'default'") + if bool(getattr(parallel, "use_hsdp", False)): + reasons.append("HSDP layouts are not supported by the final-layout BF16 consumer") + if self.quant_config is not None: + reasons.append("quantized layouts require a representation-specific HWR producer") + if getattr(self.od_config, "lora_path", None): + reasons.append("adapted weights are not reusable base-model artifacts") + + dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ()))) + if not dit_modules: + reasons.append("no DiT modules were discovered") + elif any( + getattr(dit, "host_weight_restore_contract", None) is None + or not callable(getattr(dit, "validate_restored_host_weights", None)) + for _, dit in dit_modules + ): + reasons.append("every owned DiT must declare the final-layout restore contract") + + if reasons: + message = "; ".join(reasons) + if mode is RuntimeMode.REQUIRED: + raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") + logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) + return None + return mode + + def _prepare_hwr_sources( + self, + model: nn.Module, + modules: object, + sources: Sequence[object], + ) -> tuple[PreparedWeightSource, ...]: + from vllm_omni.diffusion.model_loader.host_weights import ( + ImplementationIdentity, + PreparedWeightSource, + WeightSourceKind, + ) + + dit_prefixes = tuple(f"{name}." for name in getattr(modules, "dit_names", ())) + selected_sources = tuple( + source + for source in sources + if not source.prefix or any(prefix.startswith(source.prefix) for prefix in dit_prefixes) + ) + if not selected_sources: + raise ValueError("final-layout HWR requires canonical weight sources covering every DiT") + + prepared: list[PreparedWeightSource] = [] + for source in selected_sources: + resolved_root, weight_files, use_safetensors = self._prepare_weights( + source.model_or_path, + source.subfolder, + source.revision, + source.fall_back_to_pt, + source.allow_patterns_overrides, + ) + adapter = self._get_checkpoint_adapter(model, source, use_safetensors) + adapter_identity = None + if adapter is not None: + adapter_name = f"{type(adapter).__module__}.{type(adapter).__qualname__}" + adapter_identity = ImplementationIdentity( + implementation_id=adapter_name, + version="1", + fingerprint=self._identity_fingerprint({"adapter": adapter_name}), + ) + source_kind = ( + WeightSourceKind.LOCAL_PATH + if os.path.isdir(os.fspath(source.model_or_path)) + else WeightSourceKind.HUGGING_FACE_HUB + ) + prepared.append( + PreparedWeightSource( + model_or_path=os.fspath(source.model_or_path), + subfolder=source.subfolder, + requested_revision=source.revision, + prefix=source.prefix, + resolved_root=Path(os.fspath(resolved_root)), + weight_files=tuple(Path(os.fspath(path)) for path in weight_files), + use_safetensors=use_safetensors, + checkpoint_adapter=adapter_identity, + source_kind=source_kind, + ) + ) + return tuple(prepared) + + def _build_hwr_context( + self, + model: nn.Module, + modules: object, + *, + load_format: str, + sources: Sequence[object], + ) -> FinalLayoutIdentityContext: + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + + from vllm_omni.diffusion.model_loader.host_weights import ( + FINAL_LAYOUT_BF16_POLICY, + FinalLayoutLoaderIdentity, + FinalLayoutParallelIdentity, + FinalLayoutRequest, + ImplementationIdentity, + build_final_layout_identity, + ) + + parallel = self.parallel_config + tp_size = int(getattr(parallel, "tensor_parallel_size", 1)) + try: + tp_rank = int(get_tensor_model_parallel_rank()) if tp_size > 1 else 0 + except Exception: + tp_rank = int(getattr(parallel, "tensor_parallel_rank", 0)) + sp_size = int(getattr(parallel, "sequence_parallel_size", 1) or 1) + loader_config = { + "dtype": str(self.od_config.dtype), + "model_class": f"{type(model).__module__}.{type(model).__qualname__}", + "model_config": self._identity_value(getattr(self.od_config, "tf_model_config", None)), + "load_format": load_format, + "quantization": self._identity_value(self.quant_config), + } + contracts = [ + self._identity_value(getattr(dit, "host_weight_restore_contract")) for dit in getattr(modules, "dits", ()) + ] + loader_identity = FinalLayoutLoaderIdentity( + implementation=ImplementationIdentity( + implementation_id="vllm-omni.diffusion.diffusers-loader", + version="final-layout-v1", + fingerprint=self._identity_fingerprint( + { + "loader": "diffusers-loader-final-layout-v1", + "pipeline": f"{type(model).__module__}.{type(model).__qualname__}", + } + ), + ), + model_config_fingerprint=self._identity_fingerprint(loader_config), + weight_transform_fingerprint=self._identity_fingerprint( + { + "contracts": contracts, + "transform": "diffusion-final-layout-loader-transforms-v1", + } + ), + ) + semantic_parallel = FinalLayoutParallelIdentity( + tensor_parallel_size=tp_size, + tensor_parallel_rank=tp_rank, + sequence_parallel_size=sp_size, + ulysses_degree=int(getattr(parallel, "ulysses_degree", 1)), + ring_degree=int(getattr(parallel, "ring_degree", 1)), + allgather_degree=int(getattr(parallel, "allgather_degree", 1)), + ulysses_mode=str(getattr(parallel, "ulysses_mode", "strict")), + pipeline_parallel_size=int(getattr(parallel, "pipeline_parallel_size", 1)), + cfg_parallel_size=int(getattr(parallel, "cfg_parallel_size", 1)), + use_hsdp=bool(getattr(parallel, "use_hsdp", False)), + enable_expert_parallel=bool(getattr(parallel, "enable_expert_parallel", False)), + ) + model_id = str(getattr(self.od_config, "model", "") or "") + if not model_id: + raise ValueError("final-layout HWR requires a canonical model identifier") + request = FinalLayoutRequest( + model_id=model_id, + loader=loader_identity, + parallel=semantic_parallel, + load_format=load_format, + ) + prepared_sources = self._prepare_hwr_sources(model, modules, sources) + dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ()))) + return build_final_layout_identity( + model, + dit_modules=dit_modules, + prepared_sources=prepared_sources, + request=request, + policy=FINAL_LAYOUT_BF16_POLICY, + ) + + def _resolve_hwr( + self, + model: nn.Module, + modules: object, + *, + dist_offload: bool, + use_allgather: bool, + load_format: str, + sources: Sequence[object], + ) -> dict[str, object] | None: + """Resolve an eligible no-AllGather final-layout HWR transaction.""" + from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutTensorRestorer + from vllm_omni.host_weight_runtime import ( + HostWeightLeaseCarrier, + HostWeightRuntime, + HostWeightRuntimeConfig, + ProductionPolicy, + ResolutionOutcome, + RuntimeMode, + StorageDomainPolicy, + ) + from vllm_omni.host_weight_runtime.filesystem import detect_storage_class + + mode = self._hwr_eligibility_mode( + model, + modules, + dist_offload=dist_offload, + use_allgather=use_allgather, + load_format=load_format, + ) + if mode is None: + return None + assert isinstance(mode, RuntimeMode) + expected_prefixes = frozenset(f"{name}." for name in getattr(modules, "dit_names", ())) + available_prefixes = frozenset(getattr(source, "prefix", "") for source in sources) + if not expected_prefixes <= available_prefixes: + message = "final-layout HWR requires one dedicated source prefix per owned DiT" + if mode is RuntimeMode.REQUIRED: + raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") + logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) + return None + overlapping_prefixes = frozenset( + prefix + for prefix in available_prefixes + if any(dit_prefix.startswith(prefix) for dit_prefix in expected_prefixes) + ) + if not overlapping_prefixes <= expected_prefixes: + message = "final-layout HWR requires dedicated DiT sources and rejects mixed component sources" + if mode is RuntimeMode.REQUIRED: + raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") + logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) + return None + context = self._build_hwr_context(model, modules, load_format=load_format, sources=sources) + root_value = getattr(self.od_config, "host_weight_runtime_root", None) + if not isinstance(root_value, str) or not root_value.strip(): + raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") + root = Path(root_value).expanduser() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=mode, + domain=StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + production=ProductionPolicy( + allow_local_build=False, + allow_post_load_publish=True, + ), + ) + ) + resolution = runtime.resolve(context.identity) + state: dict[str, object] = { + "mode": mode, + "context": context, + "runtime": runtime, + "outcome": resolution.report.outcome, + } + if resolution.report.outcome is ResolutionOutcome.FAILED: + failure = next( + (attempt.failure for attempt in reversed(resolution.report.attempts) if attempt.failure is not None), + None, + ) + detail = failure.message if failure is not None else "resolution failed without a typed detail" + raise RuntimeError(f"Host Weight Runtime resolution failed: {detail}") + if resolution.report.outcome is not ResolutionOutcome.LOCAL_HIT: + return state + + lease = resolution.lease + if lease is None: + raise RuntimeError("Host Weight Runtime returned LOCAL_HIT without a lease") + restorer = FinalLayoutTensorRestorer(context) + try: + restore_plan = restorer.plan_restore(model, lease) + except Exception: + lease.close() + if mode is RuntimeMode.REQUIRED: + raise + logger.warning("HWR warm restore planning failed; falling back to canonical loading", exc_info=True) + state["outcome"] = ResolutionOutcome.CANONICAL_FALLBACK + return state + + try: + restore_plan.commit() + except Exception as exc: + lease.close() + raise _HWRCommitError( + "Host Weight Runtime restore commit failed; the partially restored model must be discarded" + ) from exc + + try: + carrier = HostWeightLeaseCarrier(lease) + warm_snapshot = self._snapshot_final_layout_tensors(model, context.tensor_names) + except Exception as exc: + lease.close() + raise _HWRCommitError( + "Host Weight Runtime committed restore could not establish its startup ownership boundary" + ) from exc + source_metadata = context.identity.source.metadata.to_value() + target_bindings = source_metadata.get("target_bindings") if isinstance(source_metadata, dict) else None + planned_prefixes = frozenset( + binding["source_prefix"] + for binding in (target_bindings or ()) + if isinstance(binding, dict) and isinstance(binding.get("source_prefix"), str) + ) + if not planned_prefixes: + lease.close() + raise _HWRCommitError("Host Weight Runtime identity did not record owned canonical source prefixes") + state["plan"] = HostWeightPlan( + backing_kind="host_weight_runtime", + bindings={name: TensorBinding(name, "") for name in context.tensor_names}, + planned_source_prefixes=planned_prefixes, + lease_carrier=carrier, + ) + state["warm_snapshot"] = warm_snapshot + return state + + def load_fresh_canonical_model(self) -> nn.Module: + """Reload a disposable startup model without HWR or checkpoint mmap.""" + if self._last_load_request is None: + raise RuntimeError("cannot construct a fresh canonical model before the initial load") + request = dict(self._last_load_request) + self._force_canonical_load = True + try: + return self.load_model(**cast(dict[str, object], request)) + finally: + self._force_canonical_load = False + + def _publish_hwr_after_load(self, model: nn.Module, modules: object, state: dict[str, object] | None) -> None: + from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutBF16Producer + from vllm_omni.host_weight_runtime import PostLoadPublicationOutcome, ResolutionOutcome + + if state is None or state.get("outcome") is not ResolutionOutcome.CANONICAL_FALLBACK: + return + context = state["context"] + runtime = state["runtime"] + dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ()))) + try: + producer = FinalLayoutBF16Producer(context, model, dit_modules) + report = runtime.publish_after_load(context.identity, producer=producer) + except Exception: + logger.warning("Host Weight Runtime post-load publication failed", exc_info=True) + return + if report.outcome is PostLoadPublicationOutcome.FAILED: + logger.warning("Host Weight Runtime post-load publication failed: %s", report.failure) + + def _attach_offload_startup_state(self, model: nn.Module) -> None: + """Publish generic offload startup state on the loaded pipeline.""" + if self.host_weight_plan is None: + return + from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state + + hwr_state = self._hwr_state + allow_retry = False + if hwr_state is not None and isinstance(hwr_state.get("plan"), HostWeightPlan): + from vllm_omni.host_weight_runtime import RuntimeMode + + allow_retry = hwr_state.get("mode") is RuntimeMode.PREFERRED + attach_offload_startup_state( + model, + OffloadStartupState( + host_weight_plan=self.host_weight_plan, + fresh_model_loader=self.load_fresh_canonical_model if allow_retry else None, + allow_fresh_retry=allow_retry, + ), + ) diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index a61d60cfb..ca755b082 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Loader-owned host-weight plans shared with diffusion offload backends.""" from __future__ import annotations @@ -9,6 +9,7 @@ import json import os from collections.abc import Callable, Sequence from dataclasses import dataclass +from typing import TYPE_CHECKING import torch from safetensors import safe_open @@ -21,6 +22,9 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( logger = init_logger(__name__) +if TYPE_CHECKING: + from vllm_omni.host_weight_runtime import HostWeightLeaseCarrier + TensorTransform = Callable[[torch.Tensor], torch.Tensor] @@ -40,6 +44,7 @@ class HostWeightPlan: backing_kind: str bindings: dict[str, TensorBinding] planned_source_prefixes: frozenset[str] = frozenset() + lease_carrier: HostWeightLeaseCarrier | None = None @dataclass(frozen=True) diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index d29daaded..f768cd14c 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import torch from vllm.logger import init_logger @@ -25,6 +25,7 @@ from .sequential_backend import ( remove_sequential_offload, sequential_offload_component, ) +from .startup import OffloadStartupState, take_offload_startup_state from .tensor_utils import ( dtype_size, is_dtensor, @@ -52,6 +53,9 @@ __all__ = [ "apply_distributed_block_hook", "remove_distributed_block_hook", "get_offload_backend", + "enable_offload_backend", + "OffloadStartupState", + "take_offload_startup_state", "get_offload_plan", "get_blocks_attr_names", "get_blocks_from_dit", @@ -137,3 +141,55 @@ def get_offload_backend( else: logger.error("Unknown offload strategy: %s", config.strategy) return None + + +def enable_offload_backend( + od_config: OmniDiffusionConfig, + pipeline: torch.nn.Module, + device: torch.device | None = None, +) -> tuple[torch.nn.Module, OffloadBackend | None]: + """Create and enable the loader-selected backend transactionally. + + The model runner only calls this generic offloader boundary. Loader-owned + host plans and optional fresh-model recovery callbacks stay inside the + startup state consumed here. + """ + + def enable_once(model: torch.nn.Module, state: OffloadStartupState | None) -> OffloadBackend | None: + backend: OffloadBackend | None = None + try: + backend = get_offload_backend( + od_config, + device=device, + host_weight_plan=state.host_weight_plan if state is not None else None, + ) + if backend is not None: + logger.info("Enabling offloader backend: %s", backend.__class__.__name__) + backend.enable(model) + elif state is not None: + state.close_loader_ownership() + return backend + except Exception: + if backend is not None: + try: + backend.disable() + except Exception: + logger.exception("Failed to clean up the offload backend after startup failure") + if state is not None: + state.close_loader_ownership() + raise + + startup_state = take_offload_startup_state(pipeline) + try: + return pipeline, enable_once(pipeline, startup_state) + except Exception: + if startup_state is None or not startup_state.allow_fresh_retry: + raise + assert startup_state.fresh_model_loader is not None + logger.warning( + "Loader-backed offload startup failed; retrying with a fresh canonical model", + exc_info=True, + ) + del pipeline + pipeline = startup_state.fresh_model_loader() + return pipeline, enable_once(pipeline, take_offload_startup_state(pipeline)) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index d184e71ea..e53b6d78e 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -30,6 +30,7 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, ) +from vllm_omni.host_weight_runtime import HostWeightLease from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig @@ -902,6 +903,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._using_mmap = False self._using_rank_local_mmap = False self.host_weight_plan = host_weight_plan + self._host_weight_lease: HostWeightLease | None = None self._mmap_transforms_by_tensor_id: dict[int, Any] = {} def load_resident_layers(self) -> None: @@ -1405,6 +1407,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): buffer.data = buffer.data.to(self.device, non_blocking=True) def enable(self, pipeline: nn.Module) -> None: + """Enable DLO and make partial startup failures transactional.""" + try: + self._enable(pipeline) + except BaseException: + try: + self.disable() + except BaseException: + logger.exception("DLO cleanup failed while handling an enable failure") + raise + + def _enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("DistributedLayerwiseOffloadBackend already enabled") return @@ -1443,16 +1456,33 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._using_mmap = self.host_weight_plan is not None self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1 if self._using_mmap: - if self.host_weight_plan.backing_kind != "checkpoint_mmap": + if self.host_weight_plan.backing_kind == "host_weight_runtime": + carrier = self.host_weight_plan.lease_carrier + if carrier is None: + raise RuntimeError("DLO received a Host Weight Runtime plan without a lease carrier") + self._host_weight_lease = carrier.take() + if self._host_weight_lease.closed: + raise RuntimeError("DLO received a closed Host Weight Runtime lease") + # The final-layout restorer has already rebound the model to + # immutable host tensors. Treat those tensors as mmap-like + # sources for the bounded two-slot rank-local staging path; + # do not reinitialize or rematerialize the restored model. + self._using_rank_local_mmap = True + logger.info( + "DLO consuming final-layout Host Weight Runtime lease %s with bounded rank-local staging", + self._host_weight_lease.provenance.resolution_id, + ) + elif self.host_weight_plan.backing_kind == "checkpoint_mmap": + self._load_weights_via_mmap( + pipeline, + modules, + self.host_weight_plan, + ) + else: raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}") - self._load_weights_via_mmap( - pipeline, - modules, - self.host_weight_plan, - ) if self._using_rank_local_mmap: logger.info( - "DLO rank-local mmap storage enabled: checkpoint pages are " + "DLO rank-local host storage enabled: source pages are " "node-shared; each worker owns only two bounded host staging slots" ) else: @@ -1722,12 +1752,19 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._cleanup_after_loading() def _release_mmap_handles(self) -> None: - """Release safetensors mmap file handles.""" + """Release source handles and the transport-owned HWR lease.""" self._mmap_transforms_by_tensor_id.clear() if hasattr(self, "_mmap_file_cache"): self._mmap_file_cache.clear() del self._mmap_file_cache logger.info("Released safetensors mmap file handles") + lease = self._host_weight_lease + self._host_weight_lease = None + if lease is not None and not lease.closed: + lease.close() + logger.info("Released Host Weight Runtime lease %s", lease.provenance.resolution_id) + if self.host_weight_plan is not None and self.host_weight_plan.lease_carrier is not None: + self.host_weight_plan.lease_carrier.close() def _cleanup_after_loading(self) -> None: """Synchronize and release freed device/CPU memory after sharding.""" @@ -1743,7 +1780,25 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): pass def disable(self) -> None: - if not self.enabled and not hasattr(self, "_mmap_file_cache"): + has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed + has_carrier = ( + self.host_weight_plan is not None + and self.host_weight_plan.lease_carrier is not None + and not self.host_weight_plan.lease_carrier.closed + ) + has_partial_hooks = bool( + self._blocks + or self._all_hook_groups + or getattr(self, "_on_demand_handles", ()) + or self._resident_layer_group is not None + ) + if ( + not self.enabled + and not hasattr(self, "_mmap_file_cache") + and not has_open_lease + and not has_carrier + and not has_partial_hooks + ): return if self._using_rank_local_mmap: diff --git a/vllm_omni/diffusion/offloader/startup.py b/vllm_omni/diffusion/offloader/startup.py new file mode 100644 index 000000000..7861ddfd2 --- /dev/null +++ b/vllm_omni/diffusion/offloader/startup.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generic startup handoff between model loaders and offload backends.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +from torch import nn + +from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan + +_STARTUP_STATE_ATTR = "_omni_offload_startup_state" + + +@dataclass +class OffloadStartupState: + """Loader-owned state consumed by the offloader startup boundary.""" + + host_weight_plan: HostWeightPlan | None = None + fresh_model_loader: Callable[[], nn.Module] | None = None + allow_fresh_retry: bool = False + + def __post_init__(self) -> None: + if self.allow_fresh_retry != (self.fresh_model_loader is not None): + raise ValueError("fresh-model retry policy requires exactly one retry callback") + + def close_loader_ownership(self) -> None: + """Release a plan that never reached a backend.""" + if self.host_weight_plan is None: + return + carrier = self.host_weight_plan.lease_carrier + if carrier is not None: + carrier.close() + + +def attach_offload_startup_state(model: nn.Module, state: OffloadStartupState) -> None: + """Attach one process-local startup handoff to a loaded pipeline.""" + previous = getattr(model, _STARTUP_STATE_ATTR, None) + if previous is not None: + previous.close_loader_ownership() + setattr(model, _STARTUP_STATE_ATTR, state) + + +def take_offload_startup_state(model: nn.Module) -> OffloadStartupState | None: + """Take and remove the loader handoff from a pipeline exactly once.""" + state = getattr(model, _STARTUP_STATE_ATTR, None) + if state is None: + return None + delattr(model, _STARTUP_STATE_ATTR) + return state + + +__all__ = [ + "OffloadStartupState", + "attach_offload_startup_state", + "take_offload_startup_state", +] diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index de10bafa4..15680d60e 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -49,7 +49,7 @@ from vllm_omni.diffusion.models.interface import ( supports_prompt_update, supports_step_execution, ) -from vllm_omni.diffusion.offloader import get_offload_backend +from vllm_omni.diffusion.offloader import enable_offload_backend from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( @@ -336,15 +336,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): f"{self.od_config.model_class_name} does not support that contract." ) - # Apply CPU offloading - self.offload_backend = get_offload_backend( + # The offloader owns loader-plan handoff and startup recovery. The + # runner only receives the pipeline/backend pair that is ready to use. + self.pipeline, self.offload_backend = enable_offload_backend( self.od_config, + self.pipeline, device=self.device, - host_weight_plan=model_loader.take_host_weight_plan(), ) - if self.offload_backend is not None: - logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}") - self.offload_backend.enable(self.pipeline) # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index f1c483d48..e1c289357 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -559,6 +559,8 @@ class OrchestratorArgs: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = 0 + host_weight_runtime_mode: str = "disabled" + host_weight_runtime_root: str | None = None boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 9d7c66e1d..e5a0e44e9 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1059,6 +1059,8 @@ class AsyncOmniEngine: "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), + "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), + "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 122cc862b..c14d306dc 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -738,6 +738,28 @@ class OmniServeCommand(CLISubcommand): help="Keep this many leading main-DiT blocks resident on the device " "while distributed layerwise offload streams the remaining blocks.", ) + omni_config_group.add_argument( + "--host-weight-runtime-mode", + choices=("disabled", "preferred", "required"), + default="disabled", + help=( + "Host Weight Runtime policy for eligible no-AllGather DLO: " + "disabled does not consult HWR; preferred restores an exact hit " + "or canonically loads and publishes on a miss; required restores " + "an exact hit or fails startup. Populate a required store with " + "preferred first." + ), + ) + omni_config_group.add_argument( + "--host-weight-runtime-root", + type=str, + default=None, + help=( + "Writable node-local Host Weight Runtime store shared by workers " + "in one storage domain. Required for preferred and required; use " + "the same persistent path for population and serving." + ), + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio", diff --git a/vllm_omni/host_weight_runtime/__init__.py b/vllm_omni/host_weight_runtime/__init__.py index d870188ef..5a504adfc 100644 --- a/vllm_omni/host_weight_runtime/__init__.py +++ b/vllm_omni/host_weight_runtime/__init__.py @@ -35,7 +35,7 @@ from .inspection import ( DomainInspection, HostWeightInspector, ) -from .lease import HostWeightLease, LeaseProvenance, MappedHostRegion +from .lease import HostWeightLease, HostWeightLeaseCarrier, LeaseProvenance, MappedHostRegion from .manifest import ( FileManifestEntry, ManifestValidationError, @@ -92,6 +92,7 @@ __all__ = [ "HostWeightError", "HostWeightFailure", "HostWeightLease", + "HostWeightLeaseCarrier", "HostWeightResolution", "HostWeightInspector", "HostWeightRuntime", diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index 9c81d2972..c5c70c490 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -863,7 +863,11 @@ class FilesystemHostWeightStore: manifest, _ = writer.finalize(identity, metadata) artifact_bytes = self._tree_bytes(staging) self._check_capacity(artifact_bytes, 0) - os.chmod(staging, 0o555) + # Keep the staging directory writable until the atomic rename. + # overlayfs rejects renaming a non-writable source directory; + # the artifact lock keeps readers out while the published + # directory is hardened below. + os.chmod(staging, 0o755) _fsync_directory(staging) except HostWeightError as exc: if writer is not None: @@ -934,6 +938,11 @@ class FilesystemHostWeightStore: self._remove_tree(staging) else: os.replace(staging, entry_path) + try: + os.chmod(entry_path, 0o555) + except OSError: + self._quarantine_locked(identity.key, reason="publication_hardening_failure") + raise _fsync_directory(self.artifacts_dir) _fsync_directory(self.tmp_dir) self._remove_deny_locked(identity.key) diff --git a/vllm_omni/host_weight_runtime/lease.py b/vllm_omni/host_weight_runtime/lease.py index 7446f0bbe..f3a62e067 100644 --- a/vllm_omni/host_weight_runtime/lease.py +++ b/vllm_omni/host_weight_runtime/lease.py @@ -140,4 +140,60 @@ class HostWeightLease: self.close() -__all__ = ["HostWeightLease", "LeaseProvenance", "MappedHostRegion"] +class HostWeightLeaseCarrier: + """Process-local, single-take ownership carrier for a host-weight lease. + + The loader creates the carrier after a restore transaction has committed. + The runner/backend takes the lease exactly once before starting any + asynchronous transport work. If setup fails before that handoff, closing + the carrier releases the loader-owned lease. + """ + + def __init__(self, lease: HostWeightLease) -> None: + if not isinstance(lease, HostWeightLease): + raise TypeError("HostWeightLeaseCarrier requires one HostWeightLease") + if lease.closed: + raise ValueError("cannot carry a closed HostWeightLease") + self._lease: HostWeightLease | None = lease + self._taken = False + self._lock = threading.RLock() + + @property + def taken(self) -> bool: + with self._lock: + return self._taken + + @property + def closed(self) -> bool: + with self._lock: + return self._lease is None or self._lease.closed + + def take(self) -> HostWeightLease: + """Transfer the lease to the transport owner exactly once.""" + with self._lock: + if self._taken: + raise RuntimeError("HostWeightLeaseCarrier was already taken") + lease = self._lease + if lease is None or lease.closed: + raise RuntimeError("HostWeightLeaseCarrier does not own an open lease") + self._lease = None + self._taken = True + return lease + + def close(self) -> None: + """Close the lease if ownership has not been transferred.""" + with self._lock: + lease = self._lease + self._lease = None + if lease is not None: + lease.close() + + def __reduce_ex__(self, _protocol: SupportsIndex) -> NoReturn: + raise TypeError("HostWeightLeaseCarrier is process-local and cannot be serialized") + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.close() + + +__all__ = ["HostWeightLease", "HostWeightLeaseCarrier", "LeaseProvenance", "MappedHostRegion"]