docs: split diffusion feature guides
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
8b0a6779d4
commit
7872663bf0
21 files changed
+1138
-801
No files matched your search
@@ -0,0 +1,65 @@
|
||||
# Layerwise Offloading
|
||||
|
||||
Layerwise, or blockwise, offloading keeps one transformer block on the
|
||||
accelerator and prefetches the next block while the current block computes.
|
||||
It is best suited to compute-heavy video DiTs whose block execution can hide
|
||||
host-to-device transfers.
|
||||
|
||||
## Execution flow
|
||||
|
||||
Each block has a pre-forward and post-forward hook. Parameters are consolidated
|
||||
in pinned host tensors and rematerialized for execution on a dedicated copy
|
||||
stream.
|
||||
|
||||
| Block | Pre-forward hook | Forward | Post-forward hook |
|
||||
| --- | --- | --- | --- |
|
||||
| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
|
||||
| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
|
||||
| ... | ... | ... | ... |
|
||||
| last block | Prefetch block 0 | Compute last block | Free last block |
|
||||
|
||||
Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
|
||||
remain device resident.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
enable_layerwise_offload=True,
|
||||
)
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--omni --enable-layerwise-offload
|
||||
```
|
||||
|
||||
## Model integration
|
||||
|
||||
Transformer classes declare containers of executable blocks:
|
||||
|
||||
```python
|
||||
class WanTransformer3DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
|
||||
class Flux2Transformer2DModel(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = [
|
||||
"transformer_blocks",
|
||||
"single_transformer_blocks",
|
||||
]
|
||||
```
|
||||
|
||||
See the [layerwise design](../../../design/feature/layerwise_offload.md) for
|
||||
the discovery and hook invariants. `OffloadPlan` is a separate declarative
|
||||
topology path for distributed layerwise offload.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Single device only.
|
||||
- Setup consolidates and pins block parameters, increasing cold-start time.
|
||||
- Performance depends on block compute time and host-to-device bandwidth;
|
||||
lightweight blocks may not hide transfers.
|
||||
Reference in new issue
Block a user