diff --git a/docs/.nav.yml b/docs/.nav.yml index 4ff9493c7..ac5ee2b28 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -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 diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md new file mode 100644 index 000000000..a52b0b8ea --- /dev/null +++ b/docs/design/feature/distributed_layerwise_offload.md @@ -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. diff --git a/docs/design/index.md b/docs/design/index.md index 0da17cce4..e8df11b77 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -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 diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 86a0e9201..02ab9f074 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -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. diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md index 27d1c77bd..8eeea5cc2 100644 --- a/recipes/cosmos3/Cosmos3-DistOffload.md +++ b/recipes/cosmos3/Cosmos3-DistOffload.md @@ -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