diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index 63382e76f..bae8cd1b5 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -309,7 +309,7 @@ For Omni or custom models, create: Required updates: 1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table -2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table) +2. `docs/user_guide/diffusion/cpu_offload.md` — if CPU offload supported (add to supported models table) 3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported 4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported 5. Offline example docs under `examples/offline_inference//` (`README.md` or category-specific `.md`) diff --git a/docs/.nav.yml b/docs/.nav.yml index 8c9b614dc..ca17d550d 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -68,12 +68,12 @@ nav: - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md - - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md + - CPU Offloading: user_guide/diffusion/cpu_offload.md - LoRA: user_guide/diffusion/lora.md - Custom Pipeline: features/custom_pipeline.md - Execution Modes: user_guide/diffusion/execution_modes.md - Quantization: - - Overview: user_guide/quantization/overview.md + - Quantization: user_guide/quantization/overview.md - Online Quantization: user_guide/quantization/online.md - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md - FP8 W8A8: user_guide/quantization/fp8.md @@ -106,7 +106,6 @@ nav: - Feature Design: - Runtime and Stage Execution: - Disaggregated Inference: design/feature/disaggregated_inference.md - - Ray-based Execution: design/feature/ray_based_execution.md - Async Chunk: design/feature/async_chunk.md - Async Diffusion Output: design/feature/async_diffusion_output.md - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md @@ -130,10 +129,11 @@ nav: - VAE Patch Parallelism: design/feature/vae_parallel.md - Attention Backends: - Skip-Softmax: design/feature/skip_softmax.md - - Quantization: user_guide/quantization/overview.md + - Quantization: design/feature/quantization.md - Cache-DiT: design/feature/cache_dit.md - TeaCache: design/feature/teacache.md - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md + - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md - Infrastructure and Performance: - Prometheus Metrics: design/metrics.md - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index d390a603a..5e65d2f48 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -1,6 +1,6 @@ **Scope** -For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/). +For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/). **Test Case Design** diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index dd7dfbd6b..678797652 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -824,7 +824,7 @@ omni = Omni(model="your-model", ### CPU Offload -See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md) +See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload.md) vLLM-Omni provides two offloading strategies to reduce GPU memory usage: diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md new file mode 100644 index 000000000..da735d60b --- /dev/null +++ b/docs/design/feature/distributed_layerwise_offload.md @@ -0,0 +1,156 @@ +# 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.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] +``` + +![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif) + +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/feature/quantization.md b/docs/design/feature/quantization.md new file mode 100644 index 000000000..37dd5e26b --- /dev/null +++ b/docs/design/feature/quantization.md @@ -0,0 +1,249 @@ +# Quantization + +This document describes the quantization architecture in vLLM-Omni and the +extension points for contributors. For user-facing configuration, supported +hardware, and method-specific instructions, see the +[quantization user guide](../../user_guide/quantization/overview.md). + +## Goals and scope + +The design has four goals: + +1. Provide one `quantization_config` entry point for diffusion-only, + multi-stage omni/TTS, and multi-stage diffusion models. +2. Reuse vLLM's quantization registry and `QuantizationConfig` contract where + a vLLM backend already provides the required behavior. +3. Allow different components of one pipeline to use different quantization + configurations. +4. Preserve the distinction between online quantization of a BF16/FP16 + checkpoint and loading a pre-quantized checkpoint with serialized scales. + +This document covers configuration resolution, component routing, model +integration, backend extension, and validation. Kernel implementations and +user-facing support matrices are outside its scope. + +## Architecture + +The quantization path is split across the following layers: + +| Layer | Responsibility | Source | +| --- | --- | --- | +| Public API | Exposes `build_quant_config`, component routing, and backend registration | `vllm_omni/quantization/__init__.py` | +| Factory | Normalizes method names, selects Omni overrides or vLLM registry backends, detects ModelOpt configs, and reconciles checkpoint metadata | `vllm_omni/quantization/factory.py` | +| Component router | Resolves a quantization config from a layer prefix using longest-prefix matching | `vllm_omni/quantization/component_config.py` | +| Method configs | Implement `QuantizationConfig` and select the platform-specific linear method | `vllm_omni/quantization/*_config.py` | +| Diffusion configuration | Stores the canonical runtime config and propagates checkpoint configuration | `vllm_omni/diffusion/data.py` | +| Worker and model integration | Passes the resolved config into vLLM's model configuration and model layers | `vllm_omni/diffusion/worker/diffusion_worker.py` and `vllm_omni/diffusion/models/` | + +The runtime contract remains vLLM's `QuantizationConfig`: + +```text +user/CLI or model config + | + v +OmniDiffusionConfig.quantization_config + | + +--> build_quant_config() + | | + | +--> Omni override + | +--> vLLM registry backend + | +--> ComponentQuantizationConfig + | + +--> checkpoint metadata reconciliation + | + +--> diffusion worker / vLLM config + | + +--> model layers call get_quant_method(layer, prefix) +``` + +The model layer receives either a method implementation or `None`. Returning +`None` leaves that layer unquantized. + +## Configuration resolution + +`build_quant_config()` accepts the following forms: + +| Input | Resolution | +| --- | --- | +| `None` or `"none"` | Disable quantization. | +| Method string, such as `"fp8"` | Resolve an Omni override first, then the vLLM quantization registry. | +| Flat dictionary with `method` or `quant_method` | Normalize the method and construct its `QuantizationConfig`. | +| Per-component dictionary | Construct a `ComponentQuantizationConfig`. | +| Existing `QuantizationConfig` | Pass through without rebuilding it. | + +Method aliases are normalized case-insensitively, with `-` and `_` treated as +equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter. +ModelOpt checkpoint metadata is detected from its method, producer metadata, or +`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration. + +The factory uses Omni-specific overrides for methods whose diffusion behavior +or platform dispatch is not provided by the vLLM registry. The current override +set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/ +AutoRound. Other methods, such as FP8, GGUF, and ModelOpt, are resolved through +vLLM's registry when their configuration is compatible. + +## Per-component routing + +Multi-stage models do not necessarily want one quantization policy for every +module. A component map expresses that policy without requiring each pipeline +to construct separate vLLM configurations: + +```python +quantization_config = { + "transformer": {"method": "fp8"}, + "vae": None, + "default": None, +} +``` + +`ComponentQuantizationConfig` applies the following rules: + +1. Component keys are interpreted as runtime layer-prefixes. +2. The longest matching prefix wins. +3. `None` disables quantization for the matching component. +4. `default`, when present, handles prefixes with no explicit match. +5. A matching child config delegates to its normal + `get_quant_method(layer, prefix)` implementation. + +Prefix matching happens after the model's weight-name mapping. New model +integrations must verify the actual runtime prefixes, especially when a +`WeightsMapper` changes checkpoint names. A prefix mismatch can silently route +a layer to the default configuration, so component-routing tests should cover +both positive matches and fall-through behavior. + +## Online and pre-quantized checkpoints + +The same method name can represent two different loading paths: + +- **Online**: the source checkpoint contains BF16/FP16 weights; the method + creates quantized weights or scales while the model is loaded. +- **Pre-quantized**: the checkpoint contains quantized weights and the metadata + required by the method; the loader consumes those serialized tensors. + +Method configuration owns the mode flag, for example +`is_checkpoint_int8_serialized` or `is_checkpoint_mxfp8_serialized`. The +configuration's `get_quant_method()` selects the corresponding linear method. + +For pipelines with transformer-specific `config.json` files, such as cascade +models, `resolve_quant_config_from_disk()` reconciles the active configuration +with each transformer's checkpoint metadata. Its invariants are: + +- If no active configuration exists, valid checkpoint metadata can be used for + auto-detection. +- A method mismatch is an error rather than an implicit conversion. +- Serialized-checkpoint flags rebuild an online configuration into the + matching offline configuration. +- Checkpoint-specific `ignored_layers` are preserved by rebuilding the active + configuration when they differ. +- AutoRound MXFP8 metadata (`data_type="mx_fp"`) selects its offline path. + +This prevents a model from interpreting serialized weights with an online +loader or silently applying the wrong method to one transformer in a +multi-transformer pipeline. + +## Model and layer integration + +Diffusion model constructors receive the resolved configuration and pass it to +quantizable vLLM layers. Each method configuration is responsible for: + +- identifying supported layers in `get_quant_method()`; +- creating the correct parameter shapes and loaders; +- processing weights after loading when conversion is required; and +- applying the platform-specific quantized operation at runtime. + +Model code must keep precision-sensitive or unsupported modules unquantized. +The shared `safe_quant_config()` helper is used by model integrations for +normalization and modulation layers where a generic FP8 configuration would be +unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers +are not implicitly quantized just because the main transformer is quantized; +they require explicit component routing and method support. + +Weight quantization and quantized KV cache are separate paths. The former is +resolved through `quantization_config`; runtime attention/KV-cache quantization +is configured by the diffusion attention configuration and must be validated +independently. + +## Platform and parallelism boundaries + +Platform checks belong in the method implementation rather than in the common +factory. This keeps configuration parsing portable while allowing CUDA, NPU, +and XPU methods to select different kernels or reject unsupported execution at +model construction time. The user guide is the source of truth for the current +hardware support matrix. + +Quantization methods must also respect the weight layout produced by the +configured parallelism strategy: + +- Tensor-parallel layers may receive partitioned weight shapes and loaders. +- Sequence-parallel execution changes activation movement but does not by + itself change the quantization configuration. +- Data-parallel replicas build their own method state; quantization metadata + must be available consistently on every rank. +- `WeightsMapper` transformations must be applied before matching ignored-layer + lists or component prefixes. + +Distributed layerwise offload has an additional boundary: its sharded mmap +AllGather path is not compatible with online quantization because the path +expects already materialized checkpoint weights. Use the rank-local +`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when +combining these features. See the [DLO design document](distributed_layerwise_offload.md) +for the full compatibility matrix. + +## Adding a quantization backend + +Contributors should follow this sequence: + +1. Check whether an existing vLLM backend can provide the required weight + format and kernel behavior. Prefer reuse over a duplicate Omni backend. +2. If Omni-specific behavior is required, implement a `QuantizationConfig` and + its `QuantizeMethodBase`/linear method in `vllm_omni/quantization/`. +3. Implement `get_name()`, supported activation dtypes, minimum capability, + checkpoint parsing, ignored-layer handling, and `get_quant_method()`. +4. Add a lazy factory override in `factory.py`. Keep optional platform + dependencies out of package import paths and normalize aliases to one + canonical method name. +5. If the method is pre-quantized, define the checkpoint metadata and the + serialized loading path before adding an online path. +6. Add model-specific component routing only when the model contains stages + that must use different policies. +7. Document user-visible configuration in the quantization user guide and keep + implementation rationale here. + +The public factory registration hook, +`register_quantization_override(method, builder)`, is intended for integrations +that need to register an Omni-compatible builder without changing the vLLM +registry. + +## Validation + +Validation should cover the complete path, not only whether a config object can +be instantiated: + +| Area | Required checks | +| --- | --- | +| Factory | Strings, aliases, flat dictionaries, invalid inputs, passthrough, and component maps | +| Method config | Online/offline selection, ignored layers, platform guards, and weight loading | +| Checkpoint metadata | Auto-detection, method mismatch errors, serialized flags, and per-transformer metadata | +| Model integration | Correct quantization scope, unquantized components, and parallel weight layouts | +| Numerical quality | Same-seed BF16 versus quantized output comparison at representative resolutions and step counts | +| Performance | Peak device memory, generation latency, and any backend-specific throughput metric | + +The unit tests live under `tests/diffusion/quantization/`. The trajectory +similarity tool in +`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` +provides a repeatable quality and memory comparison for diffusion models. +Method support and feature combinations remain model- and hardware-dependent; +passing factory tests alone does not establish production support. + +## Design invariants + +- `quantization_config` is the single runtime source of truth after config + normalization. +- A checkpoint method mismatch fails loudly instead of risking weight + corruption. +- Components without an explicit supported quantization path remain in their + original precision. +- Optional backend dependencies are imported lazily. +- Method implementations own platform-specific capability checks and kernels. +- User-facing support claims require model, hardware, numerical, and memory + validation in addition to configuration-level tests. diff --git a/docs/design/feature/ray_based_execution.md b/docs/design/feature/ray_based_execution.md deleted file mode 100644 index 4f4caa642..000000000 --- a/docs/design/feature/ray_based_execution.md +++ /dev/null @@ -1,59 +0,0 @@ -# Distributed utils - -This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends. -## 1. Installation -```bash -pip install "ray[default]" -``` -## 2. Ray Utils - -The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for: -* **Multi-node deployment**: Running pipeline stages across different physical machines. -* **Resource management**: Efficient GPU/CPU allocation. - -### 2.1 Basic Usage - -To use the Ray backend, specify `worker_backend="ray"` when initializing the engine. - -**Command Line Example:** -```bash -vllm serve Qwen/Qwen2.5-Omni-7B \ - --omni \ - --port 8091 \ - --worker-backend ray \ - --ray-address auto -``` - -### 2.2 Cluster Setup - -**Step 1: Start Head Node** -Run this on your primary machine: -```bash -ray start --head --port=6399 -``` - -**Step 2: Connect Worker Nodes** -Run this on each worker machine: -```bash -ray start --address=:6399 -``` - -> **Tip**: For a complete cluster setup script, refer to the vLLM example: -> [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh) - -### 2.3 Distributed Connector Support - -When running on Ray, the system automatically adapts its communication strategy: - -* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback). -* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma). -* **SHM payload routing**: `SharedMemoryConnector` always stores payloads in `/dev/shm`. - -### 2.4 Internal Helpers - -* **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one. - -## 3. Troubleshooting - -* **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open. -* **Version Mismatch**: Ensure all nodes run the same version of Ray and Python. diff --git a/docs/design/figures/dlo/dlo_pipeline.gif b/docs/design/figures/dlo/dlo_pipeline.gif new file mode 100644 index 000000000..cc206d03b Binary files /dev/null and b/docs/design/figures/dlo/dlo_pipeline.gif differ diff --git a/docs/design/index.md b/docs/design/index.md index d4b36e693..5434b27fb 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -13,7 +13,6 @@ while the source paths remain stable so existing links continue to work. ### Runtime and stage execution - [Disaggregated Inference](feature/disaggregated_inference.md) -- [Ray-based Execution](feature/ray_based_execution.md) - [Async Chunk](feature/async_chunk.md) - [Async Diffusion Output](feature/async_diffusion_output.md) - [Async Omni Output Materialization](feature/omni_async_output_materialization.md) @@ -48,11 +47,12 @@ while the source paths remain stable so existing links continue to work. #### Quantization -- [Quantization Overview](../user_guide/quantization/overview.md) +- [Quantization](feature/quantization.md) - [Cache-DiT](feature/cache_dit.md) - [TeaCache](feature/teacache.md) - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md) +- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md) ## Infrastructure and Performance diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload.md similarity index 88% rename from docs/user_guide/diffusion/cpu_offload_diffusion.md rename to docs/user_guide/diffusion/cpu_offload.md index 86a0e9201..02ab9f074 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload.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/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 3f631d5e5..83bc3643e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on | Method | Description | Best For | |--------|-------------|----------| -| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs | +| **[CPU Offload](diffusion/cpu_offload.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs | | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss | | **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak | @@ -104,8 +104,8 @@ The following tables show which models support each feature: > Notes: -> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md). -> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details. +> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md). +> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details. ### ImageGen @@ -274,9 +274,9 @@ Measured on NVIDIA H800: **Memory Optimization:** -- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage +- **[CPU Offload Guide](diffusion/cpu_offload.md)** - Offload model components to CPU, reduce GPU memory usage - **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos -- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models +- **[Quantization](quantization/overview.md)** - Quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models **Extensions:** diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 3a4a45573..ee36f5b60 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -5,6 +5,9 @@ path. The same configuration entrypoint is used across diffusion-only models, multi-stage omni/TTS models, and multi-stage diffusion models, but each model type has a different quantization scope. +For the internal architecture and backend extension points, see the +[quantization design document](../../design/feature/quantization.md). + ## Quantization Modes | Mode | Guide | Description | Methods | diff --git a/mkdocs.yml b/mkdocs.yml index cda933c44..fd97b9b37 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -75,6 +75,8 @@ plugins: redirect_maps: user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md + user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md + user_guide/quantization/index.md: user_guide/quantization/overview.md design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md - glightbox 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