65 lines
1.9 KiB
Markdown
65 lines
1.9 KiB
Markdown
|
# 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/offloader/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.
|