diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 542c29fa1..fa134c03c 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -44,7 +44,7 @@ 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. In no-AllGather mode, an opt-in Phase B -runtime cache can normalize those final ordinary-loader DiT tensors into an +host weight 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 host-registration budget can make the complete final mapping directly @@ -66,7 +66,7 @@ DLO is divided into four planes: configuration, loader-owned host storage, backend setup, and the per-block inference hot path. Keeping these decisions separate is intentional. In particular, an mmap cache is a host-storage choice; it does not imply a DLO AllGather, and disabling DLO AllGather does not decide -whether the host source is private memory, checkpoint mmap, or runtime-cache +whether the host source is private memory, checkpoint mmap, or host weight cache mmap. The shared control plane is described once below; the AllGather and no-AllGather startup and hot paths are then shown separately. @@ -84,14 +84,14 @@ flowchart TB Loader["DiffusersPipelineLoader"] Checkpoint["checkpoint mmap preflight"] Ordinary["ordinary load + callbacks + calibration"] - RuntimeCache["runtime_weight_cache
final-layout publication / join"] + HostWeightCache["host_weight_cache
final-layout publication / join"] Contract["HostWeightPlan
backing + bindings + transforms"] Loader --> Checkpoint Loader -->|"registered final layout requested"| Ordinary Checkpoint -->|"compatible"| Contract Checkpoint -->|"fallback"| Ordinary - Ordinary -->|"runtime cache enabled"| RuntimeCache - RuntimeCache -->|"enabled and compatible"| Contract + Ordinary -->|"host weight cache enabled"| HostWeightCache + HostWeightCache -->|"enabled and compatible"| Contract end Runner -->|"load_model"| Loader @@ -104,7 +104,7 @@ flowchart TB Discovery["ModuleDiscovery + optional OffloadPlan"] --> Backend Contract -.->|"passed once by the runner"| Backend - Backend -.->|"runtime cache + positive budget"| Registration["host_registration
platform-neutral lifecycle"] + Backend -.->|"host weight cache + positive budget"| Registration["host_registration
platform-neutral lifecycle"] Registration --> CUDA["CUDA registration backend"] Backend --> Hooks["DistributedLayerwiseOffloadHook
one circular hook per streamed block"] Forward["pipeline.forward"] --> Hooks @@ -118,8 +118,8 @@ flowchart TB | CLI, engine arguments, and `OmniDiffusionConfig` | User policy, defaults, validation, and propagation to every worker | Checkpoint inspection, cache publication, or buffer allocation | | `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, and registration budget | Host-weight compatibility decisions | | `DiffusersPipelineLoader` | Ordinary loading lifecycle and the decision to return an optional `HostWeightPlan` | DLO process groups or block-stream scheduling | -| Checkpoint mmap planner | Proving that checkpoint tensors can represent the required runtime layout before ordinary DiT materialization is skipped | Runtime-cache identity or transfer policy | -| `runtime_weight_cache` | Normalizing final ordinary-loader tensors into an immutable node-local entry and returning a validated plan | Model mutation, host registration, H2D, or collectives | +| Checkpoint mmap planner | Proving that checkpoint tensors can represent the required runtime layout before ordinary DiT materialization is skipped | Host weight cache identity or transfer policy | +| `host_weight_cache` | Normalizing final ordinary-loader tensors into an immutable node-local entry and returning a validated plan | Model mutation, host registration, H2D, or collectives | | `HostWeightPlan` | The one-way loader-to-offloader contract: backing kind, tensor bindings, source coverage, and any bounded transforms | Cache construction or runtime scheduling | | `ModuleDiscovery` and model `OffloadPlan` | Discovering DiT/components and declaring block, resident, and on-demand paths | Loader or cache behavior embedded in model pipelines | | `DistributedLayerwiseOffloadBackend` | Consuming the exact plan, selecting the existing group, mapping/repointing storage, allocating shared buffers and streams, optional host registration, installing hooks, and cleanup | Re-running loader compatibility checks or orchestrating replicas | @@ -128,14 +128,14 @@ flowchart TB | `DiffusionWorker` | Calling backend teardown before destroying the distributed environment | Managing backend-owned mappings or registrations directly | The configuration fields are routed to the component that can enforce their -contract. `dlo_enable_runtime_cache`, its cache root, and lock timeout are -consumed by the loader. `dlo_runtime_cache_pin_limit_gib` is also visible while -loading because a positive value deliberately selects the transform-complete -runtime cache; the backend later consumes the same budget for registration. -`dlo_use_allgather` controls loader plan routing because the runtime cache is -no-AllGather-only, and it selects the backend transfer protocol. This is -coordination through typed configuration and `HostWeightPlan`, not through -model-specific global state. +contract. `dlo_enable_host_weight_cache`, its cache root, and lock timeout are +consumed by the loader. `dlo_host_weight_cache_pin_limit_gib` is also visible +while loading because a positive value deliberately selects the +transform-complete host weight cache; the backend later consumes the same +budget for registration. `dlo_use_allgather` controls loader plan routing +because the host weight cache is no-AllGather-only, and it selects the backend +transfer protocol. This is coordination through typed configuration and +`HostWeightPlan`, not through model-specific global state. ### Shared startup and lifecycle @@ -162,9 +162,9 @@ sequenceDiagram ``` The loader returns a plan only after its contract is valid: checkpoint plans -have preflight-proven source coverage and metadata, while runtime-cache plans -also have full manifest, file, and mapped-content validation. The runner then -transfers ownership exactly once with +have preflight-proven source coverage and metadata, while host weight cache +plans also have full manifest, file, and mapped-content validation. The runner +then transfers ownership exactly once with `model_loader.take_host_weight_plan()`. If no plan is returned, the backend requires materialized ordinary-loader tensors. If a plan caused ordinary DiT materialization to be skipped, failure to create or enable its consumer is a @@ -206,7 +206,7 @@ uses rank-local transfer. ```mermaid sequenceDiagram participant L as DiffusersPipelineLoader - participant C as runtime_weight_cache + participant C as host_weight_cache participant B as DLO backend participant P as platform registration participant H as block hooks @@ -216,20 +216,20 @@ sequenceDiagram L-->>B: checkpoint plan via runner else ordinary loading L->>L: load, mutate, validate, and calibrate final tensors - opt runtime cache enabled + opt host weight cache enabled L->>C: publish or join final-layout entry - C-->>L: validated runtime-cache plan or fail-closed fallback + C-->>L: validated host weight cache plan or fail-closed fallback end L-->>B: plan or materialized tensors via runner end B->>B: consume complete rank-local block sources - opt runtime-cache plan and positive registration budget + opt host weight cache plan and positive registration budget B->>P: register the complete immutable mapping P-->>B: direct-H2D source or bounded-staging fallback end B->>H: install hooks and prime the first block loop each streamed block - alt registered runtime-cache mapping + alt registered host weight cache mapping H->>H: copy complete next block directly to device else pageable mmap mapping H->>H: pack into a pinned staging slot, then copy to device @@ -252,13 +252,13 @@ the source of the complete-block H2D copy. | Ordinary runtime tensors | AllGather group larger than one | Each rank retains a pinned local weight shard | shard H2D → DLO AllGather → device slot | | Checkpoint mmap | AllGather group larger than one | Each rank prepares its persistent shard, then releases the source mapping | shard H2D → DLO AllGather → device slot | | Ordinary runtime tensors | Rank-local / no-AllGather | Each rank retains the complete loader-produced local layout | complete-block H2D → device slot | -| Checkpoint or runtime-cache mmap, zero registration budget | Rank-local / no-AllGather | Immutable mapping plus two bounded pinned host-staging slots | mmap views → pinned staging slot → device slot | -| Runtime-cache mmap, positive registration budget | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot | +| Checkpoint or host weight cache mmap, zero registration budget | Rank-local / no-AllGather | Immutable mapping plus two bounded pinned host-staging slots | mmap views → pinned staging slot → device slot | +| Host weight cache mmap, positive registration budget | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot | The three axes therefore remain independent: 1. **Offload policy:** whether DLO is enabled and which layers are resident. -2. **Host backing:** ordinary tensors, checkpoint mmap, or final runtime-cache +2. **Host backing:** ordinary tensors, checkpoint mmap, or final host weight cache mmap, selected by the loader. 3. **Transfer protocol:** rank-local H2D or H2D plus DLO AllGather, selected by the backend from validated configuration and existing topology. @@ -268,8 +268,8 @@ The three axes therefore remain independent: - Invalid option combinations fail during configuration validation. - Checkpoint mmap preflight failure falls back before ordinary loading is skipped. -- Runtime-cache incompatibility, timeout, or publication failure retains the - already-valid ordinary tensors. +- Host weight cache incompatibility, timeout, or publication failure retains + the already-valid ordinary tensors. - Host-registration failure rolls back and retains bounded staging. A rollback failure aborts startup because unmapping a range still owned by the platform would be unsafe. @@ -345,7 +345,7 @@ parameter-sharding lifecycle and cannot be sharded a second time by DLO. The loader may provide ordinary runtime tensors or a compatible checkpoint plan. During backend setup, either source is normalized into persistent local -shards. The final-layout runtime cache is not selected in AllGather mode. +shards. The final-layout host weight cache is not selected in AllGather mode. With the default `dlo_use_allgather=True`, each rank stores approximately `1 / group_size` of each streamable block in pinned host memory. The next @@ -359,7 +359,7 @@ The figure uses DP2 with TP2 to expose both parallel dimensions. DLO creates one AllGather group across the DP replicas for each TP coordinate; TP0 workers reconstruct TP0 blocks together and TP1 workers reconstruct TP1 blocks together. Ordinary TP collectives remain inside each DP replica. Unlike the -no-AllGather runtime-cache path, these persistent pinned host shards are +no-AllGather host weight cache path, these persistent pinned host shards are process-owned, and every participating rank must follow the same block order. ```text @@ -398,7 +398,7 @@ domains. 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 -a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or +a loader-approved checkpoint mapping, a normalized host weight cache mapping, or ordinary runtime tensors. ![No-AllGather DLO shared host weights and independent runtime orchestration](../../figures/dlo/dlo_no_allgather_shared_host.svg) @@ -418,13 +418,13 @@ the persistent private full-model copy per pure-DP process, but each process still packs and transfers every complete block. Sharing is node-local; each node has its own page cache. -For a runtime-cache mapping, a positive registration budget replaces those +For a host weight cache mapping, a positive registration budget replaces those host slots with direct copies from the shared final tensor views. A zero budget or failed registration retains the bounded staging behavior. 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. Supported CPU layouts may then enter the opt-in runtime cache. +parameters. Supported CPU layouts may then enter the opt-in host weight cache. Unsupported layouts keep one private runtime copy per process. This mode means: @@ -436,11 +436,11 @@ 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 runtime cache can share equivalent ordinary-loader + mmap is selected. The host weight 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. -### Post-load runtime cache +### Post-load host weight 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 @@ -452,7 +452,7 @@ processing, strict validation, and calibration. The loader then: 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, +5. returns a `HostWeightPlan(backing_kind="host_weight_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 @@ -487,15 +487,16 @@ 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. -### Registered runtime-cache transfer +### Registered host weight cache transfer -`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. +`dlo_host_weight_cache_pin_limit_gib > 0` opts a worker into registration of +the complete host weight cache mapping when its platform supports it. It also +changes loader selection: no-AllGather uses the ordinary loader and +final-layout host weight 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 host weight cache 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 @@ -527,10 +528,10 @@ destroying the distributed environment and platform context. | 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; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. | -| **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. Registered transfer does not remove ordinary SP activation/attention collectives. | -| **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. Registering those final mappings avoids the recurrent staging copy but does not remove ordinary TP collectives. | -| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted only with ordinary-loader rank-local tensors. HSDP owns parameter sharding and its own gathers; runtime-cache v1 rejects HSDP/DTensor layouts. | +| **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 host weight cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. | +| **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. Host weight cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. | +| **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. Host weight cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings avoids the recurrent staging copy but does not remove ordinary TP collectives. | +| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted only with ordinary-loader rank-local tensors. HSDP owns parameter sharding and its own gathers; host weight cache v1 rejects HSDP/DTensor layouts. | ### Combined dimensions @@ -564,7 +565,7 @@ layouts are validated. ### No-AllGather constraints The no-AllGather path does not impose these DLO-specific synchronized-wave -requirements. Runtime-cache v1 is no-AllGather only and rejects all +requirements. Host weight cache v1 is no-AllGather only and rejects all quantization configurations, HSDP/DTensor, expert parallelism, CFG parallelism, PP, non-contiguous tensors, and tied/shared storage. @@ -596,8 +597,8 @@ artifacts before making hardware or performance claims. Track that work in - Use **no-AllGather** when independent replica execution is required. TP1 direct-mmap deployments can share checkpoint pages per node. Enable the - runtime cache only when independent replicas have repeated final layouts and - benchmark the target CPU/memory topology, especially for matching + host weight cache only when independent replicas have repeated final layouts + and benchmark the target CPU/memory topology, especially for matching coordinates across multiple TP engines. On CUDA, provide an explicit registration budget large enough for the complete final layout when recurrent staging cost is unacceptable; leave it at zero when page-locking @@ -605,5 +606,5 @@ artifacts before making hardware or performance claims. Track that work in ### HSDP -- Do not combine HSDP with DLO AllGather or runtime-cache v1. HSDP with +- Do not combine HSDP with DLO AllGather or host weight cache v1. HSDP with no-AllGather remains an ordinary-loader rank-local path. diff --git a/docs/design/figures/dlo/dlo_allgather_sharded_host.svg b/docs/design/figures/dlo/dlo_allgather_sharded_host.svg index 2478786ed..3214d65f5 100644 --- a/docs/design/figures/dlo/dlo_allgather_sharded_host.svg +++ b/docs/design/figures/dlo/dlo_allgather_sharded_host.svg @@ -147,7 +147,7 @@ TP stays the model-parallel group. - PROCESS-OWNED PINNED HOST SHARDS · NOT RUNTIME-CACHE MMAP + PROCESS-OWNED PINNED HOST SHARDS · NOT HOST WEIGHT CACHE MMAP DP replica 0 · private TP0 shard 0/2 · TP1 shard 0/2 diff --git a/docs/design/figures/dlo/dlo_no_allgather_shared_host.svg b/docs/design/figures/dlo/dlo_no_allgather_shared_host.svg index f7ec39116..b497a8598 100644 --- a/docs/design/figures/dlo/dlo_no_allgather_shared_host.svg +++ b/docs/design/figures/dlo/dlo_no_allgather_shared_host.svg @@ -1,6 +1,6 @@ No-AllGather DLO shared host weights and independent runtime orchestration - An external orchestrator routes requests to two independently scheduled TP2 diffusion engines on one node. Each engine contains two rank-local workers and devices, with private loaders, mmap views, DLO runtime hooks, streams, events, and rotating GPU slots. Equivalent TP0 workers across engines map one immutable TP0 runtime-cache payload, while equivalent TP1 workers map a separate TP1 payload, without a cross-engine DLO AllGather. + An external orchestrator routes requests to two independently scheduled TP2 diffusion engines on one node. Each engine contains two rank-local workers and devices, with private loaders, mmap views, DLO runtime hooks, streams, events, and rotating GPU slots. Equivalent TP0 workers across engines map one immutable TP0 entry in the host weight cache, while equivalent TP1 workers map a separate TP1 entry, without a cross-engine DLO AllGather. @@ -184,7 +184,7 @@ A/TP1 + B/TP1 → TP1 entry - NODE-LOCAL IMMUTABLE RUNTIME CACHE · ONE COPY PER TP COORDINATE + NODE-LOCAL IMMUTABLE HOST WEIGHT CACHE · ONE COPY PER TP COORDINATE TP0 physical payload shared by A/TP0 + B/TP0 diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 7d5f926ba..84733c913 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -3,7 +3,7 @@ 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. +no-AllGather path can optionally share node-local host weights through mmap. See the [shared DLO component architecture](../../../design/feature/offloader/distributed_layerwise_offload.md#architecture) for ownership and lifecycle coordination. The mode-specific internals are @@ -18,10 +18,10 @@ and | --- | --- | --- | --- | | 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 | +| no-AllGather + host weight cache | Equivalent independent workers share one node | Shared final mmap layout per TP coordinate | No DLO weight collective | +| host weight cache + host registration | The platform supports registration and recurrent staging is too expensive | Shared registered mmap layout | No DLO weight collective | -AllGather is normally the best choice for synchronized DP. Runtime-cache mode +AllGather is normally the best choice for synchronized DP. Host weight cache mode targets independently scheduled replicas, especially repeated TP engines on one node. @@ -38,13 +38,13 @@ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --dlo-no-use-allgather -# Share final runtime weights and register up to 80 GiB per worker for direct H2D +# Share final host 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-pin-limit-gib 80 + --dlo-enable-host-weight-cache \ + --dlo-host-weight-cache-pin-limit-gib 80 ``` ## Relevant options @@ -53,13 +53,14 @@ vllm serve /path/to/model --omni \ | --- | --- | --- | | `--enable-distributed-layerwise-offload` | Enable DLO | `false` | | `--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-enable-host-weight-cache` | Share final no-AllGather DiT weights through a node-local mmap cache | `false` | +| `--dlo-host-weight-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` | -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 host weight cache uses `~/.cache/vllm-omni/dlo-host-weights`. Programmatic +configuration may override `dlo_host_weight_cache_dir` and the writer lock +timeout; these advanced storage controls are intentionally not separate CLI +flags. ## Operational notes @@ -74,8 +75,8 @@ these advanced storage controls are intentionally not separate CLI flags. - 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. +- Host weight cache v1 has no automatic eviction. Stop all users of an entry + before deleting it. ## Scheduling constraints @@ -88,15 +89,15 @@ SP model collectives still synchronize ranks within each engine. ## Limitations -- Direct checkpoint mmap currently requires TP1. The runtime cache can share +- Direct checkpoint mmap currently requires TP1. The host weight 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. 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, and runtime-cache v1 rejects HSDP. -- The runtime cache changes steady-state host backing, not startup loading: +- Host weight cache v1 rejects quantized, non-contiguous, aliased/tied, + device-only, HSDP, expert-parallel, CFG-parallel, and PP layouts. +- HSDP plus DLO AllGather is unsupported, and host weight cache v1 rejects HSDP. +- The host weight cache changes steady-state host backing, not startup loading: each worker still performs ordinary loading and full content validation. - Skip-load-on-hit, cache eviction, and cross-node sharing remain follow-up work in diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 56557bc07..e1b9d8202 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -406,9 +406,9 @@ 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): +def test_dlo_host_weight_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 + import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod od_config = SimpleNamespace( dtype=torch.float32, @@ -427,9 +427,9 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, 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, + dlo_enable_host_weight_cache=True, + dlo_host_weight_cache_dir=str(tmp_path), + dlo_host_weight_cache_lock_timeout=1.0, model="unused", revision=None, model_config={}, @@ -440,7 +440,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, model.transformer.quant_method = UnquantizedLinearMethod() events: list[str] = [] runtime_plan = HostWeightPlan( - backing_kind="runtime_cache", + backing_kind="host_weight_cache", bindings={}, runtime_layout_key="layout-key", post_load_complete=True, @@ -485,16 +485,16 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, "build_checkpoint_mmap_plan", lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"), ) - monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache) + monkeypatch.setattr(cache_mod, "build_host_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_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path): +def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(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 + import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod od_config = SimpleNamespace( dtype=torch.float32, @@ -509,9 +509,9 @@ def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monke 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_pin_limit_gib=1.0, + dlo_enable_host_weight_cache=True, + dlo_host_weight_cache_dir=str(tmp_path), + dlo_host_weight_cache_pin_limit_gib=1.0, model="unused", ) loader = DiffusersPipelineLoader(LoadConfig(), od_config) @@ -519,7 +519,7 @@ def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monke model.transformer = nn.Linear(2, 2, bias=False) events: list[str] = [] runtime_plan = HostWeightPlan( - backing_kind="runtime_cache", + backing_kind="host_weight_cache", bindings={}, runtime_layout_key="layout-key", post_load_complete=True, @@ -531,23 +531,23 @@ def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monke loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] def unexpected_checkpoint_plan(*_args, **_kwargs): - pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan") + pytest.fail("registered host weight cache mode must not select the raw checkpoint mmap plan") def build_cache(_pipeline, **_kwargs): events.append("cache") return HostWeightPlanResult(runtime_plan) monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan) - monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache) + monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache) assert loader.load_model(load_device="cpu") is model assert events == ["load", "process", "cache"] assert loader.take_host_weight_plan() is runtime_plan -def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path): +def test_dlo_host_weight_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 + import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod od_config = SimpleNamespace( dtype=torch.float32, @@ -562,8 +562,8 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat 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_enable_host_weight_cache=True, + dlo_host_weight_cache_dir=str(tmp_path), model="unused", ) loader = DiffusersPipelineLoader(LoadConfig(), od_config) @@ -584,7 +584,7 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat assert kwargs["quantization_config"] == "module_quant_method" return HostWeightPlanResult(None, "writer timed out", "lock_timeout") - monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache) + monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache) original_weight = model.transformer.weight.detach().clone() assert loader.load_model(load_device="cpu") is model diff --git a/tests/diffusion/model_loader/test_runtime_weight_cache.py b/tests/diffusion/model_loader/test_host_weight_cache.py similarity index 95% rename from tests/diffusion/model_loader/test_runtime_weight_cache.py rename to tests/diffusion/model_loader/test_host_weight_cache.py index 2d0ed4339..98c6c1c06 100644 --- a/tests/diffusion/model_loader/test_runtime_weight_cache.py +++ b/tests/diffusion/model_loader/test_host_weight_cache.py @@ -10,9 +10,9 @@ 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, +from vllm_omni.diffusion.model_loader.host_weight_cache import ( + build_host_weight_cache_plan, + default_host_weight_cache_root, ) pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -42,7 +42,7 @@ class _Pipeline(nn.Module): 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") + assert default_host_weight_cache_root().endswith("/.cache/vllm-omni/dlo-host-weights") def _build(pipeline: nn.Module, cache_root: Path, **overrides): @@ -67,7 +67,7 @@ def _build(pipeline: nn.Module, cache_root: Path, **overrides): "pipeline_parallel_size": 1, } kwargs.update(overrides) - return build_runtime_weight_cache_plan(pipeline, **kwargs) + return build_host_weight_cache_plan(pipeline, **kwargs) def _multiprocess_builder(cache_root: str, queue) -> None: @@ -121,7 +121,7 @@ def test_unstable_loader_identity_fails_closed(tmp_path): assert result.plan is None assert result.fallback_code == "unstable_identity" - assert result.fallback_reason == "runtime-cache identity does not support values of type builtins.object" + assert result.fallback_reason == "host weight cache identity does not support values of type builtins.object" assert list(tmp_path.rglob("*.safetensors")) == [] diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index cf985ffa8..5aa13dc3c 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -808,7 +808,7 @@ class _MmapPostLoadPipeline(nn.Module): class TestMmapWeightLoading: - def test_runtime_cache_maps_final_weights_without_rerunning_post_load( + def test_host_weight_cache_maps_final_weights_without_rerunning_post_load( self, tmp_path, patched_offload_runtime, @@ -833,7 +833,7 @@ class TestMmapWeightLoading: 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", + backing_kind="host_weight_cache", bindings={ "transformer.proj.weight": TensorBinding( storage_key="transformer.proj.weight", @@ -854,7 +854,7 @@ class TestMmapWeightLoading: ) modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) - backend._load_weights_from_runtime_cache(pipeline, modules, plan) + backend._load_weights_from_host_weight_cache(pipeline, modules, plan) assert pipeline.transformer.post_load_calls == 0 assert pipeline.transformer.proj.weight is old_parameter @@ -862,7 +862,7 @@ class TestMmapWeightLoading: 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( + def test_host_weight_cache_mapping_failure_leaves_ordinary_tensors_intact( self, tmp_path, patched_offload_runtime, @@ -871,7 +871,7 @@ class TestMmapWeightLoading: pipeline.transformer = nn.Linear(2, 2, bias=False) old_parameter = pipeline.transformer.weight plan = HostWeightPlan( - backing_kind="runtime_cache", + backing_kind="host_weight_cache", bindings={ "transformer.weight": TensorBinding( storage_key="transformer.weight", @@ -893,7 +893,7 @@ class TestMmapWeightLoading: modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) with pytest.raises(Exception, match="missing.safetensors"): - backend._load_weights_from_runtime_cache(pipeline, modules, plan) + backend._load_weights_from_host_weight_cache(pipeline, modules, plan) assert pipeline.transformer.weight is old_parameter @@ -991,7 +991,7 @@ 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( + def test_host_weight_cache_enable_reclaims_private_weights_before_hook_setup( self, tmp_path, patched_offload_runtime, @@ -1013,7 +1013,7 @@ class TestMmapWeightLoading: weight_file = tmp_path / "runtime.safetensors" save_file(weights, str(weight_file)) plan = HostWeightPlan( - backing_kind="runtime_cache", + backing_kind="host_weight_cache", bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights}, runtime_layout_key="runtime-layout", post_load_complete=True, @@ -1038,7 +1038,7 @@ class TestMmapWeightLoading: assert len(cleanup_calls) == 2 backend.disable() - def test_runtime_cache_registration_uses_budget_and_releases_before_mmap( + def test_host_weight_cache_registration_uses_budget_and_releases_before_mmap( self, monkeypatch: pytest.MonkeyPatch, patched_offload_runtime, @@ -1067,14 +1067,14 @@ class TestMmapWeightLoading: strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=True, dlo_use_allgather=False, - dlo_runtime_cache_pin_limit_gib=1.5, + dlo_host_weight_cache_pin_limit_gib=1.5, ), torch.device("cuda"), ) sources = {"runtime.safetensors": [torch.ones(1)]} - backend._runtime_cache_mapped_sources = sources + backend._host_weight_cache_mapped_sources = sources - assert backend._try_register_runtime_cache_mmap() + assert backend._try_register_host_weight_cache_mmap() assert create_calls == [(sources, int(1.5 * 1024**3))] backend.enabled = True @@ -1085,7 +1085,7 @@ class TestMmapWeightLoading: assert events == ["unregister", "mmap"] - def test_runtime_cache_registration_falls_back_on_unsupported_platform( + def test_host_weight_cache_registration_falls_back_on_unsupported_platform( self, patched_offload_runtime, ): @@ -1094,13 +1094,13 @@ class TestMmapWeightLoading: strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=True, dlo_use_allgather=False, - dlo_runtime_cache_pin_limit_gib=1.0, + dlo_host_weight_cache_pin_limit_gib=1.0, ), torch.device("cpu"), ) - backend._runtime_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]} + backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]} - assert not backend._try_register_runtime_cache_mmap() + assert not backend._try_register_host_weight_cache_mmap() assert backend._host_registration is None diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 95cb37770..a99d13344 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -328,8 +328,8 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--omni", "--enable-distributed-layerwise-offload", "--dlo-no-use-allgather", - "--dlo-enable-runtime-cache", - "--dlo-runtime-cache-pin-limit-gib", + "--dlo-enable-host-weight-cache", + "--dlo-host-weight-cache-pin-limit-gib", "80", "--dlo-resident-layers", "20", @@ -342,15 +342,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_pin_limit_gib == 80 + assert args.dlo_enable_host_weight_cache is True + assert args.dlo_host_weight_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"] 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_enable_host_weight_cache"] is True + assert engine_args["dlo_host_weight_cache_dir"] is None + assert engine_args["dlo_host_weight_cache_lock_timeout"] == 600.0 + assert engine_args["dlo_host_weight_cache_pin_limit_gib"] == 80 assert engine_args["dlo_resident_layers"] == 20 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index f7327c5b9..3641f6d59 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -687,10 +687,10 @@ 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_runtime_cache_pin_limit_gib: float = Field(default=0.0, ge=0) + dlo_enable_host_weight_cache: bool = False + dlo_host_weight_cache_dir: str | None = None + dlo_host_weight_cache_lock_timeout: float = Field(default=600.0, gt=0) + dlo_host_weight_cache_pin_limit_gib: float = Field(default=0.0, ge=0) dlo_resident_layers: int = Field(default=0, ge=0) pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" @@ -753,12 +753,12 @@ 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.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache: - raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True") + if self.dlo_enable_host_weight_cache and not self.enable_distributed_layerwise_offload: + raise ValueError("dlo_enable_host_weight_cache requires distributed layerwise offload") + if self.dlo_enable_host_weight_cache and self.dlo_use_allgather: + raise ValueError("dlo_enable_host_weight_cache currently requires dlo_use_allgather=False") + if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_enable_host_weight_cache: + raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_enable_host_weight_cache=True") if self.tf_model_config is None: self.tf_model_config = TransformerConfig() diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 9148198e6..c97e1c23e 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -382,10 +382,10 @@ 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_runtime_cache_pin_limit_gib: float | None = None + dlo_enable_host_weight_cache: bool | None = None + dlo_host_weight_cache_dir: str | None = None + dlo_host_weight_cache_lock_timeout: float | None = None + dlo_host_weight_cache_pin_limit_gib: 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 cb6dda5d6..d86f42139 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -712,14 +712,14 @@ class OmniDiffusionConfig: # used when direct checkpoint mmap is unavailable, or when direct CUDA # registration requires a transform-complete final layout. 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 + dlo_enable_host_weight_cache: bool = False + # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-host-weights. + dlo_host_weight_cache_dir: str | None = None # Maximum wait for the per-layout POSIX writer lock. - dlo_runtime_cache_lock_timeout: float = 600.0 + dlo_host_weight_cache_lock_timeout: float = 600.0 # Maximum file-backed memory a worker may register with CUDA. Zero keeps # the bounded mmap-to-pinned staging path. - dlo_runtime_cache_pin_limit_gib: float = 0.0 + dlo_host_weight_cache_pin_limit_gib: float = 0.0 # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 @@ -948,16 +948,16 @@ 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 not math.isfinite(self.dlo_runtime_cache_pin_limit_gib) or self.dlo_runtime_cache_pin_limit_gib < 0: - raise ValueError("dlo_runtime_cache_pin_limit_gib must be non-negative") - 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.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache: - raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True") + if not math.isfinite(self.dlo_host_weight_cache_lock_timeout) or self.dlo_host_weight_cache_lock_timeout <= 0: + raise ValueError("dlo_host_weight_cache_lock_timeout must be positive") + if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib) or self.dlo_host_weight_cache_pin_limit_gib < 0: + raise ValueError("dlo_host_weight_cache_pin_limit_gib must be non-negative") + if self.dlo_enable_host_weight_cache and not self.enable_distributed_layerwise_offload: + raise ValueError("dlo_enable_host_weight_cache requires distributed layerwise offload") + if self.dlo_enable_host_weight_cache and self.dlo_use_allgather: + raise ValueError("dlo_enable_host_weight_cache currently requires dlo_use_allgather=False") + if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_enable_host_weight_cache: + raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_enable_host_weight_cache=True") 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 e0c5149fb..ec0814a74 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -441,16 +441,16 @@ class DiffusersPipelineLoader: _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1)) _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1)) _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size - _prefer_registered_runtime_cache = ( + _prefer_registered_host_weight_cache = ( _dist_offload and not _use_ag - and bool(getattr(self.od_config, "dlo_enable_runtime_cache", False)) - and float(getattr(self.od_config, "dlo_runtime_cache_pin_limit_gib", 0.0)) > 0 + and bool(getattr(self.od_config, "dlo_enable_host_weight_cache", False)) + and float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0 ) plan_result = None weight_sources = self._get_weight_sources(model) - if _dist_offload and not _prefer_registered_runtime_cache: + if _dist_offload and not _prefer_registered_host_weight_cache: modules = ModuleDiscovery.discover(model) plan_result = build_checkpoint_mmap_plan( model, @@ -462,14 +462,14 @@ class DiffusersPipelineLoader: online_quantization=_has_online_quant, ) self.host_weight_plan = plan_result.plan - elif _prefer_registered_runtime_cache: + elif _prefer_registered_host_weight_cache: # Direct checkpoint bindings can carry deferred transforms # (for example grouped-QKV reordering), so registering the # raw checkpoint mapping would not make every recurrent H2D # source directly copyable. Materialize once, apply every # post-load mutation, then publish/register the final layout. logger.info( - "DLO runtime-cache registration requested; using the ordinary loader " + "DLO host weight cache registration requested; using the ordinary loader " "once to publish a transform-complete final mmap layout" ) @@ -540,19 +540,19 @@ class DiffusersPipelineLoader: self._apply_skip_softmax_calibration(model) model.eval() - self._maybe_build_runtime_weight_cache(model, load_format) + self._maybe_build_host_weight_cache(model, load_format) return model - def _maybe_build_runtime_weight_cache(self, model: nn.Module, load_format: str) -> None: + def _maybe_build_host_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): + if not getattr(self.od_config, "dlo_enable_host_weight_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") + logger.info("DLO host weight cache is no-AllGather only; leaving AllGather storage unchanged") return parallel_config = self.parallel_config @@ -565,17 +565,18 @@ class DiffusersPipelineLoader: 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", + "DLO host weight 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 ( + from vllm_omni.diffusion.model_loader.host_weight_cache import ( DEFAULT_LOCK_TIMEOUT_SECONDS, DEFAULT_SHARD_SIZE_BYTES, - build_runtime_weight_cache_plan, + build_host_weight_cache_plan, ) loader_inputs = { @@ -602,13 +603,13 @@ class DiffusersPipelineLoader: for module in model.modules() ): runtime_quantization = "module_quant_method" - result = build_runtime_weight_cache_plan( + result = build_host_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), + cache_root=getattr(self.od_config, "dlo_host_weight_cache_dir", None), lock_timeout_seconds=float( - getattr(self.od_config, "dlo_runtime_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS) + getattr(self.od_config, "dlo_host_weight_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS) ), max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES, model_identity=str(model_identity) if model_identity is not None else None, @@ -628,7 +629,7 @@ class DiffusersPipelineLoader: 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)", + "DLO host weight 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, @@ -636,7 +637,7 @@ class DiffusersPipelineLoader: ) else: logger.warning( - "DLO runtime cache unavailable [%s]; retaining ordinary host tensors: %s", + "DLO host weight cache unavailable [%s]; retaining ordinary host tensors: %s", result.fallback_code or "unknown", result.fallback_reason or "unspecified reason", ) diff --git a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py b/vllm_omni/diffusion/model_loader/host_weight_cache.py similarity index 79% rename from vllm_omni/diffusion/model_loader/runtime_weight_cache.py rename to vllm_omni/diffusion/model_loader/host_weight_cache.py index bde659f4d..cf1b0f2d3 100644 --- a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py +++ b/vllm_omni/diffusion/model_loader/host_weight_cache.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Build node-local runtime-weight cache plans for diffusion DLO. +"""Build node-local host weight cache plans for diffusion DLO. -``build_runtime_weight_cache_plan`` is the module's single entry point. It +``build_host_weight_cache_plan`` is the module's single entry point. It validates final ordinary-loader DiT tensors, joins or atomically publishes an immutable content-addressed entry, validates the mmap result, and returns a fail-closed host-weight plan. Reusable identity, hashing, locking, and file -primitives live in :mod:`runtime_cache_utils`. +primitives live in :mod:`host_weight_cache_utils`. """ from __future__ import annotations @@ -30,14 +30,9 @@ 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, -) -from vllm_omni.diffusion.model_loader.runtime_cache_utils import ( +from vllm_omni.diffusion.model_loader.host_weight_cache_utils import ( + HostWeightCacheError, IdentityNormalizationError, - RuntimeCacheError, RuntimeTensor, canonical_json, canonicalize_existing_local_path, @@ -52,15 +47,20 @@ from vllm_omni.diffusion.model_loader.runtime_cache_utils import ( split_shards, tensor_metadata, ) +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + HostWeightPlanResult, + TensorBinding, +) 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 +HOST_WEIGHT_CACHE_SCHEMA_VERSION = 1 # Keep the default shared across local stage replicas. -DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights" +DEFAULT_HOST_WEIGHT_CACHE_ROOT = "~/.cache/vllm-omni/dlo-host-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. @@ -83,9 +83,9 @@ for _dtype_name in ("float8_e4m3fn", "float8_e5m2"): _SUPPORTED_DTYPES.add(_dtype) -def default_runtime_cache_root() -> str: +def default_host_weight_cache_root() -> str: """Return the default root, intentionally independent of VLLM_CACHE_ROOT.""" - return os.path.expanduser(DEFAULT_RUNTIME_CACHE_ROOT) + return os.path.expanduser(DEFAULT_HOST_WEIGHT_CACHE_ROOT) def _persistent_buffer(module: nn.Module, local_name: str) -> bool: @@ -135,7 +135,7 @@ def _collect_runtime_tensors( candidate = RuntimeTensor(runtime_name, tensor, "parameter") existing = records.get(runtime_name) if existing is not None and existing.tensor is not tensor: - raise RuntimeCacheError( + raise HostWeightCacheError( "ambiguous_ownership", f"multiple DiT tensors resolve to {runtime_name!r}", ) @@ -147,14 +147,14 @@ def _collect_runtime_tensors( candidate = RuntimeTensor(runtime_name, tensor, "buffer") existing = records.get(runtime_name) if existing is not None and existing.tensor is not tensor: - raise RuntimeCacheError( + raise HostWeightCacheError( "ambiguous_ownership", f"multiple DiT tensors resolve to {runtime_name!r}", ) records[runtime_name] = candidate if not records: - raise RuntimeCacheError( + raise HostWeightCacheError( "ambiguous_ownership", "no DiT parameters or persistent buffers were discovered", ) @@ -165,37 +165,37 @@ def _collect_runtime_tensors( try: pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name) except AttributeError as exc: - raise RuntimeCacheError( + raise HostWeightCacheError( "ambiguous_ownership", f"DiT tensor {record.name!r} is not owned by the pipeline", ) from exc if pipeline_tensor is not tensor: - raise RuntimeCacheError( + raise HostWeightCacheError( "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( + raise HostWeightCacheError( "unsupported_tensor", f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}", ) if tensor.layout != torch.strided: - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_tensor", f"{record.name!r} uses unsupported layout {tensor.layout}", ) if not tensor.is_contiguous(): - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_tensor", f"{record.name!r} is non-contiguous with stride {tensor.stride()}", ) if tensor.dtype not in _SUPPORTED_DTYPES: - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_tensor", f"{record.name!r} uses unsupported dtype {tensor.dtype}", ) if hasattr(tensor, "to_local"): - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_tensor", f"{record.name!r} is a distributed tensor", ) @@ -203,7 +203,7 @@ def _collect_runtime_tensors( tensor_nbytes = tensor.numel() * tensor.element_size() storage = tensor.untyped_storage() if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes: - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_tensor", f"{record.name!r} is a view into a larger storage", ) @@ -211,7 +211,7 @@ def _collect_runtime_tensors( storage_id = (storage.data_ptr(), storage.nbytes()) owner = storage_owners.setdefault(storage_id, record.name) if owner != record.name: - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_alias", f"{record.name!r} shares storage with {owner!r}", ) @@ -228,7 +228,7 @@ def _collect_runtime_tensors( storage = tensor.untyped_storage() owner = storage_owners.get((storage.data_ptr(), storage.nbytes())) if owner is not None and owner != pipeline_name: - raise RuntimeCacheError( + raise HostWeightCacheError( "unsupported_alias", f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}", ) @@ -250,9 +250,9 @@ def _publish_entry( 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( + raise HostWeightCacheError( "insufficient_disk", - f"runtime cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}", + f"host weight cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}", ) tmp_root = cache_root / ".tmp" @@ -270,7 +270,7 @@ def _publish_entry( 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)}, + metadata={"format": "pt", "dlo_host_weight_cache_schema": str(HOST_WEIGHT_CACHE_SCHEMA_VERSION)}, ) fsync_file(path) file_manifest[filename] = { @@ -285,7 +285,7 @@ def _publish_entry( } manifest = { - "schema_version": RUNTIME_CACHE_SCHEMA_VERSION, + "schema_version": HOST_WEIGHT_CACHE_SCHEMA_VERSION, "cache_key": cache_key, "layout_identity": layout_identity, "runtime_content_sha256": content_digest, @@ -328,35 +328,35 @@ def _validate_entry( 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 + raise HostWeightCacheError("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("schema_version") != HOST_WEIGHT_CACHE_SCHEMA_VERSION: + raise HostWeightCacheError("entry_rejected", "host weight cache schema mismatch") if manifest.get("cache_key") != cache_key: - raise RuntimeCacheError("entry_rejected", "runtime-cache key mismatch") + raise HostWeightCacheError("entry_rejected", "host weight cache key mismatch") if manifest.get("layout_identity") != layout_identity: - raise RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch") + raise HostWeightCacheError("entry_rejected", "host weight cache layout identity mismatch") if manifest.get("runtime_content_sha256") != content_digest: - raise RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch") + raise HostWeightCacheError("entry_rejected", "host weight 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") + raise HostWeightCacheError("entry_rejected", "host weight 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") + raise HostWeightCacheError("entry_rejected", "host weight 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}") + raise HostWeightCacheError("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 + raise HostWeightCacheError("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}") + raise HostWeightCacheError("entry_rejected", f"content digest mismatch for cache shard {path}") bindings: dict[str, TensorBinding] = {} mapped_records: list[RuntimeTensor] = [] @@ -369,43 +369,43 @@ def _validate_entry( 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}") + raise HostWeightCacheError("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}") + raise HostWeightCacheError("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}") + raise HostWeightCacheError("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( + raise HostWeightCacheError( "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}") + raise HostWeightCacheError("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}") + raise HostWeightCacheError("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") + raise HostWeightCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher") return HostWeightPlan( - backing_kind="runtime_cache", + backing_kind="host_weight_cache", bindings=bindings, runtime_layout_key=cache_key, post_load_complete=True, ) -def build_runtime_weight_cache_plan( +def build_host_weight_cache_plan( pipeline: nn.Module, *, dit_modules: Sequence[tuple[str, nn.Module]], @@ -427,7 +427,7 @@ def build_runtime_weight_cache_plan( cfg_parallel_size: int, pipeline_parallel_size: int, ) -> HostWeightPlanResult: - """Build the runtime-cache host-weight plan after ordinary model loading. + """Build a host-weight plan from the host weight cache after ordinary loading. This is the module's main entry point. It first rejects unsupported runtime layouts, then discovers and validates the final CPU DiT tensors and derives @@ -448,33 +448,37 @@ def build_runtime_weight_cache_plan( "unsupported_load_format", ) if use_hsdp: - return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp") + return HostWeightPlanResult( + None, "HSDP/DTensor layouts are not host weight cache compatible", "unsupported_hsdp" + ) if enable_expert_parallel: return HostWeightPlanResult( None, - "expert-parallel weight ownership is outside runtime-cache v1", + "expert-parallel weight ownership is outside host weight cache v1", "unsupported_expert_parallel", ) if quantization_config is not None: return HostWeightPlanResult( None, - "quantized runtime layouts are not supported by the first runtime-cache version", + "quantized runtime layouts are not supported by the first host weight 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") + return HostWeightPlanResult( + None, "pipeline-parallel ownership is outside host weight 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") + return HostWeightPlanResult(None, "host weight 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") + return HostWeightPlanResult(None, "host weight 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, + "schema_version": HOST_WEIGHT_CACHE_SCHEMA_VERSION, "model": (canonicalize_existing_local_path(model_identity) if model_identity is not None else None), "revision": revision, "runtime_dtype": str(runtime_dtype), @@ -488,8 +492,8 @@ def build_runtime_weight_cache_plan( "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}" + root = Path(cache_root or default_host_weight_cache_root()).expanduser().resolve() + entries_root = root / f"v{HOST_WEIGHT_CACHE_SCHEMA_VERSION}" entry_dir = entries_root / cache_key entries_root.mkdir(parents=True, exist_ok=True) @@ -502,10 +506,10 @@ def build_runtime_weight_cache_plan( content_digest=content_digest, expected_records=records, ) - logger.info("Reusing DLO runtime cache entry %s from %s", cache_key[:12], entry_dir) + logger.info("Reusing DLO host weight 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) + except HostWeightCacheError as exc: + logger.warning("DLO host weight 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): @@ -519,9 +523,9 @@ def build_runtime_weight_cache_plan( content_digest=content_digest, expected_records=records, ) - logger.info("Joined DLO runtime cache entry %s from %s", cache_key[:12], entry_dir) + logger.info("Joined DLO host weight cache entry %s from %s", cache_key[:12], entry_dir) return HostWeightPlanResult(plan) - except RuntimeCacheError: + except HostWeightCacheError: pass _publish_entry( @@ -541,25 +545,25 @@ def build_runtime_weight_cache_plan( expected_records=records, ) logger.info( - "Published DLO runtime cache entry %s with %d tensors at %s", + "Published DLO host weight cache entry %s with %d tensors at %s", cache_key[:12], len(records), entry_dir, ) return HostWeightPlanResult(plan) - except RuntimeCacheError as exc: + except HostWeightCacheError 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") + return HostWeightPlanResult(None, f"host weight cache operation failed: {exc}", "cache_operation_failed") __all__ = [ "DEFAULT_LOCK_TIMEOUT_SECONDS", - "DEFAULT_RUNTIME_CACHE_ROOT", + "DEFAULT_HOST_WEIGHT_CACHE_ROOT", "DEFAULT_SHARD_SIZE_BYTES", - "RUNTIME_CACHE_SCHEMA_VERSION", - "build_runtime_weight_cache_plan", - "default_runtime_cache_root", + "HOST_WEIGHT_CACHE_SCHEMA_VERSION", + "build_host_weight_cache_plan", + "default_host_weight_cache_root", ] diff --git a/vllm_omni/diffusion/model_loader/runtime_cache_utils.py b/vllm_omni/diffusion/model_loader/host_weight_cache_utils.py similarity index 93% rename from vllm_omni/diffusion/model_loader/runtime_cache_utils.py rename to vllm_omni/diffusion/model_loader/host_weight_cache_utils.py index 6859e65e2..5695a6fbb 100644 --- a/vllm_omni/diffusion/model_loader/runtime_cache_utils.py +++ b/vllm_omni/diffusion/model_loader/host_weight_cache_utils.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Reusable identity, hashing, and file helpers for the DLO runtime cache.""" +"""Reusable identity, hashing, and file helpers for the DLO host weight cache.""" from __future__ import annotations @@ -43,7 +43,7 @@ class RuntimeTensor: kind: str -class RuntimeCacheError(RuntimeError): +class HostWeightCacheError(RuntimeError): """Expected cache failure that should fall back with a stable code.""" def __init__(self, code: str, message: str): @@ -102,7 +102,7 @@ def normalize_identity(value: object) -> JsonValue: normalized: dict[str, JsonValue] = {} for key, item in value.items(): if not isinstance(key, str): - raise IdentityNormalizationError("runtime-cache identity mappings require string keys") + raise IdentityNormalizationError("host weight 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)): @@ -110,7 +110,9 @@ def normalize_identity(value: object) -> JsonValue: 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)}") + raise IdentityNormalizationError( + f"host weight cache identity does not support values of type {_type_identity(value)}" + ) def canonical_json(value: object) -> bytes: @@ -203,12 +205,12 @@ def exclusive_lock(path: Path, timeout_seconds: float) -> Iterator[None]: 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 + raise HostWeightCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc remaining = deadline - time.monotonic() if remaining <= 0: - raise RuntimeCacheError( + raise HostWeightCacheError( "lock_timeout", - f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}", + f"timed out after {timeout_seconds:g}s waiting for host weight cache writer {path.name}", ) from exc time.sleep(min(0.2, remaining)) try: @@ -268,7 +270,7 @@ def remove_stale_temps(cache_root: Path, cache_key: str) -> None: __all__ = [ "IdentityNormalizationError", "JsonValue", - "RuntimeCacheError", + "HostWeightCacheError", "RuntimeTensor", "canonical_json", "canonicalize_existing_local_path", diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 73546a721..bdeded55d 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -72,7 +72,7 @@ def get_offload_backend( od_config: OmniDiffusionConfig with offload settings device: Target device (auto-detected if None) host_weight_plan: Exact loader-produced backing plan. Checkpoint plans - may have skipped ordinary materialization; runtime-cache plans + may have skipped ordinary materialization; host weight cache plans describe already-valid ordinary tensors that can be remapped. Returns: @@ -92,7 +92,7 @@ def get_offload_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") + logger.warning("Ignoring optional host weight cache plan because distributed layerwise offload is not selected") host_weight_plan = None # Return None if no offloading requested diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 38e4c3255..433909afc 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -31,9 +31,9 @@ class OffloadConfig: # blocks from the loader-selected host backing with H2D only. dlo_use_allgather: bool = True dlo_resident_layers: int = 0 # leading DiT layers kept on device - # Per-worker budget for registering shared runtime-cache mappings with CUDA. + # Per-worker budget for registering shared host weight cache mappings with CUDA. # Zero retains bounded host staging. - dlo_runtime_cache_pin_limit_gib: float = 0.0 + dlo_host_weight_cache_pin_limit_gib: float = 0.0 @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": @@ -104,7 +104,7 @@ class OffloadConfig: # requirements (concurrent requests, dummy run skip). dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0)) - dlo_runtime_cache_pin_limit_gib = float(getattr(od_config, "dlo_runtime_cache_pin_limit_gib", 0.0)) + dlo_host_weight_cache_pin_limit_gib = float(getattr(od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) if dlo_resident_layers < 0: raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}") if dlo_resident_layers and dlo_use_allgather: @@ -140,7 +140,7 @@ class OffloadConfig: dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, dlo_resident_layers=dlo_resident_layers, - dlo_runtime_cache_pin_limit_gib=dlo_runtime_cache_pin_limit_gib, + dlo_host_weight_cache_pin_limit_gib=dlo_host_weight_cache_pin_limit_gib, ) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index db9f29591..0d4223bf8 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -966,7 +966,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._using_registered_mmap = False 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._host_weight_cache_mapped_sources: dict[str, list[torch.Tensor]] = {} self._host_registration: HostRegistration | None = None def load_resident_layers(self) -> None: @@ -1168,7 +1168,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): f"(first 5: {remaining_meta[:5]})." ) - def _load_weights_from_runtime_cache( + def _load_weights_from_host_weight_cache( self, pipeline: nn.Module, modules, @@ -1178,9 +1178,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): from safetensors import safe_open if not plan.post_load_complete: - raise RuntimeError("A runtime-cache plan must describe fully post-processed weights") + raise RuntimeError("A host weight 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") + raise RuntimeError("Host weight cache bindings cannot carry deferred checkpoint transforms") required_names: set[str] = set() for dit_name, dit_module in zip(modules.dit_names, modules.dits): @@ -1194,7 +1194,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): missing = sorted(required_names - set(plan.bindings)) unexpected = sorted(set(plan.bindings) - required_names) raise RuntimeError( - "Runtime-cache plan coverage changed after loader validation: " + "Host weight cache plan coverage changed after loader validation: " f"missing={missing[:5]}, unexpected={unexpected[:5]}" ) @@ -1209,16 +1209,16 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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 + raise RuntimeError(f"Host weight 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") + raise RuntimeError(f"Host weight 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}" + f"Host weight cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}" ) if binding.file_path not in file_cache: @@ -1230,7 +1230,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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"Host weight cache metadata changed for {runtime_name!r}: " f"cache={tuple(tensor.shape)}/{tensor.dtype}, " f"runtime={tuple(target.shape)}/{target.dtype}" ) @@ -1248,48 +1248,50 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._mmap_transforms_by_tensor_id.clear() self._mmap_file_cache = file_cache - self._runtime_cache_mapped_sources = mapped_sources + self._host_weight_cache_mapped_sources = mapped_sources logger.info( - "Remapped %d final DiT tensors from runtime cache %s", + "Remapped %d final DiT tensors from host weight cache %s", len(replacements), plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown", ) - def _try_register_runtime_cache_mmap(self) -> bool: + def _try_register_host_weight_cache_mmap(self) -> bool: """Register complete cache mappings when the operator budget permits.""" - pin_limit_gib = self.config.dlo_runtime_cache_pin_limit_gib + pin_limit_gib = self.config.dlo_host_weight_cache_pin_limit_gib if pin_limit_gib <= 0: return False if not self.config.pin_cpu_memory: - logger.warning("DLO runtime-cache registration requires pin_cpu_memory=True; using bounded host staging") + logger.warning( + "DLO host weight cache registration requires pin_cpu_memory=True; using bounded host staging" + ) return False - if not self._runtime_cache_mapped_sources: - logger.warning("DLO runtime-cache registration found no mapped sources; using bounded host staging") + if not self._host_weight_cache_mapped_sources: + logger.warning("DLO host weight cache registration found no mapped sources; using bounded host staging") return False max_bytes = int(pin_limit_gib * 1024**3) started = time.perf_counter() try: registration = register_host_mappings( - self._runtime_cache_mapped_sources, + self._host_weight_cache_mapped_sources, device=self.device, max_bytes=max_bytes, ) 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") + logger.exception("DLO host weight cache registration rollback failed") raise except HostRegistrationError as exc: logger.warning( - "DLO runtime-cache direct H2D unavailable (%s); using bounded host staging", + "DLO host weight cache direct H2D unavailable (%s); using bounded host staging", exc, ) return False self._host_registration = registration logger.info( - "Registered %.2f GiB of shared runtime-cache mappings in %d range(s) for direct H2D in %.3f s", + "Registered %.2f GiB of shared host weight cache mappings in %d range(s) for direct H2D in %.3f s", registration.total_bytes / 1024**3, registration.region_count, time.perf_counter() - started, @@ -1303,12 +1305,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): errors = registration.close() if errors: logger.error( - "DLO runtime-cache host unregistration failed; retaining mmap handles for retry/process exit: %s", + "DLO host weight cache unregistration failed; retaining mmap handles for retry/process exit: %s", errors[:3], ) return False self._host_registration = None - logger.info("Unregistered DLO runtime-cache host mappings") + logger.info("Unregistered DLO host weight cache mappings") return True def _init_dp_group(self) -> None: @@ -1657,11 +1659,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): modules, self.host_weight_plan, ) - elif self.host_weight_plan.backing_kind == "runtime_cache": + elif self.host_weight_plan.backing_kind == "host_weight_cache": if self.config.dlo_use_allgather: - raise ValueError("Runtime-cache backing is no-AllGather only in the first implementation") + raise ValueError("Host weight cache backing is no-AllGather only in the first implementation") self._using_rank_local_mmap = True - self._load_weights_from_runtime_cache( + self._load_weights_from_host_weight_cache( pipeline, modules, self.host_weight_plan, @@ -1847,9 +1849,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if not self._all_hook_groups: self.enabled = bool(self._resident_blocks) - if self.enabled and self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache": - self._using_registered_mmap = self._try_register_runtime_cache_mmap() - self._runtime_cache_mapped_sources.clear() + if ( + self.enabled + and self._using_rank_local_mmap + and self.host_weight_plan.backing_kind == "host_weight_cache" + ): + self._using_registered_mmap = self._try_register_host_weight_cache_mmap() + self._host_weight_cache_mapped_sources.clear() assert self._resident_layer_group is not None self._resident_layer_group.registered_mmap = self._using_registered_mmap if self._using_registered_mmap: @@ -1870,9 +1876,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self.dp_size > 1: unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) - if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache": - self._using_registered_mmap = self._try_register_runtime_cache_mmap() - self._runtime_cache_mapped_sources.clear() + if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "host_weight_cache": + self._using_registered_mmap = self._try_register_host_weight_cache_mmap() + self._host_weight_cache_mapped_sources.clear() for hook in all_hooks: hook.registered_mmap = self._using_registered_mmap if self._resident_layer_group is not None: @@ -1962,7 +1968,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def _release_mmap_handles(self) -> None: """Release safetensors mmap file handles.""" self._mmap_transforms_by_tensor_id.clear() - self._runtime_cache_mapped_sources.clear() + self._host_weight_cache_mapped_sources.clear() if hasattr(self, "_mmap_file_cache"): self._mmap_file_cache.clear() del self._mmap_file_cache diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 8c4c7fc43..fe4b1926a 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -545,10 +545,10 @@ 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_runtime_cache_pin_limit_gib: float = 0.0 + dlo_enable_host_weight_cache: bool = False + dlo_host_weight_cache_dir: str | None = None + dlo_host_weight_cache_lock_timeout: float = 600.0 + dlo_host_weight_cache_pin_limit_gib: float = 0.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 940c45abc..d5465f3c0 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1040,10 +1040,10 @@ 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_runtime_cache_pin_limit_gib": kwargs.get("dlo_runtime_cache_pin_limit_gib", 0.0), + "dlo_enable_host_weight_cache": kwargs.get("dlo_enable_host_weight_cache", False), + "dlo_host_weight_cache_dir": kwargs.get("dlo_host_weight_cache_dir", None), + "dlo_host_weight_cache_lock_timeout": kwargs.get("dlo_host_weight_cache_lock_timeout", 600.0), + "dlo_host_weight_cache_pin_limit_gib": kwargs.get("dlo_host_weight_cache_pin_limit_gib", 0.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 d854d3e33..88aca8f6f 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -723,7 +723,7 @@ class OmniServeCommand(CLISubcommand): ), ) omni_config_group.add_argument( - "--dlo-enable-runtime-cache", + "--dlo-enable-host-weight-cache", action="store_true", help=( "After ordinary loading, publish or join a node-local mmap cache " @@ -732,11 +732,11 @@ class OmniServeCommand(CLISubcommand): ), ) omni_config_group.add_argument( - "--dlo-runtime-cache-pin-limit-gib", + "--dlo-host-weight-cache-pin-limit-gib", type=float, default=0.0, help=( - "Maximum GiB of read-only runtime-cache mappings each worker may " + "Maximum GiB of read-only host weight cache mappings each worker may " "register for direct H2D when the platform supports it. Zero keeps " "bounded pinned staging." ),