From d30ccdc31590fc6b88d41817e23ff63ebb15baa2 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sun, 16 Aug 2026 23:15:48 +0800 Subject: [PATCH] [Diffusion] Add node-local runtime weight cache for DLO Signed-off-by: Hongsheng Liu --- .../distributed_layerwise_offload.md | 105 ++- .../distributed_layerwise_offload.md | 79 +- .../model_loader/test_diffusers_loader.py | 131 +++- .../model_loader/test_runtime_weight_cache.py | 213 ++++++ .../test_distributed_layerwise_backend.py | 142 +++- .../test_async_omni_diffusion_config.py | 11 + vllm_omni/config/omni_config.py | 8 + vllm_omni/config/stage_config.py | 3 + vllm_omni/diffusion/data.py | 13 + .../model_loader/diffusers_loader.py | 94 ++- .../model_loader/host_weight_plan.py | 29 +- .../model_loader/runtime_weight_cache.py | 693 ++++++++++++++++++ vllm_omni/diffusion/offloader/__init__.py | 19 +- .../distributed_layerwise_backend.py | 114 ++- vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/async_omni_engine.py | 3 + vllm_omni/entrypoints/cli/serve.py | 24 + 17 files changed, 1616 insertions(+), 68 deletions(-) create mode 100644 tests/diffusion/model_loader/test_runtime_weight_cache.py create mode 100644 vllm_omni/diffusion/model_loader/runtime_weight_cache.py diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index e254e3a43..383d871bf 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -43,15 +43,15 @@ DLO weight collective. Host storage is selected separately from the transfer protocol. The loader can produce a direct-checkpoint mmap plan for a proven-compatible runtime layout; -otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on -the same node can share immutable checkpoint pages when direct mmap is -selected, while the ordinary-loader fallback still keeps a private runtime -copy per process. +otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B +runtime cache can normalize those final ordinary-loader DiT tensors into an +immutable node-local mmap entry. Consequently, replicas can share either +proven-compatible checkpoint pages or equivalent final runtime layouts. The Phase A shared-mmap support boundary is TP1. TP greater than one is an ordinary-loader compatibility path: DLO can consume the resulting TP-local -tensors, but those configurations do not use checkpoint mmap and must not be -used to claim shared-mmap host-memory savings. +tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache +can share each final TP coordinate among equivalent processes. The compatibility matrix below describes the current implementation. The unit-level guards are covered, but not every parallelism combination has a @@ -105,6 +105,48 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only transforms. Model-specific direct-layout knowledge, when required, lives in a checkpoint adapter beside the ordinary loader. +### Post-load runtime cache + +The Phase B cache preserves the same ownership boundary. DLO neither derives +cache keys nor interprets model loaders. When enabled with no-AllGather, every +rank first completes the ordinary loader, custom callbacks, post-load +processing, strict validation, and calibration. The loader then: + +1. discovers final CPU DiT parameters and persistent buffers; +2. rejects unsupported non-contiguous, aliased, quantized, distributed, or + ambiguous layouts before changing the model; +3. hashes tensor names, metadata, and final bytes; +4. builds or joins one content-addressed entry under a strict POSIX lock; and +5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file, + manifest, and remapped-tensor content validation succeeds. + +The identity excludes DP size/rank and SP rank. It includes TP world size/rank +and conservative SP implementation/world-size guards. DP replicas and +equivalent SP ranks can therefore share; different TP coordinates cannot. +Final tensor content is authoritative, so an unmodeled loader/config change +that changes bytes creates a different entry instead of risking silent +sharing. + +Publication writes sharded safetensors and a JSON manifest into a same- +filesystem temporary directory, flushes them, and atomically renames the +directory. The manifest includes per-file SHA-256 digests and the final tensor +content digest. Kernel `flock` ownership distinguishes a live writer from a +stale lock file: process death releases the lock automatically. Stale temporary +entries are removed only by the next lock owner. A bounded wait or any cache +error retains the already-valid ordinary tensors. + +DLO maps this plan without moving the DiT to `meta` and without rerunning any +post-load hook. It changes tensor storage while preserving `Parameter` objects, +then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing +streaming hooks. The source mappings remain open through DLO disable and feed +the existing bounded two-slot host staging path. + +This version is deliberately opt-in and steady-state-oriented. It requires +roughly one local-disk model copy per runtime identity, has no eviction, and +still performs ordinary loading plus full content-hash passes in every rank. +Skip-load-on-hit is a separate follow-up after loader side effects can be +modeled safely. + ### AllGather path With the default `dlo_use_allgather=True`, each rank stores approximately @@ -138,11 +180,17 @@ When DP is greater than one, the engine can process one request per DP rank in the same denoising wave. Because AllGather is a collective, all participating requests must take the same execution path at every denoising step. +Fast NVLink/NVSwitch changes collective cost, not this execution contract. For +independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather +plus shared mmap storage avoids coupling their request schedules and failure +domains. + ### Rank-local path without DLO AllGather With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to one and streams complete blocks using H2D copies only. The host backing may be -either a loader-approved checkpoint mapping or ordinary runtime tensors. +a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or +ordinary runtime tensors. For direct mmap, each process retains immutable safetensors views and uses two bounded pinned host staging slots. Processes on the same node that map the same @@ -153,8 +201,8 @@ node has its own page cache. When direct mmap preflight fails, the regular model loader remains responsible for preparing each rank's weights, including TP-local tensors or HSDP-managed -parameters. In that fallback, each pure-DP process keeps a private full runtime -copy. +parameters. Supported CPU layouts may then enter the opt-in runtime cache. +Unsupported layouts keep one private runtime copy per process. This mode means: @@ -165,17 +213,17 @@ This mode means: - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only DLO's additional weight AllGather is disabled. - Pure DP deployments share one checkpoint-backed copy per node when direct - mmap is selected; the ordinary-loader fallback keeps one private runtime - copy per rank. + mmap is selected. The runtime cache can share equivalent ordinary-loader + outputs; its fallback keeps one private runtime copy per rank. - The scheduler does not require a synchronized DP request wave for DLO. ## 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. | +| **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; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. | +| **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 without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. | +| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. | | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. | ### Combined dimensions @@ -211,10 +259,9 @@ 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). +Runtime-cache v1 is no-AllGather only and rejects all quantization +configurations, HSDP/DTensor, CFG parallelism, PP, non-contiguous tensors, and +tied/shared storage. ## Validation coverage @@ -229,6 +276,11 @@ Current source-level validation includes: - exact loader-to-backend plan transfer and ordinary-loader fallback; - rank-local mmap source retention, bounded two-slot staging, and adapter transforms without parameter-side flags; +- runtime-cache content identity, sharded atomic publication, concurrent + writer election, corruption rebuild, lock timeout, and fail-closed layout + rejection; +- loader ordering through final post-load mutation and DLO remapping without + rerunning post-load hooks; - resident-layer requests requiring no-AllGather; - DP request-wave validation for denoising-step compatibility; - sharding, double-buffer, AllGather-size, and heterogeneous-block regression @@ -288,11 +340,12 @@ from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker because it counts a shared physical page in every process that maps it; summed PSS is the appropriate node-memory comparison. -The highest-value missing coverage is broader end-to-end numerical and -lifecycle comparison against ordinary layerwise offload for DP+SP, -HSDP+SP+no-AllGather, and TP greater than one across additional models and -target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not -change the Phase A direct-mmap TP1 support boundary. +The highest-value missing coverage is the Phase B hardware matrix from +[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231): cold and +warm publication, inode sharing, PSS/`Private_Dirty`, output equivalence, and +independently scheduled `TP=2` engines on target hardware. The historical +B300 rows above predate the runtime cache and must not be treated as Phase B +memory validation. ## Recommendations @@ -301,8 +354,8 @@ change the Phase A direct-mmap TP1 support boundary. - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is not the goal. - Use **no-AllGather** when independent replica execution is required. TP1 - direct-mmap deployments can share checkpoint pages per node; other layouts - retain the ordinary loader's private host memory behavior and are outside - the Phase A shared-mmap support scope. + direct-mmap deployments can share checkpoint pages per node. Enable the + runtime cache when independent replicas have repeated final layouts, + especially matching coordinates across multiple TP engines. - Prefer **HSDP alone** for production HSDP deployments until the combined HSDP + DLO no-AllGather path has broader end-to-end coverage. diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 49139e695..9044e9a1c 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -5,7 +5,9 @@ deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of the host weights and reconstructs each layer at runtime. Without AllGather, each rank streams a complete block independently. Compatible TP1 deployments can share checkpoint-backed host pages among processes on the same node; -otherwise DLO streams the ordinary loader's rank-local tensors. +otherwise DLO streams the ordinary loader's rank-local tensors. The optional +runtime-weight cache extends node-local sharing to final ordinary-loader +layouts, including matching TP coordinates across independent engines. See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md) for the implementation contract and compatibility matrix. @@ -21,8 +23,10 @@ AllGather: [AG N+1] [AG N+2] Slots: slot 0: Layer N slot 1: Layer N+1 ``` -AllGather communicates only request-independent weight shards, so data- -parallel ranks may process different requests concurrently. +AllGather communicates request-independent weight shards, but every member of +the transfer group must request the same next block in the same collective +order. A synchronized DP wave may contain different prompts only when those +requests follow the same denoising and block-execution path. ## Usage @@ -38,6 +42,14 @@ vllm serve /path/to/model --omni \ --data-parallel-size 4 \ --dlo-no-use-allgather +# Independently scheduled replicas sharing final runtime weights on one node +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --tensor-parallel-size 2 \ + --dlo-no-use-allgather \ + --dlo-enable-runtime-cache \ + --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights + # Sequence parallel deployment vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ @@ -50,7 +62,9 @@ from vllm_omni import Omni omni = Omni( model="/path/to/model", enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, + dlo_use_allgather=False, + dlo_enable_runtime_cache=True, + dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights", ) ``` @@ -62,6 +76,9 @@ omni = Omni( | `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` | | `--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-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` | +| `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` | +| `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` | | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | ## Host-weight loading @@ -78,6 +95,21 @@ DLO may still consume those TP-local tensors, but this is a compatibility path: it does not share checkpoint-backed runtime weights across DP replicas and provides no shared-mmap host-memory guarantee. +When `--dlo-enable-runtime-cache` is set with no-AllGather, that fallback gains +a second storage opportunity. Every rank still completes ordinary loading, +loader callbacks, post-load processing, validation, and calibration. It then +hashes the final CPU DiT parameters and persistent buffers. One process per +equivalent runtime layout publishes sharded safetensors plus a checksummed +manifest; peers validate and mmap the same files. DLO replaces only tensor +storage, preserves the final `Parameter` objects, and releases the private +allocator-backed copies before installing its streaming hooks. + +The runtime layout identity excludes DP rank and SP rank. It includes TP world +size/rank and conservative SP implementation/world-size guards. Therefore two +independent `TP=2` engines normally create two entries: both TP0 workers map +one entry and both TP1 workers map the other. Cache sharing does not create a +process group or an inference-time collective. + The mmap plan skips only dedicated DiT weight sources. Other component sources, such as a text encoder loaded through the shared diffusion loader, continue to use their ordinary component loader. A checkpoint source that mixes DiT and @@ -104,6 +136,27 @@ 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. +## Runtime-cache lifecycle + +Runtime-cache publication uses a per-layout POSIX file lock, a temporary +directory on the same filesystem, file and final-tensor SHA-256 validation, +`fsync`, and an atomic directory rename. A crashed writer releases its kernel +lock; the next writer removes only that layout's stale temporary directories +while holding the lock. Lock timeout, unsupported layout, insufficient disk, +or validation failure keeps the already-valid ordinary tensors. + +The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because +vLLM-Omni may isolate that variable between stage replicas. All consumers that +should share must receive the same `--dlo-runtime-cache-dir` on a local, +disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower +host-memory PSS, and do not use a cross-node directory in this version. + +There is no automatic eviction. Budget roughly one final DiT copy per distinct +TP coordinate, dtype/layout, and model version, plus temporary space while an +entry is published. Operators may remove entries or the complete cache root +only after all workers using its mmap files have stopped. The feature is +opt-in for this reason. + ## Declarative topology Models may declare an `OffloadPlan` instead of embedding offload logic: @@ -132,10 +185,10 @@ must enter each collective. ## Limitations -- 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. +- Direct checkpoint mmap currently requires TP1. TP greater than one falls + back to the ordinary TP-aware loader. With the optional runtime cache, + equivalent processes at the same TP coordinate can then share that final + rank-local layout in no-AllGather mode. - 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 @@ -143,14 +196,16 @@ must enter each collective. complete FP8 model in host memory before DLO retains only its shard. Other online quantization methods require no-AllGather until their runtime layouts are validated. +- Runtime-cache v1 rejects all quantized, non-contiguous, aliased/tied, + device-only, HSDP, CFG-parallel, and PP layouts. - Resident leading layers require `--dlo-no-use-allgather` and a model `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 -[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of -the direct-checkpoint path. +The runtime cache improves steady-state host PSS, not startup peak: every rank +loads private weights before remapping, and validation adds full content-hash +passes. Skip-load-on-cache-hit and cache eviction remain follow-up work in +[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). 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 527942961..faefa41bb 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -346,7 +346,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch): backing_kind="checkpoint_mmap", bindings={ "transformer.weight": TensorBinding( - checkpoint_key="weight", + storage_key="weight", file_path="unused", ) }, @@ -405,6 +405,135 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): assert loader.take_host_weight_plan() is None +def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=False, + tensor_parallel_size=1, + sequence_parallel_size=1, + data_parallel_size=2, + cfg_parallel_size=1, + pipeline_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + allgather_degree=1, + ulysses_mode="strict", + ), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + dlo_enable_runtime_cache=True, + dlo_runtime_cache_dir=str(tmp_path), + dlo_runtime_cache_lock_timeout=1.0, + model="unused", + revision=None, + model_config={}, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + runtime_plan = HostWeightPlan( + backing_kind="runtime_cache", + bindings={}, + runtime_layout_key="layout-key", + post_load_complete=True, + ) + + def load_weights(_model): + events.append("load") + with torch.no_grad(): + model.transformer.weight.fill_(1) + + def process_weights(_model, _device): + events.append("process") + with torch.no_grad(): + model.transformer.weight.add_(1) + + def calibrate(_model): + events.append("calibrate") + with torch.no_grad(): + model.transformer.weight.add_(1) + + original_eval = model.eval + + def eval_model(): + events.append("eval") + return original_eval() + + def build_cache(_pipeline, **kwargs): + events.append("cache") + assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0)) + assert kwargs["tensor_parallel_rank"] == 0 + assert "data_parallel_size" not in kwargs["sequence_parallel_guard"] + return HostWeightPlanResult(runtime_plan) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = load_weights # type: ignore[method-assign] + loader._process_weights_after_loading = process_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = calibrate # type: ignore[method-assign] + model.eval = eval_model # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"), + ) + monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache) + + assert loader.load_model(load_device="cpu") is model + assert events == ["load", "process", "calibrate", "eval", "cache"] + assert loader.take_host_weight_plan() is runtime_plan + + +def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=False, + tensor_parallel_size=1, + sequence_parallel_size=1, + data_parallel_size=2, + cfg_parallel_size=1, + pipeline_parallel_size=1, + ), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + dlo_enable_runtime_cache=True, + dlo_runtime_cache_dir=str(tmp_path), + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: None # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"), + ) + monkeypatch.setattr( + cache_mod, + "build_runtime_weight_cache_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(None, "writer timed out", "lock_timeout"), + ) + + original_weight = model.transformer.weight.detach().clone() + assert loader.load_model(load_device="cpu") is model + assert torch.equal(model.transformer.weight, original_weight) + assert loader.take_host_weight_plan() is None + + def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch): from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, diff --git a/tests/diffusion/model_loader/test_runtime_weight_cache.py b/tests/diffusion/model_loader/test_runtime_weight_cache.py new file mode 100644 index 000000000..f50661135 --- /dev/null +++ b/tests/diffusion/model_loader/test_runtime_weight_cache.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import fcntl +import multiprocessing as mp +import shutil +from pathlib import Path + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.model_loader.runtime_weight_cache import ( + build_runtime_weight_cache_plan, + default_runtime_cache_root, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +class _LoaderMarker: + pass + + +class _Transformer(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(3, 2) + self.register_buffer("persistent", torch.arange(2, dtype=torch.float32)) + self.register_buffer("temporary", torch.ones(1), persistent=False) + with torch.no_grad(): + self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3)) + self.proj.bias.copy_(torch.tensor([7.0, 8.0])) + + +class _Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = _Transformer() + + +def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch): + monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache") + + assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights") + + +def _build(pipeline: nn.Module, cache_root: Path, **overrides): + kwargs = { + "dit_modules": (("transformer", pipeline.transformer),), + "loader_type": _LoaderMarker, + "cache_root": cache_root, + "lock_timeout_seconds": 2.0, + "max_shard_bytes": 32, + "model_identity": "test/model", + "revision": "revision", + "runtime_dtype": torch.float32, + "load_format": "default", + "loader_inputs": {"model_config": {"hidden_size": 3}}, + "tensor_parallel_size": 1, + "tensor_parallel_rank": 0, + "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"}, + "use_hsdp": False, + "quantization_config": None, + "cfg_parallel_size": 1, + "pipeline_parallel_size": 1, + } + kwargs.update(overrides) + return build_runtime_weight_cache_plan(pipeline, **kwargs) + + +def _multiprocess_builder(cache_root: str, queue) -> None: + result = _build(_Pipeline(), Path(cache_root)) + queue.put( + ( + result.fallback_code, + result.plan.runtime_layout_key if result.plan is not None else None, + sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [], + ) + ) + + +def test_equivalent_process_layouts_reuse_one_entry(tmp_path): + first = _build(_Pipeline(), tmp_path) + second = _build(_Pipeline(), tmp_path) + + assert first.plan is not None + assert second.plan is not None + assert first.plan.runtime_layout_key == second.plan.runtime_layout_key + assert first.plan.post_load_complete + assert {binding.file_path for binding in first.plan.bindings.values()} == { + binding.file_path for binding in second.plan.bindings.values() + } + assert set(first.plan.bindings) == { + "transformer.persistent", + "transformer.proj.bias", + "transformer.proj.weight", + } + + +def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path): + pipeline = _Pipeline() + tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0) + tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1) + sp2 = _build( + pipeline, + tmp_path, + sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"}, + ) + + assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None + assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3 + + +def test_final_runtime_content_is_authoritative(tmp_path): + first_pipeline = _Pipeline() + second_pipeline = _Pipeline() + with torch.no_grad(): + second_pipeline.transformer.proj.weight[0, 0] += 1 + + first = _build(first_pipeline, tmp_path) + second = _build(second_pipeline, tmp_path) + + assert first.plan is not None and second.plan is not None + assert first.plan.runtime_layout_key != second.plan.runtime_layout_key + + +def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path): + first = _build(_Pipeline(), tmp_path) + assert first.plan is not None + shard = Path(next(iter(first.plan.bindings.values())).file_path) + old_inode = shard.stat().st_ino + with shard.open("r+b") as handle: + handle.seek(-1, 2) + byte = handle.read(1) + handle.seek(-1, 2) + handle.write(bytes([byte[0] ^ 0xFF])) + + second = _build(_Pipeline(), tmp_path) + + assert second.plan is not None + assert second.plan.runtime_layout_key == first.plan.runtime_layout_key + assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode + + +def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path): + initial = _build(_Pipeline(), tmp_path) + assert initial.plan is not None + cache_key = initial.plan.runtime_layout_key + entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent + shutil.rmtree(entry_dir) + lock_path = tmp_path / ".locks" / f"{cache_key}.lock" + + with lock_path.open("a+b") as lock_handle: + fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01) + + assert result.plan is None + assert result.fallback_code == "lock_timeout" + + +def test_next_writer_removes_stale_temp_after_process_death(tmp_path): + initial = _build(_Pipeline(), tmp_path) + assert initial.plan is not None + cache_key = initial.plan.runtime_layout_key + entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent + shutil.rmtree(entry_dir) + stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer" + stale_temp.mkdir(parents=True) + (stale_temp / "partial").write_bytes(b"partial") + + result = _build(_Pipeline(), tmp_path) + + assert result.plan is not None + assert not stale_temp.exists() + + +@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"]) +def test_unsupported_tensor_layouts_fail_closed(tmp_path, case): + pipeline = _Pipeline() + if case == "alias": + pipeline.transformer.tied = pipeline.transformer.proj.weight + expected_code = "unsupported_alias" + elif case == "external_alias": + pipeline.encoder = nn.Module() + pipeline.encoder.shared = pipeline.transformer.proj.weight + expected_code = "unsupported_alias" + else: + pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t()) + expected_code = "unsupported_tensor" + + result = _build(pipeline, tmp_path) + + assert result.plan is None + assert result.fallback_code == expected_code + assert list(tmp_path.rglob("*.safetensors")) == [] + + +def test_concurrent_processes_publish_one_immutable_entry(tmp_path): + context = mp.get_context("spawn") + queue = context.Queue() + 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] + for process in processes: + process.join(timeout=30) + + assert all(process.exitcode == 0 for process in processes) + assert all(code is None for code, _, _ in results) + assert len({key for _, key, _ in results}) == 1 + assert len({tuple(paths) for _, _, paths in results}) == 1 + assert len(list((tmp_path / "v1").iterdir())) == 1 diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index c67a2d45b..f649565ae 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -765,6 +765,95 @@ class _MmapPostLoadPipeline(nn.Module): class TestMmapWeightLoading: + def test_runtime_cache_maps_final_weights_without_rerunning_post_load( + self, + tmp_path, + patched_offload_runtime, + ): + class Transformer(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(2, 2, bias=False) + self.post_load_calls = 0 + + def post_load_weights(self): + self.post_load_calls += 1 + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + pipeline = Pipeline() + old_parameter = pipeline.transformer.proj.weight + old_parameter.runtime_marker = "preserved" + weight_file = tmp_path / "runtime.safetensors" + save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file)) + plan = HostWeightPlan( + backing_kind="runtime_cache", + bindings={ + "transformer.proj.weight": TensorBinding( + storage_key="transformer.proj.weight", + file_path=str(weight_file), + ) + }, + runtime_layout_key="runtime-layout", + post_load_complete=True, + ) + 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"]) + + backend._load_weights_from_runtime_cache(pipeline, modules, plan) + + assert pipeline.transformer.post_load_calls == 0 + assert pipeline.transformer.proj.weight is old_parameter + assert pipeline.transformer.proj.weight.runtime_marker == "preserved" + assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0)) + backend._release_mmap_handles() + + def test_runtime_cache_mapping_failure_leaves_ordinary_tensors_intact( + self, + tmp_path, + patched_offload_runtime, + ): + pipeline = nn.Module() + pipeline.transformer = nn.Linear(2, 2, bias=False) + old_parameter = pipeline.transformer.weight + plan = HostWeightPlan( + backing_kind="runtime_cache", + bindings={ + "transformer.weight": TensorBinding( + storage_key="transformer.weight", + file_path=str(tmp_path / "missing.safetensors"), + ) + }, + runtime_layout_key="runtime-layout", + post_load_complete=True, + ) + 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(Exception, match="missing.safetensors"): + backend._load_weights_from_runtime_cache(pipeline, modules, plan) + + assert pipeline.transformer.weight is old_parameter + 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()} @@ -788,7 +877,7 @@ class TestMmapWeightLoading: backing_kind="checkpoint_mmap", bindings={ name: TensorBinding( - checkpoint_key=name, + storage_key=name, file_path=str(weight_file), ) for name in weights @@ -859,6 +948,53 @@ class TestMmapWeightLoading: assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group) backend.disable() + def test_runtime_cache_enable_reclaims_private_weights_before_hook_setup( + 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() + weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()} + weight_file = tmp_path / "runtime.safetensors" + save_file(weights, str(weight_file)) + plan = HostWeightPlan( + backing_kind="runtime_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, + ) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + cleanup_calls: list[None] = [] + backend._cleanup_after_loading = lambda: cleanup_calls.append(None) # type: ignore[method-assign] + + backend.enable(pipeline) + + assert backend._using_rank_local_mmap + # One reclamation immediately after remap, plus the existing final + # post-hook cleanup. + assert len(cleanup_calls) == 2 + backend.disable() + class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): @@ -1164,7 +1300,7 @@ class TestMmapValidation: assert result.fallback_reason is None assert result.plan is not None assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding( - checkpoint_key="blocks.0.weight", + storage_key="blocks.0.weight", file_path=str(checkpoint_file), ) assert result.plan.planned_source_prefixes == frozenset({"transformer."}) @@ -1454,7 +1590,7 @@ class TestMmapValidation: assert result.plan is not None assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding( - checkpoint_key="proj_in.weight", + storage_key="proj_in.weight", file_path=str(checkpoint_file), ) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1daca1e99..1befcfe44 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -328,6 +328,11 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--omni", "--enable-distributed-layerwise-offload", "--dlo-no-use-allgather", + "--dlo-enable-runtime-cache", + "--dlo-runtime-cache-dir", + "/var/cache/vllm-omni-dlo", + "--dlo-runtime-cache-lock-timeout", + "30", "--dlo-resident-layers", "20", ] @@ -339,9 +344,15 @@ def test_serve_cli_forwards_distributed_offload_residency(): assert args.enable_distributed_layerwise_offload is True assert args.dlo_use_allgather is False + assert args.dlo_enable_runtime_cache is True + assert args.dlo_runtime_cache_dir == "/var/cache/vllm-omni-dlo" + assert args.dlo_runtime_cache_lock_timeout == 30 assert args.dlo_resident_layers == 20 assert engine_args["enable_distributed_layerwise_offload"] is True assert engine_args["dlo_use_allgather"] is False + assert engine_args["dlo_enable_runtime_cache"] is True + assert engine_args["dlo_runtime_cache_dir"] == "/var/cache/vllm-omni-dlo" + assert engine_args["dlo_runtime_cache_lock_timeout"] == 30 assert engine_args["dlo_resident_layers"] == 20 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 9cd7ed518..1ca12cd4e 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -687,6 +687,9 @@ class _DiffusionConfigProjection: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_enable_runtime_cache: bool = False + dlo_runtime_cache_dir: str | None = None + dlo_runtime_cache_lock_timeout: float = Field(default=600.0, gt=0) dlo_resident_layers: int = Field(default=0, ge=0) pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" @@ -749,6 +752,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 self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload: + raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload") + if self.dlo_enable_runtime_cache and self.dlo_use_allgather: + raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False") + if self.tf_model_config is None: self.tf_model_config = TransformerConfig() elif isinstance(self.tf_model_config, Mapping): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index cdfd35ee2..5ac625bb3 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -382,6 +382,9 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None + dlo_enable_runtime_cache: bool | None = None + dlo_runtime_cache_dir: str | None = None + dlo_runtime_cache_lock_timeout: float | None = None dlo_resident_layers: int | 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 34c8ccb96..f586367d0 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -708,6 +708,13 @@ class OmniDiffusionConfig: # This avoids AllGather synchronization, while host memory follows the # loader's existing rank-local layout instead of adding a second DP shard. dlo_use_allgather: bool = True + # Build or join a node-local mmap cache after ordinary loading when direct + # checkpoint mmap is unavailable. The first version is no-AllGather only. + dlo_enable_runtime_cache: bool = False + # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-runtime-weights. + dlo_runtime_cache_dir: str | None = None + # Maximum wait for the per-layout POSIX writer lock. + dlo_runtime_cache_lock_timeout: float = 600.0 # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 @@ -936,6 +943,12 @@ class OmniDiffusionConfig: ) def __post_init__(self): + if not math.isfinite(self.dlo_runtime_cache_lock_timeout) or self.dlo_runtime_cache_lock_timeout <= 0: + raise ValueError("dlo_runtime_cache_lock_timeout must be positive") + if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload: + raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload") + if self.dlo_enable_runtime_cache and self.dlo_use_allgather: + raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False") if self.diffusion_compile_granularity not in {"regional", "full"}: raise ValueError( "diffusion_compile_granularity must be 'regional' or 'full', " diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 62bc0525f..137f61c96 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -496,7 +496,8 @@ class DiffusersPipelineLoader: ) if _dist_offload and plan_result is not None: logger.info( - "DLO direct checkpoint mmap unavailable; using ordinary loader: %s", + "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s", + plan_result.fallback_code or "unknown", plan_result.fallback_reason, ) logger.debug("Loading weights on %s ...", load_device) @@ -521,7 +522,96 @@ class DiffusersPipelineLoader: logger.info("Quantization complete, offloaded model back to CPU") self._apply_skip_softmax_calibration(model) - return model.eval() + model.eval() + self._maybe_build_runtime_weight_cache(model, load_format) + return model + + def _maybe_build_runtime_weight_cache(self, model: nn.Module, load_format: str) -> None: + """Publish final ordinary-loader DiT tensors for no-AllGather DLO.""" + if self.host_weight_plan is not None: + return + if not getattr(self.od_config, "enable_distributed_layerwise_offload", False): + return + if not getattr(self.od_config, "dlo_enable_runtime_cache", False): + return + if getattr(self.od_config, "dlo_use_allgather", True): + logger.info("DLO runtime cache is no-AllGather only; leaving AllGather storage unchanged") + return + + parallel_config = self.parallel_config + tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1)) + tp_rank = 0 + if tp_size > 1: + try: + from vllm.distributed import get_tensor_model_parallel_rank + + tp_rank = int(get_tensor_model_parallel_rank()) + except Exception as exc: + logger.warning( + "DLO runtime cache unavailable [parallel_rank_unavailable]: cannot resolve TP rank for TP=%d: %s", + tp_size, + exc, + ) + return + + modules = ModuleDiscovery.discover(model) + from vllm_omni.diffusion.model_loader.runtime_weight_cache import ( + DEFAULT_SHARD_SIZE_BYTES, + build_runtime_weight_cache_plan, + ) + + loader_inputs = { + "model_class_name": getattr(self.od_config, "model_class_name", None), + "model_config": getattr(self.od_config, "model_config", None), + "diffusers_load_kwargs": getattr(self.od_config, "diffusers_load_kwargs", None), + "custom_pipeline_args": getattr(self.od_config, "custom_pipeline_args", None), + "lora_path": getattr(self.od_config, "lora_path", None), + "lora_scale": getattr(self.od_config, "lora_scale", None), + "lora_backend": getattr(self.od_config, "lora_backend", None), + } + sp_guard = { + "sequence_parallel_size": int(getattr(parallel_config, "sequence_parallel_size", 1)), + "ulysses_degree": int(getattr(parallel_config, "ulysses_degree", 1)), + "ring_degree": int(getattr(parallel_config, "ring_degree", 1)), + "allgather_degree": int(getattr(parallel_config, "allgather_degree", 1)), + "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"), + } + model_identity = getattr(self.od_config, "model", None) + result = build_runtime_weight_cache_plan( + model, + dit_modules=tuple(zip(modules.dit_names, modules.dits)), + loader_type=type(self), + cache_root=getattr(self.od_config, "dlo_runtime_cache_dir", None), + lock_timeout_seconds=float(getattr(self.od_config, "dlo_runtime_cache_lock_timeout", 600.0)), + max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES, + model_identity=str(model_identity) if model_identity is not None else None, + revision=getattr(self.od_config, "revision", None), + runtime_dtype=getattr(self.od_config, "dtype", None), + load_format=load_format, + loader_inputs=loader_inputs, + tensor_parallel_size=tp_size, + tensor_parallel_rank=tp_rank, + sequence_parallel_guard=sp_guard, + use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)), + quantization_config=self.quant_config, + cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)), + pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)), + ) + if result.plan is not None: + self.host_weight_plan = result.plan + logger.info( + "DLO runtime cache selected (layout=%s, TP=%d/%d, SP=%d)", + result.plan.runtime_layout_key[:12] if result.plan.runtime_layout_key else "unknown", + tp_rank, + tp_size, + sp_guard["sequence_parallel_size"], + ) + else: + logger.warning( + "DLO runtime cache unavailable [%s]; retaining ordinary host tensors: %s", + result.fallback_code or "unknown", + result.fallback_reason or "unspecified reason", + ) @staticmethod def _request_offload_after_quant(model: nn.Module) -> int: diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index a61d60cfb..849f05679 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -28,7 +28,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor] class TensorBinding: """One runtime tensor backed by one safetensors entry.""" - checkpoint_key: str + storage_key: str file_path: str transform: TensorTransform | None = None @@ -40,6 +40,8 @@ class HostWeightPlan: backing_kind: str bindings: dict[str, TensorBinding] planned_source_prefixes: frozenset[str] = frozenset() + runtime_layout_key: str | None = None + post_load_complete: bool = False @dataclass(frozen=True) @@ -48,6 +50,7 @@ class HostWeightPlanResult: plan: HostWeightPlan | None fallback_reason: str | None = None + fallback_code: str | None = None class _PlanIncompatibleError(RuntimeError): @@ -258,11 +261,11 @@ def _validate_source_metadata( with safe_open(file_path, framework="pt", device="cpu") as handle: available = set(handle.keys()) for runtime_name, target, binding in entries: - if binding.checkpoint_key not in available: + if binding.storage_key not in available: raise _PlanIncompatibleError( - f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}" + f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}" ) - source = handle.get_slice(binding.checkpoint_key) + source = handle.get_slice(binding.storage_key) source_shape = tuple(source.get_shape()) source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype()) if source_shape != tuple(target.shape): @@ -287,11 +290,19 @@ def build_checkpoint_mmap_plan( ) -> HostWeightPlanResult: """Build a complete direct-checkpoint plan or return a fallback reason.""" if tensor_parallel_size != 1: - return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader") + return HostWeightPlanResult( + None, + f"TP={tensor_parallel_size} requires the ordinary loader", + "unsupported_tp", + ) if use_hsdp: - return HostWeightPlanResult(None, "HSDP requires the ordinary loader") + return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp") if online_quantization: - return HostWeightPlanResult(None, "online quantization requires the ordinary loader") + return HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + "unsupported_online_quantization", + ) remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) if not callable(remap_fn): @@ -328,14 +339,14 @@ def build_checkpoint_mmap_plan( ) checkpoint_key, file_path = model_to_ckpt[runtime_name] bindings[runtime_name] = TensorBinding( - checkpoint_key=checkpoint_key, + storage_key=checkpoint_key, file_path=file_path, transform=policy.transform if policy is not None else None, ) _validate_source_metadata(required, bindings) except (OSError, ValueError, _PlanIncompatibleError) as exc: - return HostWeightPlanResult(None, str(exc)) + return HostWeightPlanResult(None, str(exc), "checkpoint_plan_incompatible") return HostWeightPlanResult( HostWeightPlan( diff --git a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py new file mode 100644 index 000000000..537615e09 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py @@ -0,0 +1,693 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Node-local, content-addressed runtime-weight cache for diffusion DLO.""" + +from __future__ import annotations + +import contextlib +import dataclasses +import errno +import fcntl +import hashlib +import inspect +import json +import math +import os +import shutil +import time +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + HostWeightPlanResult, + TensorBinding, +) + +logger = init_logger(__name__) + +RUNTIME_CACHE_SCHEMA_VERSION = 1 +DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights" +DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0 +DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3 +_HASH_CHUNK_BYTES = 64 * 1024**2 +_FILE_HASH_CHUNK_BYTES = 8 * 1024**2 + +_SUPPORTED_DTYPES = { + torch.bool, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, +} +for _dtype_name in ("float8_e4m3fn", "float8_e5m2"): + if (_dtype := getattr(torch, _dtype_name, None)) is not None: + _SUPPORTED_DTYPES.add(_dtype) + + +@dataclass(frozen=True) +class _RuntimeTensor: + name: str + tensor: torch.Tensor + kind: str + + +class _RuntimeCacheError(RuntimeError): + def __init__(self, code: str, message: str): + super().__init__(message) + self.code = code + + +def default_runtime_cache_root() -> str: + """Return the default root, intentionally independent of VLLM_CACHE_ROOT.""" + return os.path.expanduser(DEFAULT_RUNTIME_CACHE_ROOT) + + +def _persistent_buffer(module: nn.Module, local_name: str) -> bool: + parent_path, _, leaf_name = local_name.rpartition(".") + owner = module.get_submodule(parent_path) + return leaf_name not in owner._non_persistent_buffers_set + + +def _named_parameters_with_duplicates(module: nn.Module): + try: + return module.named_parameters(remove_duplicate=False) + except TypeError: # pragma: no cover - compatibility with old torch + return module.named_parameters() + + +def _named_buffers_with_duplicates(module: nn.Module): + try: + return module.named_buffers(remove_duplicate=False) + except TypeError: # pragma: no cover - compatibility with old torch + return module.named_buffers() + + +def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None: + parent_path, _, leaf_name = runtime_name.rpartition(".") + parent = pipeline.get_submodule(parent_path) + tensor = parent._parameters.get(leaf_name) + if tensor is None: + tensor = parent._buffers.get(leaf_name) + return tensor + + +def _collect_runtime_tensors( + pipeline: nn.Module, + dit_modules: Sequence[tuple[str, nn.Module]], +) -> list[_RuntimeTensor]: + records: dict[str, _RuntimeTensor] = {} + for dit_name, dit_module in dit_modules: + for local_name, tensor in _named_parameters_with_duplicates(dit_module): + runtime_name = f"{dit_name}.{local_name}" + candidate = _RuntimeTensor(runtime_name, tensor, "parameter") + existing = records.get(runtime_name) + if existing is not None and existing.tensor is not tensor: + raise _RuntimeCacheError( + "ambiguous_ownership", + f"multiple DiT tensors resolve to {runtime_name!r}", + ) + records[runtime_name] = candidate + for local_name, tensor in _named_buffers_with_duplicates(dit_module): + if not _persistent_buffer(dit_module, local_name): + continue + runtime_name = f"{dit_name}.{local_name}" + candidate = _RuntimeTensor(runtime_name, tensor, "buffer") + existing = records.get(runtime_name) + if existing is not None and existing.tensor is not tensor: + raise _RuntimeCacheError( + "ambiguous_ownership", + f"multiple DiT tensors resolve to {runtime_name!r}", + ) + records[runtime_name] = candidate + + if not records: + raise _RuntimeCacheError( + "ambiguous_ownership", + "no DiT parameters or persistent buffers were discovered", + ) + + storage_owners: dict[tuple[int, int], str] = {} + for record in records.values(): + tensor = record.tensor + try: + pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name) + except AttributeError as exc: + raise _RuntimeCacheError( + "ambiguous_ownership", + f"DiT tensor {record.name!r} is not owned by the pipeline", + ) from exc + if pipeline_tensor is not tensor: + raise _RuntimeCacheError( + "ambiguous_ownership", + f"DiT tensor {record.name!r} does not resolve to the discovered object", + ) + if tensor.device.type != "cpu" or tensor.is_meta: + raise _RuntimeCacheError( + "unsupported_tensor", + f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}", + ) + if tensor.layout != torch.strided: + raise _RuntimeCacheError( + "unsupported_tensor", + f"{record.name!r} uses unsupported layout {tensor.layout}", + ) + if not tensor.is_contiguous(): + raise _RuntimeCacheError( + "unsupported_tensor", + f"{record.name!r} is non-contiguous with stride {tensor.stride()}", + ) + if tensor.dtype not in _SUPPORTED_DTYPES: + raise _RuntimeCacheError( + "unsupported_tensor", + f"{record.name!r} uses unsupported dtype {tensor.dtype}", + ) + if hasattr(tensor, "to_local"): + raise _RuntimeCacheError( + "unsupported_tensor", + f"{record.name!r} is a distributed tensor", + ) + + tensor_nbytes = tensor.numel() * tensor.element_size() + storage = tensor.untyped_storage() + if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes: + raise _RuntimeCacheError( + "unsupported_tensor", + f"{record.name!r} is a view into a larger storage", + ) + if tensor_nbytes: + storage_id = (storage.data_ptr(), storage.nbytes()) + owner = storage_owners.setdefault(storage_id, record.name) + if owner != record.name: + raise _RuntimeCacheError( + "unsupported_alias", + f"{record.name!r} shares storage with {owner!r}", + ) + + # The ownership boundary is the complete pipeline, not only the discovered + # DiT traversal. Reject a DiT storage that is also registered by an encoder, + # VAE, resident component, or non-persistent buffer. + for pipeline_name, tensor in chain( + _named_parameters_with_duplicates(pipeline), + _named_buffers_with_duplicates(pipeline), + ): + if tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0: + continue + storage = tensor.untyped_storage() + owner = storage_owners.get((storage.data_ptr(), storage.nbytes())) + if owner is not None and owner != pipeline_name: + raise _RuntimeCacheError( + "unsupported_alias", + f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}", + ) + + return [records[name] for name in sorted(records)] + + +def _json_normalize(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, os.PathLike): + return os.fspath(value) + if dataclasses.is_dataclass(value): + return _json_normalize(dataclasses.asdict(value)) + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + try: + dumped = model_dump(mode="json") + except (TypeError, ValueError): + dumped = model_dump() + return _json_normalize(dumped) + if isinstance(value, Mapping): + return {str(key): _json_normalize(item) for key, item in sorted(value.items(), key=lambda item: str(item[0]))} + if isinstance(value, (set, frozenset)): + return sorted((_json_normalize(item) for item in value), key=lambda item: _canonical_json(item)) + if isinstance(value, (list, tuple)): + return [_json_normalize(item) for item in value] + if inspect.isclass(value): + return f"{value.__module__}.{value.__qualname__}" + return { + "type": f"{type(value).__module__}.{type(value).__qualname__}", + "value": str(value), + } + + +def _canonical_json(value: Any) -> bytes: + return json.dumps( + _json_normalize(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode() + + +def _implementation_fingerprint( + loader_type: type, pipeline: nn.Module, dit_modules: Sequence[tuple[str, nn.Module]] +) -> str: + objects: list[Any] = [loader_type, type(pipeline)] + load_weights = getattr(type(pipeline), "load_weights", None) + if load_weights is not None: + objects.append(load_weights) + for _, dit_module in dit_modules: + objects.append(type(dit_module)) + post_load = getattr(type(dit_module), "post_load_weights", None) + if post_load is not None: + objects.append(post_load) + + digest = hashlib.sha256() + identities: set[str] = set() + for obj in objects: + identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}" + if identity in identities: + continue + identities.add(identity) + digest.update(identity.encode()) + try: + digest.update(inspect.getsource(obj).encode()) + except (OSError, TypeError): + pass + return digest.hexdigest() + + +def _tensor_metadata(record: _RuntimeTensor) -> dict[str, Any]: + tensor = record.tensor + return { + "kind": record.kind, + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "nbytes": tensor.numel() * tensor.element_size(), + "layout": "contiguous", + } + + +def _update_hash_with_tensor(digest: Any, tensor: torch.Tensor) -> None: + byte_view = tensor.detach().reshape(-1).view(torch.uint8) + for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES): + chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES] + digest.update(memoryview(chunk.numpy())) + + +def _runtime_content_digest(records: Sequence[_RuntimeTensor]) -> str: + digest = hashlib.sha256() + for record in records: + digest.update(_canonical_json({"name": record.name, **_tensor_metadata(record)})) + _update_hash_with_tensor(digest, record.tensor) + return digest.hexdigest() + + +def _file_digest(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(_FILE_HASH_CHUNK_BYTES): + digest.update(chunk) + return digest.hexdigest() + + +@contextlib.contextmanager +def _exclusive_lock(path: Path, timeout_seconds: float): + path.parent.mkdir(parents=True, exist_ok=True) + deadline = time.monotonic() + timeout_seconds + with path.open("a+b") as handle: + while True: + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except OSError as exc: + if exc.errno not in (errno.EACCES, errno.EAGAIN): + raise _RuntimeCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc + remaining = deadline - time.monotonic() + if remaining <= 0: + raise _RuntimeCacheError( + "lock_timeout", + f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}", + ) from exc + time.sleep(min(0.2, remaining)) + try: + yield + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + +def _fsync_file(path: Path) -> None: + with path.open("rb") as handle: + os.fsync(handle.fileno()) + + +def _fsync_directory(path: Path) -> None: + fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _split_shards(records: Sequence[_RuntimeTensor], max_shard_bytes: int) -> list[list[_RuntimeTensor]]: + shards: list[list[_RuntimeTensor]] = [] + current: list[_RuntimeTensor] = [] + current_bytes = 0 + for record in records: + nbytes = record.tensor.numel() * record.tensor.element_size() + if current and current_bytes + nbytes > max_shard_bytes: + shards.append(current) + current = [] + current_bytes = 0 + current.append(record) + current_bytes += nbytes + if current: + shards.append(current) + return shards + + +def _publish_entry( + *, + cache_root: Path, + entry_dir: Path, + cache_key: str, + layout_identity: dict[str, Any], + content_digest: str, + records: Sequence[_RuntimeTensor], + max_shard_bytes: int, +) -> None: + total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records) + free_bytes = shutil.disk_usage(cache_root).free + required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100) + if free_bytes < required_bytes: + raise _RuntimeCacheError( + "insufficient_disk", + f"runtime cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}", + ) + + tmp_root = cache_root / ".tmp" + tmp_root.mkdir(parents=True, exist_ok=True) + temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}" + temp_dir.mkdir() + try: + shards = _split_shards(records, max_shard_bytes) + tensor_manifest: dict[str, dict[str, Any]] = {} + file_manifest: dict[str, dict[str, Any]] = {} + shard_count = len(shards) + for index, shard in enumerate(shards, start=1): + filename = f"model-{index:05d}-of-{shard_count:05d}.safetensors" + path = temp_dir / filename + save_file( + {record.name: record.tensor.detach() for record in shard}, + str(path), + metadata={"format": "pt", "dlo_runtime_cache_schema": str(RUNTIME_CACHE_SCHEMA_VERSION)}, + ) + _fsync_file(path) + file_manifest[filename] = { + "size": path.stat().st_size, + "sha256": _file_digest(path), + } + for record in shard: + tensor_manifest[record.name] = { + **_tensor_metadata(record), + "file": filename, + "storage_key": record.name, + } + + manifest = { + "schema_version": RUNTIME_CACHE_SCHEMA_VERSION, + "cache_key": cache_key, + "layout_identity": layout_identity, + "runtime_content_sha256": content_digest, + "total_tensor_bytes": total_nbytes, + "created_at_unix_seconds": time.time(), + "writer_pid": os.getpid(), + "files": file_manifest, + "tensors": tensor_manifest, + } + manifest_path = temp_dir / "manifest.json" + with manifest_path.open("w", encoding="utf-8") as handle: + json.dump(manifest, handle, sort_keys=True, separators=(",", ":")) + handle.flush() + os.fsync(handle.fileno()) + _fsync_directory(temp_dir) + + invalid_dir: Path | None = None + if entry_dir.exists(): + invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}" + os.replace(entry_dir, invalid_dir) + os.replace(temp_dir, entry_dir) + _fsync_directory(entry_dir.parent) + if invalid_dir is not None: + shutil.rmtree(invalid_dir, ignore_errors=True) + except Exception: + shutil.rmtree(temp_dir, ignore_errors=True) + raise + + +def _validate_entry( + *, + entry_dir: Path, + cache_key: str, + layout_identity: dict[str, Any], + content_digest: str, + expected_records: Sequence[_RuntimeTensor], +) -> HostWeightPlan: + manifest_path = entry_dir / "manifest.json" + try: + with manifest_path.open(encoding="utf-8") as handle: + manifest = json.load(handle) + except (OSError, ValueError) as exc: + raise _RuntimeCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc + + if manifest.get("schema_version") != RUNTIME_CACHE_SCHEMA_VERSION: + raise _RuntimeCacheError("entry_rejected", "runtime-cache schema mismatch") + if manifest.get("cache_key") != cache_key: + raise _RuntimeCacheError("entry_rejected", "runtime-cache key mismatch") + if manifest.get("layout_identity") != layout_identity: + raise _RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch") + if manifest.get("runtime_content_sha256") != content_digest: + raise _RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch") + + expected_metadata = {record.name: _tensor_metadata(record) for record in expected_records} + tensor_manifest = manifest.get("tensors") + file_manifest = manifest.get("files") + if not isinstance(tensor_manifest, dict) or not isinstance(file_manifest, dict): + raise _RuntimeCacheError("entry_rejected", "runtime-cache manifest has invalid sections") + if set(tensor_manifest) != set(expected_metadata): + raise _RuntimeCacheError("entry_rejected", "runtime-cache tensor names differ from the loaded DiT") + + for filename, metadata in file_manifest.items(): + if Path(filename).name != filename or not isinstance(metadata, dict): + raise _RuntimeCacheError("entry_rejected", f"invalid cache filename {filename!r}") + path = entry_dir / filename + try: + stat = path.stat() + except OSError as exc: + raise _RuntimeCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc + if stat.st_size != metadata.get("size") or _file_digest(path) != metadata.get("sha256"): + raise _RuntimeCacheError("entry_rejected", f"content digest mismatch for cache shard {path}") + + bindings: dict[str, TensorBinding] = {} + mapped_records: list[_RuntimeTensor] = [] + with contextlib.ExitStack() as stack: + handles = { + filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu")) + for filename in file_manifest + } + expected_keys_by_file: dict[str, set[str]] = {filename: set() for filename in file_manifest} + for record in expected_records: + stored = tensor_manifest.get(record.name) + if not isinstance(stored, dict): + raise _RuntimeCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}") + filename = stored.get("file") + storage_key = stored.get("storage_key") + if filename not in handles or not isinstance(storage_key, str): + raise _RuntimeCacheError("entry_rejected", f"invalid storage binding for {record.name!r}") + actual_metadata = {key: stored.get(key) for key in expected_metadata[record.name]} + if actual_metadata != expected_metadata[record.name]: + raise _RuntimeCacheError("entry_rejected", f"metadata mismatch for {record.name!r}") + expected_keys_by_file[filename].add(storage_key) + try: + tensor = handles[filename].get_tensor(storage_key) + except Exception as exc: + raise _RuntimeCacheError( + "entry_rejected", + f"cannot map {record.name!r} from {filename}: {exc}", + ) from exc + if tuple(tensor.shape) != tuple(record.tensor.shape) or tensor.dtype != record.tensor.dtype: + raise _RuntimeCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}") + mapped_records.append(_RuntimeTensor(record.name, tensor, record.kind)) + bindings[record.name] = TensorBinding(storage_key=storage_key, file_path=str(entry_dir / filename)) + + for filename, expected_keys in expected_keys_by_file.items(): + if set(handles[filename].keys()) != expected_keys: + raise _RuntimeCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}") + + if _runtime_content_digest(mapped_records) != content_digest: + raise _RuntimeCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher") + + return HostWeightPlan( + backing_kind="runtime_cache", + bindings=bindings, + runtime_layout_key=cache_key, + post_load_complete=True, + ) + + +def _remove_stale_temps(cache_root: Path, cache_key: str) -> None: + temp_root = cache_root / ".tmp" + if not temp_root.is_dir(): + return + for path in temp_root.glob(f"{cache_key}.*"): + if path.is_dir(): + shutil.rmtree(path, ignore_errors=True) + else: + path.unlink(missing_ok=True) + + +def build_runtime_weight_cache_plan( + pipeline: nn.Module, + *, + dit_modules: Sequence[tuple[str, nn.Module]], + loader_type: type, + cache_root: str | os.PathLike[str] | None, + lock_timeout_seconds: float, + max_shard_bytes: int, + model_identity: str | None, + revision: str | None, + runtime_dtype: Any, + load_format: str, + loader_inputs: Mapping[str, Any], + tensor_parallel_size: int, + tensor_parallel_rank: int, + sequence_parallel_guard: Mapping[str, Any], + use_hsdp: bool, + quantization_config: Any, + cfg_parallel_size: int, + pipeline_parallel_size: int, +) -> HostWeightPlanResult: + """Publish or join one immutable cache entry for final CPU DiT tensors.""" + if load_format == "dummy": + return HostWeightPlanResult( + None, + "dummy/random weights do not have a reusable runtime identity", + "unsupported_load_format", + ) + if use_hsdp: + return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp") + if quantization_config is not None: + return HostWeightPlanResult( + None, + "quantized runtime layouts are not supported by the first runtime-cache version", + "unsupported_quantization", + ) + if cfg_parallel_size != 1: + return HostWeightPlanResult(None, "CFG-parallel component ownership is not proven", "unsupported_cfg_parallel") + if pipeline_parallel_size != 1: + return HostWeightPlanResult(None, "pipeline-parallel ownership is outside runtime-cache v1", "unsupported_pp") + if not math.isfinite(lock_timeout_seconds) or lock_timeout_seconds <= 0: + return HostWeightPlanResult(None, "runtime-cache lock timeout must be positive", "invalid_configuration") + if max_shard_bytes <= 0: + return HostWeightPlanResult(None, "runtime-cache shard size must be positive", "invalid_configuration") + + try: + records = _collect_runtime_tensors(pipeline, dit_modules) + content_digest = _runtime_content_digest(records) + layout_identity = { + "schema_version": RUNTIME_CACHE_SCHEMA_VERSION, + "model": model_identity, + "revision": revision, + "runtime_dtype": str(runtime_dtype), + "load_format": load_format, + "loader_inputs": _json_normalize(loader_inputs), + "loader_implementation_sha256": _implementation_fingerprint(loader_type, pipeline, dit_modules), + "components": [name for name, _ in dit_modules], + "tensor_parallel_size": tensor_parallel_size, + "tensor_parallel_rank": tensor_parallel_rank, + "sequence_parallel_guard": _json_normalize(sequence_parallel_guard), + "runtime_content_sha256": content_digest, + } + cache_key = hashlib.sha256(_canonical_json(layout_identity)).hexdigest() + root = Path(cache_root or default_runtime_cache_root()).expanduser().resolve() + entries_root = root / f"v{RUNTIME_CACHE_SCHEMA_VERSION}" + entry_dir = entries_root / cache_key + entries_root.mkdir(parents=True, exist_ok=True) + + if entry_dir.is_dir(): + try: + plan = _validate_entry( + entry_dir=entry_dir, + cache_key=cache_key, + layout_identity=layout_identity, + content_digest=content_digest, + expected_records=records, + ) + logger.info("Reusing DLO runtime cache entry %s from %s", cache_key[:12], entry_dir) + return HostWeightPlanResult(plan) + except _RuntimeCacheError as exc: + logger.warning("DLO runtime cache entry %s will be rebuilt: %s", cache_key[:12], exc) + + lock_path = root / ".locks" / f"{cache_key}.lock" + with _exclusive_lock(lock_path, lock_timeout_seconds): + _remove_stale_temps(root, cache_key) + if entry_dir.is_dir(): + try: + plan = _validate_entry( + entry_dir=entry_dir, + cache_key=cache_key, + layout_identity=layout_identity, + content_digest=content_digest, + expected_records=records, + ) + logger.info("Joined DLO runtime cache entry %s from %s", cache_key[:12], entry_dir) + return HostWeightPlanResult(plan) + except _RuntimeCacheError: + pass + + _publish_entry( + cache_root=root, + entry_dir=entry_dir, + cache_key=cache_key, + layout_identity=layout_identity, + content_digest=content_digest, + records=records, + max_shard_bytes=max_shard_bytes, + ) + plan = _validate_entry( + entry_dir=entry_dir, + cache_key=cache_key, + layout_identity=layout_identity, + content_digest=content_digest, + expected_records=records, + ) + logger.info( + "Published DLO runtime cache entry %s with %d tensors at %s", + cache_key[:12], + len(records), + entry_dir, + ) + return HostWeightPlanResult(plan) + except _RuntimeCacheError as exc: + return HostWeightPlanResult(None, str(exc), exc.code) + except Exception as exc: + return HostWeightPlanResult(None, f"runtime-cache operation failed: {exc}", "cache_operation_failed") + + +__all__ = [ + "DEFAULT_LOCK_TIMEOUT_SECONDS", + "DEFAULT_RUNTIME_CACHE_ROOT", + "DEFAULT_SHARD_SIZE_BYTES", + "RUNTIME_CACHE_SCHEMA_VERSION", + "build_runtime_weight_cache_plan", + "default_runtime_cache_root", +] diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 2502fa76b..73546a721 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -71,8 +71,9 @@ def get_offload_backend( Args: od_config: OmniDiffusionConfig with offload settings device: Target device (auto-detected if None) - host_weight_plan: Exact loader-produced backing plan, if ordinary - weight materialization was skipped. + host_weight_plan: Exact loader-produced backing plan. Checkpoint plans + may have skipped ordinary materialization; runtime-cache plans + describe already-valid ordinary tensors that can be remapped. Returns: OffloadBackend instance or None if offloading disabled @@ -85,10 +86,14 @@ def get_offload_backend( # Extract and validate configuration config = OffloadConfig.from_od_config(od_config) + plan_requires_backend = host_weight_plan is not None and not host_weight_plan.post_load_complete if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE: - raise RuntimeError( - "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected" - ) + if plan_requires_backend: + raise RuntimeError( + "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected" + ) + logger.warning("Ignoring optional runtime-cache plan because distributed layerwise offload is not selected") + host_weight_plan = None # Return None if no offloading requested if config.strategy == OffloadStrategy.NONE: @@ -96,7 +101,7 @@ def get_offload_backend( # Validate platform (CUDA required for now) if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1: - if host_weight_plan is not None: + if plan_requires_backend: raise RuntimeError( "The loader skipped ordinary weight materialization for DLO, " "but this platform cannot create the required offload backend" @@ -112,7 +117,7 @@ def get_offload_backend( try: device = current_omni_platform.get_torch_device() except (NotImplementedError, AttributeError) as exc: - if host_weight_plan is not None: + if plan_requires_backend: raise RuntimeError( "The loader skipped ordinary weight materialization for DLO, " "but the target device could not be resolved" diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index d184e71ea..597f478dd 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1023,7 +1023,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): framework="pt", device="cpu", ) - tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key) + tensor = file_cache[binding.file_path].get_tensor(binding.storage_key) if is_parameter: replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad) parent._parameters[leaf_name] = replacement @@ -1103,6 +1103,90 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): f"(first 5: {remaining_meta[:5]})." ) + def _load_weights_from_runtime_cache( + self, + pipeline: nn.Module, + modules, + plan: HostWeightPlan, + ) -> None: + """Replace final ordinary-loader tensors with validated cache views.""" + from safetensors import safe_open + + if not plan.post_load_complete: + raise RuntimeError("A runtime-cache plan must describe fully post-processed weights") + if any(binding.transform is not None for binding in plan.bindings.values()): + raise RuntimeError("Runtime-cache bindings cannot carry deferred checkpoint transforms") + + 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()) + for name, _ in dit_module.named_buffers(): + parent_path, _, leaf_name = name.rpartition(".") + owner = dit_module.get_submodule(parent_path) + if leaf_name not in owner._non_persistent_buffers_set: + required_names.add(f"{dit_name}.{name}") + if required_names != set(plan.bindings): + missing = sorted(required_names - set(plan.bindings)) + unexpected = sorted(set(plan.bindings) - required_names) + raise RuntimeError( + "Runtime-cache plan coverage changed after loader validation: " + f"missing={missing[:5]}, unexpected={unexpected[:5]}" + ) + + file_cache: dict[str, Any] = {} + replacements: list[tuple[torch.Tensor, torch.Tensor]] = [] + try: + # Validate and map every tensor before mutating the pipeline. A + # missing/corrupt shard therefore leaves the ordinary tensors intact. + for runtime_name, binding in plan.bindings.items(): + parent_path, _, leaf_name = runtime_name.rpartition(".") + try: + parent = pipeline.get_submodule(parent_path) + except AttributeError as exc: + raise RuntimeError(f"Runtime-cache target module {parent_path!r} no longer exists") from exc + + target = parent._parameters.get(leaf_name) + if target is None: + target = parent._buffers.get(leaf_name) + if target is None: + raise RuntimeError(f"Runtime-cache target tensor {runtime_name!r} no longer exists") + if target.device.type != "cpu" or target.is_meta: + raise RuntimeError( + f"Runtime-cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}" + ) + + if binding.file_path not in file_cache: + file_cache[binding.file_path] = safe_open( + binding.file_path, + framework="pt", + device="cpu", + ) + tensor = file_cache[binding.file_path].get_tensor(binding.storage_key) + if tensor.shape != target.shape or tensor.dtype != target.dtype: + raise RuntimeError( + f"Runtime-cache metadata changed for {runtime_name!r}: " + f"cache={tuple(tensor.shape)}/{tensor.dtype}, " + f"runtime={tuple(target.shape)}/{target.dtype}" + ) + replacements.append((target, tensor)) + except Exception: + file_cache.clear() + raise + + # Preserve Parameter objects and their loader/runtime attributes; only + # their storage changes. This also keeps any references held by model + # code valid while dropping the private allocator-backed storage. + for target, tensor in replacements: + set_tensor_storage(target, tensor) + + self._mmap_transforms_by_tensor_id.clear() + self._mmap_file_cache = file_cache + logger.info( + "Remapped %d final DiT tensors from runtime cache %s", + len(replacements), + plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown", + ) + def _init_dp_group(self) -> None: """Reuse the process group initialized by parallel_state. @@ -1443,16 +1527,30 @@ 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 == "checkpoint_mmap": + self._load_weights_via_mmap( + pipeline, + modules, + self.host_weight_plan, + ) + elif self.host_weight_plan.backing_kind == "runtime_cache": + if self.config.dlo_use_allgather: + raise ValueError("Runtime-cache backing is no-AllGather only in the first implementation") + self._using_rank_local_mmap = True + self._load_weights_from_runtime_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() + 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 mmap storage enabled: file-backed pages are " "node-shared; each worker owns only two bounded host staging slots" ) else: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fcd11b173..db79c2a34 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -545,6 +545,9 @@ class OrchestratorArgs: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_enable_runtime_cache: bool = False + dlo_runtime_cache_dir: str | None = None + dlo_runtime_cache_lock_timeout: float = 600.0 dlo_resident_layers: int = 0 boundary_ratio: float | None = None flow_shift: float | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e35315f7e..cc1d5c469 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1040,6 +1040,9 @@ class AsyncOmniEngine: "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), + "dlo_enable_runtime_cache": kwargs.get("dlo_enable_runtime_cache", False), + "dlo_runtime_cache_dir": kwargs.get("dlo_runtime_cache_dir", None), + "dlo_runtime_cache_lock_timeout": kwargs.get("dlo_runtime_cache_lock_timeout", 600.0), "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 62ee10d94..1e74cb2a4 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -722,6 +722,30 @@ class OmniServeCommand(CLISubcommand): "independently (including existing TP shards)." ), ) + omni_config_group.add_argument( + "--dlo-enable-runtime-cache", + action="store_true", + help=( + "After ordinary loading, publish or join a node-local mmap cache " + "for final DiT weights. This first version requires " + "--dlo-no-use-allgather." + ), + ) + omni_config_group.add_argument( + "--dlo-runtime-cache-dir", + type=str, + default=None, + help=( + "Shared local-disk directory for DLO runtime weights. Defaults " + "to ~/.cache/vllm-omni/dlo-runtime-weights." + ), + ) + omni_config_group.add_argument( + "--dlo-runtime-cache-lock-timeout", + type=float, + default=600.0, + help="Seconds to wait for another process publishing the same runtime layout.", + ) omni_config_group.add_argument( "--dlo-resident-layers", type=int,