diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 5f00487e1..59a252d45 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -47,8 +47,9 @@ 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. An -opt-in CUDA registration budget can make the complete final mapping directly -H2D-copyable, avoiding the recurrent host packing otherwise required by mmap. +opt-in host-registration budget can make the complete final mapping directly +H2D-copyable when the platform provides a backend, avoiding the recurrent host +packing otherwise required by mmap. CUDA is the first implementation. 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 @@ -156,10 +157,10 @@ modeled safely. ### Registered runtime-cache transfer -`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete -runtime-cache registration. It also changes loader selection: no-AllGather -uses the ordinary loader and final runtime cache even when a TP1 direct- -checkpoint plan is available. A checkpoint binding can carry a deferred +`dlo_runtime_cache_pin_limit_gib > 0` opts a worker into complete runtime-cache +registration when its platform supports it. It also changes loader selection: +no-AllGather uses the ordinary loader and final runtime cache even when a TP1 +direct-checkpoint plan is available. A checkpoint binding can carry a deferred adapter such as grouped-QKV reordering. Registering that raw source would not remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache contains the transform-complete final tensors. @@ -167,13 +168,13 @@ contains the transform-complete final tensors. After mapping and validating every final tensor, DLO groups source storages by backing file, page-aligns their address ranges, and coalesces overlapping or adjacent ranges within that mapping. The complete byte count is checked -against the per-worker budget before any CUDA call. It then registers every -range with the read-only CUDA host-registration flag and verifies that PyTorch -recognizes every mapped tensor as pinned. Failure rolls back any ranges already -registered and preserves the existing two-slot staging path. Partial direct -transfer is deliberately unsupported. A rollback failure aborts worker -initialization so a still-registered range is retained until process/context -teardown rather than being unsafely unmapped. +against the per-worker budget before any platform mutation. The CUDA backend +then registers every range with the read-only host-registration flag and +verifies that PyTorch recognizes every mapped tensor as pinned. Failure rolls +back any ranges already registered and preserves the existing two-slot staging +path. Partial direct transfer is deliberately unsupported. A rollback failure +aborts worker initialization so a still-registered range is retained until +process/context teardown rather than being unsafely unmapped. On success, the hooks copy each immutable tensor view directly into its offset in the existing flattened rotating device buffer on the copy stream. Parameter @@ -183,13 +184,13 @@ but avoids the much larger recurrent CPU copy. A packed cache schema should be considered only if launch-count measurements justify its added format and publication complexity. -Registration is per process and CUDA context; physical file-cache pages remain -shared and are not copied into one private host allocation per worker. The -operator must nevertheless budget page-locked memory and satisfy OS/CUDA -registration limits for every worker. During shutdown, DLO first synchronizes -outstanding device work, unregisters all ranges, and then closes the safetensors -mappings. Worker teardown invokes this cleanup before destroying the -distributed environment and CUDA context. +Registration is per process and platform context; physical file-cache pages +remain shared and are not copied into one private host allocation per worker. +The operator must nevertheless budget page-locked memory and satisfy the +platform and OS registration limits for every worker. During shutdown, DLO +first synchronizes outstanding device work, unregisters all ranges, and then +closes the safetensors mappings. Worker teardown invokes this cleanup before +destroying the distributed environment and platform context. ### AllGather path diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index eff153fe1..9a7adda6a 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -1,265 +1,103 @@ # Distributed Layerwise Offloading -Distributed layerwise offloading (DLO) extends block streaming to multi-device -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. The optional -runtime-weight cache extends node-local sharing to final ordinary-loader -layouts, including matching TP coordinates across independent engines. On -CUDA, an explicit registration budget can make those shared mappings directly -H2D-copyable and avoid recurrent mmap-to-pinned host packing. +Distributed layerwise offloading (DLO) streams DiT blocks through two device +buffers instead of keeping the complete DiT in HBM. It supports sharded host +weights with AllGather or independently streamed rank-local weights. The +no-AllGather path can optionally share node-local runtime weights through mmap. See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md) -for the implementation contract and compatibility matrix. +for loader contracts, compatibility details, failure handling, and benchmarks. -## Execution model +## Choose a mode -DLO overlaps three operations with a fixed two-block device buffer: +| Mode | Use when | Host weights | Runtime synchronization | +| --- | --- | --- | --- | +| DLO AllGather (default) | DP ranks execute the same block path in lockstep | About `1 / dp_size` per rank | DLO weight AllGather | +| no-AllGather | Ranks or engines must schedule independently | Complete rank-local layout | No DLO weight collective | +| no-AllGather + runtime cache | Equivalent independent workers share one node | Shared final mmap layout per TP coordinate | No DLO weight collective | +| runtime cache + host registration | The platform supports registration and recurrent staging is too expensive | Shared registered mmap layout | No DLO weight collective | -```text -Compute stream: [Layer N] [Layer N+1] [Layer N+2] -H2D stream: [H2D shard N+1] [H2D shard N+2] -AllGather: [AG N+1] [AG N+2] -Slots: slot 0: Layer N slot 1: Layer N+1 -``` - -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. +AllGather is normally the best choice for synchronized DP. Runtime-cache mode +targets independently scheduled replicas, especially repeated TP engines on +one node. ## Usage ```bash -# Four ranks with sharded host weights and AllGather +# Sharded host weights with DLO AllGather vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --data-parallel-size 4 -# Standard-loader rank-local weights, without DLO AllGather +# Independently streamed rank-local weights vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ - --data-parallel-size 4 \ --dlo-no-use-allgather -# Independently scheduled replicas sharing final runtime weights on one node +# Share final runtime weights and register up to 80 GiB per worker for direct H2D 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 \ --dlo-runtime-cache-pin-limit-gib 80 - -# Sequence parallel deployment -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --usp 4 ``` -```python -from vllm_omni import Omni - -omni = Omni( - model="/path/to/model", - enable_distributed_layerwise_offload=True, - dlo_use_allgather=False, - dlo_enable_runtime_cache=True, - dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights", - dlo_runtime_cache_pin_limit_gib=80, -) -``` - -## Flags +## Relevant options | Flag | Meaning | Default | | --- | --- | --- | | `--enable-distributed-layerwise-offload` | Enable DLO | `false` | -| `--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-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` | -| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | +| `--dlo-no-use-allgather` | Stream complete rank-local blocks independently | `false` | +| `--dlo-enable-runtime-cache` | Share final no-AllGather DiT weights through a node-local mmap cache | `false` | +| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker registration budget; zero uses bounded host staging | `0` | +| `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` | -## Host-weight loading +The runtime cache uses `~/.cache/vllm-omni/dlo-runtime-weights`. Programmatic +configuration may override `dlo_runtime_cache_dir` and the writer lock timeout; +these advanced storage controls are intentionally not separate CLI flags. -The diffusion loader chooses host storage before DLO is enabled. It first -attempts to build a complete, validated direct-checkpoint mmap plan. If names, -coverage, shape, dtype, topology, or loader-callback compatibility cannot be -proven, it runs the ordinary model loader instead. DLO consumes that result and -does not make a second checkpoint-compatibility decision. +## Operational notes -The shared-mmap optimization in this phase is supported only with TP1. TP -greater than one falls back before model mutation to ordinary TP-aware loading. -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. +- All workers that should share must see the same local, disk-backed cache + directory. Do not use tmpfs or a cross-node filesystem for this version. +- Cache entries are immutable and validated before use. A cache or registration + failure keeps the ordinary loader weights or falls back to bounded staging. +- A positive registration budget must cover the complete page-aligned mapping + reported in the worker log. Registration is all-or-nothing. +- CUDA is the first registration backend. Platforms without an equivalent + implementation continue to use bounded staging. +- Registration is process-local but does not duplicate the underlying file + pages. Each worker must still satisfy its platform and OS page-locking limits. +- Shutdown unregisters host ranges before closing their mmap handles. +- Runtime-cache v1 has no automatic eviction. Stop all users of an entry before + deleting it. -When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading -gains a second storage opportunity. Every rank completes 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. +## Scheduling constraints -A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final -runtime layout even when a direct TP1 checkpoint plan is available. Raw -checkpoint mappings may require deferred loader transforms such as grouped-QKV -reordering; registering them would still leave recurrent host work. Publishing -the fully transformed layout once makes every streamed source directly -copyable. +AllGather ranks must request blocks in the same collective order. Concurrent DP +requests may use different prompts, but they must follow the same denoising and +block-execution path and set the same explicit `num_inference_steps`. -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. - -With a zero registration budget, cached tensors remain pageable mmap views. -DLO packs each streamed block into one of two bounded pinned host slots before -the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent -host-to-host copies; competing TP ranks or engines can turn DDR contention and -staging skew into collective wait time. - -With a positive budget on CUDA, DLO instead attempts to register the complete -page-aligned cache mapping as read-only pinned host memory. Registration is -all-or-nothing: the total mapped range must fit the per-worker budget, and any -CUDA registration or verification failure rolls back partial registrations and -uses the bounded staging path. If CUDA cannot roll back a partial registration, -worker initialization aborts instead of unmapping memory that CUDA may still -own. On success, each immutable tensor view is copied directly into the -existing rotating device buffers on the H2D stream, so no private host staging -slots are allocated. - -CUDA registration is process-local, but it does not clone file contents. The -same physical page-cache pages remain shared among equivalent workers; each -worker does, however, hold its own registration reference and must satisfy the -host's page-locking limits. Shutdown synchronizes outstanding GPU work, -unregisters every range, and only then closes the mmap handles. - -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 -non-DiT weights falls back completely rather than leaving an unplanned -component uninitialized. - -With direct checkpoint mmap, the loader: - -1. saves non-persistent buffers such as RoPE frequencies; -2. moves the normally created transformer to the meta device; -3. loads checkpoint tensors as mmap views backed by the shared OS page cache; -4. applies any loader-owned bounded layout adapters while packing blocks; -5. restores saved non-persistent buffers; and -6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle - hooks. - -For AllGather with a group larger than one, each process copies only its -persistent shard and then releases the source mapping. For no-AllGather, each -process keeps the mapping open and packs complete blocks through two bounded -pinned staging slots. Processes mapping the same files on one node share the -immutable pages; no-AllGather still performs a complete-block H2D copy in each -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. - -For direct H2D, set the registration limit above the complete page-aligned -layout size reported in the worker log. Partial registration is intentionally -unsupported because mixing direct and staged tensors would complicate buffer -lifetime and make performance topology-dependent. Registration is CUDA-only; -an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or -registration error is reported and falls back to bounded staging. - -## Declarative topology - -Models may declare an `OffloadPlan` instead of embedding offload logic: - -```python -from vllm_omni.diffusion.offloader import OffloadPlan - - -class MyPipeline(nn.Module): - _dit_modules = ["transformer"] - _offload_plan = OffloadPlan( - block_attrs={"transformer": ("blocks",)}, - offload_submodules={"context_encoder": "layers"}, - ) -``` - -When no plan exists, discovery falls back to -`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup. - -## Data-parallel concurrency - -With `data_parallel_size > 1` and AllGather enabled, the scheduler can process -up to `dp_size` requests per denoising step. Every concurrent request must set -the same explicit `num_inference_steps`; `None` is rejected because every rank -must enter each collective. +No-AllGather workers do not have this DLO lockstep requirement. Ordinary TP or +SP model collectives still synchronize ranks within each engine. ## Limitations -- 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. A positive registration budget also - chooses the final runtime cache at TP1 so deferred checkpoint transforms are - not repeated during inference. -- HSDP plus AllGather is rejected to avoid double sharding. HSDP without - AllGather has limited end-to-end validation. +- Direct checkpoint mmap currently requires TP1. The runtime cache can share + final ordinary-loader layouts between matching TP coordinates at TP1 or TP>1. - Per-tensor online FP8 linears use the ordinary loader and can run with either - DLO transfer path. With AllGather, every rank temporarily materializes the - 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, expert-parallel, 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. - -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. Direct H2D removes recurrent staging but does not change that startup -contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in -[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). - -Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3 -smoke documented in the L20X runtime-cache matrix in the -[feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix), -two independent TP2 engines reduced request-time host PSS by 36.6% with the -cache, but recurrent mmap-to-pinned staging reduced combined throughput by -69.3%. The registered path recovered near-private per-engine latency while -retaining the shared-memory benefit; a matched TP1 run improved mean wave -latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the -feature design for the full smoke-test caveats and metrics. + DLO transfer path. Other online quantization methods require no-AllGather + until their runtime layouts are validated. +- Runtime-cache v1 rejects quantized, non-contiguous, aliased/tied, device-only, + HSDP, expert-parallel, CFG-parallel, and PP layouts. +- HSDP plus DLO AllGather is unsupported. HSDP without AllGather has limited + end-to-end validation. +- The runtime cache lowers steady-state host PSS, not startup peak: each worker + still performs ordinary loading and full content validation. +- Skip-load-on-hit, cache eviction, cross-node sharing, and broader SP hardware + validation 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/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py index 37bdff116..01a1f763e 100644 --- a/tests/diffusion/offloader/test_cuda_host_registration.py +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -6,6 +6,7 @@ from __future__ import annotations from types import SimpleNamespace import pytest +import torch import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module from vllm_omni.diffusion.offloader.cuda_host_registration import ( @@ -15,6 +16,10 @@ from vllm_omni.diffusion.offloader.cuda_host_registration import ( CudaHostRegistrationError, _coalesce_ranges, ) +from vllm_omni.diffusion.offloader.host_registration import ( + HostRegistrationError, + register_host_mappings, +) pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -78,6 +83,25 @@ class _FakeRuntime: return f"error-{error}" +def test_platform_factory_reports_unsupported_registration() -> None: + with pytest.raises(HostRegistrationError, match="not supported on cpu"): + register_host_mappings({}, device=torch.device("cpu"), max_bytes=4096) + + +def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + sources = {"weights": []} + + def create(actual_sources, *, max_bytes): + assert actual_sources is sources + assert max_bytes == 4096 + return sentinel + + monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create)) + + assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel + + def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None: assert _coalesce_ranges( [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)], diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index e49375683..cf985ffa8 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -1056,11 +1056,12 @@ class TestMmapWeightLoading: create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = [] - def create(sources, *, max_bytes): + def register(sources, *, device, max_bytes): + assert device == torch.device("cuda") create_calls.append((sources, max_bytes)) return Registration() - monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create)) + monkeypatch.setattr(dist_backend_module, "register_host_mappings", register) backend = DistributedLayerwiseOffloadBackend( OffloadConfig( strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, @@ -1084,6 +1085,24 @@ class TestMmapWeightLoading: assert events == ["unregister", "mmap"] + def test_runtime_cache_registration_falls_back_on_unsupported_platform( + self, + patched_offload_runtime, + ): + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dlo_use_allgather=False, + dlo_runtime_cache_pin_limit_gib=1.0, + ), + torch.device("cpu"), + ) + backend._runtime_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]} + + assert not backend._try_register_runtime_cache_mmap() + assert backend._host_registration is None + class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 565eeaf84..95cb37770 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -329,10 +329,6 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--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-runtime-cache-pin-limit-gib", "80", "--dlo-resident-layers", @@ -347,15 +343,13 @@ 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_runtime_cache_pin_limit_gib == 80 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_runtime_cache_dir"] is None + assert engine_args["dlo_runtime_cache_lock_timeout"] == 600.0 assert engine_args["dlo_runtime_cache_pin_limit_gib"] == 80 assert engine_args["dlo_resident_layers"] == 20 diff --git a/vllm_omni/diffusion/model_loader/runtime_cache_utils.py b/vllm_omni/diffusion/model_loader/runtime_cache_utils.py new file mode 100644 index 000000000..431343137 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/runtime_cache_utils.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deterministic identity helpers for the DLO runtime-weight cache.""" + +from __future__ import annotations + +import dataclasses +import json +import os +from collections.abc import Mapping +from enum import Enum +from pathlib import Path +from typing import TypeAlias + +import torch + +JsonScalar: TypeAlias = None | bool | int | float | str +JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"] +PathIdentity: TypeAlias = str | os.PathLike[str] | os.PathLike[bytes] + + +class IdentityNormalizationError(ValueError): + """Raised when a cache identity contains a process-unstable value.""" + + +def _type_identity(value: object) -> str: + value_type = value if isinstance(value, type) else type(value) + return f"{value_type.__module__}.{value_type.__qualname__}" + + +def canonicalize_existing_local_path(value: PathIdentity) -> str: + """Resolve equivalent local paths while leaving model repository IDs alone.""" + original = os.fsdecode(os.fspath(value)) + candidate = Path(original).expanduser() + try: + if candidate.exists(): + return str(candidate.resolve()) + except OSError: + pass + return original + + +def normalize_identity(value: object) -> JsonValue: + """Convert supported loader inputs into deterministic JSON-compatible data.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, Enum): + return {"enum": _type_identity(value), "name": value.name} + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, os.PathLike): + return canonicalize_existing_local_path(value) + if isinstance(value, type): + return _type_identity(value) + if dataclasses.is_dataclass(value): + return normalize_identity(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 normalize_identity(dumped) + + if isinstance(value, Mapping): + normalized: dict[str, JsonValue] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise IdentityNormalizationError("runtime-cache identity mappings require string keys") + normalized[key] = normalize_identity(item) + return {key: normalized[key] for key in sorted(normalized)} + if isinstance(value, (set, frozenset)): + return sorted((normalize_identity(item) for item in value), key=canonical_json) + if isinstance(value, (list, tuple)): + return [normalize_identity(item) for item in value] + + raise IdentityNormalizationError(f"runtime-cache identity does not support values of type {_type_identity(value)}") + + +def canonical_json(value: object) -> bytes: + """Serialize a supported identity with stable ordering and separators.""" + return json.dumps( + normalize_identity(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode() + + +__all__ = [ + "IdentityNormalizationError", + "JsonValue", + "canonical_json", + "canonicalize_existing_local_path", + "normalize_identity", +] diff --git a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py index 3fc9031d2..96dfc4d7f 100644 --- a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py +++ b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py @@ -5,7 +5,6 @@ from __future__ import annotations import contextlib -import dataclasses import errno import fcntl import hashlib @@ -18,7 +17,6 @@ import time import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass -from enum import Enum from itertools import chain from pathlib import Path from typing import Any @@ -34,12 +32,24 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlanResult, TensorBinding, ) +from vllm_omni.diffusion.model_loader.runtime_cache_utils import ( + IdentityNormalizationError, + canonical_json, + canonicalize_existing_local_path, + normalize_identity, +) logger = init_logger(__name__) +# Internal cache-format and I/O defaults, not environment variables. Public +# configuration may override the root and lock timeout; chunk sizes stay private. +# Bump the schema to isolate incompatible manifest layouts. RUNTIME_CACHE_SCHEMA_VERSION = 1 +# Keep the default shared across local stage replicas. DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights" +# Bound waits for another publisher instead of hanging initialization forever. DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0 +# Avoid one model-sized output file and bound hashing/read working sets. DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3 _HASH_CHUNK_BYTES = 64 * 1024**2 _FILE_HASH_CHUNK_BYTES = 8 * 1024**2 @@ -112,6 +122,13 @@ def _collect_runtime_tensors( pipeline: nn.Module, dit_modules: Sequence[tuple[str, nn.Module]], ) -> list[_RuntimeTensor]: + """Collect a complete, safely remappable view of final CPU DiT tensors. + + Parameters and persistent buffers must resolve to the same objects through + the owning pipeline, occupy complete contiguous CPU storages, use supported + dtypes, and have no aliases inside or outside the discovered DiT modules. + The stable name ordering is reused by content hashing and publication. + """ records: dict[str, _RuntimeTensor] = {} for dit_name, dit_module in dit_modules: for local_name, tensor in _named_parameters_with_duplicates(dit_module): @@ -220,68 +237,6 @@ def _collect_runtime_tensors( return [records[name] for name in sorted(records)] -def _type_identity(value: Any) -> str: - value_type = value if inspect.isclass(value) else type(value) - return f"{value_type.__module__}.{value_type.__qualname__}" - - -def _canonicalize_existing_local_path(value: str | os.PathLike[str]) -> str: - original = os.fsdecode(os.fspath(value)) - candidate = Path(original).expanduser() - try: - if candidate.exists(): - return str(candidate.resolve()) - except OSError: - pass - return original - - -def _json_normalize(value: Any) -> Any: - if value is None or isinstance(value, (bool, int, float, str)): - return value - if isinstance(value, Enum): - return {"enum": _type_identity(value), "name": value.name} - if isinstance(value, torch.dtype): - return str(value) - if isinstance(value, os.PathLike): - return _canonicalize_existing_local_path(value) - if inspect.isclass(value): - return _type_identity(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): - if not all(isinstance(key, str) for key in value): - raise _RuntimeCacheError( - "unstable_identity", - "runtime-cache identity mappings require string keys", - ) - return {key: _json_normalize(value[key]) for key in sorted(value)} - 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] - raise _RuntimeCacheError( - "unstable_identity", - f"runtime-cache identity does not support values of type {_type_identity(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: @@ -331,7 +286,7 @@ def _update_hash_with_tensor(digest: Any, tensor: torch.Tensor) -> None: 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)})) + digest.update(canonical_json({"name": record.name, **_tensor_metadata(record)})) _update_hash_with_tensor(digest, record.tensor) return digest.hexdigest() @@ -636,19 +591,19 @@ def build_runtime_weight_cache_plan( content_digest = _runtime_content_digest(records) layout_identity = { "schema_version": RUNTIME_CACHE_SCHEMA_VERSION, - "model": (_canonicalize_existing_local_path(model_identity) if model_identity is not None else None), + "model": (canonicalize_existing_local_path(model_identity) if model_identity is not None else None), "revision": revision, "runtime_dtype": str(runtime_dtype), "load_format": load_format, - "loader_inputs": _json_normalize(loader_inputs), + "loader_inputs": normalize_identity(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), + "sequence_parallel_guard": normalize_identity(sequence_parallel_guard), "runtime_content_sha256": content_digest, } - cache_key = hashlib.sha256(_canonical_json(layout_identity)).hexdigest() + 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 @@ -710,6 +665,8 @@ def build_runtime_weight_cache_plan( return HostWeightPlanResult(plan) except _RuntimeCacheError as exc: return HostWeightPlanResult(None, str(exc), exc.code) + except IdentityNormalizationError as exc: + return HostWeightPlanResult(None, str(exc), "unstable_identity") except Exception as exc: return HostWeightPlanResult(None, f"runtime-cache operation failed: {exc}", "cache_operation_failed") diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py index e4a4ed305..8b9e6314e 100644 --- a/vllm_omni/diffusion/offloader/cuda_host_registration.py +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -1,29 +1,35 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Bounded CUDA registration for existing read-only host mappings.""" +"""CUDA implementation of read-only host-mapping registration.""" from __future__ import annotations import mmap from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Any +from typing import Protocol import torch +from .host_registration import ( + HostRegistrationBudgetError, + HostRegistrationCleanupError, + HostRegistrationError, +) + _CUDA_HOST_REGISTER_READ_ONLY = 0x08 - -class CudaHostRegistrationError(RuntimeError): - """Raised when an existing host range cannot be registered safely.""" +CudaHostRegistrationError = HostRegistrationError +CudaHostRegistrationBudgetError = HostRegistrationBudgetError +CudaHostRegistrationCleanupError = HostRegistrationCleanupError -class CudaHostRegistrationBudgetError(CudaHostRegistrationError): - """Raised before mutation when the complete mapping exceeds its budget.""" +class _CudaRuntime(Protocol): + def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ... + def cudaHostUnregister(self, address: int) -> int: ... -class CudaHostRegistrationCleanupError(CudaHostRegistrationError): - """Raised when partial registration could not be rolled back safely.""" + def cudaGetErrorString(self, error: int) -> str | bytes: ... @dataclass(frozen=True) @@ -74,7 +80,7 @@ def _tensor_regions( ranges: list[tuple[int, int]] = [] for tensor in tensors: if tensor.device.type != "cpu": - raise CudaHostRegistrationError( + raise HostRegistrationError( f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}" ) storage = tensor.untyped_storage() @@ -83,7 +89,7 @@ def _tensor_regions( return tuple(regions) -def _error_message(runtime: Any, error: Any) -> str: +def _error_message(runtime: _CudaRuntime, error: int) -> str: try: message = runtime.cudaGetErrorString(error) except Exception: @@ -98,7 +104,7 @@ class CudaHostRegistration: def __init__( self, - runtime: Any, + runtime: _CudaRuntime, regions: tuple[_AddressRange, ...], ) -> None: self._runtime = runtime @@ -113,22 +119,22 @@ class CudaHostRegistration: max_bytes: int, ) -> CudaHostRegistration: if max_bytes <= 0: - raise CudaHostRegistrationBudgetError("CUDA host-registration budget is disabled") + raise HostRegistrationBudgetError("CUDA host-registration budget is disabled") regions = _tensor_regions(sources_by_mapping) total_bytes = sum(region.size for region in regions) if total_bytes > max_bytes: - raise CudaHostRegistrationBudgetError( + raise HostRegistrationBudgetError( f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget" ) if not regions: - raise CudaHostRegistrationError("no non-empty host ranges were available for registration") + raise HostRegistrationError("no non-empty host ranges were available for registration") if not torch.cuda.is_available(): - raise CudaHostRegistrationError("CUDA is not available") + raise HostRegistrationError("CUDA is not available") try: runtime = torch.cuda.cudart() except Exception as exc: - raise CudaHostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc + raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc registered: list[_AddressRange] = [] try: @@ -139,7 +145,7 @@ class CudaHostRegistration: _CUDA_HOST_REGISTER_READ_ONLY, ) if int(error) != 0: - raise CudaHostRegistrationError( + raise HostRegistrationError( "cudaHostRegister(read-only) failed for " f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}" ) @@ -151,7 +157,7 @@ class CudaHostRegistration: if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors) ] if unpinned: - raise CudaHostRegistrationError( + raise HostRegistrationError( f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}" ) except Exception as exc: @@ -166,12 +172,12 @@ class CudaHostRegistration: except Exception as rollback_exc: rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}") if rollback_errors: - raise CudaHostRegistrationCleanupError( + raise HostRegistrationCleanupError( f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}" ) from exc - if isinstance(exc, CudaHostRegistrationError): + if isinstance(exc, HostRegistrationError): raise - raise CudaHostRegistrationError(f"CUDA host registration raised: {exc}") from exc + raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc return cls(runtime, regions) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 42541537d..db9f29591 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -35,10 +35,11 @@ from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig from .block_discovery import get_blocks_from_dit -from .cuda_host_registration import ( - CudaHostRegistration, - CudaHostRegistrationCleanupError, - CudaHostRegistrationError, +from .host_registration import ( + HostRegistration, + HostRegistrationCleanupError, + HostRegistrationError, + register_host_mappings, ) from .module_collector import ModuleDiscovery from .offload_plan import ( @@ -966,7 +967,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.host_weight_plan = host_weight_plan self._mmap_transforms_by_tensor_id: dict[int, Any] = {} self._runtime_cache_mapped_sources: dict[str, list[torch.Tensor]] = {} - self._cuda_host_registration: CudaHostRegistration | None = None + self._host_registration: HostRegistration | None = None def load_resident_layers(self) -> None: """Load the model-declared leading blocks for the denoise stage.""" @@ -1259,12 +1260,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): pin_limit_gib = self.config.dlo_runtime_cache_pin_limit_gib if pin_limit_gib <= 0: return False - if self.device.type != "cuda": - logger.warning( - "DLO runtime-cache registration is CUDA-only; using bounded host staging on %s", - self.device, - ) - return False if not self.config.pin_cpu_memory: logger.warning("DLO runtime-cache registration requires pin_cpu_memory=True; using bounded host staging") return False @@ -1275,23 +1270,24 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): max_bytes = int(pin_limit_gib * 1024**3) started = time.perf_counter() try: - registration = CudaHostRegistration.create( + registration = register_host_mappings( self._runtime_cache_mapped_sources, + device=self.device, max_bytes=max_bytes, ) - except CudaHostRegistrationCleanupError: - # Falling back could unmap a range that CUDA still owns. Abort the - # worker so process/context teardown provides the final cleanup. + except HostRegistrationCleanupError: + # Falling back could unmap a range the platform still owns. Abort + # the worker so process/context teardown provides final cleanup. logger.exception("DLO runtime-cache registration rollback failed") raise - except CudaHostRegistrationError as exc: + except HostRegistrationError as exc: logger.warning( "DLO runtime-cache direct H2D unavailable (%s); using bounded host staging", exc, ) return False - self._cuda_host_registration = registration + self._host_registration = registration logger.info( "Registered %.2f GiB of shared runtime-cache mappings in %d range(s) for direct H2D in %.3f s", registration.total_bytes / 1024**3, @@ -1301,7 +1297,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): return True def _release_registered_mmap(self) -> bool: - registration = self._cuda_host_registration + registration = self._host_registration if registration is None: return True errors = registration.close() @@ -1311,7 +1307,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): errors[:3], ) return False - self._cuda_host_registration = None + self._host_registration = None logger.info("Unregistered DLO runtime-cache host mappings") return True diff --git a/vllm_omni/diffusion/offloader/host_registration.py b/vllm_omni/diffusion/offloader/host_registration.py new file mode 100644 index 000000000..72f44f218 --- /dev/null +++ b/vllm_omni/diffusion/offloader/host_registration.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Platform-neutral contract for registering existing host mappings.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Protocol + +import torch + + +class HostRegistrationError(RuntimeError): + """Raised when existing host ranges cannot be registered safely.""" + + +class HostRegistrationBudgetError(HostRegistrationError): + """Raised before mutation when the complete mapping exceeds its budget.""" + + +class HostRegistrationCleanupError(HostRegistrationError): + """Raised when partial registration cannot be rolled back safely.""" + + +class HostRegistration(Protocol): + """Lifecycle owned by DLO for one platform's registered host ranges.""" + + @property + def total_bytes(self) -> int: ... + + @property + def region_count(self) -> int: ... + + def close(self) -> list[str]: ... + + +def register_host_mappings( + sources_by_mapping: Mapping[str, Sequence[torch.Tensor]], + *, + device: torch.device, + max_bytes: int, +) -> HostRegistration: + """Register mappings with the active platform or report unsupported use. + + CUDA is the first implementation. Other platforms retain DLO's bounded + host-staging path until they provide an equivalent registration backend. + """ + if device.type == "cuda": + from .cuda_host_registration import CudaHostRegistration + + return CudaHostRegistration.create(sources_by_mapping, max_bytes=max_bytes) + raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}") + + +__all__ = [ + "HostRegistration", + "HostRegistrationBudgetError", + "HostRegistrationCleanupError", + "HostRegistrationError", + "register_host_mappings", +] diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 154374852..d854d3e33 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -731,28 +731,14 @@ class OmniServeCommand(CLISubcommand): "--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-runtime-cache-pin-limit-gib", type=float, default=0.0, help=( - "Maximum GiB of read-only runtime-cache mappings each CUDA worker " - "may register for direct H2D. Zero keeps bounded pinned staging." + "Maximum GiB of read-only runtime-cache mappings each worker may " + "register for direct H2D when the platform supports it. Zero keeps " + "bounded pinned staging." ), ) omni_config_group.add_argument(