Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 58eebee14b [Diffusion] Simplify DLO host plan setup
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 23:00:31 +08:00
Hongsheng Liu 6bcd350301 [Diffusion] Prefer host-cache registration when pinned
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 18:13:44 +08:00
Hongsheng Liu 859003fb97 [Diffusion] Gate CUDA read-only host registration
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu 42c6ae9007 [Diffusion] Document TP host layout identity
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu 16642dc2e8 [Diffusion] Make DLO host weight cache opt-in
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu 75af33a76f [Diffusion] Address DLO host cache review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu 44d3245a6d [Diffusion] Clear handled CUDA registration errors
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu 010b3c1d74 [Diffusion] Make no-AllGather host weight cache automatic
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu 3bb7547a0f [Diffusion] Rename runtime cache to host weight cache
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:36 +08:00
Hongsheng Liu ee665fcfc3 [Docs] Add AllGather DLO architecture diagram
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:16 +08:00
Hongsheng Liu 649558e1a1 [Docs] Show TP2 devices in DLO architecture
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:16 +08:00
Hongsheng Liu 3b9727c1e1 [Docs] Reorder no-AllGather architecture diagram
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:16 +08:00
Hongsheng Liu 47153d86e6 [Docs] Visualize no-AllGather shared host weights
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:16 +08:00
Hongsheng Liu 1ccdc2e80e [Docs] Separate DLO transfer mode design
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:23:16 +08:00
Hongsheng Liu 8f93a74951 [Docs] Remove stale DLO validation results
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:22:52 +08:00
Hongsheng Liu 4acb3b1387 [Docs] Add DLO component architecture
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:59 +08:00
Hongsheng Liu 6b7b5368f2 [Diffusion] Clarify runtime cache module boundary
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:59 +08:00
Hongsheng Liu 6751c1f404 [Diffusion] Address runtime cache review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:59 +08:00
Hongsheng Liu c1196216a3 [Diffusion] Stabilize runtime cache identities
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:28 +08:00
Hongsheng Liu 23ab14a15b [Diffusion] Register runtime cache for direct DLO H2D
Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup.

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:28 +08:00
Hongsheng Liu 67f15e73cd [Diffusion] Document runtime cache hardware tradeoffs
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:28 +08:00
Hongsheng Liu 59549968cf [Diffusion] Tighten runtime cache compatibility guards
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:28 +08:00
Hongsheng Liu d30ccdc315 [Diffusion] Add node-local runtime weight cache for DLO
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-19 11:21:01 +08:00
27 changed files with 4021 additions and 288 deletions

No files matched your search

