[diffusion] feat: add explicit snapshot-offload component residency (#38535)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Mick
2026-09-09 13:44:54 +08:00
committed by GitHub
co-authored by Mick Qian
parent 0ee8e41a4e
commit 00a9028e87
20 changed files with 835 additions and 29 deletions
+29
View File
@@ -313,6 +313,7 @@ The available modes are:
- `resident`: keep the complete component on the accelerator.
- `component-offload`: keep the complete component on CPU between uses, moving it to the accelerator before each declared use and back to CPU afterward.
- `snapshot-offload`: retain CPU parameter storage while a complete component runs on GPU. At the end of a use, restore the CPU parameters instead of copying GPU weights back. Mutable buffers still move back to CPU.
- `layerwise-offload`: keep component weights on CPU and stream its declared layers during execution.
Selectors match exact loaded component keys from `model_index.json`, including names such as `transformer_2`, `audio_vae`, and `connectors`. The group selectors `dit`, `text_encoder`, `image_encoder`, and `vae` are also available, together with `all`. An exact key overrides a matching group, and a group overrides `all`. Components without a matching canonical selector retain their explicit legacy setting or automatic/model default.
@@ -321,6 +322,34 @@ The existing `--dit-cpu-offload`, `--text-encoder-cpu-offload`, `--image-encoder
Layerwise selection is strict. A native weighted component selected for `layerwise-offload` must declare its layer structure; otherwise startup fails with the unsupported component name instead of silently changing modes. Explicit non-resident placement also requires a request-time component-use declaration, so it cannot silently select a module that the pipeline does not manage. FSDP applies only to resident components. The Diffusers backend supports only pipeline-wide `all=resident` and `all=component-offload`.
#### Snapshot offload
On NVIDIA CUDA GPUs with separate host and device memory, explicitly select
`--component-residency vae=snapshot-offload` to avoid repeated weight D2H transfers.
Existing defaults and `component-offload` behavior are unchanged.
Snapshot offload retains a full host copy during GPU execution. With
`--pin-cpu-memory` enabled (the default), it shares a per-worker host pin budget
with layerwise offload. Layerwise initialization claims its allowance first;
snapshot weights use the remainder. Shared parameter storage is pinned once and
reused across requests; its allowance is returned when the storage is released.
Pinning also checks current host headroom, including the cgroup limit.
Weights that do not fit the budget, or have pinning disabled, retain their
existing CPU storage. Checkpoint mmap stays file-backed unless pinning, dtype
conversion, or weight mutation materializes a copy. Pageable H2D can be slower,
and pinning temporarily needs space for both the original and pinned storage.
This budget limits new pinned weight allocations, not total process RAM or the
CUDA host allocator's cache. Benchmark repeated requests with the actual host
storage and workload; avoiding D2H does not guarantee an end-to-end speedup.
The whole component must still fit on GPU. Unlike layerwise offload with all layers
retained across denoising steps, snapshot offload does not require a layer-structure
declaration and also covers parameters outside transformer blocks. Weight updates,
LoRA merge/unmerge, and sleep restore the CPU weights before mutation or release.
FSDP-managed components, shared-memory GPUs, and snapshot-offloaded DiTs with
breakable CUDA graphs are not supported.
### Layerwise Offload Tuning
Use layerwise offload when a component does not fit comfortably in GPU memory. The compatibility options `--dit-layerwise-offload` and `--layerwise-offload-components` remain available (`--layerwise-offload-modules` is an alias), while new deployments can select the mode directly:
@@ -132,7 +132,7 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served
The modes tune native pipeline components declared to the component residency manager. DiTs, text/image encoders, VAEs, vocoders, adapters, and upsamplers can use layerwise offload when their native module declares its executable layer structure. Explicitly selecting an unsupported component fails at startup instead of falling back to another residency mode.
For direct control, assign one of `resident`, `component-offload`, or `layerwise-offload` with `--component-residency COMPONENT=MODE`:
For direct control, assign one of `resident`, `component-offload`, `snapshot-offload`, or `layerwise-offload` with `--component-residency COMPONENT=MODE`:
```bash
sglang generate \