anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

docs: document distributed layerwise offload compatibility

Browse Source
This commit is contained in:
Hongsheng Liu committed 2026-08-06 12:17:29 +08:00
1 parent f879ad4983
commit d4e0a5a5fd
5 files changed
+203 -22
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/.nav.yml
+1
View File
@@ -120,6 +120,7 @@ nav:
- design/feature/omni_async_output_materialization.md
- design/feature/vae_parallel.md
- design/feature/diffusion_continuous_batching.md
- design/feature/distributed_layerwise_offload.md
- Module Design:
- design/module/ar_module.md
- design/module/dit_module.md
docs/design/feature/distributed_layerwise_offload.md
+154
View File
@@ -0,0 +1,154 @@
# Distributed Layerwise Offload
This document describes distributed layerwise offload (DLO) for diffusion
models. DLO keeps only a small number of DiT blocks on the accelerator and
streams the remaining blocks from host memory. The distributed backend can
either shard those host-side weights across an existing parallel group or keep
the standard loader's rank-local weights and avoid an additional collective.
For user-facing commands, see the
[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload_diffusion.md)
and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
## Status
DLO is implemented for multi-device diffusion execution. The default
AllGather path is the primary path for DP and SP deployments. The
`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
for standard-loader sharding, workstation bring-up, and systems where an
additional DLO collective is undesirable, but it does not reduce host weight
storage across ranks.
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.
## Design
### DLO consumes the existing parallel topology
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 path
With the default `dlo_use_allgather=True`, each rank stores approximately
`1 / group_size` of each streamable block in pinned host memory. The next
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.
```text
Compute: [Block N] [Block N+1] [Block N+2]
H2D: [shard N+1] [shard N+2]
AllGather: [full N+1] [full N+2]
Buffers: [current slot] [prefetch slot] [current slot]
```
The backend uses two shared device buffers, so accelerator weight residency is
bounded by the largest streamed blocks rather than the complete model.
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
With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
one. The regular model loader remains responsible for preparing each rank's
weights, including TP-local tensors or HSDP-managed parameters. DLO then
streams those rank-local tensors block by block using H2D copies only.
This mode means:
- DP still provides independent replicas, but DLO does not shard weights
across DP ranks.
- SP still performs its normal activation/attention collectives, but DLO does
not shard weights across SP ranks.
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
DLO's additional weight AllGather is disabled.
- Pure DP deployments keep a full host-side model copy per rank, subject to
shared OS page-cache behavior.
- The scheduler does not require a synchronized DP request wave for DLO.
This path is intentionally not implemented by reusing the AllGather mmap
loader. It relies on the standard loader so model-specific TP/HSDP transforms
remain intact.
## 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. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
| **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** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
### Combined dimensions
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
than one; SP continues to use its own sequence-parallel group. If DP is one,
the SP group becomes DLO's sharding group in AllGather mode.
- **DP + TP/SP without AllGather:** standard model loading defines the
rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
collective.
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
materialization and synchronization.
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
configuration.
## Request and loading constraints
AllGather DP multi-concurrency requires:
- explicit `num_inference_steps`;
- the same `num_inference_steps` for all requests in a wave; and
- identical request arguments that affect the collective execution path.
The no-AllGather path does not impose these DLO-specific synchronized-wave
requirements.
The mmap loader is used only by the supported DLO+AllGather path when the
model has a compatible checkpoint layout. Online quantization is incompatible
with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
quantization. Models that do not meet the mmap requirements use the regular
loader path.
## Validation coverage
Current source-level validation includes:
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- TP rejection in the DLO+AllGather mmap path;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
tests.
The highest-value missing coverage is end-to-end numerical comparison against
ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
on the target CUDA/NCCL or CANN/HCCL hardware.
## Recommendations
- Use **DP + DLO AllGather** for the supported throughput and host-memory
scaling path.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
not the goal.
- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
the expectation of higher host-memory use and lower validation confidence.
- Prefer **HSDP alone** for production HSDP deployments until the combined
HSDP + DLO no-AllGather path has broader end-to-end coverage.
docs/design/index.md
+1
View File
@@ -12,6 +12,7 @@ This section contains design documents and architecture specifications for vLLM-
- [Ray-based Execution](feature/ray_based_execution.md)
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
## Infrastructure Design Documents
docs/user_guide/diffusion/cpu_offload_diffusion.md
+29 -12
View File
@@ -184,10 +184,15 @@ class Flux2Transformer2DModel(nn.Module):
### How It Works
Distributed layerwise offloading extends single-GPU layerwise offloading to
multi-device deployments. Each DP rank stores only **1/dp_size** of the model
weights in host memory; full layer weights are reconstructed at runtime via
**AllGather** on a dedicated communication stream, overlapped with computation
via a fixed double-buffer scheme.
multi-device deployments. In the default mode, each rank stores only
**1/dp_size** of the model weights in host memory; full layer weights are
reconstructed at runtime via **AllGather** on a dedicated communication stream,
overlapped with computation via a fixed double-buffer scheme. The
`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
weights independently and does not shard weights across ranks.
For the implementation design and compatibility matrix, see
[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
**Key features:**
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
@@ -219,7 +224,7 @@ vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4
# Without AllGather (each rank loads full weights, no sharding)
# Without DLO AllGather (each rank streams standard-loader rank-local weights)
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--data-parallel-size 4 \
@@ -247,9 +252,9 @@ m = Omni(
| Flag | Description | Default |
|------|-------------|---------|
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
### How model weights are loaded (mmap path)
@@ -303,9 +308,15 @@ request while AllGather synchronizes only weight shards (request-independent).
### Limitations
- Online quantization (FP8) is incompatible with mmap loading — falls back to
regular `load_weights()` automatically
- Tensor Parallel is not supported (DLO uses DP-based sharding)
- HSDP + AllGather is rejected (would double-shard weights)
regular loading only when the model is outside the DLO+AllGather mmap path;
online quantization with DLO+AllGather is rejected. Use
`--dlo-no-use-allgather` or disable online quantization.
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
weight-loader callbacks are bypassed. The no-AllGather path retains the
standard TP loader and is experimental.
- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
with no-AllGather is accepted at configuration level, but has limited
end-to-end validation.
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
**Module Discovery**
@@ -338,7 +349,7 @@ OffloadBackend (base class)
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
│ (single-GPU, full weights on host)
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
(multi-GPU, 1/dp_size sharded weights + AllGather)
(multi-GPU, sharded AllGather or rank-local mode)
```
Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -369,4 +380,10 @@ reasoner/generator components inside the model forward pass.
**Notes:**
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
- Distributed Layerwise Offloading reuses the layerwise block topology, but the
model and parallelism combination still needs validation. AllGather mode
additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
retains the standard loader's rank-local tensors. See the
[DLO design](../../design/feature/distributed_layerwise_offload.md) and
[Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
for usage examples.
recipes/cosmos3/Cosmos3-DistOffload.md
+18 -10
View File
@@ -56,10 +56,14 @@ vllm serve /path/to/Cosmos3-Super \
--data-parallel-size 4
```
### Without AllGather (full weights per rank, no sharding)
### Without DLO AllGather (rank-local weights, no DLO sharding)
Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
no AllGather synchronization, no concurrent request requirement, but N× host memory.
Use `--dlo-no-use-allgather` when you want each rank to stream the standard
loader's rank-local weights independently — no DLO AllGather synchronization,
no concurrent request requirement, but no DLO weight sharding and higher host
memory use. In pure DP, this is equivalent to independent full-weight
layerwise-offload replicas; with TP or HSDP, each rank retains the standard
loader's local tensor layout.
```bash
vllm serve /path/to/Cosmos3-Nano \
@@ -154,9 +158,9 @@ wait
| Flag | Description | Default |
|------|-------------|---------|
| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
| `--data-parallel-size N` | Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled | `1` |
| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
| `--dlo-no-use-allgather` | Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather) | `false` |
| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
## Mode comparison
@@ -166,7 +170,7 @@ wait
| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Standard-loader rank-local weights | 2 blocks | 0.4× HSDP | No DLO AllGather overhead, high CPU |
| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
## Memory expectations
@@ -185,7 +189,11 @@ wait
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
`load_weights()` automatically
- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
- HSDP + AllGather is rejected (would double-shard weights)
- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
use `--dlo-no-use-allgather` or disable online quantization
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
`--dlo-no-use-allgather` retains the standard TP loader and should be treated
as an experimental compatibility path
- HSDP + AllGather is rejected (would double-shard weights). HSDP with
`--dlo-no-use-allgather` is accepted at configuration level, but is not yet a
broadly validated production combination