docs: organize offloader feature designs
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed
1 parent
8b61499662
commit
a52c19ea8b
14 files changed
+29
-25
No files matched your search
@@ -0,0 +1,80 @@
|
||||
# CPU Offloading
|
||||
|
||||
This document defines the shared architecture for diffusion CPU offloading.
|
||||
The three strategies have separate user and design pages:
|
||||
|
||||
| Strategy | User guide | Design contract |
|
||||
| --- | --- | --- |
|
||||
| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
|
||||
| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
|
||||
| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
|
||||
|
||||
## Strategy selection
|
||||
|
||||
`OffloadConfig.from_od_config()` converts diffusion configuration into one
|
||||
`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
|
||||
|
||||
1. `DISTRIBUTED_LAYER_WISE`;
|
||||
2. `LAYER_WISE`;
|
||||
3. `MODEL_LEVEL`; and
|
||||
4. `NONE`.
|
||||
|
||||
The factory logs when a higher-priority strategy disables another. It derives
|
||||
parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
|
||||
a separate offload group size.
|
||||
|
||||
`get_offload_backend()` then validates platform offload support, resolves the
|
||||
device, and creates exactly one backend. Returning `None` means offloading is
|
||||
disabled or unsupported and must not leave partially installed hooks.
|
||||
|
||||
## Shared lifecycle
|
||||
|
||||
Every backend implements `OffloadBackend`:
|
||||
|
||||
- `enable(pipeline)` discovers modules, establishes initial residency, and
|
||||
installs hooks;
|
||||
- `disable()` removes owned hooks and resources; and
|
||||
- `is_enabled()` reports lifecycle state.
|
||||
|
||||
`disable()` does not promise to restore the pipeline's original device
|
||||
placement. The caller owns any subsequent rematerialization.
|
||||
|
||||
Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
|
||||
must use distinct hook names and remove only hooks they own.
|
||||
|
||||
## Discovery and topology
|
||||
|
||||
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
|
||||
|
||||
- `_dit_modules`;
|
||||
- `_encoder_modules`;
|
||||
- `_vae_modules`; and
|
||||
- `_resident_modules`.
|
||||
|
||||
Dotted paths are supported. Legacy pipelines may use the fallback scan of
|
||||
well-known attribute names, but new integrations should declare components
|
||||
explicitly.
|
||||
|
||||
Ordinary layerwise offload obtains topology from declared
|
||||
`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
|
||||
compatibility path). Distributed layerwise offload may instead use a pipeline
|
||||
`OffloadPlan`, then falls back to declared block attributes and its documented
|
||||
compatibility discovery. Discovery metadata describes structure only; the
|
||||
backend remains responsible for transfer, synchronization, and storage
|
||||
ownership.
|
||||
|
||||
## Cross-strategy invariants
|
||||
|
||||
- At most one strategy owns offload hooks for a pipeline.
|
||||
- A parameter has one authoritative host representation while offloaded.
|
||||
- Device storage is not freed until dependent compute or transfer work is
|
||||
complete.
|
||||
- Non-persistent and model-specific buffers remain correct after movement or
|
||||
checkpoint rematerialization.
|
||||
- Unsupported parallel or loading combinations fail before hooks mutate the
|
||||
model.
|
||||
- Platform streams, events, synchronization, and cache management go through
|
||||
the vLLM-Omni platform abstraction.
|
||||
|
||||
The diffusion [Offloader module design](../../module/diffusion/offloader.md)
|
||||
describes how these feature contracts fit into the larger diffusion runtime.
|
||||
@@ -0,0 +1,56 @@
|
||||
# Model-Level Offload
|
||||
|
||||
Model-level offload enforces mutual exclusion between pipeline component
|
||||
groups. For commands and model support, see the
|
||||
[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
|
||||
|
||||
## Generic hook path
|
||||
|
||||
`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It
|
||||
makes encoders and VAEs device resident for initialization, then installs
|
||||
`SequentialOffloadHook` instances:
|
||||
|
||||
- each DiT hook offloads all encoders and other DiTs before loading itself;
|
||||
- each encoder hook offloads every DiT before loading itself; and
|
||||
- VAEs are not part of the mutual-exclusion hook set.
|
||||
|
||||
The hook moves parameters and buffers without recursively calling
|
||||
`module.to()`. This avoids recursion through modules that retain references to
|
||||
other transformer blocks. Host tensors are pinned when configured; DTensors
|
||||
are not pinned through the ordinary tensor path.
|
||||
|
||||
Before forward returns control to the model, the hook synchronizes the active
|
||||
platform so compute cannot race the transfer. XPU uses blocking device-to-host
|
||||
copies because its allocator cannot safely evict storage while an asynchronous
|
||||
copy is pending.
|
||||
|
||||
## Split-model delegation
|
||||
|
||||
If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend
|
||||
delegates instead of installing generic encoder/DiT hooks. This supports models
|
||||
whose mutually exclusive components are nested inside one transformer, such
|
||||
as a reasoner and generator.
|
||||
|
||||
The delegated pipeline owns its internal contexts, but must preserve the same
|
||||
contract:
|
||||
|
||||
- only the executing component is device resident;
|
||||
- enable and disable are idempotent;
|
||||
- transfers use the supplied device, pinning, and HSDP settings; and
|
||||
- disabling removes all model-local offload hooks or contexts.
|
||||
|
||||
The delegation is currently duck typed. If another independent implementation
|
||||
is added, formalize the enable/disable pair as a protocol rather than expanding
|
||||
attribute-name heuristics.
|
||||
|
||||
## Extension invariants
|
||||
|
||||
- New ordinary pipelines declare both DiT and encoder groups through
|
||||
`SupportsComponentDiscovery`.
|
||||
- Models with multiple DiTs offload the inactive DiTs as well as encoders.
|
||||
- Empty or parameterless modules do not trigger transfers.
|
||||
- Hook removal does not move modules; callers must not assume original
|
||||
residency is restored.
|
||||
|
||||
Shared selection and lifecycle behavior is defined in the
|
||||
[CPU Offloading design](README.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/offloader/distributed_layerwise_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]
|
||||
```
|
||||
|
||||
|
||||
|
||||
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.
|
||||
@@ -0,0 +1,57 @@
|
||||
# Layerwise Offload
|
||||
|
||||
Layerwise offload streams transformer blocks through one accelerator. For
|
||||
commands and model support, see the
|
||||
[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md).
|
||||
|
||||
## Hook ring
|
||||
|
||||
`LayerWiseOffloadBackend` discovers every streamable block and registers one
|
||||
`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned
|
||||
host representation of its registered block and a reference to the next block
|
||||
it must prefetch; the final hook prefetches the first block for the next
|
||||
denoising iteration.
|
||||
|
||||
During initialization, parameters and buffers are grouped by dtype and copied
|
||||
into contiguous pinned CPU tensors. Original tensor storage is replaced with
|
||||
an empty or meta placeholder while identity and module structure are retained.
|
||||
Metadata records names, offsets, shapes, and dtypes for rematerialization.
|
||||
|
||||
## Forward lifecycle
|
||||
|
||||
For each block:
|
||||
|
||||
1. the previous hook's completion event ensures this block is materialized;
|
||||
2. the copy stream asynchronously allocates and fills the next block's device
|
||||
tensors;
|
||||
3. parameter and buffer storage is rebound to views of those tensors;
|
||||
4. the current block computes on the normal stream; and
|
||||
5. its post-forward hook replaces device storage with placeholders after all
|
||||
dependent work is safe.
|
||||
|
||||
The copy stream waits on the compute stream before reusing storage, and the
|
||||
compute stream waits on the per-block prefetch event before forward. These
|
||||
dependencies, not incidental global synchronization, define correctness.
|
||||
|
||||
## Topology contract
|
||||
|
||||
A DiT declares ordered block-container attributes using
|
||||
`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
|
||||
be an executable `nn.Module`; ordering must match forward execution. Multiple
|
||||
containers are concatenated in execution order.
|
||||
|
||||
Non-block DiT modules, encoders, VAEs, and declared resident modules are not
|
||||
streamed by this backend.
|
||||
|
||||
## Invariants and limitations
|
||||
|
||||
- Each offloaded tensor has one pinned host backing store.
|
||||
- Parameter identity is preserved when storage is rebound.
|
||||
- Blocks with heterogeneous shapes and dtypes use their recorded metadata;
|
||||
implementations must not assume identical blocks.
|
||||
- The ring is initialized before the first forward so block zero is available.
|
||||
- This backend is single-device; distributed weight reconstruction belongs to
|
||||
[Distributed Layerwise Offload](distributed_layerwise_offload.md).
|
||||
|
||||
Shared selection and lifecycle behavior is defined in the
|
||||
[CPU Offloading design](README.md).
|
||||
Reference in new issue
Block a user