KyleandClaude Opus 5
f247a307a8
Review of the first version found that the flag was unsafe as a global
default. Releasing a component and loading it again is only correct when
nothing outside the loader has changed it, and two habits in this tree break
that without raising:
* mutating a component after load. LongCatPipeline.initialize_pipeline turns
on block-sparse attention and writes parameters into every transformer
block. That runs once, so a re-materialized component silently comes back
with the feature off and the second generation quietly degrades.
* reading a component's attributes while stages are built. The shared
DenoisingStage.__init__ derives the attention backend from
transformer.hidden_size, which materializes the DiT during post_init and
defeats the deferral it was meant to gain.
_lazy_module_names is now empty in the base class, so an unchecked pipeline
gets no deferral and says so. MiniMax-H3 opts in to the four components this
PR measured, and nothing else changes behaviour.
Also from review:
Releasing on the failure path. A stage's hook frees only what that stage is
the last user of, and it ran only after a successful forward. Both the stage
and the whole run now release on the way out, so the retry a memory
constrained caller attempts does not start from a worse position than the
request that just failed. A failing release cannot replace the exception
being propagated.
Walking into nested stages. Cosmos25AutoDenoisingStage keeps the transformer
inside child stages, so a one-level scan called it unreferenced and never
freed it. The scan now recurses through stages and containers with cycle
protection.
Keeping the proxy when torch.compile is skipped. The FSDP check ran after the
proxy had already been replaced by the real module, so an FSDP-wrapped
component lost both the compile and its release hook.
Not attaching the activation trace to a deferred component, since the hook
manager pins every module it wraps.
test_parser.py asserted on a whole serialized config dict, so adding a field
to OffloadConfig broke it. Grepping the field name could not find that; only
running the suite could.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
194 lines
7.2 KiB
Markdown
194 lines
7.2 KiB
Markdown
# Offloading
|
|
|
|
This page describes how to use offloading techniques for inference to reduce GPU memory usage while maintaining acceptable performance.
|
|
|
|
## Default Behavior
|
|
|
|
```python
|
|
dit_cpu_offload: bool = True
|
|
use_fsdp_inference: bool = False
|
|
dit_layerwise_offload: bool = True
|
|
text_encoder_cpu_offload: bool = True
|
|
image_encoder_cpu_offload: bool = True
|
|
vae_cpu_offload: bool = True
|
|
pin_cpu_memory: bool = True
|
|
lazy_module_load: bool = False
|
|
```
|
|
|
|
On unified-memory accelerators such as NVIDIA GB10 and Apple silicon, FastVideo
|
|
detects the selected device inside each worker and disables all five host-offload
|
|
modes before loading modules. Host and accelerator allocations share one physical
|
|
pool there, so offload adds transfers and duplicate residency instead of freeing
|
|
memory. CUDA FSDP sharding remains enabled when requested; MPS continues to
|
|
disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged.
|
|
|
|
## Behavior Explanation
|
|
|
|
!!! note
|
|
For CLI usage, replace underscores (`_`) with hyphens (`-`).
|
|
|
|
### `use_fsdp_inference`
|
|
|
|
Enables [FSDP](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) for inference. The model weights are sharded across multiple GPUs to reduce memory usage per GPU, and weights are broadcast to all GPUs layer by layer during inference.
|
|
|
|
#### Performance Impact
|
|
|
|
FSDP inference introduces negligible performance overhead due to weight prefetching. Performance overhead may be visible when GPU interconnect is slow (e.g., multiple consumer-level GPUs connected by slow PCIe without GPU P2P support).
|
|
|
|
#### Usage Recommendation
|
|
|
|
We recommend enabling this option when multiple GPUs are available.
|
|
|
|
### `dit_cpu_offload`
|
|
|
|
Enables CPU offloading for FSDP inference. When enabled, the model weights are offloaded to CPU memory, and the weight of each layer is moved to GPU memory only when that layer is being computed.
|
|
|
|
#### Performance Impact
|
|
|
|
The PyTorch FSDP implementation does not overlap computation and data transfer perfectly for inference, so enabling this option will harm performance.
|
|
|
|
#### Usage Recommendation
|
|
|
|
This option only takes effect when FSDP is enabled. For single GPU usage, we recommend using `dit_layerwise_offload` instead.
|
|
|
|
### `dit_layerwise_offload`
|
|
|
|
This option is similar to `dit_cpu_offload`, but with two key differences:
|
|
|
|
1. It overlaps computation and PCIe data transfer
|
|
2. It only works for single GPU inference
|
|
|
|
#### Performance Impact
|
|
|
|
This option introduces negligible performance overhead.
|
|
|
|
#### Usage Recommendation
|
|
|
|
We recommend enabling this option for single GPU usage. This option is not compatible with FSDP.
|
|
|
|
### `text_encoder_cpu_offload`
|
|
|
|
When enabled, the text encoder model weights are offloaded to CPU memory, and text encoding is computed on CPU.
|
|
|
|
#### Performance Impact
|
|
|
|
This option significantly slows down text encoding computation, but text encoding is usually not the bottleneck.
|
|
|
|
#### Usage Recommendation
|
|
|
|
We recommend enabling this option only when OOM happens.
|
|
|
|
### `image_encoder_cpu_offload` and `vae_cpu_offload`
|
|
|
|
When enabled, the weights are stored in CPU memory and moved to GPU memory when the corresponding module is being computed. After computation, the weights are moved back to CPU memory.
|
|
|
|
#### Performance Impact
|
|
|
|
These options introduce performance overhead due to PCIe data transfer.
|
|
|
|
#### Usage Recommendation
|
|
|
|
We recommend enabling these options when OOM happens.
|
|
|
|
### `lazy_module_load`
|
|
|
|
Every option above moves weights between host and device. This one changes
|
|
whether they are in memory at all.
|
|
|
|
By default a pipeline loads every component before the first stage runs, so
|
|
peak memory is the sum of all of them even though no two are needed at the same
|
|
moment. With `lazy_module_load` enabled, each heavy component loads on first use
|
|
and is freed once the last stage that needs it has returned, so peak memory
|
|
becomes the largest overlapping set instead of the sum. For a text-to-video
|
|
pipeline that is roughly `max(text encoder, DiT + VAE)` rather than
|
|
`text encoder + DiT + VAE`.
|
|
|
|
#### Performance Impact
|
|
|
|
A freed component is read from disk again on the next generation, so a
|
|
multi-prompt run pays one reload per component per request. For a large text
|
|
encoder that is tens of seconds.
|
|
|
|
#### Usage Recommendation
|
|
|
|
Enable this when a model does not fit at load time, which the CPU offload
|
|
options above cannot help with because they act after loading. It is
|
|
particularly relevant on unified-memory devices, where host and device draw on
|
|
the same pool and moving weights to the host frees nothing. Leave it off when
|
|
the model already fits.
|
|
|
|
This option applies to inference only. Training keeps every component resident
|
|
and logs a warning if the flag is set.
|
|
|
|
Deferral is opt-in per pipeline. Releasing a component and loading it again is
|
|
only safe when nothing outside the loader has changed it, and two common habits
|
|
break that without raising: mutating a component after load, as LongCat does
|
|
when it enables block-sparse attention, and reading a component's attributes
|
|
while stages are built, as the shared denoising stage does to pick an attention
|
|
backend. A pipeline therefore lists the components it has checked in
|
|
`_lazy_module_names`, which is empty in the base class. MiniMax-H3 opts in. On
|
|
a pipeline that has not, the flag logs a warning and changes nothing.
|
|
|
|
## General Recommendations
|
|
|
|
### Single GPU Inference
|
|
|
|
We recommend enabling `dit_layerwise_offload`. If OOM happens, also enable `image_encoder_cpu_offload` and `vae_cpu_offload`. If OOM still happens, consider enabling `text_encoder_cpu_offload`.
|
|
|
|
### Multi-GPU Inference
|
|
|
|
We recommend enabling `use_fsdp_inference` and disabling both `dit_layerwise_offload` and `dit_cpu_offload`. If OOM happens, consider enabling `text_encoder_cpu_offload`, `image_encoder_cpu_offload`, and `vae_cpu_offload`. If OOM still happens, consider enabling `dit_cpu_offload`.
|
|
|
|
### When the Model Does Not Fit at Load Time
|
|
|
|
The offload options only help once loading has finished. If the run dies while
|
|
components are still being placed, or if the machine has unified memory so
|
|
there is no separate host pool to offload into, enable `lazy_module_load`.
|
|
|
|
## Examples
|
|
|
|
### Single GPU with Layerwise Offloading
|
|
|
|
```python
|
|
from fastvideo import VideoGenerator
|
|
|
|
generator = VideoGenerator.from_pretrained(
|
|
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
|
num_gpus=1,
|
|
# Recommended for single GPU
|
|
dit_layerwise_offload=True,
|
|
# Enable if OOM happens
|
|
vae_cpu_offload=True,
|
|
image_encoder_cpu_offload=True,
|
|
text_encoder_cpu_offload=True,
|
|
# Speeds up CPU-GPU transfer
|
|
pin_cpu_memory=True,
|
|
)
|
|
|
|
prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers."
|
|
video = generator.generate_video(prompt, output_path="output/", save_video=True)
|
|
```
|
|
|
|
### Multi-GPU with FSDP
|
|
|
|
```python
|
|
from fastvideo import VideoGenerator
|
|
|
|
generator = VideoGenerator.from_pretrained(
|
|
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
|
num_gpus=2,
|
|
# Recommended for multi-GPU
|
|
use_fsdp_inference=True,
|
|
dit_layerwise_offload=False,
|
|
dit_cpu_offload=False,
|
|
# Enable if OOM happens
|
|
vae_cpu_offload=True,
|
|
image_encoder_cpu_offload=True,
|
|
text_encoder_cpu_offload=True,
|
|
pin_cpu_memory=True,
|
|
)
|
|
|
|
prompt = "A majestic lion strides across the golden savanna."
|
|
video = generator.generate_video(prompt, output_path="output/", save_video=True)
|
|
```
|