@@ -43,21 +43,291 @@ DLO weight collective.
Host storage is selected separately from the transfer protocol. The loader can
produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
the same node can share immutable checkpoint pages when direct mmap is
selected, while the ordinary-loader fallback still keeps a private runtime
copy per process.
otherwise it uses the ordinary loader. In no-AllGather mode, the Phase B host
weight cache can explicitly normalize supported final ordinary-loader DiT
tensors into an immutable node-local mmap entry. Ordinary pinned weights remain
the default after checkpoint mmap fallback when the cache is not selected; cache
publication and validation add startup and disk lifecycle costs. Replicas can
share either proven-compatible checkpoint pages or, when selected, equivalent
final runtime layouts. Under the existing
pinned-memory policy, the backend attempts to make the complete final mapping
directly H2D-copyable when the platform provides registration, avoiding the
recurrent host packing otherwise required by mmap. An optional positive limit
adds a per-worker safety ceiling. 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
tensors, but those configurations do not use checkpoint mmap and must not be
used to claim shared-mmap host-memory savings.
tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
can share each final TP coordinate among equivalent processes.
The compatibility matrix below describes the current implementation. The
unit-level guards are covered, but not every parallelism combination has a
full model-and-hardware end-to-end test.
The compatibility matrix below defines the current implementation boundary.
Hardware and performance qualification is tracked separately from this
long-lived design contract.
## Design
## Architecture
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 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.
### Component ownership
The map includes optional mode-specific components, but the ownership rules
and loader-to-backend boundary are shared.
```mermaid
flowchart TB
CLI["CLI / EngineArgs"] --> Config["OmniDiffusionConfig<br/>validation and operator policy"]
Config --> Runner["DiffusionModelRunner<br/>startup coordinator"]
subgraph LoaderPlane["Loader-owned host-storage plane"]
Loader["DiffusersPipelineLoader"]
Checkpoint["checkpoint mmap preflight"]
Ordinary["ordinary load + callbacks + calibration"]
HostWeightCache["host_weight_cache<br/>final-layout publication / join"]
Contract["HostWeightPlan<br/>backing + bindings + transforms"]
Loader --> Checkpoint
Loader -->|"registered final layout requested"| Ordinary
Checkpoint -->|"compatible"| Contract
Checkpoint -->|"fallback"| Ordinary
Ordinary -->|"host weight cache enabled"| HostWeightCache
HostWeightCache -->|"enabled and compatible"| Contract
end
Runner -->|"load_model"| Loader
Loader -->|"pipeline + optional HostWeightPlan"| Runner
Config --> Factory["get_offload_backend<br/>OffloadConfig normalization"]
Runner -->|"device + optional plan"| Factory
Factory --> Backend["DistributedLayerwiseOffloadBackend<br/>resource and transfer owner"]
Topology["existing DP / SP process groups"] --> Backend
Discovery["ModuleDiscovery + optional OffloadPlan"] --> Backend
Contract -.->|"passed once by the runner"| Backend
Backend -.->|"host weight cache + pinned-memory policy"| Registration["host_registration<br/>platform-neutral lifecycle"]
Registration --> CUDA["CUDA registration backend"]
Backend --> Hooks["DistributedLayerwiseOffloadHook<br/>one circular hook per streamed block"]
Forward["pipeline.forward"] --> Hooks
Hooks --> DataPath["two device slots<br/>H2D only or H2D + AllGather"]
Shutdown["DiffusionWorker.shutdown"] -->|"disable before process-group teardown"| Backend
```
| Component | Owns | Must not own |
| --- | --- | --- |
| 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, pinned-memory policy, and optional registration ceiling | 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 | 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 |
| `DistributedLayerwiseOffloadHook` | The block hot path: prefetch, optional DLO AllGather, readiness events, parameter repointing, and replacement with offload placeholders | Cache keys, manifest validation, or request routing |
| `host_registration` and platform backend | All-or-nothing registration lifecycle for existing immutable mappings | Choosing which weights are cached or changing the two-slot device-buffer contract |
| `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. The cache root and lock timeout are consumed by the loader.
`dlo_use_host_weight_cache` explicitly selects final-layout cache publication
instead of direct-checkpoint mmap. It defaults to false so the ordinary
pinned-loader path remains available without requiring an opt-out.
With `pin_cpu_memory=True`, selecting the host weight cache deliberately chooses
the transform-complete layout even when direct checkpoint mmap is available,
because only that final layout can be registered without deferred transforms.
The backend consumes `dlo_host_weight_cache_pin_limit_gib` as an optional
registration ceiling; zero applies no additional ceiling and does not select
the cache by itself.
`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.
### Runtime host-layout identity (v1)
Host weight cache v1 identifies the loader-produced host representation, not
the topology used later to move weights to devices. These are separate
contracts:
- `RuntimeHostLayout` describes tensor ownership and representation after the
ordinary loader has applied its parallel-aware transforms. This is the
layout stored in the host weight cache.
- `DLOTransferLayout` describes whether each worker copies its local weights
directly or participates in the existing DLO AllGather protocol. It is not
part of the host weight cache in v1; AllGather continues to use its existing
process-owned host shards.
The v1 physical-sharing invariant is one cache entry per TP coordinate,
shared by equivalent DP and SP consumers on the same node. The identity
therefore includes TP world size and TP rank, because different TP coordinates
own different tensors. It excludes DP world size/rank and SP rank, because
those dimensions do not change the final tensor ownership supported by v1.
SP world size and implementation details remain conservative compatibility
guards until equivalent layouts across those configurations are proven.
PP, HSDP/DTensor, expert parallelism, and CFG parallelism are rejected rather
than assigned speculative cache identities. Supporting one of these layouts
requires first defining its loader-produced ownership, representation, and
sharing domain; it does not follow automatically from supporting its runtime
collectives. Consequently, this phase implements a TP-coordinate runtime host
layout, not a general parallel-layout cache.
### Shared startup and lifecycle
The shared control path selects host backing and transfers its ownership to
the backend before either transfer mode installs hooks:
```mermaid
sequenceDiagram
participant C as OmniDiffusionConfig
participant R as DiffusionModelRunner
participant L as DiffusersPipelineLoader
participant B as DLO backend
participant W as DiffusionWorker
C->>R: validated offload policy
R->>L: load_model(config)
L->>L: select and validate host backing
L-->>R: pipeline + optional HostWeightPlan
R->>B: enable(config, device, optional plan)
B->>B: discover modules and consume host backing
B-->>R: mode-specific backend is ready
W->>B: shutdown calls disable()
B->>B: synchronize, unregister, remove hooks, close mappings
```
The loader returns a plan only after its contract is valid: checkpoint plans
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
hard startup error rather than a fallback to uninitialized parameters.
At inference time, the model pipeline remains the caller: normal block forward
calls activate the installed hooks. Both modes use one copy stream and two
shared device slots sized for the largest streamed block. The hooks do not
call the loader or cache.
### AllGather startup and hot path
```mermaid
sequenceDiagram
participant B as DLO backend
participant H as block hooks
participant G as existing DP or SP group
B->>B: choose existing group and prepare persistent local shards
B->>B: allocate two device slots and a communication stream
B->>H: install hooks and prime the first block
loop each streamed block
H->>H: copy the next local shard to the alternate slot
H->>G: all_gather_into_tensor(next block)
G-->>H: complete next block is ready
H->>H: execute current block and rotate slots
end
```
AllGather uses the existing DP group when DP is greater than one, otherwise
the existing SP group when SP is greater than one. TP is never selected as
DLO's weight-sharding group. Every participating rank must request blocks in
the same order; fast P2P reduces communication cost but does not remove that
lockstep contract. An effective group size of one issues no collective and
uses rank-local transfer.
### No-AllGather startup and hot path
```mermaid
sequenceDiagram
participant L as DiffusersPipelineLoader
participant C as host_weight_cache
participant B as DLO backend
participant P as platform registration
participant H as block hooks
alt host weight cache explicitly selected
L->>L: load, mutate, validate, and calibrate final tensors
L->>C: publish or join final-layout entry
C-->>L: validated host weight cache plan or fail-closed fallback
L-->>B: cache plan or materialized fallback via runner
else automatic checkpoint selection
L->>L: checkpoint mmap preflight
alt loader selects a compatible checkpoint plan
L-->>B: checkpoint plan via runner
else ordinary loading
L->>L: load, mutate, validate, and calibrate final tensors
L-->>B: materialized tensors via runner
end
end
B->>B: consume complete rank-local block sources
opt host weight cache plan and pinned-memory policy
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 host weight cache mapping
H->>H: copy complete next block directly to device
else pageable mmap mapping
H->>H: pack into a bounded staging slot, then copy to device
else ordinary runtime tensors
H->>H: copy complete next block to device
end
H->>H: execute current block and rotate slots
end
```
No-AllGather creates no DLO weight group and issues no DLO collective.
Workers may therefore schedule requests independently. Ordinary TP and SP
collectives remain inside model execution; host registration changes only
the source of the complete-block H2D copy.
### Host backing and transfer composition
| Loader-selected host backing | Backend transfer mode | Setup result | Per-block hot path |
| --- | --- | --- | --- |
| 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 mmap, or host weight cache with pinned memory disabled/registration unavailable | Rank-local / no-AllGather | Immutable mapping plus two bounded host-staging slots | mmap views → staging slot → device slot |
| Host weight cache mmap with pinned memory enabled and registration available | 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 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.
### Failure and lifecycle boundaries
- Invalid option combinations fail during configuration validation.
- Checkpoint mmap preflight failure falls back before ordinary loading is
skipped.
- 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.
- During shutdown, DLO synchronizes outstanding device work, unregisters host
mappings, removes hooks, and closes mmap handles only after safe unregister;
the worker destroys its process groups afterward.
DLO deliberately does not own DP replica placement, independent-engine request
routing, cross-engine admission control, or cross-node cache distribution.
Those belong to an external router/orchestrator. It also does not replace TP or
SP model collectives; it only adds an optional weight collective over an
already-existing DP or SP group.
## Shared design
The following topology and loader contracts apply to both transfer modes.
### DLO consumes the existing parallel topology
@@ -65,15 +335,10 @@ DLO does not create a new DP, TP, or SP topology. It reads the configured
`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
standard distributed groups have been initialized.
The DLO weight-sharding group is selected as follows:
1. Use the existing DP group when `data_parallel_size > 1`.
2. When DP is one and SP is greater than one, use the SP group.
3. Otherwise, run rank-locally without a DLO process group.
TP is deliberately not used as DLO's AllGather group. HSDP has its own
parameter-sharding lifecycle and is not allowed to be sharded a second time by
DLO's AllGather path.
AllGather mode selects one existing group for DLO weight sharding; no-AllGather
mode creates no DLO group. Neither mode changes ordinary TP or SP model
collectives. The mode-specific group and scheduling contracts are described
below.
### The loader owns host-weight planning
@@ -105,7 +370,24 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
transforms. Model-specific direct-layout knowledge, when required, lives in a
checkpoint adapter beside the ordinary loader.
### AllGather path
## AllGather design
### Group selection
The DLO weight-sharding group is selected in this order:
1. Use the existing DP group when `data_parallel_size > 1`.
2. When DP is one and SP is greater than one, use the existing SP group.
3. Otherwise, use an effective group size of one and issue no collective.
TP is deliberately not used as DLO's AllGather group. HSDP has its own
parameter-sharding lifecycle and cannot be sharded a second time by DLO.
### Sharded transfer and double buffering
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 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
@@ -113,6 +395,15 @@ block's shard is copied to a device buffer and reconstructed with
`all_gather_into_tensor` on a communication stream while the current block is
executing.
![AllGather DLO sharded host weights and lockstep reconstruction](../../figures/dlo/dlo_allgather_sharded_host.svg)
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 host weight cache path, these persistent pinned host shards are
process-owned, and every participating rank must follow the same block order.
```text
Compute: [Block N] [Block N+1] [Block N+2]
H2D: [shard N+1] [shard N+2]
@@ -138,11 +429,29 @@ When DP is greater than one, the engine can process one request per DP rank in
the same denoising wave. Because AllGather is a collective, all participating
requests must take the same execution path at every denoising step.
### Rank-local path without DLO AllGather
Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
plus shared mmap storage avoids coupling their request schedules and failure
domains.
## No-AllGather design
### Rank-local transfer
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
one and streams complete blocks using H2D copies only. The host backing may be
either a loader-approved checkpoint mapping or ordinary runtime tensors.
a loader-approved checkpoint mapping, a normalized 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)
The figure uses TP2 to make per-coordinate physical sharing explicit. Each
engine contains two rank-local workers and devices. The cache keeps one
immutable entry per TP coordinate, so equivalent TP0 workers across engines
share the TP0 entry and equivalent TP1 workers share the separate TP1 entry.
An external router or orchestrator chooses each engine's requests; DLO
coordinates only the worker-local mapping, hooks, transfer streams, events,
and device slots.
For direct mmap, each process retains immutable safetensors views and uses two
bounded pinned host staging slots. Processes on the same node that map the same
@@ -151,10 +460,17 @@ 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.
When direct mmap preflight fails, the regular model loader remains responsible
for preparing each rank's weights, including TP-local tensors or HSDP-managed
parameters. In that fallback, each pure-DP process keeps a private full runtime
copy.
For a host weight cache mapping, the default pinned-memory policy attempts to
replace those host slots with direct copies from the shared final tensor views.
Disabling pinned memory or a failed registration retains staging behavior.
When direct mmap preflight fails or the host weight cache is selected, the
regular model loader remains responsible for preparing each rank's weights,
including TP-local tensors or HSDP-managed parameters. Ordinary pinned tensors
remain the default without cache selection. When the host weight cache is
explicitly selected,
supported CPU layouts enter the cache; unsupported layouts keep one private
runtime copy per process.
This mode means:
@@ -165,18 +481,109 @@ This mode means:
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
DLO's additional weight AllGather is disabled.
- Pure DP deployments share one checkpoint-backed copy per node when direct
mmap is selected; the ordinary-loader fallback keeps one private runtime
copy per rank.
mmap is selected. The 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 host weight cache
The Phase B cache preserves the same ownership boundary. DLO neither derives
cache keys nor interprets model loaders. No-AllGather DLO first attempts a
compatible direct-checkpoint mmap plan. When that plan is unavailable, every
rank completes the ordinary loader, custom callbacks, post-load processing,
strict validation, and calibration. When cache use or registration is
explicitly selected, the loader:
1. discovers final CPU DiT parameters and persistent buffers;
2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
ambiguous layouts before changing the model;
3. hashes tensor names, metadata, and final bytes;
4. builds or joins one content-addressed entry under a strict POSIX lock; and
5. returns a `HostWeightPlan(backing_kind="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
and conservative SP implementation/world-size guards. DP replicas and
equivalent SP ranks can therefore share; different TP coordinates cannot.
Final tensor content is authoritative, so an unmodeled loader/config change
that changes bytes creates a different entry instead of risking silent
sharing.
Publication writes sharded safetensors and a JSON manifest into a same-
filesystem temporary directory, flushes them, and atomically renames the
directory. The manifest includes per-file SHA-256 digests and the final tensor
content digest. Kernel `flock` ownership distinguishes a live writer from a
stale lock file: process death releases the lock automatically. Stale temporary
entries are removed only by the next lock owner. A bounded wait or any cache
error retains the already-valid ordinary tensors.
DLO maps this plan without moving the DiT to `meta` and without rerunning any
post-load hook. It changes tensor storage while preserving `Parameter` objects,
then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
streaming hooks. The source mappings remain open through DLO disable. With
pinned memory disabled, or when registration is unavailable, they feed the
existing bounded two-slot host-staging path. Every streamed block is first
packed from pageable mmap storage into a staging slot and is then copied to the
device. This removes model-sized private pinned storage, but adds a recurrent
host-to-host copy and can amplify TP synchronization waits when several ranks
compete for host bandwidth.
This version is deliberately steady-state-oriented. It requires roughly one
local-disk model copy per runtime identity, has no eviction, and still performs
ordinary loading plus full content-hash passes in every rank. Cache publication
or validation failure falls back to the already-valid ordinary host tensors.
Skip-load-on-hit is a separate follow-up after loader side effects can be
modeled safely.
### Registered host weight cache transfer
Selecting the host weight cache with `pin_cpu_memory=True` opts a worker into
registration of the complete 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
adjacent ranges within that mapping. When a positive per-worker ceiling is
configured, the complete byte count is checked against it before any platform
mutation. The CUDA backend
first requires `cudaDevAttrHostRegisterReadOnlySupported`, then registers every
range with the read-only host-registration flag and verifies that PyTorch
recognizes every mapped tensor as pinned. A device without that capability
retains bounded staging; vLLM-Omni does not replace immutable mappings with a
writable or copy-on-write CUDA-specific layout. Other registration failures
roll back any ranges already registered and preserve 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
repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
direct path trades recurrent host packing for finer-grained H2D operations. A
packed cache schema should be considered only if that additional format and
publication complexity becomes necessary.
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.
## Parallelism compatibility
| Parallelism | DLO + AllGather | DLO without AllGather |
|---|---|---|
| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
| **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
@@ -192,7 +599,9 @@ This mode means:
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
configuration.
## Request and loading constraints
## Mode constraints
### AllGather constraints
AllGather DP multi-concurrency requires:
@@ -200,109 +609,54 @@ AllGather DP multi-concurrency requires:
- the same `num_inference_steps` for all requests in a wave; and
- identical request arguments that affect the collective execution path.
Per-tensor online FP8 linears can use DLO AllGather after the ordinary loader
finalizes their runtime weights and scales. Other online methods must use
`--dlo-no-use-allgather` or disable online quantization until their runtime
layouts are validated.
### No-AllGather constraints
The no-AllGather path does not impose these DLO-specific synchronized-wave
requirements.
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.
### Shared loading constraint
Direct checkpoint mmap can back either transfer path. It is currently limited
to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
ordinary loader finalizes their runtime weights and scales; DLO then shards and
reconstructs those tensors with their recorded layouts. Other online methods
must use `--dlo-no-use-allgather` or disable online quantization until their
runtime layouts are validated.
ordinary loader.
A normalized runtime mmap cache, built through the ordinary loader, is the
proposed general mechanism for sharing transformed TP or quantized layouts.
That cache and its publication/lifecycle protocol are intentionally outside
this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
## Validation status
## Validation coverage
Current source-level validation includes:
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
loaders, missing keys, and shape/dtype mismatches;
- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
sharding of finalized weights and scales;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- rank-local mmap source retention, bounded two-slot staging, and adapter
transforms without parameter-side flags;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
### B300 parallel-topology smoke matrix
A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
CUDNN attention backend, 256x256 output, two denoising steps, and
`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
or shared-mmap host-memory savings.
| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
|---|---:|---:|---:|---:|
| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
Within each topology, the AllGather and no-AllGather video and audio outputs
were byte-identical. All four runs completed without an `ERROR` or traceback
and released their device allocations. For DP4xTP1, no-AllGather direct mmap
reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
AllGather, because DP replicas did not share checkpoint-backed runtime
weights. This is a functional and memory smoke test, not a production-quality
performance or output-quality benchmark.
### Host-memory measurement
A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
ordinary-loader fallback with direct mmap. Both runs used
DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
256x256 four-second request. The ordinary-loader workers were sampled after
initialization. The mmap workers were sampled after one completed request, so
the checkpoint working set had been faulted into the page cache; this is the
more conservative point for mmap.
The values below come from `/proc/<worker>/smaps_rollup` and include the whole
worker, not only the DiT. The stable rank-to-rank difference comes from other
pipeline components, so each worker should be compared with the same worker in
the other storage mode.
| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
|---|---:|---:|---:|---:|---:|
| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
31.20 GiB `Pss_File`, which is the proportional charge expected when the same
resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
because it counts a shared physical page in every process that maps it; summed
PSS is the appropriate node-memory comparison.
The highest-value missing coverage is broader end-to-end numerical and
lifecycle comparison against ordinary layerwise offload for DP+SP,
HSDP+SP+no-AllGather, and TP greater than one across additional models and
target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
change the Phase A direct-mmap TP1 support boundary.
Historical test counts and hardware measurements are intentionally not kept
in this design document because they can drift as the implementation changes.
Revalidation must run against the final commit and publish commit-pinned
artifacts before making hardware or performance claims. Track that work in
[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
## Recommendations
- Use **DP + DLO AllGather** for the supported throughput and host-memory
scaling path.
### AllGather
- Use **DP + DLO AllGather** when ranks already execute the same block path in
lockstep and belong to the same fast P2P domain. It shards persistent host
weights across that group.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
not the goal.
### No-AllGather
- Use **no-AllGather** when independent replica execution is required. TP1
direct-mmap deployments can share checkpoint pages per node; other layouts
retain the ordinary loader's private host memory behavior and are outside
the Phase A shared-mmap support scope.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
direct-mmap deployments can share checkpoint pages per node. Enable the
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. With the default pinned-memory
policy, supported CUDA workers prefer registered direct H2D. Configure a
positive ceiling when an explicit page-locking bound is required, or disable
pinned memory when pageable staging is operationally preferable.
### HSDP
- Do not combine HSDP with DLO AllGather or host weight cache v1. HSDP with
no-AllGather remains an ordinary-loader rank-local path.
@@ -0,0 +1,165 @@
<svg xmlns="http://www.w3.org/2000/svg" width="1280" height="850" viewBox="0 0 1280 850" role="img" aria-labelledby="title desc">
<title id="title">AllGather DLO sharded host weights and lockstep device reconstruction</title>
<desc id="desc">An external orchestrator submits one synchronized wave to a DP2 by TP2 diffusion engine. Each DP replica contains TP0 and TP1 workers with private pinned host shards and two rotating GPU weight slots. DLO AllGather groups connect matching TP coordinates across DP replicas, reconstructing full rank-local blocks on each device in lockstep. Ordinary TP groups remain within each DP replica.</desc>
<defs>
<marker id="request-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#7c3aed"/>
</marker>
<marker id="h2d-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#d97706"/>
</marker>
<marker id="collective-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#7c3aed"/>
</marker>
<style>
.canvas { fill: #ffffff; stroke: #e2e8f0; }
.node { fill: #f8fafc; stroke: #cbd5e1; }
.engine { fill: #ffffff; stroke: #cbd5e1; }
.replica { fill: #f8fafc; stroke: #cbd5e1; }
.router { fill: #f5f3ff; stroke: #c4b5fd; }
.scope { fill: #f8fafc; stroke: #cbd5e1; }
.host { fill: #fff7ed; stroke: #fdba74; }
.shard { fill: #d97706; }
.gpu { fill: #f0fdf4; stroke: #86efac; }
.slot-current { fill: #16a34a; }
.slot-next { fill: #bbf7d0; stroke: #22c55e; }
.lockstep { fill: #f5f3ff; stroke: #c4b5fd; }
.title { font: 700 25px system-ui, -apple-system, Segoe UI, sans-serif; fill: #0f172a; }
.subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.section { font: 700 15px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
.section-sub { font: 500 11.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.eyebrow { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .8px; fill: #64748b; }
.box-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
.box-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
.shard-title { font: 700 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.shard-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.slot-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.slot-text { font: 500 9px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.slot-next-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
.slot-next-text { font: 500 9px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
.request-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
.h2d-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b45309; }
.collective-label { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
.lock-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
.lock-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #5b21b6; }
.request-path { fill: none; stroke: #7c3aed; stroke-width: 2; marker-end: url(#request-arrow); }
.h2d-path { fill: none; stroke: #d97706; stroke-width: 2; marker-end: url(#h2d-arrow); }
.collective-path { fill: none; stroke: #7c3aed; stroke-width: 2.5; marker-start: url(#collective-arrow); marker-end: url(#collective-arrow); }
</style>
</defs>
<rect x="12" y="12" width="1256" height="826" rx="14" class="canvas"/>
<text id="title-text" x="40" y="50" class="title">AllGather DLO: sharded host weights, lockstep reconstruction</text>
<text x="40" y="74" class="subtitle">DP2 &#215; TP2 shown: DLO groups span DP replicas per TP coordinate; ordinary TP stays within each replica</text>
<!-- The external control plane submits one synchronized DP wave. -->
<rect x="390" y="94" width="500" height="68" rx="10" class="router"/>
<text x="640" y="121" text-anchor="middle" class="section">External router / orchestrator</text>
<text x="640" y="144" text-anchor="middle" class="section-sub">submits one synchronized DP wave &#183; outside vLLM-Omni DLO</text>
<rect x="30" y="184" width="1220" height="634" rx="14" class="node"/>
<text x="52" y="211" class="eyebrow">ONE NODE &#183; ONE SYNCHRONIZED ENGINE, PROCESS-OWNED HOST SHARDS</text>
<path d="M 640 162 L 640 232" class="request-path"/>
<text x="654" y="202" class="request-label">request wave</text>
<!-- One engine contains two DP replicas that execute in lockstep. -->
<rect x="45" y="232" width="1190" height="262" rx="12" class="engine"/>
<text x="68" y="261" class="section">One synchronized diffusion engine &#183; DP2 &#215; TP2</text>
<text x="68" y="281" class="section-sub">two DP requests follow the same block order; DLO reuses existing DP groups at each TP coordinate</text>
<!-- DP replica 0. -->
<rect x="65" y="296" width="515" height="186" rx="10" class="replica"/>
<text x="82" y="318" class="box-title">DP replica 0 &#183; request 0</text>
<text x="563" y="318" text-anchor="end" class="section-sub">ordinary TP group</text>
<rect x="82" y="332" width="473" height="64" rx="8" class="gpu"/>
<text x="98" y="351" class="eyebrow">TP0 &#183; GPU 0</text>
<text x="98" y="376" class="box-text">DLO backend &#183; local shard H2D</text>
<rect x="250" y="340" width="132" height="44" rx="6" class="slot-current"/>
<text x="316" y="357" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
<text x="316" y="374" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
<rect x="397" y="340" width="142" height="44" rx="6" class="slot-next"/>
<text x="468" y="357" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
<text x="468" y="374" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
<rect x="82" y="410" width="473" height="64" rx="8" class="gpu"/>
<text x="98" y="429" class="eyebrow">TP1 &#183; GPU 1</text>
<text x="98" y="454" class="box-text">DLO backend &#183; local shard H2D</text>
<rect x="250" y="418" width="132" height="44" rx="6" class="slot-current"/>
<text x="316" y="435" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
<text x="316" y="452" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
<rect x="397" y="418" width="142" height="44" rx="6" class="slot-next"/>
<text x="468" y="435" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
<text x="468" y="452" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
<!-- DP replica 1. -->
<rect x="700" y="296" width="515" height="186" rx="10" class="replica"/>
<text x="717" y="318" class="box-title">DP replica 1 &#183; request 1</text>
<text x="1198" y="318" text-anchor="end" class="section-sub">ordinary TP group</text>
<rect x="725" y="332" width="473" height="64" rx="8" class="gpu"/>
<text x="741" y="351" class="eyebrow">TP0 &#183; GPU 2</text>
<text x="741" y="376" class="box-text">DLO backend &#183; local shard H2D</text>
<rect x="893" y="340" width="132" height="44" rx="6" class="slot-current"/>
<text x="959" y="357" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
<text x="959" y="374" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
<rect x="1040" y="340" width="142" height="44" rx="6" class="slot-next"/>
<text x="1111" y="357" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
<text x="1111" y="374" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
<rect x="725" y="410" width="473" height="64" rx="8" class="gpu"/>
<text x="741" y="429" class="eyebrow">TP1 &#183; GPU 3</text>
<text x="741" y="454" class="box-text">DLO backend &#183; local shard H2D</text>
<rect x="893" y="418" width="132" height="44" rx="6" class="slot-current"/>
<text x="959" y="435" text-anchor="middle" class="slot-title">SLOT 0 &#183; CURRENT</text>
<text x="959" y="452" text-anchor="middle" class="slot-text">full TP-local N &#183; compute</text>
<rect x="1040" y="418" width="142" height="44" rx="6" class="slot-next"/>
<text x="1111" y="435" text-anchor="middle" class="slot-next-title">SLOT 1 &#183; PREFETCH</text>
<text x="1111" y="452" text-anchor="middle" class="slot-next-text">full TP-local N+1 after AG</text>
<!-- DLO collectives connect matching TP coordinates across DP replicas. -->
<text x="640" y="350" text-anchor="middle" class="collective-label">DLO AG &#183; TP0 GROUP</text>
<path d="M 555 364 L 725 364" class="collective-path"/>
<text x="640" y="428" text-anchor="middle" class="collective-label">DLO AG &#183; TP1 GROUP</text>
<path d="M 555 442 L 725 442" class="collective-path"/>
<!-- All participants must reach the same collective in the same order. -->
<rect x="365" y="516" width="550" height="48" rx="8" class="lockstep"/>
<circle cx="395" cy="540" r="12" fill="#7c3aed"/>
<text x="395" y="545" text-anchor="middle" font-family="system-ui, sans-serif" font-size="15" font-weight="700" fill="#ffffff">&#8644;</text>
<text x="420" y="536" class="lock-title">DLO AllGather requires lockstep block order</text>
<text x="420" y="553" class="lock-text">every participant enters each block collective in the same sequence</text>
<!-- Each replica copies two process-owned host shards before reconstruction. -->
<path d="M 455 650 C 412 625, 340 610, 340 575 L 340 540 C 340 516, 318 500, 300 482" class="h2d-path"/>
<text x="276" y="618" text-anchor="middle" class="h2d-label">local shard H2D &#215;2</text>
<path d="M 825 650 C 868 625, 940 610, 940 575 L 940 540 C 940 516, 962 500, 980 482" class="h2d-path"/>
<text x="1004" y="618" text-anchor="middle" class="h2d-label">local shard H2D &#215;2</text>
<!-- Persistent host shards are private to workers, not mmap-shared. -->
<rect x="58" y="650" width="242" height="140" rx="10" class="scope"/>
<text x="78" y="676" class="eyebrow">GROUP TOPOLOGY</text>
<text x="78" y="702" class="box-title">DP2 &#215; TP2</text>
<text x="78" y="725" class="box-text">DLO groups: same TP coordinate</text>
<text x="78" y="746" class="box-text">across DP replicas</text>
<text x="78" y="767" class="box-text">TP stays the model-parallel group.</text>
<rect x="330" y="650" width="620" height="140" rx="10" class="host"/>
<text x="350" y="676" class="eyebrow">PROCESS-OWNED PINNED HOST SHARDS &#183; NOT HOST WEIGHT CACHE MMAP</text>
<rect x="370" y="692" width="260" height="60" rx="8" class="shard"/>
<text x="500" y="716" text-anchor="middle" class="shard-title">DP replica 0 &#183; private</text>
<text x="500" y="738" text-anchor="middle" class="shard-text">TP0 shard 0/2 &#183; TP1 shard 0/2</text>
<rect x="650" y="692" width="260" height="60" rx="8" class="shard"/>
<text x="780" y="716" text-anchor="middle" class="shard-title">DP replica 1 &#183; private</text>
<text x="780" y="738" text-anchor="middle" class="shard-text">TP0 shard 1/2 &#183; TP1 shard 1/2</text>
<text x="640" y="773" text-anchor="middle" class="box-text">Across each DLO group, shards total approximately one rank-local runtime block copy</text>
<rect x="980" y="650" width="242" height="140" rx="10" class="scope"/>
<text x="1000" y="676" class="eyebrow">HOT PATH &#183; PER WORKER</text>
<text x="1000" y="702" class="box-title">Prefetch next block</text>
<text x="1000" y="725" class="box-text">local pinned shard &#8594; GPU</text>
<text x="1000" y="746" class="box-text">DLO AllGather &#8594; full block</text>
<text x="1000" y="767" class="box-text">rotate slots; overlap compute</text>
</svg>

After

Width:  |  Height:  |  Size: 12 KiB

@@ -0,0 +1,202 @@
<svg xmlns="http://www.w3.org/2000/svg" width="1280" height="850" viewBox="0 0 1280 850" role="img" aria-labelledby="title desc">
<title id="title">No-AllGather DLO shared host weights and independent runtime orchestration</title>
<desc id="desc">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.</desc>
<defs>
<marker id="request-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#7c3aed"/>
</marker>
<marker id="mapping-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#d97706"/>
</marker>
<marker id="runtime-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
<path d="M 0 0 L 10 5 L 0 10 z" fill="#2563eb"/>
</marker>
<style>
.canvas { fill: #ffffff; stroke: #e2e8f0; }
.node { fill: #f8fafc; stroke: #cbd5e1; }
.engine { fill: #ffffff; stroke: #cbd5e1; }
.router { fill: #f5f3ff; stroke: #c4b5fd; }
.scope { fill: #f8fafc; stroke: #cbd5e1; }
.cache { fill: #fff7ed; stroke: #fdba74; }
.payload { fill: #d97706; }
.loader { fill: #f1f5f9; stroke: #cbd5e1; }
.mapping { fill: #fffbeb; stroke: #f59e0b; }
.hook { fill: #eff6ff; stroke: #60a5fa; }
.gpu { fill: #f0fdf4; stroke: #86efac; }
.device { fill: #ffffff; stroke: #86efac; }
.slot-current { fill: #16a34a; }
.slot-next { fill: #bbf7d0; stroke: #22c55e; }
.no-collective { fill: #fef2f2; stroke: #fca5a5; }
.title { font: 700 25px system-ui, -apple-system, Segoe UI, sans-serif; fill: #0f172a; }
.subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.section { font: 700 15px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
.section-sub { font: 500 11.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
.eyebrow { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .8px; fill: #64748b; }
.box-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e293b; }
.box-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
.payload-title { font: 700 14px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.payload-text { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.device-title { font: 700 9px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
.slot-text { font: 700 8px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
.slot-next-text { font: 700 8px system-ui, -apple-system, Segoe UI, sans-serif; fill: #166534; }
.arrow-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
.request-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #6d28d9; }
.mapping-label { font: 600 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b45309; }
.no-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b91c1c; }
.no-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #7f1d1d; }
.request-path { fill: none; stroke: #7c3aed; stroke-width: 2; marker-end: url(#request-arrow); }
.mapping-path { fill: none; stroke: #d97706; stroke-width: 2; marker-end: url(#mapping-arrow); }
.runtime-path { fill: none; stroke: #2563eb; stroke-width: 2; marker-end: url(#runtime-arrow); }
.startup-path { fill: none; stroke: #94a3b8; stroke-width: 1.5; stroke-dasharray: 4 3; marker-end: url(#runtime-arrow); }
</style>
</defs>
<rect x="12" y="12" width="1256" height="826" rx="14" class="canvas"/>
<text id="title-text" x="40" y="50" class="title">No-AllGather DLO: TP-coordinate host sharing, independent engines</text>
<text x="40" y="74" class="subtitle">TP2 shown: each engine owns two devices; equivalent workers share one host payload per TP coordinate</text>
<!-- The control plane remains outside DLO. -->
<rect x="390" y="94" width="500" height="68" rx="10" class="router"/>
<text x="640" y="121" text-anchor="middle" class="section">External router / orchestrator</text>
<text x="640" y="144" text-anchor="middle" class="section-sub">independent request placement &#183; outside vLLM-Omni DLO</text>
<!-- One node contains independent engine runtimes over one shared physical backing. -->
<rect x="30" y="184" width="1220" height="634" rx="14" class="node"/>
<text x="52" y="211" class="eyebrow">ONE NODE &#183; PROCESS-LOCAL ENGINES OVER SHARED PHYSICAL HOST MEMORY</text>
<!-- Independent request routing from the external orchestrator. -->
<path d="M 500 162 C 456 181, 400 204, 332 232" class="request-path"/>
<text x="416" y="194" class="request-label">request A</text>
<path d="M 780 162 C 824 181, 880 204, 948 232" class="request-path"/>
<text x="864" y="194" class="request-label">request B</text>
<!-- Independent engine A. -->
<rect x="45" y="232" width="575" height="250" rx="12" class="engine"/>
<text x="68" y="261" class="section">Independent engine A</text>
<text x="68" y="281" class="section-sub">TP2 &#183; two rank-local workers &#183; own request schedule and DLO state</text>
<rect x="68" y="310" width="98" height="130" rx="8" class="loader"/>
<text x="117" y="335" text-anchor="middle" class="eyebrow">STARTUP &#215;2</text>
<text x="117" y="359" text-anchor="middle" class="box-title">rank-local</text>
<text x="117" y="377" text-anchor="middle" class="box-title">loaders</text>
<text x="117" y="402" text-anchor="middle" class="box-text">final transforms</text>
<text x="117" y="421" text-anchor="middle" class="box-text">publish / join</text>
<rect x="180" y="310" width="112" height="130" rx="8" class="mapping"/>
<text x="236" y="335" text-anchor="middle" class="eyebrow">HOST VIEWS &#215;2</text>
<text x="236" y="359" text-anchor="middle" class="box-title">TP0 mmap</text>
<text x="236" y="377" text-anchor="middle" class="box-title">TP1 mmap</text>
<text x="236" y="402" text-anchor="middle" class="box-text">process-local</text>
<text x="236" y="421" text-anchor="middle" class="box-text">optional registration</text>
<rect x="306" y="310" width="120" height="130" rx="8" class="hook"/>
<text x="366" y="335" text-anchor="middle" class="eyebrow">RUNTIME &#215;2</text>
<text x="366" y="359" text-anchor="middle" class="box-title">DLO backends</text>
<text x="366" y="377" text-anchor="middle" class="box-title">+ block hooks</text>
<text x="366" y="402" text-anchor="middle" class="box-text">private streams</text>
<text x="366" y="421" text-anchor="middle" class="box-text">events + state</text>
<rect x="440" y="310" width="157" height="130" rx="8" class="gpu"/>
<text x="518" y="330" text-anchor="middle" class="eyebrow">TWO DEVICES</text>
<rect x="450" y="340" width="137" height="42" rx="6" class="device"/>
<text x="518" y="353" text-anchor="middle" class="device-title">TP0 &#183; GPU A0</text>
<rect x="457" y="359" width="58" height="16" rx="4" class="slot-current"/>
<text x="486" y="370" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
<rect x="521" y="359" width="58" height="16" rx="4" class="slot-next"/>
<text x="550" y="370" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
<rect x="450" y="389" width="137" height="42" rx="6" class="device"/>
<text x="518" y="402" text-anchor="middle" class="device-title">TP1 &#183; GPU A1</text>
<rect x="457" y="408" width="58" height="16" rx="4" class="slot-current"/>
<text x="486" y="419" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
<rect x="521" y="408" width="58" height="16" rx="4" class="slot-next"/>
<text x="550" y="419" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
<path d="M 166 375 L 180 375" class="startup-path"/>
<path d="M 292 375 L 306 375" class="runtime-path"/>
<path d="M 426 375 L 440 375" class="runtime-path"/>
<!-- Independent engine B. -->
<rect x="660" y="232" width="575" height="250" rx="12" class="engine"/>
<text x="683" y="261" class="section">Independent engine B</text>
<text x="683" y="281" class="section-sub">TP2 &#183; two rank-local workers &#183; own request schedule and DLO state</text>
<rect x="683" y="310" width="98" height="130" rx="8" class="loader"/>
<text x="732" y="335" text-anchor="middle" class="eyebrow">STARTUP &#215;2</text>
<text x="732" y="359" text-anchor="middle" class="box-title">rank-local</text>
<text x="732" y="377" text-anchor="middle" class="box-title">loaders</text>
<text x="732" y="402" text-anchor="middle" class="box-text">final transforms</text>
<text x="732" y="421" text-anchor="middle" class="box-text">publish / join</text>
<rect x="795" y="310" width="112" height="130" rx="8" class="mapping"/>
<text x="851" y="335" text-anchor="middle" class="eyebrow">HOST VIEWS &#215;2</text>
<text x="851" y="359" text-anchor="middle" class="box-title">TP0 mmap</text>
<text x="851" y="377" text-anchor="middle" class="box-title">TP1 mmap</text>
<text x="851" y="402" text-anchor="middle" class="box-text">process-local</text>
<text x="851" y="421" text-anchor="middle" class="box-text">optional registration</text>
<rect x="921" y="310" width="120" height="130" rx="8" class="hook"/>
<text x="981" y="335" text-anchor="middle" class="eyebrow">RUNTIME &#215;2</text>
<text x="981" y="359" text-anchor="middle" class="box-title">DLO backends</text>
<text x="981" y="377" text-anchor="middle" class="box-title">+ block hooks</text>
<text x="981" y="402" text-anchor="middle" class="box-text">private streams</text>
<text x="981" y="421" text-anchor="middle" class="box-text">events + state</text>
<rect x="1055" y="310" width="157" height="130" rx="8" class="gpu"/>
<text x="1133" y="330" text-anchor="middle" class="eyebrow">TWO DEVICES</text>
<rect x="1065" y="340" width="137" height="42" rx="6" class="device"/>
<text x="1133" y="353" text-anchor="middle" class="device-title">TP0 &#183; GPU B0</text>
<rect x="1072" y="359" width="58" height="16" rx="4" class="slot-current"/>
<text x="1101" y="370" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
<rect x="1136" y="359" width="58" height="16" rx="4" class="slot-next"/>
<text x="1165" y="370" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
<rect x="1065" y="389" width="137" height="42" rx="6" class="device"/>
<text x="1133" y="402" text-anchor="middle" class="device-title">TP1 &#183; GPU B1</text>
<rect x="1072" y="408" width="58" height="16" rx="4" class="slot-current"/>
<text x="1101" y="419" text-anchor="middle" class="slot-text">SLOT 0 NOW</text>
<rect x="1136" y="408" width="58" height="16" rx="4" class="slot-next"/>
<text x="1165" y="419" text-anchor="middle" class="slot-next-text">SLOT 1 NEXT</text>
<path d="M 781 375 L 795 375" class="startup-path"/>
<path d="M 907 375 L 921 375" class="runtime-path"/>
<path d="M 1041 375 L 1055 375" class="runtime-path"/>
<!-- The runtimes are deliberately not coupled by a DLO collective. -->
<rect x="365" y="516" width="550" height="44" rx="8" class="no-collective"/>
<circle cx="395" cy="538" r="12" fill="#dc2626"/>
<path d="M 389 532 L 401 544 M 401 532 L 389 544" stroke="#ffffff" stroke-width="2.2"/>
<text x="420" y="534" class="no-title">No cross-engine DLO AllGather</text>
<text x="420" y="551" class="no-text">shared host pages do not couple request order, progress, or failure domains</text>
<text x="640" y="592" text-anchor="middle" class="section-sub">Ordinary TP / SP model collectives, when configured, remain inside each engine</text>
<!-- Each engine maps both TP-coordinate payloads from the shared host cache. -->
<path d="M 455 650 C 412 625, 340 610, 340 570 L 340 525 C 340 478, 278 454, 236 440" class="mapping-path"/>
<text x="286" y="617" text-anchor="middle" class="mapping-label">TP0 + TP1 mappings</text>
<path d="M 825 650 C 868 625, 940 610, 940 570 L 940 525 C 940 478, 893 454, 851 440" class="mapping-path"/>
<text x="994" y="617" text-anchor="middle" class="mapping-label">TP0 + TP1 mappings</text>
<!-- One node-local physical backing, visually beneath its consumers. -->
<rect x="58" y="650" width="242" height="140" rx="10" class="scope"/>
<text x="78" y="676" class="eyebrow">SHARING SCOPE</text>
<text x="78" y="702" class="box-title">TP2: two coordinate entries</text>
<text x="78" y="725" class="box-text">one physical copy per coordinate</text>
<text x="78" y="746" class="box-text">A/TP0 + B/TP0 &#8594; TP0 entry</text>
<text x="78" y="767" class="box-text">A/TP1 + B/TP1 &#8594; TP1 entry</text>
<rect x="330" y="650" width="620" height="140" rx="10" class="cache"/>
<text x="350" y="676" class="eyebrow">NODE-LOCAL IMMUTABLE HOST WEIGHT CACHE &#183; ONE COPY PER TP COORDINATE</text>
<rect x="370" y="692" width="260" height="60" rx="8" class="payload"/>
<text x="500" y="717" text-anchor="middle" class="payload-title">TP0 physical payload</text>
<text x="500" y="738" text-anchor="middle" class="payload-text">shared by A/TP0 + B/TP0</text>
<rect x="650" y="692" width="260" height="60" rx="8" class="payload"/>
<text x="780" y="717" text-anchor="middle" class="payload-title">TP1 physical payload</text>
<text x="780" y="738" text-anchor="middle" class="payload-text">shared by A/TP1 + B/TP1</text>
<text x="640" y="773" text-anchor="middle" class="box-text">Identity includes final layout + TP coordinate; engine and DP replica are excluded</text>
<rect x="980" y="650" width="242" height="140" rx="10" class="scope"/>
<text x="1000" y="676" class="eyebrow">WHAT STAYS PRIVATE</text>
<text x="1000" y="702" class="box-title">Per worker / process</text>
<text x="1000" y="725" class="box-text">mmap object and validation state</text>
<text x="1000" y="746" class="box-text">optional host registration</text>
<text x="1000" y="767" class="box-text">hooks, streams, events, GPU slots</text>
</svg>

After

Width:  |  Height:  |  Size: 14 KiB

@@ -1,156 +1,127 @@
# 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.
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 host weights through mmap.
See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
for the implementation contract and compatibility matrix.
See the [shared DLO component architecture](../../../design/feature/offloader/distributed_layerwise_offload.md#architecture)
for ownership and lifecycle coordination. The mode-specific internals are
documented separately in the
[AllGather design](../../../design/feature/offloader/distributed_layerwise_offload.md#allgather-design)
and
[no-AllGather design](../../../design/feature/offloader/distributed_layerwise_offload.md#no-allgather-design).
## 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 + 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 |
```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 only request-independent weight shards, so data-
parallel ranks may process different requests concurrently.
AllGather is normally the best choice for synchronized DP. Host weight 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
# Sequence parallel deployment
# Share final host weights; prefer registered direct H2D under the pinned-memory policy
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--usp 4
--dlo-no-use-allgather \
--dlo-use-host-weight-cache
# Apply an 80 GiB per-worker safety ceiling to the same registration attempt
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--tensor-parallel-size 2 \
--dlo-no-use-allgather \
--dlo-use-host-weight-cache \
--dlo-host-weight-cache-pin-limit-gib 80
```
```python
from vllm_omni import Omni
omni = Omni(
model="/path/to/model",
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
```
## 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-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-use-host-weight-cache` | Use the final-layout cache and prefer registration when pinned memory is enabled | `false` |
| `--dlo-host-weight-cache-pin-limit-gib GIB` | Optional per-worker registration ceiling; zero means no additional ceiling | `0` |
| `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` |
## Host-weight loading
No-AllGather DLO automatically uses a compatible checkpoint mmap plan. When
checkpoint mmap is unavailable, it preserves ordinary pinned loader weights by
default. `--dlo-use-host-weight-cache` explicitly selects the transform-complete
final-layout cache. Under the default `pin_cpu_memory=True` policy, DLO attempts
to register the complete immutable mapping for direct H2D. Setting
`pin_cpu_memory=False` through structured or programmatic configuration retains
the staged mmap path instead. A positive pin limit is only an optional safety
ceiling and does not select the cache. The 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.
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.
- The bounded-staging fallback is a memory-first mode: recurrent pageable-mmap to
pinned-slot copies can reduce throughput, especially across concurrent TP
engines. Use it intentionally by disabling pinned host memory, or when the
active platform cannot register the immutable mapping.
- When a positive registration ceiling is configured, it must cover the
complete page-aligned mapping reported in the worker log. Registration is
all-or-nothing. Zero applies no additional ceiling.
- CUDA is the first registration backend and requires read-only host
registration support for immutable cache mappings. Platforms or devices
without that capability continue to use bounded staging; pinned-memory policy
requests direct H2D but does not guarantee it.
- 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.
- Host weight cache v1 has no automatic eviction. Stop all users of an entry
before deleting it.
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.
## Scheduling constraints
With direct checkpoint mmap, the loader:
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`.
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.
## 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 is
outside the Phase A shared-mmap support scope and falls back before model
mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
but it provides no shared-mmap host-memory benefit or guarantee.
- 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 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. 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.
- Resident leading layers require `--dlo-no-use-allgather` and a model
`OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
Sharing transformed TP or quantized runtime layouts through a normalized mmap
cache is a follow-up design in
[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
the direct-checkpoint path.
DLO transfer path. Other online quantization methods require no-AllGather
until their runtime layouts are validated.
- 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
[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.
+46
View File
@@ -908,6 +908,52 @@ def test_structured_diffusion_config_rejects_invalid_compile_granularity():
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
@pytest.mark.parametrize(
("enable_dlo", "use_allgather"),
[(False, False), (True, True)],
)
@pytest.mark.parametrize(
("field_name", "field_value"),
[
("dlo_use_host_weight_cache", True),
("dlo_host_weight_cache_pin_limit_gib", 1.0),
],
)
def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather_dlo(
enable_dlo,
use_allgather,
field_name,
field_value,
):
with pytest.raises(ValueError, match="requires no-AllGather distributed layerwise offload"):
omni_config_module._DiffusionConfigProjection(
enable_distributed_layerwise_offload=enable_dlo,
dlo_use_allgather=use_allgather,
**{field_name: field_value},
)
def test_structured_diffusion_config_requires_cache_for_registration_ceiling():
with pytest.raises(ValueError, match="requires dlo_use_host_weight_cache=True"):
omni_config_module._DiffusionConfigProjection(
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_host_weight_cache_pin_limit_gib=1.0,
)
@pytest.mark.parametrize(
"field_name",
[
"dlo_host_weight_cache_lock_timeout",
"dlo_host_weight_cache_pin_limit_gib",
],
)
def test_structured_diffusion_config_rejects_non_finite_host_weight_cache_values(field_name):
with pytest.raises(ValueError, match=rf"{field_name} must be finite"):
omni_config_module._DiffusionConfigProjection(**{field_name: float("inf")})
def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]
@@ -12,6 +12,7 @@ import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
from vllm.config.load import LoadConfig
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
@@ -309,6 +310,7 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch
assert loader.load_model(load_device="cpu") is model
assert not loaded_ordinary_weights
assert plan.requires_backend
assert loader.take_host_weight_plan() is plan
assert loader.take_host_weight_plan() is None
@@ -346,7 +348,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
backing_kind="checkpoint_mmap",
bindings={
"transformer.weight": TensorBinding(
checkpoint_key="weight",
storage_key="weight",
file_path="unused",
)
},
@@ -374,8 +376,9 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
assert loader.take_host_weight_plan() is plan
def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
def test_dlo_plan_fallback_retains_ordinary_loader_weights_by_default(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
@@ -400,11 +403,202 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
)
def build_cache(*_args, **_kwargs):
pytest.fail("host weight cache must require explicit selection")
monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache)
assert loader.load_model(load_device="cpu") is model
assert calls == ["load", "process"]
assert loader.take_host_weight_plan() is None
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.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(
use_hsdp=False,
tensor_parallel_size=1,
sequence_parallel_size=1,
data_parallel_size=2,
cfg_parallel_size=1,
pipeline_parallel_size=1,
ulysses_degree=1,
ring_degree=1,
allgather_degree=1,
ulysses_mode="strict",
),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_use_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={},
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = UnquantizedLinearMethod()
events: list[str] = []
runtime_plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={},
runtime_layout_key="layout-key",
)
def load_weights(_model):
events.append("load")
with torch.no_grad():
model.transformer.weight.fill_(1)
def process_weights(_model, _device):
events.append("process")
with torch.no_grad():
model.transformer.weight.add_(1)
def calibrate(_model):
events.append("calibrate")
with torch.no_grad():
model.transformer.weight.add_(1)
original_eval = model.eval
def eval_model():
events.append("eval")
return original_eval()
def build_cache(_pipeline, **kwargs):
events.append("cache")
assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
assert kwargs["tensor_parallel_rank"] == 0
assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
assert kwargs["quantization_config"] is None
return HostWeightPlanResult(runtime_plan)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = load_weights # type: ignore[method-assign]
loader._process_weights_after_loading = process_weights # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = calibrate # type: ignore[method-assign]
model.eval = eval_model # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
)
monkeypatch.setattr(cache_mod, "build_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
@pytest.mark.parametrize("pin_cpu_memory", [True, False])
def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeypatch, tmp_path, pin_cpu_memory):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(
use_hsdp=False,
tensor_parallel_size=1,
sequence_parallel_size=1,
data_parallel_size=1,
cfg_parallel_size=1,
pipeline_parallel_size=1,
),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_use_host_weight_cache=True,
pin_cpu_memory=pin_cpu_memory,
dlo_host_weight_cache_dir=str(tmp_path),
dlo_host_weight_cache_pin_limit_gib=0.0,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
events: list[str] = []
runtime_plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={},
runtime_layout_key="layout-key",
)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
loader._process_weights_after_loading = lambda *_args: events.append("process") # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
def unexpected_checkpoint_plan(*_args, **_kwargs):
pytest.fail("selected 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_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_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.host_weight_cache as cache_mod
od_config = SimpleNamespace(
dtype=torch.float32,
parallel_config=SimpleNamespace(
use_hsdp=False,
tensor_parallel_size=1,
sequence_parallel_size=1,
data_parallel_size=2,
cfg_parallel_size=1,
pipeline_parallel_size=1,
),
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
dlo_use_host_weight_cache=True,
dlo_host_weight_cache_dir=str(tmp_path),
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
model = nn.Module()
model.transformer = nn.Linear(2, 2, bias=False)
model.transformer.quant_method = object()
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
loader.load_weights = lambda _model: None # type: ignore[method-assign]
loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign]
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
monkeypatch.setattr(
loader_mod,
"build_checkpoint_mmap_plan",
lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
)
def build_cache(_pipeline, **kwargs):
assert kwargs["quantization_config"] == "module_quant_method"
return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
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
assert torch.equal(model.transformer.weight, original_weight)
assert loader.take_host_weight_plan() is None
def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
from vllm.model_executor.layers.quantization.online.fp8 import (
Fp8PerTensorOnlineLinearMethod,
@@ -0,0 +1,260 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import fcntl
import multiprocessing as mp
import shutil
from pathlib import Path
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.model_loader.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]
_SPAWN_TIMEOUT_SECONDS = 120
class _LoaderMarker:
pass
class _Transformer(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(3, 2)
self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
self.register_buffer("temporary", torch.ones(1), persistent=False)
with torch.no_grad():
self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
class _Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = _Transformer()
def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
assert default_host_weight_cache_root().endswith("/.cache/vllm-omni/dlo-host-weights")
def _build(pipeline: nn.Module, cache_root: Path, **overrides):
kwargs = {
"dit_modules": (("transformer", pipeline.transformer),),
"loader_type": _LoaderMarker,
"cache_root": cache_root,
"lock_timeout_seconds": 2.0,
"max_shard_bytes": 32,
"model_identity": "test/model",
"revision": "revision",
"runtime_dtype": torch.float32,
"load_format": "default",
"loader_inputs": {"model_config": {"hidden_size": 3}},
"tensor_parallel_size": 1,
"tensor_parallel_rank": 0,
"sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
"use_hsdp": False,
"enable_expert_parallel": False,
"quantization_config": None,
"cfg_parallel_size": 1,
"pipeline_parallel_size": 1,
}
kwargs.update(overrides)
return build_host_weight_cache_plan(pipeline, **kwargs)
def _multiprocess_builder(cache_root: str, queue) -> None:
result = _build(_Pipeline(), Path(cache_root))
queue.put(
(
result.fallback_code,
result.plan.runtime_layout_key if result.plan is not None else None,
sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
)
)
def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
first = _build(_Pipeline(), tmp_path)
second = _build(_Pipeline(), tmp_path)
assert first.plan is not None
assert second.plan is not None
assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
assert not first.plan.requires_backend
assert {binding.file_path for binding in first.plan.bindings.values()} == {
binding.file_path for binding in second.plan.bindings.values()
}
assert set(first.plan.bindings) == {
"transformer.persistent",
"transformer.proj.bias",
"transformer.proj.weight",
}
def test_equivalent_local_model_paths_reuse_one_entry(tmp_path, monkeypatch):
model_dir = tmp_path / "model"
model_dir.mkdir()
model_link = tmp_path / "model-link"
model_link.symlink_to(model_dir, target_is_directory=True)
cache_root = tmp_path / "cache"
monkeypatch.setenv("HOME", str(tmp_path))
monkeypatch.chdir(tmp_path)
identities = (str(model_dir), "~/model", "model", str(model_link))
results = [_build(_Pipeline(), cache_root, model_identity=identity) for identity in identities]
assert all(result.plan is not None for result in results)
assert len({result.plan.runtime_layout_key for result in results if result.plan is not None}) == 1
assert len(list((cache_root / "v1").iterdir())) == 1
def test_unstable_loader_identity_fails_closed(tmp_path):
result = _build(_Pipeline(), tmp_path, loader_inputs={"opaque": object()})
assert result.plan is None
assert result.fallback_code == "unstable_identity"
assert result.fallback_reason == "host weight cache identity does not support values of type builtins.object"
assert list(tmp_path.rglob("*.safetensors")) == []
def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
pipeline = _Pipeline()
tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
sp2 = _build(
pipeline,
tmp_path,
sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
)
assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
def test_final_runtime_content_is_authoritative(tmp_path):
first_pipeline = _Pipeline()
second_pipeline = _Pipeline()
with torch.no_grad():
second_pipeline.transformer.proj.weight[0, 0] += 1
first = _build(first_pipeline, tmp_path)
second = _build(second_pipeline, tmp_path)
assert first.plan is not None and second.plan is not None
assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
def test_expert_parallel_ownership_fails_closed(tmp_path):
result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
assert result.plan is None
assert result.fallback_code == "unsupported_expert_parallel"
assert list(tmp_path.rglob("*.safetensors")) == []
def test_pipeline_parallel_ownership_fails_closed(tmp_path):
result = _build(_Pipeline(), tmp_path, pipeline_parallel_size=2)
assert result.plan is None
assert result.fallback_code == "unsupported_pp"
assert list(tmp_path.rglob("*.safetensors")) == []
def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
first = _build(_Pipeline(), tmp_path)
assert first.plan is not None
shard = Path(next(iter(first.plan.bindings.values())).file_path)
old_inode = shard.stat().st_ino
with shard.open("r+b") as handle:
handle.seek(-1, 2)
byte = handle.read(1)
handle.seek(-1, 2)
handle.write(bytes([byte[0] ^ 0xFF]))
second = _build(_Pipeline(), tmp_path)
assert second.plan is not None
assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
initial = _build(_Pipeline(), tmp_path)
assert initial.plan is not None
cache_key = initial.plan.runtime_layout_key
entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
shutil.rmtree(entry_dir)
lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
with lock_path.open("a+b") as lock_handle:
fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
assert result.plan is None
assert result.fallback_code == "lock_timeout"
def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
initial = _build(_Pipeline(), tmp_path)
assert initial.plan is not None
cache_key = initial.plan.runtime_layout_key
entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
shutil.rmtree(entry_dir)
stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
stale_temp.mkdir(parents=True)
(stale_temp / "partial").write_bytes(b"partial")
result = _build(_Pipeline(), tmp_path)
assert result.plan is not None
assert not stale_temp.exists()
@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
pipeline = _Pipeline()
if case == "alias":
pipeline.transformer.tied = pipeline.transformer.proj.weight
expected_code = "unsupported_alias"
elif case == "external_alias":
pipeline.encoder = nn.Module()
pipeline.encoder.shared = pipeline.transformer.proj.weight
expected_code = "unsupported_alias"
else:
pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
expected_code = "unsupported_tensor"
result = _build(pipeline, tmp_path)
assert result.plan is None
assert result.fallback_code == expected_code
assert list(tmp_path.rglob("*.safetensors")) == []
def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
context = mp.get_context("spawn")
queue = context.Queue()
processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
for process in processes:
process.start()
# Spawn workers import the full vLLM stack before entering the target;
# loaded CI hosts can spend more than 30 seconds in that startup path.
results = [queue.get(timeout=_SPAWN_TIMEOUT_SECONDS) for _ in processes]
for process in processes:
process.join(timeout=_SPAWN_TIMEOUT_SECONDS)
assert all(process.exitcode == 0 for process in processes)
assert all(code is None for code, _, _ in results)
assert len({key for _, key, _ in results}) == 1
assert len({tuple(paths) for _, _, paths in results}) == 1
assert len(list((tmp_path / "v1").iterdir())) == 1
@@ -0,0 +1,346 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import ctypes
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 (
CudaHostRegistration,
CudaHostRegistrationBudgetError,
CudaHostRegistrationCleanupError,
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]
class _FakeStorage:
def __init__(self, pointer: int, size: int) -> None:
self._pointer = pointer
self._size = size
def data_ptr(self) -> int:
return self._pointer
def nbytes(self) -> int:
return self._size
class _FakeTensor:
device = SimpleNamespace(type="cpu")
def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
self._storage = _FakeStorage(pointer, size)
self._pinned = pinned
def untyped_storage(self) -> _FakeStorage:
return self._storage
def numel(self) -> int:
return self._storage.nbytes()
def is_pinned(self) -> bool:
return self._pinned
class _FakeRuntime:
def __init__(
self,
register_results: list[int | Exception],
unregister_results: list[int | Exception] | None = None,
*,
read_only_registration_supported: bool = True,
attribute_result: int = 0,
) -> None:
self._register_results = iter(register_results)
self._unregister_results = iter(unregister_results or [])
self._read_only_registration_supported = read_only_registration_supported
self._attribute_result = attribute_result
self.attribute_queries: list[tuple[int, int]] = []
self.registered: list[tuple[int, int, int]] = []
self.unregistered: list[int] = []
def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int:
self.attribute_queries.append((attribute, device))
ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_registration_supported)
return self._attribute_result
def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
self.registered.append((pointer, size, flags))
result = next(self._register_results)
if isinstance(result, Exception):
raise result
return result
def cudaHostUnregister(self, pointer: int) -> int:
self.unregistered.append(pointer)
result = next(self._unregister_results, 0)
if isinstance(result, Exception):
raise result
return result
@staticmethod
def cudaGetErrorString(error: int) -> str:
return f"error-{error}"
@pytest.fixture(autouse=True)
def _current_cuda_device(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
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_platform_factory_allows_registration_without_additional_budget(monkeypatch: pytest.MonkeyPatch) -> None:
sentinel = object()
sources = {"weights": []}
def create(actual_sources, *, max_bytes):
assert actual_sources is sources
assert max_bytes is None
return sentinel
monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=None) is sentinel
def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
assert _coalesce_ranges(
[(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
page_size=4096,
) == (
registration_module._AddressRange(0x1000, 0x3000),
registration_module._AddressRange(0x9000, 0xA000),
)
def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
CudaHostRegistration.create(
{"weights": [_FakeTensor(0x1003, 4096)]}, # type: ignore[list-item]
max_bytes=4096,
)
assert runtime.registered == []
def test_registration_without_additional_budget_registers_complete_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
registration = CudaHostRegistration.create(
{"weights": [_FakeTensor(0x1003, 4096)]}, # type: ignore[list-item]
max_bytes=None,
)
assert registration.total_bytes == 8192
assert len(runtime.registered) == 1
assert registration.close() == []
def test_registration_rejects_unsupported_read_only_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0], read_only_registration_supported=False)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
with pytest.raises(CudaHostRegistrationError, match="does not support read-only host registration"):
CudaHostRegistration.create(
{"weights": [_FakeTensor(0x1003, 1)]}, # type: ignore[list-item]
max_bytes=4096,
)
assert runtime.attribute_queries == [
(registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
]
assert runtime.registered == []
def test_registration_consumes_capability_query_error(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0], attribute_result=7)
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
consumed_errors: list[int] = []
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
with pytest.raises(CudaHostRegistrationError, match="error-7"):
CudaHostRegistration.create(
{"weights": [_FakeTensor(0x1003, 1)]}, # type: ignore[list-item]
max_bytes=4096,
)
assert consumed_errors == [7]
assert runtime.registered == []
def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0, 7])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
consumed_errors: list[int] = []
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
with pytest.raises(CudaHostRegistrationError, match="error-7"):
CudaHostRegistration.create(
{
"first": [_FakeTensor(0x1003, 1)], # type: ignore[list-item]
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
},
max_bytes=8192,
)
assert runtime.unregistered == [0x1000]
assert consumed_errors == [7]
def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
CudaHostRegistration.create(
{
"first": [_FakeTensor(0x1003, 1)], # type: ignore[list-item]
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
},
max_bytes=8192,
)
assert runtime.unregistered == [0x1000]
def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0, 7], unregister_results=[9])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
consumed_errors: list[int] = []
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
CudaHostRegistration.create(
{
"first": [_FakeTensor(0x1003, 1)], # type: ignore[list-item]
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
},
max_bytes=8192,
)
assert consumed_errors == [7, 9]
@pytest.mark.parametrize("pending_error", [0, 801])
def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
monkeypatch: pytest.MonkeyPatch,
pending_error: int,
) -> None:
class _GetLastError:
def __call__(self) -> int:
return pending_error
runtime = SimpleNamespace(cudaGetLastError=_GetLastError())
monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime)
registration_module._consume_last_cuda_error(801)
def test_consume_last_cuda_error_rejects_unrelated_pending_error(monkeypatch: pytest.MonkeyPatch) -> None:
class _GetLastError:
def __call__(self) -> int:
return 700
runtime = SimpleNamespace(cudaGetLastError=_GetLastError())
monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime)
with pytest.raises(CudaHostRegistrationCleanupError, match="cudaGetLastError reported 700"):
registration_module._consume_last_cuda_error(801)
def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0, 0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
registration = CudaHostRegistration.create(
{
"first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)], # type: ignore[list-item]
"second": [_FakeTensor(0x9003, 1)], # type: ignore[list-item]
},
max_bytes=12288,
)
assert registration.total_bytes == 12288
assert registration.region_count == 2
assert registration.close() == []
assert runtime.unregistered == [0x9000, 0x1000]
assert registration.close() == []
def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0, 0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
registration = CudaHostRegistration.create(
{
"first.safetensors": [_FakeTensor(0x1000, 4096)], # type: ignore[list-item]
"second.safetensors": [_FakeTensor(0x2000, 4096)], # type: ignore[list-item]
},
max_bytes=8192,
)
assert registration.region_count == 2
assert runtime.registered == [
(0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
(0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
]
assert registration.close() == []
def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0], unregister_results=[9, 0])
monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
consumed_errors: list[int] = []
monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append)
registration = CudaHostRegistration.create(
{"weights": [_FakeTensor(0x1003, 1)]}, # type: ignore[list-item]
max_bytes=4096,
)
assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
assert registration.close() == []
assert runtime.unregistered == [0x1000, 0x1000]
assert consumed_errors == [9]
@@ -18,6 +18,7 @@ from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
from vllm_omni.diffusion.model_loader.host_weight_cache_utils import file_digest
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
TensorBinding,
@@ -546,6 +547,31 @@ class TestDistributedLayerwiseOffloadHook:
assert staging[0][torch.float32].numel() == 12
assert staging[1][torch.float32].numel() == 12
def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
current_block = nn.Linear(2, 2, bias=False)
next_block = nn.Linear(2, 2, bias=False)
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
next_block.weight.data.copy_(expected)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
dp_group=None,
dp_size=1,
rank=0,
pin_memory=False,
rank_local_mmap=True,
)
hook.initialize_hook(current_block)
hook.registered_mmap = True
def fail_staging(_slot):
raise AssertionError("registered mmap must bypass host staging")
hook._stage_mmap_sources = fail_staging # type: ignore[method-assign]
hook.prefetch_layer(slot=0, non_blocking=True)
assert torch.equal(next_block.weight, expected)
class TestPinnedResidentLayerGroup:
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -636,6 +662,24 @@ class TestPinnedResidentLayerGroup:
group.offload()
assert all(block.weight.numel() == 0 for block in blocks)
def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
block = nn.Linear(2, 2, bias=False)
expected = torch.arange(4, dtype=torch.float32).view(2, 2)
block.weight.data.copy_(expected)
group = PinnedResidentLayerGroup(
[block],
device=torch.device("cpu"),
copy_stream=DummyStream(),
pin_memory=False,
rank_local_mmap=True,
)
group.registered_mmap = True
group._cpu_staging_buffers.clear()
group.load()
assert torch.equal(block.weight, expected)
def test_shared_allgather_output_is_narrowed_to_current_block(
self, dist_group, patched_offload_runtime, monkeypatch
):
@@ -765,6 +809,140 @@ class _MmapPostLoadPipeline(nn.Module):
class TestMmapWeightLoading:
def test_host_weight_cache_maps_final_weights_without_rerunning_post_load(
self,
tmp_path,
patched_offload_runtime,
):
class Transformer(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(2, 2, bias=False)
self.post_load_calls = 0
def post_load_weights(self):
self.post_load_calls += 1
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
pipeline = Pipeline()
old_parameter = pipeline.transformer.proj.weight
old_parameter.runtime_marker = "preserved"
weight_file = tmp_path / "runtime.safetensors"
save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={
"transformer.proj.weight": TensorBinding(
storage_key="transformer.proj.weight",
file_path=str(weight_file),
)
},
runtime_layout_key="runtime-layout",
expected_file_digests={str(weight_file): file_digest(weight_file)},
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
),
torch.device("cpu"),
host_weight_plan=plan,
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
assert pipeline.transformer.post_load_calls == 0
assert pipeline.transformer.proj.weight is old_parameter
assert pipeline.transformer.proj.weight.runtime_marker == "preserved"
assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
backend._release_mmap_handles()
def test_host_weight_cache_mapping_failure_leaves_ordinary_tensors_intact(
self,
tmp_path,
patched_offload_runtime,
):
pipeline = nn.Module()
pipeline.transformer = nn.Linear(2, 2, bias=False)
old_parameter = pipeline.transformer.weight
plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={
"transformer.weight": TensorBinding(
storage_key="transformer.weight",
file_path=str(tmp_path / "missing.safetensors"),
)
},
runtime_layout_key="runtime-layout",
expected_file_digests={str(tmp_path / "missing.safetensors"): "unused"},
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
),
torch.device("cpu"),
host_weight_plan=plan,
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
with pytest.raises(Exception, match="missing.safetensors"):
backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
assert pipeline.transformer.weight is old_parameter
def test_host_weight_cache_digest_change_leaves_ordinary_tensors_intact(
self,
tmp_path,
patched_offload_runtime,
):
pipeline = nn.Module()
pipeline.transformer = nn.Linear(2, 2, bias=False)
old_parameter = pipeline.transformer.weight
old_value = old_parameter.detach().clone()
weight_file = tmp_path / "runtime.safetensors"
save_file({"transformer.weight": torch.full((2, 2), 9.0)}, str(weight_file))
expected_digest = file_digest(weight_file)
with weight_file.open("r+b") as handle:
handle.seek(-1, os.SEEK_END)
original = handle.read(1)
handle.seek(-1, os.SEEK_END)
handle.write(bytes([original[0] ^ 0xFF]))
plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={
"transformer.weight": TensorBinding(
storage_key="transformer.weight",
file_path=str(weight_file),
)
},
runtime_layout_key="runtime-layout",
expected_file_digests={str(weight_file): expected_digest},
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
),
torch.device("cpu"),
host_weight_plan=plan,
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
with pytest.raises(RuntimeError, match="changed after loader validation"):
backend._load_weights_from_host_weight_cache(pipeline, modules, plan)
assert pipeline.transformer.weight is old_parameter
assert torch.equal(pipeline.transformer.weight, old_value)
def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
pipeline = _MmapPostLoadPipeline()
weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -788,7 +966,7 @@ class TestMmapWeightLoading:
backing_kind="checkpoint_mmap",
bindings={
name: TensorBinding(
checkpoint_key=name,
storage_key=name,
file_path=str(weight_file),
)
for name in weights
@@ -859,6 +1037,292 @@ class TestMmapWeightLoading:
assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
backend.disable()
def test_host_weight_cache_enable_reclaims_private_weights_before_hook_setup(
self,
tmp_path,
patched_offload_runtime,
):
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
pipeline = Pipeline()
weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
weight_file = tmp_path / "runtime.safetensors"
save_file(weights, str(weight_file))
plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
runtime_layout_key="runtime-layout",
expected_file_digests={str(weight_file): file_digest(weight_file)},
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
),
torch.device("cpu"),
host_weight_plan=plan,
)
cleanup_calls: list[None] = []
backend._cleanup_after_loading = lambda: cleanup_calls.append(None) # type: ignore[method-assign]
backend.enable(pipeline)
assert backend._using_rank_local_mmap
# One reclamation immediately after remap, plus the existing final
# post-hook cleanup.
assert len(cleanup_calls) == 2
backend.disable()
def test_host_weight_cache_enable_falls_back_to_ordinary_tensors(
self,
tmp_path,
patched_offload_runtime,
):
class Transformer(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
def __init__(self):
super().__init__()
self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
class Pipeline(nn.Module):
def __init__(self):
super().__init__()
self.transformer = Transformer()
pipeline = Pipeline()
weight_file = tmp_path / "runtime.safetensors"
weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
save_file(weights, str(weight_file))
plan = HostWeightPlan(
backing_kind="host_weight_cache",
bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
runtime_layout_key="runtime-layout",
expected_file_digests={str(weight_file): "stale-digest"},
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
),
torch.device("cpu"),
host_weight_plan=plan,
)
backend.enable(pipeline)
assert backend.enabled
assert backend.host_weight_plan is None
assert not backend._using_mmap
assert not backend._using_rank_local_mmap
backend.disable()
def test_disable_synchronizes_allgather_device_work(
self,
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
synchronize_calls: list[None] = []
monkeypatch.setattr(
dist_backend_module.current_omni_platform,
"synchronize",
lambda: synchronize_calls.append(None),
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=True,
),
torch.device("cpu"),
)
backend.enabled = True
backend._using_rank_local_mmap = False
backend.disable()
assert synchronize_calls == [None]
def test_host_weight_cache_registration_uses_budget_and_releases_before_mmap(
self,
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
events: list[str] = []
class Registration:
total_bytes = 4096
region_count = 1
@staticmethod
def close() -> list[str]:
events.append("unregister")
return []
create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
def register(sources, *, device, max_bytes):
assert device == torch.device("cuda")
create_calls.append((sources, max_bytes))
return Registration()
monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dlo_use_allgather=False,
dlo_host_weight_cache_pin_limit_gib=1.5,
),
torch.device("cuda"),
)
sources = {"runtime.safetensors": [torch.ones(1)]}
backend._host_weight_cache_mapped_sources = sources
assert backend._try_register_host_weight_cache_mmap()
assert create_calls == [(sources, int(1.5 * 1024**3))]
backend.enabled = True
backend._using_rank_local_mmap = True
backend._mmap_file_cache = {}
backend._release_mmap_handles = lambda: events.append("mmap") # type: ignore[method-assign]
backend.disable()
assert events == ["unregister", "mmap"]
def test_host_weight_cache_registration_defaults_to_pinned_memory_policy(
self,
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
class Registration:
total_bytes = 4096
region_count = 1
@staticmethod
def close() -> list[str]:
return []
registration = Registration()
create_calls: list[tuple[dict[str, list[torch.Tensor]], int | None]] = []
def register(sources, *, device, max_bytes):
assert device == torch.device("cuda")
create_calls.append((sources, max_bytes))
return registration
monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dlo_use_allgather=False,
dlo_host_weight_cache_pin_limit_gib=0.0,
),
torch.device("cuda"),
)
sources = {"runtime.safetensors": [torch.ones(1)]}
backend._host_weight_cache_mapped_sources = sources
assert backend._try_register_host_weight_cache_mmap()
assert create_calls == [(sources, None)]
assert backend._host_registration is registration
def test_host_weight_cache_registration_respects_disabled_pinned_memory(
self,
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
monkeypatch.setattr(
dist_backend_module,
"register_host_mappings",
lambda *_args, **_kwargs: pytest.fail("registration must follow pin_cpu_memory=False"),
)
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
dlo_use_allgather=False,
dlo_host_weight_cache_pin_limit_gib=0.0,
),
torch.device("cuda"),
)
backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
assert not backend._try_register_host_weight_cache_mmap()
assert backend._host_registration is None
def test_host_weight_cache_transfer_is_configured_once_for_all_consumers(
self,
monkeypatch: pytest.MonkeyPatch,
patched_offload_runtime,
):
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=True,
dlo_use_allgather=False,
),
torch.device("cuda"),
host_weight_plan=HostWeightPlan(
backing_kind="host_weight_cache",
bindings={},
),
)
backend._using_rank_local_mmap = True
backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
streamed_hook = SimpleNamespace(registered_mmap=False)
resident_group = SimpleNamespace(
registered_mmap=False,
_cpu_staging_buffers=[{"buffer": torch.ones(1)}],
)
backend._resident_layer_group = resident_group
register_calls: list[None] = []
def register() -> bool:
register_calls.append(None)
return True
monkeypatch.setattr(backend, "_try_register_host_weight_cache_mmap", register)
backend._configure_host_weight_cache_transfer([streamed_hook]) # type: ignore[list-item]
assert register_calls == [None]
assert streamed_hook.registered_mmap
assert resident_group.registered_mmap
assert resident_group._cpu_staging_buffers == []
assert backend._host_weight_cache_mapped_sources == {}
def test_host_weight_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_host_weight_cache_pin_limit_gib=1.0,
),
torch.device("cpu"),
)
backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
assert not backend._try_register_host_weight_cache_mmap()
assert backend._host_registration is None
class TestGetBlocksFromDit:
def test_get_blocks_from_dit_single_block_attr(self):
@@ -1164,7 +1628,7 @@ class TestMmapValidation:
assert result.fallback_reason is None
assert result.plan is not None
assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
checkpoint_key="blocks.0.weight",
storage_key="blocks.0.weight",
file_path=str(checkpoint_file),
)
assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1454,7 +1918,7 @@ class TestMmapValidation:
assert result.plan is not None
assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
checkpoint_key="proj_in.weight",
storage_key="proj_in.weight",
file_path=str(checkpoint_file),
)
@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
- Dynamic worker class extension
"""
from types import SimpleNamespace
from typing import Any
import pytest
from pytest_mock import MockerFixture
import vllm_omni.diffusion.worker.diffusion_worker as worker_module
from vllm_omni.diffusion.worker.diffusion_worker import (
CustomPipelineWorkerExtension,
DiffusionWorker,
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
wrapper.worker.shutdown.assert_called_once()
assert result is None
def test_worker_shutdown_disables_offloader_before_distributed_teardown(
self,
mocker: MockerFixture,
monkeypatch: pytest.MonkeyPatch,
):
worker = object.__new__(DiffusionWorker)
offload_backend = mocker.Mock()
kv_transfer_manager = mocker.Mock()
worker.model_runner = SimpleNamespace(
offload_backend=offload_backend,
kv_transfer_manager=kv_transfer_manager,
)
events: list[str] = []
offload_backend.disable.side_effect = lambda: events.append("offloader")
kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
worker.shutdown()
offload_backend.disable.assert_called_once_with()
kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
assert events == ["offloader", "kv_transfer", "distributed"]
# -------------------------------------------------------------------------
# Tests: execute_method
@@ -328,6 +328,9 @@ def test_serve_cli_forwards_distributed_offload_residency():
"--omni",
"--enable-distributed-layerwise-offload",
"--dlo-no-use-allgather",
"--dlo-use-host-weight-cache",
"--dlo-host-weight-cache-pin-limit-gib",
"80",
"--dlo-resident-layers",
"20",
]
@@ -339,9 +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_use_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_use_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
+19
View File
@@ -11,6 +11,7 @@ later PRs cut consumers over to these classes.
from __future__ import annotations
import copy
import math
from collections.abc import Mapping
from dataclasses import InitVar, dataclass, field, fields
from functools import wraps
@@ -687,6 +688,10 @@ class _DiffusionConfigProjection:
enable_layerwise_offload: bool = False
enable_distributed_layerwise_offload: bool = False
dlo_use_allgather: bool = True
dlo_use_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"
@@ -749,6 +754,20 @@ class _DiffusionConfigProjection:
from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
from vllm_omni.quantization import build_quant_config
if not math.isfinite(self.dlo_host_weight_cache_lock_timeout):
raise ValueError("dlo_host_weight_cache_lock_timeout must be finite")
if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib):
raise ValueError("dlo_host_weight_cache_pin_limit_gib must be finite")
if self.dlo_host_weight_cache_pin_limit_gib and (
not self.enable_distributed_layerwise_offload or self.dlo_use_allgather
):
raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True")
if self.tf_model_config is None:
self.tf_model_config = TransformerConfig()
elif isinstance(self.tf_model_config, Mapping):
+4
View File
@@ -382,6 +382,10 @@ class StageDeployConfig:
enable_distributed_layerwise_offload: bool | None = None
dlo_use_allgather: bool | None = None
dlo_use_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
+22
View File
@@ -708,6 +708,16 @@ class OmniDiffusionConfig:
# This avoids AllGather synchronization, while host memory follows the
# loader's existing rank-local layout instead of adding a second DP shard.
dlo_use_allgather: bool = True
# Opt into publishing/joining the transform-complete final-layout host
# weight cache instead of direct-checkpoint mmap.
dlo_use_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_host_weight_cache_lock_timeout: float = 600.0
# Optional maximum file-backed memory a worker may register with CUDA.
# Zero applies no additional ceiling; pin_cpu_memory controls registration.
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
@@ -936,6 +946,18 @@ class OmniDiffusionConfig:
)
def __post_init__(self):
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_host_weight_cache_pin_limit_gib and (
not self.enable_distributed_layerwise_offload or self.dlo_use_allgather
):
raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True")
if self.diffusion_compile_granularity not in {"regional", "full"}:
raise ValueError(
"diffusion_compile_granularity must be 'regional' or 'full', "
@@ -14,6 +14,7 @@ import torch
from torch import nn
from vllm.config.load import LoadConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.model_loader.weight_utils import (
download_safetensors_index_file_from_hf,
@@ -440,10 +441,13 @@ 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
_select_host_weight_cache = (
_dist_offload and not _use_ag and bool(getattr(self.od_config, "dlo_use_host_weight_cache", False))
)
plan_result = None
weight_sources = self._get_weight_sources(model)
if _dist_offload:
if _dist_offload and not _select_host_weight_cache:
modules = ModuleDiscovery.discover(model)
plan_result = build_checkpoint_mmap_plan(
model,
@@ -455,6 +459,16 @@ class DiffusersPipelineLoader:
online_quantization=_has_online_quant,
)
self.host_weight_plan = plan_result.plan
elif _select_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 host weight cache selected; using the ordinary loader "
"once to publish a transform-complete final mmap layout"
)
_skip_load = self.host_weight_plan is not None
@@ -496,7 +510,8 @@ class DiffusersPipelineLoader:
)
if _dist_offload and plan_result is not None:
logger.info(
"DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
"DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
plan_result.fallback_code or "unknown",
plan_result.fallback_reason,
)
logger.debug("Loading weights on %s ...", load_device)
@@ -521,7 +536,110 @@ class DiffusersPipelineLoader:
logger.info("Quantization complete, offloaded model back to CPU")
self._apply_skip_softmax_calibration(model)
return model.eval()
model.eval()
self._maybe_build_host_weight_cache(model, load_format)
return model
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 getattr(self.od_config, "dlo_use_allgather", True):
return
if not getattr(self.od_config, "dlo_use_host_weight_cache", False):
logger.info(
"DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback"
)
return
parallel_config = self.parallel_config
tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1))
tp_rank = 0
if tp_size > 1:
try:
from vllm.distributed import get_tensor_model_parallel_rank
tp_rank = int(get_tensor_model_parallel_rank())
except Exception as exc:
logger.warning(
"DLO 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.host_weight_cache import (
DEFAULT_LOCK_TIMEOUT_SECONDS,
DEFAULT_SHARD_SIZE_BYTES,
build_host_weight_cache_plan,
)
loader_inputs = {
"model_class_name": getattr(self.od_config, "model_class_name", None),
"model_config": getattr(self.od_config, "model_config", None),
"diffusers_load_kwargs": getattr(self.od_config, "diffusers_load_kwargs", None),
"custom_pipeline_args": getattr(self.od_config, "custom_pipeline_args", None),
"lora_path": getattr(self.od_config, "lora_path", None),
"lora_scale": getattr(self.od_config, "lora_scale", None),
"lora_backend": getattr(self.od_config, "lora_backend", None),
}
sp_guard = {
"sequence_parallel_size": int(getattr(parallel_config, "sequence_parallel_size", 1)),
"ulysses_degree": int(getattr(parallel_config, "ulysses_degree", 1)),
"ring_degree": int(getattr(parallel_config, "ring_degree", 1)),
"allgather_degree": int(getattr(parallel_config, "allgather_degree", 1)),
"ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"),
}
model_identity = getattr(self.od_config, "model", None)
runtime_quantization = self.quant_config
if runtime_quantization is None and any(
(quant_method := getattr(module, "quant_method", None)) is not None
and not isinstance(quant_method, UnquantizedLinearMethod)
for module in model.modules()
):
runtime_quantization = "module_quant_method"
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_host_weight_cache_dir", None),
lock_timeout_seconds=float(
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,
revision=getattr(self.od_config, "revision", None),
runtime_dtype=getattr(self.od_config, "dtype", None),
load_format=load_format,
loader_inputs=loader_inputs,
tensor_parallel_size=tp_size,
tensor_parallel_rank=tp_rank,
sequence_parallel_guard=sp_guard,
use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)),
enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)),
quantization_config=runtime_quantization,
cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)),
pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)),
)
if result.plan is not None:
self.host_weight_plan = result.plan
logger.info(
"DLO 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,
sp_guard["sequence_parallel_size"],
)
else:
logger.warning(
"DLO host weight cache unavailable [%s]; retaining ordinary host tensors: %s",
result.fallback_code or "unknown",
result.fallback_reason or "unspecified reason",
)
@staticmethod
def _request_offload_after_quant(model: nn.Module) -> int:
@@ -0,0 +1,574 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Build node-local host weight cache plans for diffusion DLO.
``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:`host_weight_cache_utils`.
"""
from __future__ import annotations
import contextlib
import hashlib
import json
import math
import os
import shutil
import time
import uuid
from collections.abc import Mapping, Sequence
from itertools import chain
from pathlib import Path
from typing import Any
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.model_loader.host_weight_cache_utils import (
HostWeightCacheError,
IdentityNormalizationError,
RuntimeTensor,
canonical_json,
canonicalize_existing_local_path,
exclusive_lock,
file_digest,
fsync_directory,
fsync_file,
implementation_fingerprint,
normalize_identity,
remove_stale_temps,
runtime_content_digest,
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.
HOST_WEIGHT_CACHE_SCHEMA_VERSION = 1
# Keep the default shared across local stage replicas.
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.
DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
_SUPPORTED_DTYPES = {
torch.bool,
torch.uint8,
torch.int8,
torch.int16,
torch.int32,
torch.int64,
torch.float16,
torch.bfloat16,
torch.float32,
torch.float64,
}
for _dtype_name in ("float8_e4m3fn", "float8_e5m2"):
if (_dtype := getattr(torch, _dtype_name, None)) is not None:
_SUPPORTED_DTYPES.add(_dtype)
def default_host_weight_cache_root() -> str:
"""Return the default root, intentionally independent of VLLM_CACHE_ROOT."""
return os.path.expanduser(DEFAULT_HOST_WEIGHT_CACHE_ROOT)
def _persistent_buffer(module: nn.Module, local_name: str) -> bool:
parent_path, _, leaf_name = local_name.rpartition(".")
owner = module.get_submodule(parent_path)
return leaf_name not in owner._non_persistent_buffers_set
def _named_parameters_with_duplicates(module: nn.Module):
try:
return module.named_parameters(remove_duplicate=False)
except TypeError: # pragma: no cover - compatibility with old torch
return module.named_parameters()
def _named_buffers_with_duplicates(module: nn.Module):
try:
return module.named_buffers(remove_duplicate=False)
except TypeError: # pragma: no cover - compatibility with old torch
return module.named_buffers()
def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None:
parent_path, _, leaf_name = runtime_name.rpartition(".")
parent = pipeline.get_submodule(parent_path)
tensor = parent._parameters.get(leaf_name)
if tensor is None:
tensor = parent._buffers.get(leaf_name)
return tensor
def _collect_runtime_tensors(
pipeline: nn.Module,
dit_modules: Sequence[tuple[str, nn.Module]],
) -> list[RuntimeTensor]:
"""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):
runtime_name = f"{dit_name}.{local_name}"
candidate = RuntimeTensor(runtime_name, tensor, "parameter")
existing = records.get(runtime_name)
if existing is not None and existing.tensor is not tensor:
raise HostWeightCacheError(
"ambiguous_ownership",
f"multiple DiT tensors resolve to {runtime_name!r}",
)
records[runtime_name] = candidate
for local_name, tensor in _named_buffers_with_duplicates(dit_module):
if not _persistent_buffer(dit_module, local_name):
continue
runtime_name = f"{dit_name}.{local_name}"
candidate = RuntimeTensor(runtime_name, tensor, "buffer")
existing = records.get(runtime_name)
if existing is not None and existing.tensor is not tensor:
raise HostWeightCacheError(
"ambiguous_ownership",
f"multiple DiT tensors resolve to {runtime_name!r}",
)
records[runtime_name] = candidate
if not records:
raise HostWeightCacheError(
"ambiguous_ownership",
"no DiT parameters or persistent buffers were discovered",
)
storage_owners: dict[tuple[int, int], str] = {}
for record in records.values():
tensor = record.tensor
try:
pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
except AttributeError as exc:
raise HostWeightCacheError(
"ambiguous_ownership",
f"DiT tensor {record.name!r} is not owned by the pipeline",
) from exc
if pipeline_tensor is not tensor:
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 HostWeightCacheError(
"unsupported_tensor",
f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
)
if tensor.layout != torch.strided:
raise HostWeightCacheError(
"unsupported_tensor",
f"{record.name!r} uses unsupported layout {tensor.layout}",
)
if not tensor.is_contiguous():
raise HostWeightCacheError(
"unsupported_tensor",
f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
)
if tensor.dtype not in _SUPPORTED_DTYPES:
raise HostWeightCacheError(
"unsupported_tensor",
f"{record.name!r} uses unsupported dtype {tensor.dtype}",
)
if hasattr(tensor, "to_local"):
raise HostWeightCacheError(
"unsupported_tensor",
f"{record.name!r} is a distributed tensor",
)
tensor_nbytes = tensor.numel() * tensor.element_size()
storage = tensor.untyped_storage()
if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes:
raise HostWeightCacheError(
"unsupported_tensor",
f"{record.name!r} is a view into a larger storage",
)
if tensor_nbytes:
storage_id = (storage.data_ptr(), storage.nbytes())
owner = storage_owners.setdefault(storage_id, record.name)
if owner != record.name:
raise HostWeightCacheError(
"unsupported_alias",
f"{record.name!r} shares storage with {owner!r}",
)
# The ownership boundary is the complete pipeline, not only the discovered
# DiT traversal. Reject a DiT storage that is also registered by an encoder,
# VAE, resident component, or non-persistent buffer.
for pipeline_name, tensor in chain(
_named_parameters_with_duplicates(pipeline),
_named_buffers_with_duplicates(pipeline),
):
if tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0:
continue
storage = tensor.untyped_storage()
owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
if owner is not None and owner != pipeline_name:
raise HostWeightCacheError(
"unsupported_alias",
f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}",
)
return [records[name] for name in sorted(records)]
def _publish_entry(
*,
cache_root: Path,
entry_dir: Path,
cache_key: str,
layout_identity: dict[str, Any],
content_digest: str,
records: Sequence[RuntimeTensor],
max_shard_bytes: int,
) -> None:
total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records)
free_bytes = shutil.disk_usage(cache_root).free
required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100)
if free_bytes < required_bytes:
raise HostWeightCacheError(
"insufficient_disk",
f"host weight cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}",
)
tmp_root = cache_root / ".tmp"
tmp_root.mkdir(parents=True, exist_ok=True)
temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}"
temp_dir.mkdir()
try:
shards = split_shards(records, max_shard_bytes)
tensor_manifest: dict[str, dict[str, Any]] = {}
file_manifest: dict[str, dict[str, Any]] = {}
shard_count = len(shards)
for index, shard in enumerate(shards, start=1):
filename = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
path = temp_dir / filename
save_file(
{record.name: record.tensor.detach() for record in shard},
str(path),
metadata={"format": "pt", "dlo_host_weight_cache_schema": str(HOST_WEIGHT_CACHE_SCHEMA_VERSION)},
)
fsync_file(path)
file_manifest[filename] = {
"size": path.stat().st_size,
"sha256": file_digest(path),
}
for record in shard:
tensor_manifest[record.name] = {
**tensor_metadata(record),
"file": filename,
"storage_key": record.name,
}
manifest = {
"schema_version": HOST_WEIGHT_CACHE_SCHEMA_VERSION,
"cache_key": cache_key,
"layout_identity": layout_identity,
"runtime_content_sha256": content_digest,
"total_tensor_bytes": total_nbytes,
"created_at_unix_seconds": time.time(),
"writer_pid": os.getpid(),
"files": file_manifest,
"tensors": tensor_manifest,
}
manifest_path = temp_dir / "manifest.json"
with manifest_path.open("w", encoding="utf-8") as handle:
json.dump(manifest, handle, sort_keys=True, separators=(",", ":"))
handle.flush()
os.fsync(handle.fileno())
fsync_directory(temp_dir)
invalid_dir: Path | None = None
if entry_dir.exists():
invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}"
os.replace(entry_dir, invalid_dir)
os.replace(temp_dir, entry_dir)
fsync_directory(entry_dir.parent)
if invalid_dir is not None:
shutil.rmtree(invalid_dir, ignore_errors=True)
except Exception:
shutil.rmtree(temp_dir, ignore_errors=True)
raise
def _validate_entry(
*,
entry_dir: Path,
cache_key: str,
layout_identity: dict[str, Any],
content_digest: str,
expected_records: Sequence[RuntimeTensor],
) -> HostWeightPlan:
manifest_path = entry_dir / "manifest.json"
try:
with manifest_path.open(encoding="utf-8") as handle:
manifest = json.load(handle)
except (OSError, ValueError) as exc:
raise HostWeightCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
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 HostWeightCacheError("entry_rejected", "host weight cache key mismatch")
if manifest.get("layout_identity") != layout_identity:
raise HostWeightCacheError("entry_rejected", "host weight cache layout identity mismatch")
if manifest.get("runtime_content_sha256") != content_digest:
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 HostWeightCacheError("entry_rejected", "host weight cache manifest has invalid sections")
if set(tensor_manifest) != set(expected_metadata):
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 HostWeightCacheError("entry_rejected", f"invalid cache filename {filename!r}")
path = entry_dir / filename
try:
stat = path.stat()
except OSError as 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 HostWeightCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
bindings: dict[str, TensorBinding] = {}
mapped_records: list[RuntimeTensor] = []
with contextlib.ExitStack() as stack:
handles = {
filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu"))
for filename in file_manifest
}
expected_keys_by_file: dict[str, set[str]] = {filename: set() for filename in file_manifest}
for record in expected_records:
stored = tensor_manifest.get(record.name)
if not isinstance(stored, dict):
raise 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 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 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 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 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 HostWeightCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
if runtime_content_digest(mapped_records) != content_digest:
raise HostWeightCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
return HostWeightPlan(
backing_kind="host_weight_cache",
bindings=bindings,
runtime_layout_key=cache_key,
expected_file_digests={
str(entry_dir / filename): metadata["sha256"] for filename, metadata in file_manifest.items()
},
)
def build_host_weight_cache_plan(
pipeline: nn.Module,
*,
dit_modules: Sequence[tuple[str, nn.Module]],
loader_type: type,
cache_root: str | os.PathLike[str] | None,
lock_timeout_seconds: float,
max_shard_bytes: int,
model_identity: str | None,
revision: str | None,
runtime_dtype: Any,
load_format: str,
loader_inputs: Mapping[str, Any],
tensor_parallel_size: int,
tensor_parallel_rank: int,
sequence_parallel_guard: Mapping[str, Any],
use_hsdp: bool,
enable_expert_parallel: bool,
quantization_config: Any,
cfg_parallel_size: int,
pipeline_parallel_size: int,
) -> HostWeightPlanResult:
"""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
a content-addressed identity from their metadata, bytes, loader inputs, and
parallelism guards. A valid existing entry is joined immediately; otherwise
one process publishes the entry under a per-identity file lock and peers
validate it after the lock is released.
Success returns bindings to immutable safetensors mmap files. Any expected
incompatibility, lock, publication, or validation failure returns a
fail-closed result with no plan and a stable fallback code, leaving the
already-loaded pipeline tensors intact.
"""
if load_format == "dummy":
return HostWeightPlanResult(
None,
"dummy/random weights do not have a reusable runtime identity",
"unsupported_load_format",
)
if use_hsdp:
return HostWeightPlanResult(
None, "HSDP/DTensor layouts are not host weight cache compatible", "unsupported_hsdp"
)
if enable_expert_parallel:
return HostWeightPlanResult(
None,
"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 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 host weight cache v1", "unsupported_pp"
)
if not math.isfinite(lock_timeout_seconds) or lock_timeout_seconds <= 0:
return HostWeightPlanResult(None, "host weight cache lock timeout must be positive", "invalid_configuration")
if max_shard_bytes <= 0:
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)
# The cache stores loader-produced ownership: TP coordinates differ,
# while equivalent DP replicas and SP ranks may share one entry. The
# conservative SP guard prevents reuse across unproven implementations.
layout_identity = {
"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),
"load_format": load_format,
"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": normalize_identity(sequence_parallel_guard),
"runtime_content_sha256": content_digest,
}
cache_key = hashlib.sha256(canonical_json(layout_identity)).hexdigest()
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)
if entry_dir.is_dir():
try:
plan = _validate_entry(
entry_dir=entry_dir,
cache_key=cache_key,
layout_identity=layout_identity,
content_digest=content_digest,
expected_records=records,
)
logger.info("Reusing DLO host weight cache entry %s from %s", cache_key[:12], entry_dir)
return HostWeightPlanResult(plan)
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):
remove_stale_temps(root, cache_key)
if entry_dir.is_dir():
try:
plan = _validate_entry(
entry_dir=entry_dir,
cache_key=cache_key,
layout_identity=layout_identity,
content_digest=content_digest,
expected_records=records,
)
logger.info("Joined DLO host weight cache entry %s from %s", cache_key[:12], entry_dir)
return HostWeightPlanResult(plan)
except HostWeightCacheError:
pass
_publish_entry(
cache_root=root,
entry_dir=entry_dir,
cache_key=cache_key,
layout_identity=layout_identity,
content_digest=content_digest,
records=records,
max_shard_bytes=max_shard_bytes,
)
plan = _validate_entry(
entry_dir=entry_dir,
cache_key=cache_key,
layout_identity=layout_identity,
content_digest=content_digest,
expected_records=records,
)
logger.info(
"Published DLO host weight cache entry %s with %d tensors at %s",
cache_key[:12],
len(records),
entry_dir,
)
return HostWeightPlanResult(plan)
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"host weight cache operation failed: {exc}", "cache_operation_failed")
__all__ = [
"DEFAULT_LOCK_TIMEOUT_SECONDS",
"DEFAULT_HOST_WEIGHT_CACHE_ROOT",
"DEFAULT_SHARD_SIZE_BYTES",
"HOST_WEIGHT_CACHE_SCHEMA_VERSION",
"build_host_weight_cache_plan",
"default_host_weight_cache_root",
]
@@ -0,0 +1,287 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reusable identity, hashing, and file helpers for the DLO host weight cache."""
from __future__ import annotations
import contextlib
import dataclasses
import errno
import fcntl
import hashlib
import inspect
import json
import os
import shutil
import time
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Protocol, TypeAlias, TypeVar
import torch
from torch import nn
JsonScalar: TypeAlias = None | bool | int | float | str
JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]
PathIdentity: TypeAlias = str | os.PathLike[str] | os.PathLike[bytes]
_HASH_CHUNK_BYTES = 64 * 1024**2
_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
class IdentityNormalizationError(ValueError):
"""Raised when a cache identity contains a process-unstable value."""
@dataclass(frozen=True)
class RuntimeTensor:
"""One final runtime tensor and its manifest identity."""
name: str
tensor: torch.Tensor
kind: str
class HostWeightCacheError(RuntimeError):
"""Expected cache failure that should fall back with a stable code."""
def __init__(self, code: str, message: str):
super().__init__(message)
self.code = code
class _HashUpdater(Protocol):
def update(self, data: bytes | bytearray | memoryview) -> object: ...
_RuntimeTensorT = TypeVar("_RuntimeTensorT", bound=RuntimeTensor)
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("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)):
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"host weight 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()
def implementation_fingerprint(
loader_type: type,
pipeline: nn.Module,
dit_modules: Sequence[tuple[str, nn.Module]],
) -> str:
"""Hash relevant loader and module implementations for cache identity."""
objects: list[object] = [loader_type, type(pipeline)]
load_weights = getattr(type(pipeline), "load_weights", None)
if load_weights is not None:
objects.append(load_weights)
for _, dit_module in dit_modules:
objects.append(type(dit_module))
post_load = getattr(type(dit_module), "post_load_weights", None)
if post_load is not None:
objects.append(post_load)
digest = hashlib.sha256()
identities: set[str] = set()
for obj in objects:
identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
if identity in identities:
continue
identities.add(identity)
digest.update(identity.encode())
try:
digest.update(inspect.getsource(obj).encode())
except (OSError, TypeError):
pass
return digest.hexdigest()
def tensor_metadata(record: RuntimeTensor) -> dict[str, JsonValue]:
"""Return stable manifest metadata for one runtime tensor."""
tensor = record.tensor
return {
"kind": record.kind,
"shape": list(tensor.shape),
"dtype": str(tensor.dtype),
"nbytes": tensor.numel() * tensor.element_size(),
"layout": "contiguous",
}
def _update_hash_with_tensor(digest: _HashUpdater, tensor: torch.Tensor) -> None:
byte_view = tensor.detach().reshape(-1).view(torch.uint8)
for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
digest.update(memoryview(chunk.numpy()))
def runtime_content_digest(records: Sequence[RuntimeTensor]) -> str:
"""Hash ordered runtime tensor metadata and contents."""
digest = hashlib.sha256()
for record in records:
digest.update(canonical_json({"name": record.name, **tensor_metadata(record)}))
_update_hash_with_tensor(digest, record.tensor)
return digest.hexdigest()
def file_digest(path: Path) -> str:
"""Hash one cache file with a bounded read working set."""
digest = hashlib.sha256()
with path.open("rb") as handle:
while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
digest.update(chunk)
return digest.hexdigest()
@contextlib.contextmanager
def exclusive_lock(path: Path, timeout_seconds: float) -> Iterator[None]:
"""Acquire one non-blocking file lock within a bounded wait."""
path.parent.mkdir(parents=True, exist_ok=True)
deadline = time.monotonic() + timeout_seconds
with path.open("a+b") as handle:
while True:
try:
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
break
except OSError as exc:
if exc.errno not in (errno.EACCES, errno.EAGAIN):
raise HostWeightCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
remaining = deadline - time.monotonic()
if remaining <= 0:
raise HostWeightCacheError(
"lock_timeout",
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:
yield
finally:
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
def fsync_file(path: Path) -> None:
"""Persist one cache file before publication."""
with path.open("rb") as handle:
os.fsync(handle.fileno())
def fsync_directory(path: Path) -> None:
"""Persist directory entries needed for atomic publication."""
fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
try:
os.fsync(fd)
finally:
os.close(fd)
def split_shards(
records: Sequence[_RuntimeTensorT],
max_shard_bytes: int,
) -> list[list[_RuntimeTensorT]]:
"""Split ordered tensors without exceeding the target when possible."""
shards: list[list[_RuntimeTensorT]] = []
current: list[_RuntimeTensorT] = []
current_bytes = 0
for record in records:
nbytes = record.tensor.numel() * record.tensor.element_size()
if current and current_bytes + nbytes > max_shard_bytes:
shards.append(current)
current = []
current_bytes = 0
current.append(record)
current_bytes += nbytes
if current:
shards.append(current)
return shards
def remove_stale_temps(cache_root: Path, cache_key: str) -> None:
"""Remove abandoned temporary entries for one cache identity."""
temp_root = cache_root / ".tmp"
if not temp_root.is_dir():
return
for path in temp_root.glob(f"{cache_key}.*"):
if path.is_dir():
shutil.rmtree(path, ignore_errors=True)
else:
path.unlink(missing_ok=True)
__all__ = [
"IdentityNormalizationError",
"JsonValue",
"HostWeightCacheError",
"RuntimeTensor",
"canonical_json",
"canonicalize_existing_local_path",
"exclusive_lock",
"file_digest",
"fsync_directory",
"fsync_file",
"implementation_fingerprint",
"normalize_identity",
"remove_stale_temps",
"runtime_content_digest",
"split_shards",
"tensor_metadata",
]
@@ -9,6 +9,7 @@ import json
import os
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import Literal
import torch
from safetensors import safe_open
@@ -28,7 +29,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor]
class TensorBinding:
"""One runtime tensor backed by one safetensors entry."""
checkpoint_key: str
storage_key: str
file_path: str
transform: TensorTransform | None = None
@@ -37,9 +38,16 @@ class TensorBinding:
class HostWeightPlan:
"""Complete, prevalidated host backing consumed by an offload backend."""
backing_kind: str
backing_kind: Literal["checkpoint_mmap", "host_weight_cache"]
bindings: dict[str, TensorBinding]
planned_source_prefixes: frozenset[str] = frozenset()
runtime_layout_key: str | None = None
expected_file_digests: dict[str, str] | None = None
@property
def requires_backend(self) -> bool:
"""Whether ordinary loading was skipped and this plan must be consumed."""
return self.backing_kind == "checkpoint_mmap"
@dataclass(frozen=True)
@@ -48,6 +56,7 @@ class HostWeightPlanResult:
plan: HostWeightPlan | None
fallback_reason: str | None = None
fallback_code: str | None = None
class _PlanIncompatibleError(RuntimeError):
@@ -258,11 +267,11 @@ def _validate_source_metadata(
with safe_open(file_path, framework="pt", device="cpu") as handle:
available = set(handle.keys())
for runtime_name, target, binding in entries:
if binding.checkpoint_key not in available:
if binding.storage_key not in available:
raise _PlanIncompatibleError(
f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
)
source = handle.get_slice(binding.checkpoint_key)
source = handle.get_slice(binding.storage_key)
source_shape = tuple(source.get_shape())
source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
if source_shape != tuple(target.shape):
@@ -287,11 +296,19 @@ def build_checkpoint_mmap_plan(
) -> HostWeightPlanResult:
"""Build a complete direct-checkpoint plan or return a fallback reason."""
if tensor_parallel_size != 1:
return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
return HostWeightPlanResult(
None,
f"TP={tensor_parallel_size} requires the ordinary loader",
"unsupported_tp",
)
if use_hsdp:
return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
if online_quantization:
return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
return HostWeightPlanResult(
None,
"online quantization requires the ordinary loader",
"unsupported_online_quantization",
)
remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
if not callable(remap_fn):
@@ -328,14 +345,14 @@ def build_checkpoint_mmap_plan(
)
checkpoint_key, file_path = model_to_ckpt[runtime_name]
bindings[runtime_name] = TensorBinding(
checkpoint_key=checkpoint_key,
storage_key=checkpoint_key,
file_path=file_path,
transform=policy.transform if policy is not None else None,
)
_validate_source_metadata(required, bindings)
except (OSError, ValueError, _PlanIncompatibleError) as exc:
return HostWeightPlanResult(None, str(exc))
return HostWeightPlanResult(None, str(exc), "checkpoint_plan_incompatible")
return HostWeightPlanResult(
HostWeightPlan(
+9 -4
View File
@@ -71,8 +71,9 @@ def get_offload_backend(
Args:
od_config: OmniDiffusionConfig with offload settings
device: Target device (auto-detected if None)
host_weight_plan: Exact loader-produced backing plan, if ordinary
weight materialization was skipped.
host_weight_plan: Exact loader-produced backing plan. Checkpoint plans
may have skipped ordinary materialization; host weight cache plans
describe already-valid ordinary tensors that can be remapped.
Returns:
OffloadBackend instance or None if offloading disabled
@@ -85,10 +86,14 @@ def get_offload_backend(
# Extract and validate configuration
config = OffloadConfig.from_od_config(od_config)
plan_requires_backend = host_weight_plan is not None and host_weight_plan.requires_backend
if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
if plan_requires_backend:
raise RuntimeError(
"A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
)
logger.warning("Ignoring optional host weight cache plan because distributed layerwise offload is not selected")
host_weight_plan = None
# Return None if no offloading requested
if config.strategy == OffloadStrategy.NONE:
@@ -96,7 +101,7 @@ def get_offload_backend(
# Validate platform (CUDA required for now)
if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
if host_weight_plan is not None:
if plan_requires_backend:
raise RuntimeError(
"The loader skipped ordinary weight materialization for DLO, "
"but this platform cannot create the required offload backend"
@@ -112,7 +117,7 @@ def get_offload_backend(
try:
device = current_omni_platform.get_torch_device()
except (NotImplementedError, AttributeError) as exc:
if host_weight_plan is not None:
if plan_requires_backend:
raise RuntimeError(
"The loader skipped ordinary weight materialization for DLO, "
"but the target device could not be resolved"
+5
View File
@@ -31,6 +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
# Optional per-worker ceiling for registering shared host weight cache mappings.
# Zero applies no additional ceiling; pin_cpu_memory controls registration.
dlo_host_weight_cache_pin_limit_gib: float = 0.0
@classmethod
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -101,6 +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_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:
@@ -136,6 +140,7 @@ class OffloadConfig:
dp_size=dp_size,
dlo_use_allgather=dlo_use_allgather,
dlo_resident_layers=dlo_resident_layers,
dlo_host_weight_cache_pin_limit_gib=dlo_host_weight_cache_pin_limit_gib,
)
@@ -0,0 +1,280 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CUDA implementation of read-only host-mapping registration."""
from __future__ import annotations
import ctypes
import mmap
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Protocol
import torch
from .host_registration import (
HostRegistrationBudgetError,
HostRegistrationCleanupError,
HostRegistrationError,
)
_CUDA_HOST_REGISTER_READ_ONLY = 0x08
_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113
CudaHostRegistrationError = HostRegistrationError
CudaHostRegistrationBudgetError = HostRegistrationBudgetError
CudaHostRegistrationCleanupError = HostRegistrationCleanupError
class _CudaRuntime(Protocol):
def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ...
def cudaHostUnregister(self, address: int) -> int: ...
def cudaGetErrorString(self, error: int) -> str | bytes: ...
@dataclass(frozen=True)
class _AddressRange:
start: int
end: int
@property
def size(self) -> int:
return self.end - self.start
def _coalesce_ranges(
ranges: Sequence[tuple[int, int]],
page_size: int = mmap.PAGESIZE,
) -> tuple[_AddressRange, ...]:
"""Page-align and merge overlapping ranges from one backing mapping."""
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
aligned: list[_AddressRange] = []
for start, size in ranges:
if start <= 0 or size < 0:
raise ValueError(f"invalid host range start={start}, size={size}")
if size == 0:
continue
aligned_start = start - start % page_size
end = start + size
aligned_end = ((end + page_size - 1) // page_size) * page_size
aligned.append(_AddressRange(aligned_start, aligned_end))
merged: list[_AddressRange] = []
for region in sorted(aligned, key=lambda item: (item.start, item.end)):
if merged and region.start <= merged[-1].end:
previous = merged[-1]
merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
else:
merged.append(region)
return tuple(merged)
def _tensor_regions(
sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
) -> tuple[_AddressRange, ...]:
"""Resolve storage spans without merging unrelated file mappings."""
regions: list[_AddressRange] = []
for mapping_name, tensors in sources_by_mapping.items():
ranges: list[tuple[int, int]] = []
for tensor in tensors:
if tensor.device.type != "cpu":
raise HostRegistrationError(
f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}"
)
storage = tensor.untyped_storage()
ranges.append((storage.data_ptr(), storage.nbytes()))
regions.extend(_coalesce_ranges(ranges))
return tuple(regions)
def _error_message(runtime: _CudaRuntime, error: int) -> str:
try:
message = runtime.cudaGetErrorString(error)
except Exception:
return str(error)
if isinstance(message, bytes):
return message.decode(errors="replace")
return str(message)
def _consume_last_cuda_error(expected_error: int) -> None:
"""Clear a handled CUDA Runtime error before returning to PyTorch.
PyTorch's ``torch.cuda.cudart()`` binding exposes host registration but not
``cudaGetLastError``. Resolve the already-loaded process symbol so this call
clears the same CUDA Runtime instance that produced the registration error.
A different pending error is not safe to hide behind the staging fallback.
"""
try:
get_last_error = ctypes.CDLL(None).cudaGetLastError
get_last_error.argtypes = []
get_last_error.restype = ctypes.c_int
pending_error = int(get_last_error())
except (AttributeError, OSError) as exc:
raise HostRegistrationCleanupError("cannot clear CUDA's pending error after host-registration failure") from exc
if pending_error not in (0, expected_error):
raise HostRegistrationCleanupError(
f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}"
)
def _handled_error_message(runtime: _CudaRuntime, error: int) -> str:
"""Format and consume one nonzero CUDA Runtime return code."""
error_code = int(error)
message = _error_message(runtime, error)
try:
_consume_last_cuda_error(error_code)
except HostRegistrationCleanupError as exc:
raise HostRegistrationCleanupError(f"{message}; {exc}") from exc
return message
def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool:
"""Query support required by immutable file-backed cache mappings."""
try:
get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None)
if get_attribute is None:
get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute
get_attribute.argtypes = [
ctypes.POINTER(ctypes.c_int),
ctypes.c_int,
ctypes.c_int,
]
get_attribute.restype = ctypes.c_int
except (AttributeError, OSError) as exc:
raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc
supported = ctypes.c_int()
try:
error = get_attribute(
ctypes.byref(supported),
_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED,
torch.accelerator.current_device_index(),
)
except Exception as exc:
raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc
if int(error) != 0:
raise HostRegistrationError(
f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}"
)
return bool(supported.value)
class CudaHostRegistration:
"""Own CUDA registrations for already-existing file-backed CPU tensors."""
def __init__(
self,
runtime: _CudaRuntime,
regions: tuple[_AddressRange, ...],
) -> None:
self._runtime = runtime
self._regions = regions
self._closed = False
@classmethod
def create(
cls,
sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
*,
max_bytes: int | None,
) -> CudaHostRegistration:
regions = _tensor_regions(sources_by_mapping)
total_bytes = sum(region.size for region in regions)
if max_bytes is not None and total_bytes > max_bytes:
raise HostRegistrationBudgetError(
f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
)
if not regions:
raise HostRegistrationError("no non-empty host ranges were available for registration")
if not torch.cuda.is_available():
raise HostRegistrationError("CUDA is not available")
try:
runtime = torch.cuda.cudart()
except Exception as exc:
raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
if not _supports_read_only_host_registration(runtime):
raise HostRegistrationError(
"CUDA device does not support read-only host registration required by immutable "
"host weight cache mappings"
)
registered: list[_AddressRange] = []
try:
for region in regions:
error = runtime.cudaHostRegister(
region.start,
region.size,
_CUDA_HOST_REGISTER_READ_ONLY,
)
if int(error) != 0:
raise HostRegistrationError(
"cudaHostRegister(read-only) failed for "
f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}"
)
registered.append(region)
unpinned = [
mapping_name
for mapping_name, tensors in sources_by_mapping.items()
if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors)
]
if unpinned:
raise HostRegistrationError(
f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}"
)
except Exception as exc:
rollback_errors: list[str] = []
for region in reversed(registered):
try:
error = runtime.cudaHostUnregister(region.start)
if int(error) != 0:
rollback_errors.append(
f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}"
)
except Exception as rollback_exc:
rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
if rollback_errors:
raise HostRegistrationCleanupError(
f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}"
) from exc
if isinstance(exc, HostRegistrationError):
raise
raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc
return cls(runtime, regions)
@property
def total_bytes(self) -> int:
return sum(region.size for region in self._regions)
@property
def region_count(self) -> int:
return len(self._regions)
def close(self) -> list[str]:
"""Unregister every range, returning errors after best-effort cleanup."""
if self._closed:
return []
errors: list[str] = []
failed: list[_AddressRange] = []
for region in reversed(self._regions):
try:
error = self._runtime.cudaHostUnregister(region.start)
if int(error) != 0:
errors.append(
f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}"
)
failed.append(region)
except Exception as exc:
errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
failed.append(region)
self._regions = tuple(reversed(failed))
self._closed = not self._regions
return errors
@@ -18,7 +18,9 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
from __future__ import annotations
import time
from itertools import chain
from pathlib import Path
from typing import Any
import torch
@@ -27,6 +29,7 @@ from torch import nn
from vllm.logger import init_logger
from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
from vllm_omni.diffusion.model_loader.host_weight_cache_utils import file_digest
from vllm_omni.diffusion.model_loader.host_weight_plan import (
HostWeightPlan,
)
@@ -34,6 +37,12 @@ from vllm_omni.platforms import current_omni_platform
from .base import OffloadBackend, OffloadConfig
from .block_discovery import get_blocks_from_dit
from .host_registration import (
HostRegistration,
HostRegistrationCleanupError,
HostRegistrationError,
register_host_mappings,
)
from .module_collector import ModuleDiscovery
from .offload_plan import (
OffloadPlan,
@@ -93,6 +102,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
self.rank = rank
self.pin_memory = pin_memory
self.rank_local_mmap = rank_local_mmap
self.registered_mmap = False
self.tensor_transforms = tensor_transforms or {}
self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -445,17 +455,11 @@ class DistributedLayerwiseOffloadHook(ModelHook):
# ------------------------------------------------------------------ #
@staticmethod
def _pack_mmap_sources(
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
metadata: dict[torch.dtype, list[dict[str, Any]]],
slot_buffers: dict[torch.dtype, torch.Tensor],
) -> dict[torch.dtype, torch.Tensor]:
staged: dict[torch.dtype, torch.Tensor] = {}
for dtype, metas in metadata.items():
total_numel = sum(meta["numel"] for meta in metas)
destination = slot_buffers[dtype][:total_numel]
sources = cpu_sources[dtype]
for source_info, meta in zip(sources, metas, strict=True):
def _resolve_mmap_source(
source_info: dict[str, Any],
meta: dict[str, Any],
dtype: torch.dtype,
) -> torch.Tensor:
source = source_info["tensor"]
transform = source_info["transform"]
if callable(transform):
@@ -468,6 +472,21 @@ class DistributedLayerwiseOffloadHook(ModelHook):
f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
f"stride={source.stride()}"
)
return source
@staticmethod
def _pack_mmap_sources(
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
metadata: dict[torch.dtype, list[dict[str, Any]]],
slot_buffers: dict[torch.dtype, torch.Tensor],
) -> dict[torch.dtype, torch.Tensor]:
staged: dict[torch.dtype, torch.Tensor] = {}
for dtype, metas in metadata.items():
total_numel = sum(meta["numel"] for meta in metas)
destination = slot_buffers[dtype][:total_numel]
sources = cpu_sources[dtype]
for source_info, meta in zip(sources, metas, strict=True):
source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
start = meta["offset"]
physical_storage = destination[start : start + meta["numel"]]
if source.is_contiguous():
@@ -481,6 +500,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
staged[dtype] = destination
return staged
@staticmethod
def _copy_mmap_sources_to_device(
cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
metadata: dict[torch.dtype, list[dict[str, Any]]],
gpu_buffers: dict[torch.dtype, torch.Tensor],
*,
non_blocking: bool,
) -> None:
"""Copy registered source views directly into flattened device buffers."""
for dtype, metas in metadata.items():
destination = gpu_buffers[dtype]
sources = cpu_sources[dtype]
for source_info, meta in zip(sources, metas, strict=True):
source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
start = meta["offset"]
physical_storage = destination[start : start + meta["numel"]]
async_copy = non_blocking and source.is_pinned()
if source.is_contiguous():
physical_storage.copy_(source.flatten(), non_blocking=async_copy)
else:
torch.as_strided(
physical_storage,
size=source.shape,
stride=source.stride(),
).copy_(source, non_blocking=async_copy)
def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
"""Pack this block's mmap views into one bounded host staging slot."""
previous_copy = self.cpu_staging_events[slot]
@@ -514,6 +559,16 @@ class DistributedLayerwiseOffloadHook(ModelHook):
assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
if self.dp_size <= 1 or self.dp_group is None:
if self.rank_local_mmap and self.registered_mmap:
with current_omni_platform.stream(self.copy_stream):
self._copy_mmap_sources_to_device(
self.cpu_sources,
self.metadata,
gpu_weights,
non_blocking=non_blocking,
)
evt.record(self.copy_stream)
else:
cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
with current_omni_platform.stream(self.copy_stream):
for dtype, cpu_shard in cpu_weights.items():
@@ -734,6 +789,7 @@ class PinnedResidentLayerGroup:
self.copy_stream = copy_stream
self.loaded = False
self.rank_local_mmap = rank_local_mmap
self.registered_mmap = False
self.pin_memory = pin_memory
self._states: list[dict[str, Any]] = []
self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -804,6 +860,14 @@ class PinnedResidentLayerGroup:
ready = current_omni_platform.Event()
with current_omni_platform.stream(self.copy_stream):
for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
if self.rank_local_mmap and self.registered_mmap:
DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
state["cpu_sources"],
state["metadata"],
block_buffers,
non_blocking=True,
)
continue
if self.rank_local_mmap:
slot = index % 2
previous_copy = self._cpu_staging_events[slot]
@@ -901,8 +965,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
self._resident_layer_group: PinnedResidentLayerGroup | None = None
self._using_mmap = False
self._using_rank_local_mmap = False
self._using_registered_mmap = False
self.host_weight_plan = host_weight_plan
self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
self._host_weight_cache_mapped_sources: dict[str, list[torch.Tensor]] = {}
self._host_registration: HostRegistration | None = None
def load_resident_layers(self) -> None:
"""Load the model-declared leading blocks for the denoise stage."""
@@ -1023,7 +1090,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
framework="pt",
device="cpu",
)
tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
if is_parameter:
replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
parent._parameters[leaf_name] = replacement
@@ -1103,6 +1170,179 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
f"(first 5: {remaining_meta[:5]})."
)
def _load_weights_from_host_weight_cache(
self,
pipeline: nn.Module,
modules,
plan: HostWeightPlan,
) -> None:
"""Replace final ordinary-loader tensors with validated cache views."""
from safetensors import safe_open
if any(binding.transform is not None for binding in plan.bindings.values()):
raise RuntimeError("Host weight cache bindings cannot carry deferred checkpoint transforms")
cache_files = {binding.file_path for binding in plan.bindings.values()}
if plan.expected_file_digests is None or set(plan.expected_file_digests) != cache_files:
raise RuntimeError("Host weight cache plan does not cover every shard digest")
for file_path in sorted(cache_files):
try:
actual_digest = file_digest(Path(file_path))
except OSError as exc:
raise RuntimeError(f"Cannot revalidate host weight cache shard {file_path}: {exc}") from exc
if actual_digest != plan.expected_file_digests[file_path]:
raise RuntimeError(f"Host weight cache shard changed after loader validation: {file_path}")
required_names: set[str] = set()
for dit_name, dit_module in zip(modules.dit_names, modules.dits):
required_names.update(f"{dit_name}.{name}" for name, _ in dit_module.named_parameters())
for name, _ in dit_module.named_buffers():
parent_path, _, leaf_name = name.rpartition(".")
owner = dit_module.get_submodule(parent_path)
if leaf_name not in owner._non_persistent_buffers_set:
required_names.add(f"{dit_name}.{name}")
if required_names != set(plan.bindings):
missing = sorted(required_names - set(plan.bindings))
unexpected = sorted(set(plan.bindings) - required_names)
raise RuntimeError(
"Host weight cache plan coverage changed after loader validation: "
f"missing={missing[:5]}, unexpected={unexpected[:5]}"
)
file_cache: dict[str, Any] = {}
replacements: list[tuple[torch.Tensor, torch.Tensor]] = []
mapped_sources: dict[str, list[torch.Tensor]] = {}
try:
# Validate and map every tensor before mutating the pipeline. A
# missing/corrupt shard therefore leaves the ordinary tensors intact.
for runtime_name, binding in plan.bindings.items():
parent_path, _, leaf_name = runtime_name.rpartition(".")
try:
parent = pipeline.get_submodule(parent_path)
except AttributeError as exc:
raise RuntimeError(f"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"Host weight cache target tensor {runtime_name!r} no longer exists")
if target.device.type != "cpu" or target.is_meta:
raise RuntimeError(
f"Host weight cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}"
)
if binding.file_path not in file_cache:
file_cache[binding.file_path] = safe_open(
binding.file_path,
framework="pt",
device="cpu",
)
tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
if tensor.shape != target.shape or tensor.dtype != target.dtype:
raise RuntimeError(
f"Host weight cache metadata changed for {runtime_name!r}: "
f"cache={tuple(tensor.shape)}/{tensor.dtype}, "
f"runtime={tuple(target.shape)}/{target.dtype}"
)
replacements.append((target, tensor))
mapped_sources.setdefault(binding.file_path, []).append(tensor)
except Exception:
file_cache.clear()
raise
# Preserve Parameter objects and their loader/runtime attributes; only
# their storage changes. This also keeps any references held by model
# code valid while dropping the private allocator-backed storage.
for target, tensor in replacements:
set_tensor_storage(target, tensor)
self._mmap_transforms_by_tensor_id.clear()
self._mmap_file_cache = file_cache
self._host_weight_cache_mapped_sources = mapped_sources
logger.info(
"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_host_weight_cache_mmap(self) -> bool:
"""Register complete cache mappings under the pinned-memory policy."""
pin_limit_gib = self.config.dlo_host_weight_cache_pin_limit_gib
if not self.config.pin_cpu_memory:
logger.info("DLO host weight cache registration disabled by pin_cpu_memory=False; using host staging")
return False
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) if pin_limit_gib > 0 else None
started = time.perf_counter()
try:
registration = register_host_mappings(
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 host weight cache registration rollback failed")
raise
except HostRegistrationError as exc:
logger.warning(
"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 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,
)
return True
def _configure_host_weight_cache_transfer(
self,
hooks: list[DistributedLayerwiseOffloadHook],
) -> None:
"""Select registered mmap or staging once for every cache consumer."""
plan = self.host_weight_plan
if (
(not hooks and self._resident_layer_group is None)
or not self._using_rank_local_mmap
or plan is None
or plan.backing_kind != "host_weight_cache"
):
return
self._using_registered_mmap = self._try_register_host_weight_cache_mmap()
self._host_weight_cache_mapped_sources.clear()
for hook in hooks:
hook.registered_mmap = self._using_registered_mmap
if self._resident_layer_group is not None:
self._resident_layer_group.registered_mmap = self._using_registered_mmap
if self._using_registered_mmap:
self._resident_layer_group._cpu_staging_buffers.clear()
def _release_registered_mmap(self) -> bool:
registration = self._host_registration
if registration is None:
return True
errors = registration.close()
if errors:
logger.error(
"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 host weight cache mappings")
return True
def _init_dp_group(self) -> None:
"""Reuse the process group initialized by parallel_state.
@@ -1443,17 +1683,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
self._using_mmap = self.host_weight_plan is not None
self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
if self._using_mmap:
if self.host_weight_plan.backing_kind != "checkpoint_mmap":
raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
if self.host_weight_plan.backing_kind == "checkpoint_mmap":
self._load_weights_via_mmap(
pipeline,
modules,
self.host_weight_plan,
)
elif self.host_weight_plan.backing_kind == "host_weight_cache":
if self.config.dlo_use_allgather:
raise ValueError("Host weight cache backing is no-AllGather only in the first implementation")
self._using_rank_local_mmap = True
try:
self._load_weights_from_host_weight_cache(
pipeline,
modules,
self.host_weight_plan,
)
except Exception as exc:
logger.warning(
"DLO host weight cache remapping failed; retaining ordinary host tensors: %s",
exc,
)
self.host_weight_plan = None
self._using_mmap = False
self._using_rank_local_mmap = False
else:
# Assignments above dropped the model's private tensor
# references. Reclaim those allocations before hook setup
# so steady-state PSS reflects the shared mmap backing.
self._cleanup_after_loading()
else:
raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
if self._using_rank_local_mmap:
logger.info(
"DLO rank-local mmap storage enabled: checkpoint pages are "
"node-shared; each worker owns only two bounded host staging slots"
"DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
"transfer setup will select registered direct H2D or bounded host staging"
)
else:
remaining_meta = [
@@ -1623,6 +1887,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
)
pipeline._dlo_residency_controller = self
all_hooks = [hook for group in self._all_hook_groups for hook in group]
self._configure_host_weight_cache_transfer(all_hooks)
if not self._all_hook_groups:
self.enabled = bool(self._resident_blocks)
if self._using_mmap and not self.enabled:
@@ -1632,17 +1899,14 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
# Unified allocation: 2 shared output buffers + 2 shared shard buffers
# sized to the max block across ALL module groups (gen_layers +
# language_model). Groups execute sequentially, so 2 buffers suffice.
all_hooks: list[DistributedLayerwiseOffloadHook] = []
for group in self._all_hook_groups:
all_hooks.extend(group)
unified_buffers = self._allocate_shared_buffers(all_hooks)
unified_shard_buffers = None
if self.dp_size > 1:
unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
unified_cpu_staging = None
cpu_staging_events = None
if self._using_rank_local_mmap:
if self._using_rank_local_mmap and not self._using_registered_mmap:
unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
all_hooks,
self._resident_layer_group,
@@ -1724,6 +1988,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
def _release_mmap_handles(self) -> None:
"""Release safetensors mmap file handles."""
self._mmap_transforms_by_tensor_id.clear()
self._host_weight_cache_mapped_sources.clear()
if hasattr(self, "_mmap_file_cache"):
self._mmap_file_cache.clear()
del self._mmap_file_cache
@@ -1746,8 +2011,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
if not self.enabled and not hasattr(self, "_mmap_file_cache"):
return
if self._using_rank_local_mmap:
# Copy and communication streams are asynchronous in both AllGather
# and no-AllGather modes. Drain all device work before hooks, process
# groups, registrations, or mapped storage can be torn down.
current_omni_platform.synchronize()
registration_released = self._release_registered_mmap()
for blocks in self._blocks:
for block in blocks:
@@ -1763,9 +2031,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
self._all_hook_groups.clear()
self._resident_blocks.clear()
self._resident_layer_group = None
if registration_released:
self._release_mmap_handles()
self._using_mmap = False
self._using_rank_local_mmap = False
self._using_registered_mmap = False
self.enabled = False
logger.info("Distributed layer-wise offloading disabled")
@@ -0,0 +1,62 @@
# 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 | None,
) -> 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.
``max_bytes=None`` applies no additional ceiling to the complete mapping.
"""
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",
]
@@ -859,10 +859,17 @@ class DiffusionWorker:
def shutdown(self) -> None:
"""Shutdown the worker and cleanup distributed environment."""
try:
if self.model_runner is not None:
mgr = getattr(self.model_runner, "kv_transfer_manager", None)
try:
offload_backend = getattr(self.model_runner, "offload_backend", None)
if offload_backend is not None:
offload_backend.disable()
finally:
if mgr is not None:
mgr.shutdown_prefetch()
finally:
destroy_distributed_env()
+4
View File
@@ -545,6 +545,10 @@ class OrchestratorArgs:
enable_layerwise_offload: bool = False
enable_distributed_layerwise_offload: bool = False
dlo_use_allgather: bool = True
dlo_use_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
+4
View File
@@ -1040,6 +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_use_host_weight_cache": kwargs.get("dlo_use_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": (
+19
View File
@@ -722,6 +722,25 @@ class OmniServeCommand(CLISubcommand):
"independently (including existing TP shards)."
),
)
omni_config_group.add_argument(
"--dlo-use-host-weight-cache",
action="store_true",
help=(
"Use the node-local final-layout host weight cache for no-AllGather DLO. "
"With the default pinned-memory policy, supported platforms register the immutable mapping "
"for direct H2D; otherwise DLO uses bounded host staging."
),
)
omni_config_group.add_argument(
"--dlo-host-weight-cache-pin-limit-gib",
type=float,
default=0.0,
help=(
"Optional maximum GiB of read-only host weight cache mappings each worker may "
"register for direct H2D. Zero applies no additional registration ceiling; "
"this option requires --dlo-use-host-weight-cache."
),
)
omni_config_group.add_argument(
"--dlo-resident-layers",
type=int,