[diffusion] refactor: unify component residency controls (#34736)
This commit is contained in:
@@ -245,16 +245,18 @@ SGLANG_CACHE_DIT_MC=3 \
|
||||
sglang serve --model-path krea/Krea-2-Raw
|
||||
```
|
||||
|
||||
#### 4.2.2 Memory & CPU Offload
|
||||
#### 4.2.2 Memory and Component Residency
|
||||
|
||||
Krea-2's DiT is ~24 GB in bf16 (the bulk of the model). On memory-constrained GPUs you can keep less of it resident:
|
||||
|
||||
- `--dit-layerwise-offload`: stream the DiT's transformer blocks layer-by-layer with async host-to-device prefetch overlap, so only a small working set stays on the GPU. This is the primary way to fit Krea-2 on a single consumer / 32 GB-class card, at a modest latency cost. Tune the memory/latency trade-off with `--dit-offload-prefetch-size` (`0.0` prefetches one layer for the lowest memory; larger values prefetch more layers -- faster but more memory).
|
||||
- `--dit-cpu-offload`: keep the whole DiT in host memory. Combine it with `--dit-layerwise-offload` for the lowest peak GPU memory (weights stay on host and only the layers needed for the current step are brought on-device).
|
||||
- `--text-encoder-cpu-offload`: offload the Qwen3-VL text encoder (it is idle during the denoise loop).
|
||||
- `--vae-cpu-offload`: offload the VAE.
|
||||
- `--component-residency dit=layerwise-offload`: stream the DiT's transformer blocks layer-by-layer with async host-to-device prefetch overlap, so only a small working set stays on the GPU. This is the primary way to fit Krea-2 on a single consumer / 32 GB-class card, at a modest latency cost. Tune the memory/latency trade-off with `--dit-offload-prefetch-size` (`0.0` prefetches one layer for the lowest memory; larger values prefetch more layers -- faster but more memory).
|
||||
- `--component-residency dit=component-offload`: keep the complete DiT on CPU between denoising uses. This and layerwise offload are distinct modes; do not combine them for the same component.
|
||||
- `--component-residency text_encoder=component-offload`: offload the Qwen3-VL text encoder while it is idle during denoising.
|
||||
- `--component-residency vae=component-offload`: offload the VAE between uses.
|
||||
- `--pin-cpu-memory`: pin host memory for offload. Add only as a temporary workaround if you hit `CUDA error: invalid argument`.
|
||||
|
||||
The legacy `--dit-layerwise-offload`, `--dit-cpu-offload`, `--text-encoder-cpu-offload`, and `--vae-cpu-offload` forms remain accepted. If both legacy DiT offload flags are enabled, layerwise offload is the effective DiT mode.
|
||||
|
||||
On large-VRAM GPUs (e.g. H200), keep everything resident (offloads off) for the fastest latency.
|
||||
|
||||
## 5. Benchmark
|
||||
|
||||
@@ -32,6 +32,14 @@ Each recipe provides step-by-step instructions to help you quickly implement SGL
|
||||
3. Adapt configurations to your specific hardware and requirements
|
||||
4. Join our community to share feedback and improvements
|
||||
|
||||
For memory placement, prefer the unified
|
||||
[`--component-residency`](/docs/sglang-diffusion/api/cli#component-residency)
|
||||
selector. Each component resolves to exactly one of `resident`,
|
||||
`component-offload`, or `layerwise-offload`. Existing options such as
|
||||
`--dit-cpu-offload`, `--text-encoder-cpu-offload`,
|
||||
`--image-encoder-cpu-offload`, and `--vae-cpu-offload` remain supported by all
|
||||
recipes that already use them.
|
||||
|
||||
The sglang diffusion cookbook directory structure are shown below:
|
||||
|
||||
```text Example
|
||||
|
||||
@@ -200,33 +200,42 @@ For supported native pipelines, set `SGLANG_CACHE_DIT_ENABLED=true` to enable Ca
|
||||
|
||||
For supported image pipelines, breakable CUDA graph can be enabled with `--enable-breakable-cuda-graph`, but you must declare every served resolution in `--warmup-resolutions` so warmup captures matching graph signatures.
|
||||
|
||||
### Component CPU Offload
|
||||
### Component Residency
|
||||
|
||||
Use `--cpu-offload-components` to explicitly select components for coarse CPU offload:
|
||||
Use `--component-residency COMPONENT=MODE` to assign one runtime residency mode to each native pipeline component:
|
||||
|
||||
```bash Command
|
||||
sglang generate \
|
||||
--model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--cpu-offload-components dit text_encoder \
|
||||
--component-residency all=resident text_encoder=layerwise-offload vae=component-offload \
|
||||
--prompt "A quiet city street after rain"
|
||||
```
|
||||
|
||||
Component names are matched against the loaded pipeline's component keys from `model_index.json`, so names such as `transformer_2`, `audio_vae`, and `connectors` are supported without adding them to a registry. The group aliases `dit`, `text_encoder`, `image_encoder`, and `vae` remain available. Use `all` to offload every loaded `torch.nn.Module`, or `none` to disable coarse component CPU offload. Non-module components such as tokenizers and schedulers are unaffected. This unified option cannot be combined with the legacy per-component CPU offload flags. Layerwise offload remains independently controlled by `--layerwise-offload-components`.
|
||||
The available modes are:
|
||||
|
||||
### Layerwise Offload
|
||||
- `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.
|
||||
- `layerwise-offload`: keep component weights on CPU and stream its declared layers during execution.
|
||||
|
||||
Use layerwise offload when a large component does not fit comfortably in GPU memory. By default, `--dit-layerwise-offload` only applies to legacy DiT components. Use `--layerwise-offload-components` to select pipeline component names explicitly (`--layerwise-offload-modules` is accepted as an alias):
|
||||
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.
|
||||
|
||||
The existing `--dit-cpu-offload`, `--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`, `--vae-cpu-offload`, and `--cpu-offload-components` options remain supported. New and legacy options may be mixed: `--component-residency` wins only for components it matches, while unmatched legacy settings remain effective. Legacy layerwise selectors take precedence over legacy component-offload selectors for the same component. Explicit `--dit-layerwise-offload false` makes the DiT resident unless another explicit DiT selector, such as `--dit-cpu-offload true` or `--component-residency dit=component-offload`, selects a different mode.
|
||||
|
||||
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. FSDP applies only to resident components. The Diffusers backend supports only pipeline-wide `all=resident` and `all=component-offload`.
|
||||
|
||||
### 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:
|
||||
|
||||
```bash Command
|
||||
sglang generate \
|
||||
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--dit-layerwise-offload \
|
||||
--layerwise-offload-components transformer text_encoder \
|
||||
--component-residency transformer=layerwise-offload text_encoder=layerwise-offload \
|
||||
--dit-offload-prefetch-size 0 \
|
||||
--prompt "A quiet city street after rain"
|
||||
```
|
||||
|
||||
The values must match keys in the selected pipeline's `pipeline.modules`, such as `transformer`, `text_encoder`, `image_encoder`, `vae`, `condition_image_encoder`, `spatial_upsampler`, or `vocoder`. Use `all` to select every layerwise-offloadable component. Prefer the smallest component set that solves the memory issue because layerwise offload can increase latency.
|
||||
Layerwise tuning options such as `--dit-offload-prefetch-size`, `--dit-layerwise-resident-layers`, and `--dit-layerwise-residency-policy` continue to control the streamed layer working set. Prefer the smallest component set that solves the memory issue because layerwise offload can increase latency.
|
||||
|
||||
## Serve
|
||||
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
---
|
||||
title: "Deployment and Performance Modes"
|
||||
description: "Choose CPU offload, FSDP, CFG parallelism, SP, TP, and performance-mode presets in SGLang Diffusion."
|
||||
description: "Choose component residency, FSDP, CFG parallelism, SP, TP, and performance-mode presets in SGLang Diffusion."
|
||||
tag: "preserve"
|
||||
---
|
||||
|
||||
This page gives practical defaults for choosing `--performance-mode`, CPU offload, FSDP, CFG parallelism, SP, and TP.
|
||||
This page gives practical defaults for choosing `--performance-mode`, component residency, FSDP, CFG parallelism, SP, and TP.
|
||||
|
||||
## Quick Rule
|
||||
|
||||
@@ -24,15 +24,15 @@ Use the simplest setting that fits your memory target:
|
||||
<tbody>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Fastest single-GPU run when the model fits</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable CPU offload and do not use FSDP.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use resident components and do not use FSDP.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Lower single-GPU memory usage</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use component CPU offload, or layerwise DiT offload for supported Wan/MOVA models.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use component offload, then layerwise offload when a complete component still does not fit comfortably.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Faster multi-GPU Qwen/Wan CFG generation</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use FSDP with CFG parallelism and disable CPU offload.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use FSDP with CFG parallelism and keep the sharded component resident.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>Sequence length or video-shape scaling</td>
|
||||
@@ -48,7 +48,7 @@ Use the simplest setting that fits your memory target:
|
||||
Base the decision on available memory on the selected GPU(s).
|
||||
|
||||
- For multi-GPU deployment: the least-free selected GPU is the bottleneck. A busy 80GiB GPU can behave like a much smaller GPU.
|
||||
- For single-GPU deployment: FSDP shards DiT weights across multiple GPUs. It is not useful for keeping a single-GPU deployment on one GPU; for that case use CPU offload.
|
||||
- For single-GPU deployment: FSDP shards weights across multiple GPUs. It is not useful for keeping a single-GPU deployment on one GPU; use component or layerwise offload instead.
|
||||
|
||||
## Health Probes
|
||||
|
||||
@@ -130,7 +130,17 @@ See [OpenAI API: Served model name](/docs/sglang-diffusion/api/openai_api#served
|
||||
|
||||
`auto` checks selected GPU memory before applying FSDP. In multi-GPU runs it uses the least available memory across selected GPUs, and only turns on FSDP automatically when doing so can replace DiT offload. For image workloads with at least 45 GiB available per selected GPU, it keeps the repeatedly reused DiT resident and uses layerwise offload for large auxiliary encoders; below that threshold it keeps the DiT offloaded. Video DiT residency remains model- and workload-specific because frame count and resolution change its peak memory substantially. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `speed` intentionally does not check memory; it is the mode for users who prefer latency/throughput and accept OOM risk. It keeps `torch.compile` disabled by default because its effect varies by model and workload. A model-specific deployment config may enable a validated compile path, and `--enable-torch-compile true` always opts in explicitly.
|
||||
|
||||
The modes tune residency for native pipeline components declared to the component residency manager. Today this covers the major DiT, text/image encoder, VAE, vocoder, and upsampler components; DiT can use layerwise offload when supported, while text encoders use either resident execution or component CPU offload. Do not assume text-encoder layerwise offload unless a model implements and validates it.
|
||||
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`:
|
||||
|
||||
```bash
|
||||
sglang generate \
|
||||
--model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--component-residency dit=layerwise-offload text_encoder=component-offload vae=resident
|
||||
```
|
||||
|
||||
Existing per-component CPU-offload and layerwise flags remain supported. Canonical selectors override matching legacy settings only; unmatched legacy settings and automatic defaults remain effective. See [Component Residency](/docs/sglang-diffusion/api/cli#component-residency) for the complete precedence rules.
|
||||
|
||||
When `torch.compile` is enabled, `--offload-during-compile` stays on by default. During compile warmup it temporarily offloads the DiT and evicts resident non-DiT components so `max-autotune` fits on tighter-memory GPUs, then restores the configured serving residency before real traffic.
|
||||
|
||||
@@ -168,11 +178,11 @@ In this example, `auto` will not re-enable FSDP. The same applies to parallelism
|
||||
|
||||
## Interpreting The Levers
|
||||
|
||||
**No offload** keeps model components resident on GPU. It is usually fastest when memory is sufficient.
|
||||
**Resident** keeps the complete component on the accelerator. It is usually fastest when memory is sufficient.
|
||||
|
||||
**Component CPU offload** lowers GPU memory by moving large components to CPU. It is simple and robust, but it usually trades latency for memory.
|
||||
**Component offload** keeps a complete component on CPU between declared uses. It is simple and robust, but each use pays a whole-component transfer.
|
||||
|
||||
**Layerwise DiT offload** lowers DiT memory further for supported Wan/MOVA models by moving DiT layers between CPU and GPU. It can be the best single-GPU memory mode, but may increase latency and lower throughput.
|
||||
**Layerwise offload** streams the declared layers of any supported native weighted component. It lowers peak accelerator memory further, but may increase latency and lower throughput.
|
||||
|
||||
**FSDP** shards DiT weights across multiple GPUs and all-gathers weights during forward. It can reduce DiT CPU offload cost on multi-GPU deployments, especially for validated Wan I2V workloads.
|
||||
|
||||
|
||||
@@ -72,11 +72,28 @@ Or, more simply, with the CLI:
|
||||
|
||||
```bash
|
||||
sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--text-encoder-cpu-offload --pin-cpu-memory \
|
||||
--component-residency text_encoder=component-offload --pin-cpu-memory \
|
||||
--prompt "A curious raccoon" \
|
||||
--save-output
|
||||
```
|
||||
|
||||
### Component residency
|
||||
|
||||
Use `--component-residency COMPONENT=MODE` to choose one runtime mode for each
|
||||
loaded component:
|
||||
|
||||
- `resident` keeps the complete component on the accelerator.
|
||||
- `component-offload` stores the complete component on CPU between uses.
|
||||
- `layerwise-offload` streams the component's declared layers from CPU.
|
||||
|
||||
`COMPONENT` can be an exact `model_index.json` key or one of `all`, `dit`,
|
||||
`text_encoder`, `image_encoder`, and `vae`. Exact keys override groups, and
|
||||
groups override `all`. Existing options such as `--dit-cpu-offload`,
|
||||
`--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`,
|
||||
`--vae-cpu-offload`, and `--cpu-offload-components` remain supported. See the
|
||||
[CLI reference](https://docs.sglang.io/docs/sglang-diffusion/api/cli#component-residency)
|
||||
for precedence and compatibility details.
|
||||
|
||||
### LoRA support
|
||||
|
||||
Apply LoRA adapters via `--lora-path`:
|
||||
|
||||
@@ -451,6 +451,7 @@ def launch_pool_disagg_server(
|
||||
base_dict = {
|
||||
f.name: getattr(server_args, f.name)
|
||||
for f in dataclasses.fields(server_args)
|
||||
if f.init
|
||||
}
|
||||
base_dict.update(role_overrides)
|
||||
base_dict.pop("pipeline_config", None)
|
||||
@@ -727,7 +728,9 @@ def launch_disagg_role(server_args: ServerArgs):
|
||||
}
|
||||
|
||||
base_dict = {
|
||||
f.name: getattr(server_args, f.name) for f in dataclasses.fields(server_args)
|
||||
f.name: getattr(server_args, f.name)
|
||||
for f in dataclasses.fields(server_args)
|
||||
if f.init
|
||||
}
|
||||
base_dict.update(role_overrides)
|
||||
base_dict.pop("pipeline_config", None)
|
||||
|
||||
@@ -65,7 +65,7 @@ class AdapterLoader(ComponentLoader):
|
||||
# Not a fixed name: connectors follow DiT offload, while the duration
|
||||
# head stays resident unless selected explicitly.
|
||||
target_device = self.target_device(
|
||||
server_args.should_cpu_offload_component(component_name)
|
||||
server_args.should_start_component_on_cpu(component_name)
|
||||
)
|
||||
default_dtype = resolve_precision(
|
||||
server_args, component_name, precision_attr="dit_precision"
|
||||
|
||||
@@ -9,6 +9,9 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
|
||||
from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
|
||||
from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
|
||||
from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
RESIDENT,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
|
||||
@@ -74,12 +77,17 @@ class BridgeLoader(ComponentLoader):
|
||||
default_dtype,
|
||||
)
|
||||
|
||||
component_cpu_offload = server_args.should_cpu_offload_component(component_name)
|
||||
use_fsdp = server_args.should_use_fsdp_for_component(component_name)
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
component_name
|
||||
)
|
||||
|
||||
# Use the FSDP loader when FSDP is requested or shard rules are declared.
|
||||
fsdp_shard_conditions = getattr(model_cls, "_fsdp_shard_conditions", None)
|
||||
if server_args.use_fsdp_inference or (
|
||||
server_args.hsdp_shard_dim is not None and fsdp_shard_conditions
|
||||
if use_fsdp or (
|
||||
server_args.residency_mode(component_name) == RESIDENT
|
||||
and server_args.hsdp_shard_dim is not None
|
||||
and fsdp_shard_conditions
|
||||
):
|
||||
local_torch_device = get_local_torch_device()
|
||||
# Load with FSDP support
|
||||
@@ -90,9 +98,9 @@ class BridgeLoader(ComponentLoader):
|
||||
device=local_torch_device,
|
||||
hsdp_replicate_dim=server_args.hsdp_replicate_dim,
|
||||
hsdp_shard_dim=server_args.hsdp_shard_dim,
|
||||
cpu_offload=component_cpu_offload,
|
||||
component_starts_on_cpu=component_starts_on_cpu,
|
||||
pin_cpu_memory=server_args.pin_cpu_memory,
|
||||
fsdp_inference=server_args.use_fsdp_inference,
|
||||
fsdp_inference=use_fsdp,
|
||||
param_dtype=default_dtype,
|
||||
reduce_dtype=torch.float32,
|
||||
output_dtype=None,
|
||||
@@ -106,7 +114,7 @@ class BridgeLoader(ComponentLoader):
|
||||
model = model_cls.from_pretrained(
|
||||
component_model_path, torch_dtype=default_dtype
|
||||
)
|
||||
target_device = self.target_device(component_cpu_offload)
|
||||
target_device = self.target_device(component_starts_on_cpu)
|
||||
model = model.to(device=target_device, dtype=default_dtype)
|
||||
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
|
||||
@@ -14,7 +14,6 @@ from diffusers import AutoModel
|
||||
from torch import nn
|
||||
from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
|
||||
|
||||
from sglang.multimodal_gen.configs.models import ModelConfig
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.layers.attention.selector import (
|
||||
component_attn_backend_context_manager,
|
||||
@@ -25,19 +24,12 @@ from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
component_name_to_loader_cls,
|
||||
get_memory_usage_of_component,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
ComponentResidencyError,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
is_fsdp_managed_module,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
configure_layerwise_offload_modules,
|
||||
is_layerwise_offloaded_module,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS,
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
layerwise_component_matches_any_selection,
|
||||
normalize_layerwise_offload_components,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
|
||||
@@ -94,25 +86,15 @@ class ComponentLoader(ABC):
|
||||
self.device = device
|
||||
self.component_architecture: str | None = None
|
||||
|
||||
def should_offload(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
model_config: ModelConfig | None = None,
|
||||
component_name: str | None = None,
|
||||
):
|
||||
return component_name is not None and server_args.should_cpu_offload_component(
|
||||
component_name
|
||||
)
|
||||
|
||||
def target_device(self, should_offload):
|
||||
if should_offload:
|
||||
@staticmethod
|
||||
def target_device(component_starts_on_cpu: bool) -> torch.device:
|
||||
if component_starts_on_cpu:
|
||||
return (
|
||||
torch.device("mps")
|
||||
if current_platform.is_mps()
|
||||
else torch.device("cpu")
|
||||
)
|
||||
else:
|
||||
return get_local_torch_device()
|
||||
return get_local_torch_device()
|
||||
|
||||
def customized_load_kwargs_for_component(
|
||||
self, _server_args: ServerArgs, _component_name: str
|
||||
@@ -127,62 +109,6 @@ class ComponentLoader(ABC):
|
||||
)
|
||||
return component_name in native_only_components
|
||||
|
||||
@staticmethod
|
||||
def _is_component_set_as_layerwise_load(
|
||||
server_args: ServerArgs, component_name: str
|
||||
) -> bool:
|
||||
"""if a component should be loaded in a layerwise-fashion"""
|
||||
selected_component_names = normalize_layerwise_offload_components(
|
||||
server_args.layerwise_offload_components
|
||||
)
|
||||
if selected_component_names is None:
|
||||
return False
|
||||
selected_component_names = set(selected_component_names)
|
||||
if LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names:
|
||||
return True
|
||||
explicit_component_names = selected_component_names - {
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP
|
||||
}
|
||||
return layerwise_component_matches_any_selection(
|
||||
component_name, explicit_component_names
|
||||
)
|
||||
|
||||
def _maybe_configure_layerwise_after_startup_cpu_staging(
|
||||
self,
|
||||
component: AutoModel,
|
||||
server_args: ServerArgs,
|
||||
component_name: str,
|
||||
load_kwargs: dict[str, Any],
|
||||
) -> AutoModel:
|
||||
if not load_kwargs.get("cpu_offload_flag"):
|
||||
return component
|
||||
if not isinstance(component, nn.Module):
|
||||
return component
|
||||
|
||||
# try to configure layerwise-offload with the component
|
||||
configured_components = configure_layerwise_offload_modules(
|
||||
{component_name: component},
|
||||
server_args,
|
||||
component_names=server_args.layerwise_offload_components,
|
||||
warn_missing=False,
|
||||
)
|
||||
if is_layerwise_offloaded_module(component):
|
||||
logger.info(
|
||||
"Configured layerwise offload for %s immediately after startup CPU staging",
|
||||
component_name,
|
||||
)
|
||||
return component
|
||||
|
||||
logger.warning(
|
||||
"Layerwise startup CPU staging was requested for %s, but the loaded "
|
||||
"module did not enable layerwise offload. Moving it to GPU.",
|
||||
component_name,
|
||||
)
|
||||
# ensures the module is on GPU
|
||||
if component_name in configured_components:
|
||||
return component
|
||||
return component.to(get_local_torch_device())
|
||||
|
||||
def _load_customized_with_context(
|
||||
self,
|
||||
component_model_path: str,
|
||||
@@ -197,12 +123,9 @@ class ComponentLoader(ABC):
|
||||
load_kwargs = self.customized_load_kwargs_for_component(
|
||||
server_args, component_name
|
||||
)
|
||||
component = self.load_customized(
|
||||
return self.load_customized(
|
||||
component_model_path, server_args, component_name, **load_kwargs
|
||||
)
|
||||
return self._maybe_configure_layerwise_after_startup_cpu_staging(
|
||||
component, server_args, component_name, load_kwargs
|
||||
)
|
||||
|
||||
def _load_native_with_context(
|
||||
self,
|
||||
@@ -222,9 +145,7 @@ class ComponentLoader(ABC):
|
||||
transformers_or_diffusers,
|
||||
component_name,
|
||||
)
|
||||
should_offload = self.should_offload(server_args, component_name=component_name)
|
||||
target_device = self.target_device(should_offload)
|
||||
return component.to(device=target_device)
|
||||
return component
|
||||
|
||||
def load(
|
||||
self,
|
||||
@@ -270,6 +191,8 @@ class ComponentLoader(ABC):
|
||||
component_attn_name,
|
||||
)
|
||||
source = "sgl-diffusion"
|
||||
except ComponentResidencyError:
|
||||
raise
|
||||
except Exception as e:
|
||||
if self.should_raise_customized_load_error(server_args, component_name):
|
||||
traceback.print_exc()
|
||||
@@ -308,12 +231,12 @@ class ComponentLoader(ABC):
|
||||
else:
|
||||
if isinstance(component, nn.Module):
|
||||
component = component.eval()
|
||||
if (
|
||||
server_args.cpu_offload_components is not None
|
||||
and server_args.should_cpu_offload_component(component_name)
|
||||
and not is_fsdp_managed_module(component)
|
||||
):
|
||||
component = component.to("cpu")
|
||||
if not is_fsdp_managed_module(component):
|
||||
component = component.to(
|
||||
self.target_device(
|
||||
server_args.should_start_component_on_cpu(component_name)
|
||||
)
|
||||
)
|
||||
current_gpu_mem = current_platform.get_available_gpu_memory()
|
||||
model_size = get_memory_usage_of_component(component) or "NA"
|
||||
consumed = gpu_mem_before_loading - current_gpu_mem
|
||||
|
||||
+1
-1
@@ -45,7 +45,7 @@ class DiffusionDecoderLoader(ComponentLoader):
|
||||
server_args.model_paths[component_name] = component_model_path
|
||||
model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
target_device = self.target_device(
|
||||
server_args.should_cpu_offload_component(component_name)
|
||||
server_args.should_start_component_on_cpu(component_name)
|
||||
)
|
||||
dtype = resolve_precision(
|
||||
server_args, component_name, precision_attr="vae_precision"
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from sglang.multimodal_gen.configs.models import ModelConfig
|
||||
from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
|
||||
TextEncoderLoader,
|
||||
)
|
||||
@@ -16,32 +15,11 @@ class ImageEncoderLoader(TextEncoderLoader):
|
||||
component_names = ["image_encoder"]
|
||||
expected_library = "transformers"
|
||||
|
||||
def should_offload(
|
||||
self,
|
||||
server_args,
|
||||
model_config: ModelConfig | None = None,
|
||||
component_name: str | None = None,
|
||||
):
|
||||
component_name = component_name or "image_encoder"
|
||||
should_offload = server_args.should_cpu_offload_component(component_name)
|
||||
if not should_offload:
|
||||
return False
|
||||
# _fsdp_shard_conditions is in arch_config, not directly on model_config
|
||||
arch_config = (
|
||||
getattr(model_config, "arch_config", model_config) if model_config else None
|
||||
)
|
||||
fsdp_shard_conditions = (
|
||||
getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
|
||||
)
|
||||
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
|
||||
return use_cpu_offload
|
||||
|
||||
def load_customized(
|
||||
self,
|
||||
component_model_path: str,
|
||||
server_args: ServerArgs,
|
||||
component_name: str = "image_encoder",
|
||||
cpu_offload_flag: bool | None = None,
|
||||
):
|
||||
"""Load the text encoders based on the model path, and inference args."""
|
||||
# model_config: PretrainedConfig = get_hf_config(
|
||||
@@ -69,10 +47,5 @@ class ImageEncoderLoader(TextEncoderLoader):
|
||||
encoder_config,
|
||||
server_args,
|
||||
server_args.pipeline_config.image_encoder_precision,
|
||||
cpu_offload_flag=(
|
||||
cpu_offload_flag
|
||||
if cpu_offload_flag is not None
|
||||
else server_args.should_cpu_offload_component(component_name)
|
||||
),
|
||||
component_name=component_name,
|
||||
)
|
||||
|
||||
@@ -174,7 +174,11 @@ class PELoader(ComponentLoader):
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
)
|
||||
|
||||
device = get_local_torch_device()
|
||||
device = (
|
||||
torch.device("cpu")
|
||||
if server_args.should_start_component_on_cpu(component_name)
|
||||
else get_local_torch_device()
|
||||
)
|
||||
model = model.to(device).eval()
|
||||
|
||||
logger.info(
|
||||
|
||||
+1
-1
@@ -41,7 +41,7 @@ class SoundTokenizerLoader(ComponentLoader):
|
||||
precision = "bf16"
|
||||
dtype = PRECISION_TO_TYPE[precision]
|
||||
target_device = self.target_device(
|
||||
server_args.should_cpu_offload_component(component_name)
|
||||
server_args.should_start_component_on_cpu(component_name)
|
||||
)
|
||||
|
||||
with set_default_torch_dtype(dtype), skip_init_modules():
|
||||
|
||||
+17
-77
@@ -7,12 +7,10 @@ from contextlib import nullcontext
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch import nn
|
||||
from torch.distributed import init_device_mesh
|
||||
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
|
||||
|
||||
from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig
|
||||
from sglang.multimodal_gen.configs.models import EncoderConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
|
||||
QwenImageEditPipelineConfig,
|
||||
)
|
||||
@@ -27,10 +25,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
|
||||
ComponentLoader,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.fsdp_load import (
|
||||
register_fsdp_entrypoints,
|
||||
shard_model,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
set_default_torch_dtype,
|
||||
skip_init_modules,
|
||||
@@ -84,39 +78,6 @@ class TextEncoderLoader(ComponentLoader):
|
||||
allow_patterns_overrides: list[str] | None = None
|
||||
"""If defined, weights will load exclusively using these patterns."""
|
||||
|
||||
def should_offload(
|
||||
self,
|
||||
server_args,
|
||||
model_config: ModelConfig | None = None,
|
||||
component_name: str | None = None,
|
||||
):
|
||||
component_name = component_name or "text_encoder"
|
||||
should_offload = server_args.should_cpu_offload_component(component_name)
|
||||
if not should_offload:
|
||||
return False
|
||||
# _fsdp_shard_conditions is in arch_config, not directly on model_config
|
||||
arch_config = (
|
||||
getattr(model_config, "arch_config", model_config) if model_config else None
|
||||
)
|
||||
fsdp_shard_conditions = (
|
||||
getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
|
||||
)
|
||||
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
|
||||
return use_cpu_offload
|
||||
|
||||
def customized_load_kwargs_for_component(
|
||||
self, server_args: ServerArgs, component_name: str
|
||||
) -> dict[str, bool]:
|
||||
if ComponentLoader._is_component_set_as_layerwise_load(
|
||||
server_args, component_name
|
||||
):
|
||||
logger.info(
|
||||
"Loading %s on CPU first because it is selected for layerwise offload",
|
||||
component_name,
|
||||
)
|
||||
return {"cpu_offload_flag": True}
|
||||
return {}
|
||||
|
||||
def load_native(
|
||||
self,
|
||||
component_model_path: str,
|
||||
@@ -327,7 +288,7 @@ class TextEncoderLoader(ComponentLoader):
|
||||
component_model_path: str,
|
||||
server_args: ServerArgs,
|
||||
component_name: str,
|
||||
cpu_offload_flag: bool | None = None,
|
||||
component_starts_on_cpu: bool | None = None,
|
||||
):
|
||||
"""Load the text encoders based on the model path, and inference args."""
|
||||
diffusers_pretrained_config = get_config(
|
||||
@@ -381,7 +342,7 @@ class TextEncoderLoader(ComponentLoader):
|
||||
encoder_config,
|
||||
server_args,
|
||||
encoder_dtype,
|
||||
cpu_offload_flag=cpu_offload_flag,
|
||||
component_starts_on_cpu=component_starts_on_cpu,
|
||||
component_name=component_name,
|
||||
)
|
||||
|
||||
@@ -413,37 +374,36 @@ class TextEncoderLoader(ComponentLoader):
|
||||
model_config: EncoderConfig,
|
||||
server_args: ServerArgs,
|
||||
dtype: str = "fp16",
|
||||
cpu_offload_flag: bool | None = None,
|
||||
component_starts_on_cpu: bool | None = None,
|
||||
component_name: str = "text_encoder",
|
||||
):
|
||||
# Determine CPU offload behavior and target device
|
||||
|
||||
local_torch_device = get_local_torch_device()
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
fsdp_cpu_offload = self.should_offload(
|
||||
server_args, model_config, component_name
|
||||
)
|
||||
should_offload = (
|
||||
cpu_offload_flag if cpu_offload_flag is not None else fsdp_cpu_offload
|
||||
component_starts_on_cpu = (
|
||||
component_starts_on_cpu
|
||||
if component_starts_on_cpu is not None
|
||||
else server_args.should_start_component_on_cpu(component_name)
|
||||
)
|
||||
else:
|
||||
fsdp_cpu_offload = False
|
||||
should_offload = False
|
||||
component_starts_on_cpu = False
|
||||
|
||||
if (
|
||||
getattr(
|
||||
model_config.arch_config, "requires_gpu_resident_text_encoder", False
|
||||
)
|
||||
and should_offload
|
||||
and component_starts_on_cpu
|
||||
):
|
||||
server_args.require_component_resident(
|
||||
component_name, feature_name="bitsandbytes 4-bit text encoder"
|
||||
)
|
||||
logger.warning(
|
||||
"Keeping bitsandbytes 4-bit text encoder GPU-resident; CUDA "
|
||||
"weights and quant states are required for this checkpoint."
|
||||
)
|
||||
should_offload = False
|
||||
component_starts_on_cpu = False
|
||||
|
||||
if should_offload and not current_platform.is_mps():
|
||||
if component_starts_on_cpu and not current_platform.is_mps():
|
||||
model_device = torch.device("cpu")
|
||||
else:
|
||||
model_device = local_torch_device
|
||||
@@ -480,33 +440,13 @@ class TextEncoderLoader(ComponentLoader):
|
||||
self._get_all_weights(
|
||||
model,
|
||||
model_path,
|
||||
to_cpu=should_offload,
|
||||
to_cpu=component_starts_on_cpu,
|
||||
)
|
||||
)
|
||||
|
||||
if should_offload:
|
||||
# Disable FSDP for MPS as it's not compatible
|
||||
if component_starts_on_cpu:
|
||||
if current_platform.is_mps():
|
||||
logger.info(
|
||||
"Disabling FSDP sharding for MPS platform as it's not compatible"
|
||||
)
|
||||
model = model.to(local_torch_device)
|
||||
elif fsdp_cpu_offload:
|
||||
mesh = init_device_mesh(
|
||||
current_platform.device_type,
|
||||
mesh_shape=(1, dist.get_world_size()),
|
||||
mesh_dim_names=("offload", "replicate"),
|
||||
)
|
||||
shard_model(
|
||||
model,
|
||||
cpu_offload=True,
|
||||
reshard_after_forward=True,
|
||||
mesh=mesh["offload"],
|
||||
fsdp_shard_conditions=model_config.arch_config._fsdp_shard_conditions
|
||||
or getattr(model, "_fsdp_shard_conditions", None),
|
||||
pin_cpu_memory=server_args.pin_cpu_memory,
|
||||
)
|
||||
register_fsdp_entrypoints(model)
|
||||
else:
|
||||
model = model.to("cpu")
|
||||
else:
|
||||
|
||||
+20
-14
@@ -43,10 +43,10 @@ logger = init_logger(__name__)
|
||||
def _resolve_checkpoint_load_device(
|
||||
runtime_device: torch.device,
|
||||
*,
|
||||
component_cpu_offload: bool,
|
||||
component_starts_on_cpu: bool,
|
||||
runtime_quant_config: object | None,
|
||||
) -> torch.device:
|
||||
if component_cpu_offload and runtime_quant_config is None:
|
||||
if component_starts_on_cpu and runtime_quant_config is None:
|
||||
return torch.device("cpu")
|
||||
return runtime_device
|
||||
|
||||
@@ -157,11 +157,6 @@ class TransformerLoader(ComponentLoader):
|
||||
component_server_args = _server_args_for_transformer_component(
|
||||
server_args, component_name
|
||||
)
|
||||
if server_args.cpu_offload_components is not None:
|
||||
component_server_args = copy.copy(component_server_args)
|
||||
component_server_args.dit_cpu_offload = (
|
||||
server_args.should_cpu_offload_component(component_name)
|
||||
)
|
||||
|
||||
# 1. hf config
|
||||
config = get_diffusers_component_config(component_path=component_model_path)
|
||||
@@ -172,11 +167,15 @@ class TransformerLoader(ComponentLoader):
|
||||
|
||||
# 2. dit config
|
||||
# Config from Diffusers supersedes sgl_diffusion's model config
|
||||
component_name = _normalize_component_type(component_name)
|
||||
component_type = _normalize_component_type(component_name)
|
||||
server_args.model_paths[component_name] = component_model_path
|
||||
if component_name in ("transformer", "unconditional_transformer", "video_dit"):
|
||||
if component_type in (
|
||||
"transformer",
|
||||
"unconditional_transformer",
|
||||
"video_dit",
|
||||
):
|
||||
pipeline_dit_config_attr = "dit_config"
|
||||
elif component_name in ("audio_dit",):
|
||||
elif component_type == "audio_dit":
|
||||
pipeline_dit_config_attr = "audio_dit_config"
|
||||
else:
|
||||
raise ValueError(f"Invalid module name: {component_name}")
|
||||
@@ -193,7 +192,14 @@ class TransformerLoader(ComponentLoader):
|
||||
component_model_path=component_model_path,
|
||||
model_cls=model_cls,
|
||||
cls_name=cls_name,
|
||||
component_name=component_name,
|
||||
)
|
||||
# Quantization adapters may require resident weights, so placement must
|
||||
# be resolved after they have validated the component configuration.
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
component_name
|
||||
)
|
||||
use_fsdp = server_args.should_use_fsdp_for_component(component_name)
|
||||
|
||||
logger.info(
|
||||
"Loading %s from %s safetensors file(s) %s, param_dtype: %s",
|
||||
@@ -255,7 +261,7 @@ class TransformerLoader(ComponentLoader):
|
||||
local_torch_device = get_local_torch_device()
|
||||
checkpoint_load_device = _resolve_checkpoint_load_device(
|
||||
local_torch_device,
|
||||
component_cpu_offload=bool(component_server_args.dit_cpu_offload),
|
||||
component_starts_on_cpu=component_starts_on_cpu,
|
||||
runtime_quant_config=quant_spec.runtime_quant_config,
|
||||
)
|
||||
direct_gpu_weight_loading = bool(
|
||||
@@ -268,7 +274,7 @@ class TransformerLoader(ComponentLoader):
|
||||
weight_load_plan = WeightLoadPlan.for_component(
|
||||
checkpoint_load_device=checkpoint_load_device,
|
||||
needs_device_weight_postprocess=quant_spec.needs_device_weight_postprocess,
|
||||
component_cpu_offload=bool(component_server_args.dit_cpu_offload),
|
||||
component_starts_on_cpu=component_starts_on_cpu,
|
||||
load_full_state_dict_on_device=direct_gpu_weight_loading,
|
||||
)
|
||||
if direct_gpu_weight_loading:
|
||||
@@ -304,9 +310,9 @@ class TransformerLoader(ComponentLoader):
|
||||
device=local_torch_device,
|
||||
hsdp_replicate_dim=server_args.hsdp_replicate_dim,
|
||||
hsdp_shard_dim=server_args.hsdp_shard_dim,
|
||||
cpu_offload=component_server_args.dit_cpu_offload,
|
||||
component_starts_on_cpu=component_starts_on_cpu,
|
||||
pin_cpu_memory=component_server_args.pin_cpu_memory,
|
||||
fsdp_inference=component_server_args.use_fsdp_inference,
|
||||
fsdp_inference=use_fsdp,
|
||||
param_dtype=quant_spec.param_dtype,
|
||||
reduce_dtype=torch.float32,
|
||||
output_dtype=None,
|
||||
|
||||
@@ -213,8 +213,10 @@ class UpsamplerLoader(ComponentLoader):
|
||||
|
||||
logger.info("Loading LatentUpsampler with config: %s", config)
|
||||
|
||||
should_offload = server_args.should_cpu_offload_component(component_name)
|
||||
target_device = self.target_device(should_offload)
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
component_name
|
||||
)
|
||||
target_device = self.target_device(component_starts_on_cpu)
|
||||
|
||||
with torch.device("meta"):
|
||||
model = LatentUpsampler(**config)
|
||||
|
||||
@@ -133,8 +133,10 @@ class VAELoader(ComponentLoader):
|
||||
# NOTE: some post init logics are only available after updated with config
|
||||
vae_config.post_init()
|
||||
|
||||
should_offload = server_args.should_cpu_offload_component(component_name)
|
||||
target_device = self.target_device(should_offload)
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
component_name
|
||||
)
|
||||
target_device = self.target_device(component_starts_on_cpu)
|
||||
|
||||
native_only = component_name in getattr(
|
||||
server_args.pipeline_config, "native_only_components", ()
|
||||
|
||||
@@ -59,7 +59,7 @@ class VisionLanguageEncoderLoader(ComponentLoader):
|
||||
revision=server_args.revision,
|
||||
)
|
||||
target_device = self.target_device(
|
||||
server_args.should_cpu_offload_component("vision_language_encoder")
|
||||
server_args.should_start_component_on_cpu("vision_language_encoder")
|
||||
)
|
||||
model = GlmImageForConditionalGeneration.from_pretrained(
|
||||
component_model_path,
|
||||
|
||||
@@ -51,8 +51,10 @@ class VocoderLoader(ComponentLoader):
|
||||
else PRECISION_TO_TYPE["fp32"]
|
||||
)
|
||||
|
||||
should_offload = server_args.should_cpu_offload_component(component_name)
|
||||
target_device = self.target_device(should_offload)
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
component_name
|
||||
)
|
||||
target_device = self.target_device(component_starts_on_cpu)
|
||||
|
||||
with set_default_torch_dtype(vocoder_dtype), skip_init_modules():
|
||||
vocoder_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
|
||||
@@ -231,7 +231,7 @@ def maybe_load_fsdp_model(
|
||||
hsdp_shard_dim: int,
|
||||
param_dtype: torch.dtype,
|
||||
reduce_dtype: torch.dtype,
|
||||
cpu_offload: bool = False,
|
||||
component_starts_on_cpu: bool = False,
|
||||
fsdp_inference: bool = False,
|
||||
output_dtype: torch.dtype | None = None,
|
||||
pin_cpu_memory: bool = True,
|
||||
@@ -251,6 +251,8 @@ def maybe_load_fsdp_model(
|
||||
original parameter dtypes
|
||||
- Weight loading and casting
|
||||
reduce_dtype: Data type for gradient reduction in FSDP mixed precision.
|
||||
component_starts_on_cpu: Load a non-FSDP component onto CPU initially.
|
||||
Runtime residency strategies move it to the compute device before use.
|
||||
strict: If True, enforce strict state dict loading (all keys must match).
|
||||
weight_load_plan: Optional checkpoint/postprocess device plan for this load.
|
||||
"""
|
||||
@@ -293,16 +295,12 @@ def maybe_load_fsdp_model(
|
||||
logger.info("Disabling FSDP for MPS platform as it's not compatible")
|
||||
|
||||
weight_load_plan = weight_load_plan or WeightLoadPlan(checkpoint_load_device=device)
|
||||
defer_cpu_offload = bool(
|
||||
cpu_offload and weight_load_plan.defer_component_cpu_offload
|
||||
defer_cpu_placement = bool(
|
||||
component_starts_on_cpu
|
||||
and weight_load_plan.defer_cpu_placement
|
||||
and not use_fsdp
|
||||
)
|
||||
if defer_cpu_offload and use_fsdp:
|
||||
logger.warning(
|
||||
"Ignoring deferred CPU offload for FSDP loading; keeping the existing "
|
||||
"FSDP offload policy."
|
||||
)
|
||||
defer_cpu_offload = False
|
||||
load_cpu_offload = bool(cpu_offload and not defer_cpu_offload)
|
||||
load_on_cpu = bool(component_starts_on_cpu and not defer_cpu_placement)
|
||||
weight_postprocess_device = weight_load_plan.weight_postprocess_device
|
||||
if use_fsdp and weight_postprocess_device is not None:
|
||||
logger.warning("Ignoring weight postprocess device override for FSDP loading.")
|
||||
@@ -327,7 +325,7 @@ def maybe_load_fsdp_model(
|
||||
)
|
||||
shard_model(
|
||||
model,
|
||||
cpu_offload=load_cpu_offload,
|
||||
cpu_offload=False,
|
||||
reshard_after_forward=True,
|
||||
mp_policy=mp_policy,
|
||||
mesh=device_mesh,
|
||||
@@ -408,7 +406,7 @@ def maybe_load_fsdp_model(
|
||||
weight_load_plan.checkpoint_load_device,
|
||||
param_dtype,
|
||||
strict=strict,
|
||||
cpu_offload=load_cpu_offload,
|
||||
cpu_offload=load_on_cpu,
|
||||
param_names_mapping=param_names_mapping_fn,
|
||||
preconverted_state_dict=preconverted_state_dict,
|
||||
)
|
||||
@@ -444,7 +442,7 @@ def maybe_load_fsdp_model(
|
||||
p.requires_grad = False
|
||||
|
||||
# 4. deferred cpu offload
|
||||
if defer_cpu_offload:
|
||||
if defer_cpu_placement:
|
||||
model.to("cpu")
|
||||
|
||||
return model
|
||||
|
||||
@@ -21,6 +21,10 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
|
||||
_patch_nunchaku_scales,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_OFFLOAD,
|
||||
ComponentResidencyError,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
|
||||
maybe_download_model,
|
||||
@@ -146,6 +150,35 @@ class _TransformerQuantAdapter:
|
||||
return []
|
||||
|
||||
|
||||
def _uses_component_offload(
|
||||
server_args: ServerArgs,
|
||||
component_name: str | None,
|
||||
*,
|
||||
legacy_enabled: bool,
|
||||
) -> bool:
|
||||
if component_name is None:
|
||||
return legacy_enabled
|
||||
return server_args.residency_mode(component_name) == COMPONENT_OFFLOAD
|
||||
|
||||
|
||||
def _reject_explicit_component_selector(
|
||||
server_args: ServerArgs,
|
||||
component_name: str | None,
|
||||
*,
|
||||
feature_name: str,
|
||||
) -> None:
|
||||
if component_name is None:
|
||||
return
|
||||
selected_by_component_residency = (
|
||||
server_args.canonical_residency_mode(component_name) == COMPONENT_OFFLOAD
|
||||
)
|
||||
if selected_by_component_residency:
|
||||
raise ComponentResidencyError(
|
||||
f"{feature_name} does not support component-offload for "
|
||||
f"{component_name!r}; select resident or layerwise-offload"
|
||||
)
|
||||
|
||||
|
||||
class _NunchakuQuantAdapter(_TransformerQuantAdapter):
|
||||
"""Adapter for Nunchaku checkpoints"""
|
||||
|
||||
@@ -196,16 +229,19 @@ class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter):
|
||||
cls_name: str,
|
||||
server_args: ServerArgs,
|
||||
quant_config: Optional[QuantizationConfig],
|
||||
component_name: str | None,
|
||||
) -> None:
|
||||
self.cls_name = cls_name
|
||||
self.server_args = server_args
|
||||
self.quant_config = quant_config
|
||||
self.component_name = component_name
|
||||
|
||||
@staticmethod
|
||||
def _maybe_adjust_flux2_nvfp4_fallback_defaults(
|
||||
cls_name: str,
|
||||
server_args: ServerArgs,
|
||||
quant_config: Optional[QuantizationConfig],
|
||||
component_name: str | None = None,
|
||||
) -> None:
|
||||
if cls_name != "Flux2Transformer2DModel" or quant_config is None:
|
||||
return
|
||||
@@ -219,14 +255,47 @@ class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter):
|
||||
if not weights_path.endswith("-mixed.safetensors") or server_args.tp_size <= 1:
|
||||
return
|
||||
|
||||
if server_args.dit_cpu_offload or server_args.text_encoder_cpu_offload:
|
||||
server_args.dit_cpu_offload = False
|
||||
server_args.text_encoder_cpu_offload = False
|
||||
dit_component_offload = _uses_component_offload(
|
||||
server_args,
|
||||
component_name,
|
||||
legacy_enabled=bool(server_args.dit_cpu_offload),
|
||||
)
|
||||
text_encoder_component_offload = _uses_component_offload(
|
||||
server_args,
|
||||
"text_encoder" if component_name is not None else None,
|
||||
legacy_enabled=bool(server_args.text_encoder_cpu_offload),
|
||||
)
|
||||
if dit_component_offload:
|
||||
_reject_explicit_component_selector(
|
||||
server_args,
|
||||
component_name,
|
||||
feature_name="FLUX.2 mixed NVFP4 with tensor parallelism",
|
||||
)
|
||||
if text_encoder_component_offload:
|
||||
_reject_explicit_component_selector(
|
||||
server_args,
|
||||
"text_encoder" if component_name is not None else None,
|
||||
feature_name="FLUX.2 mixed NVFP4 with tensor parallelism",
|
||||
)
|
||||
if dit_component_offload or text_encoder_component_offload:
|
||||
if component_name is None:
|
||||
server_args.dit_cpu_offload = False
|
||||
server_args.text_encoder_cpu_offload = False
|
||||
else:
|
||||
if dit_component_offload:
|
||||
server_args.require_component_resident(
|
||||
component_name,
|
||||
feature_name="FLUX.2 mixed NVFP4 with tensor parallelism",
|
||||
)
|
||||
if text_encoder_component_offload:
|
||||
server_args.require_component_resident(
|
||||
"text_encoder",
|
||||
feature_name="FLUX.2 mixed NVFP4 with tensor parallelism",
|
||||
)
|
||||
logger.warning(
|
||||
"FLUX.2 mixed NVFP4 is using the ModelOpt FP4 path with tp_size=%d; "
|
||||
"disabling dit/text-encoder CPU offload to avoid TP all-gather "
|
||||
"launch failures. Override the offload flags explicitly if you need "
|
||||
"the old behavior.",
|
||||
"keeping the DiT and text encoder resident to avoid TP all-gather "
|
||||
"launch failures.",
|
||||
server_args.tp_size,
|
||||
)
|
||||
|
||||
@@ -235,6 +304,7 @@ class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter):
|
||||
cls_name=self.cls_name,
|
||||
server_args=self.server_args,
|
||||
quant_config=self.quant_config,
|
||||
component_name=self.component_name,
|
||||
)
|
||||
|
||||
|
||||
@@ -246,14 +316,17 @@ class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter):
|
||||
*,
|
||||
server_args: ServerArgs,
|
||||
quant_config: Optional[QuantizationConfig],
|
||||
component_name: str | None,
|
||||
) -> None:
|
||||
self.server_args = server_args
|
||||
self.quant_config = quant_config
|
||||
self.component_name = component_name
|
||||
|
||||
@staticmethod
|
||||
def _maybe_disable_incompatible_dit_offload_modes(
|
||||
server_args: ServerArgs,
|
||||
quant_config: Optional[QuantizationConfig],
|
||||
component_name: str | None = None,
|
||||
) -> None:
|
||||
if quant_config is None:
|
||||
return
|
||||
@@ -264,18 +337,34 @@ class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter):
|
||||
if quant_name != "modelopt_fp8":
|
||||
return
|
||||
|
||||
if server_args.dit_cpu_offload:
|
||||
server_args.dit_cpu_offload = False
|
||||
component_offload = _uses_component_offload(
|
||||
server_args,
|
||||
component_name,
|
||||
legacy_enabled=bool(server_args.dit_cpu_offload),
|
||||
)
|
||||
if component_offload:
|
||||
_reject_explicit_component_selector(
|
||||
server_args,
|
||||
component_name,
|
||||
feature_name="ModelOpt FP8 diffusion checkpoints",
|
||||
)
|
||||
if component_name is None:
|
||||
server_args.dit_cpu_offload = False
|
||||
else:
|
||||
server_args.require_component_resident(
|
||||
component_name,
|
||||
feature_name="ModelOpt FP8 diffusion checkpoints",
|
||||
)
|
||||
logger.warning(
|
||||
"ModelOpt FP8 diffusion checkpoints currently keep dit_cpu_offload "
|
||||
"disabled. Layerwise DiT offload stays enabled because the runtime "
|
||||
"now preserves the restored FP8 tensor strides.",
|
||||
"ModelOpt FP8 diffusion checkpoints keep the DiT resident instead "
|
||||
"of using component offload. Layerwise offload remains supported.",
|
||||
)
|
||||
|
||||
def prepare(self) -> None:
|
||||
_ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes(
|
||||
server_args=self.server_args,
|
||||
quant_config=self.quant_config,
|
||||
component_name=self.component_name,
|
||||
)
|
||||
|
||||
|
||||
@@ -287,25 +376,52 @@ class _BitsAndBytes4BitAdapter(_TransformerQuantAdapter):
|
||||
*,
|
||||
server_args: ServerArgs,
|
||||
quant_config: Optional[QuantizationConfig],
|
||||
component_name: str | None,
|
||||
) -> None:
|
||||
self.server_args = server_args
|
||||
self.quant_config = quant_config
|
||||
self.component_name = component_name
|
||||
|
||||
@staticmethod
|
||||
def _maybe_disable_incompatible_offload_modes(
|
||||
server_args: ServerArgs,
|
||||
quant_config: Optional[QuantizationConfig],
|
||||
component_name: str | None = None,
|
||||
) -> None:
|
||||
if _get_quant_config_name(quant_config) != "bitsandbytes":
|
||||
return
|
||||
|
||||
changed = []
|
||||
if server_args.dit_cpu_offload:
|
||||
server_args.dit_cpu_offload = False
|
||||
changed.append("dit_cpu_offload=False")
|
||||
if server_args.use_fsdp_inference:
|
||||
server_args.use_fsdp_inference = False
|
||||
changed.append("use_fsdp_inference=False")
|
||||
component_offload = _uses_component_offload(
|
||||
server_args,
|
||||
component_name,
|
||||
legacy_enabled=bool(server_args.dit_cpu_offload),
|
||||
)
|
||||
if component_offload:
|
||||
_reject_explicit_component_selector(
|
||||
server_args,
|
||||
component_name,
|
||||
feature_name="bitsandbytes 4-bit transformer checkpoints",
|
||||
)
|
||||
if component_name is None:
|
||||
server_args.dit_cpu_offload = False
|
||||
else:
|
||||
server_args.require_component_resident(
|
||||
component_name,
|
||||
feature_name="bitsandbytes 4-bit transformer checkpoints",
|
||||
)
|
||||
changed.append(
|
||||
"dit_cpu_offload=False"
|
||||
if component_name is None
|
||||
else f"{component_name}=resident"
|
||||
)
|
||||
if component_name is None:
|
||||
if server_args.use_fsdp_inference:
|
||||
server_args.use_fsdp_inference = False
|
||||
changed.append("use_fsdp_inference=False")
|
||||
elif server_args.should_use_fsdp_for_component(component_name):
|
||||
server_args.disable_fsdp_for_component(component_name)
|
||||
changed.append(f"{component_name}.fsdp=False")
|
||||
if changed:
|
||||
logger.warning(
|
||||
"Keeping bitsandbytes 4-bit transformer GPU-resident: %s",
|
||||
@@ -316,6 +432,7 @@ class _BitsAndBytes4BitAdapter(_TransformerQuantAdapter):
|
||||
_BitsAndBytes4BitAdapter._maybe_disable_incompatible_offload_modes(
|
||||
server_args=self.server_args,
|
||||
quant_config=self.quant_config,
|
||||
component_name=self.component_name,
|
||||
)
|
||||
|
||||
|
||||
@@ -440,6 +557,7 @@ def resolve_transformer_quant_load_spec(
|
||||
component_model_path: str,
|
||||
model_cls: type[nn.Module],
|
||||
cls_name: str,
|
||||
component_name: str | None = None,
|
||||
) -> TransformerQuantLoadSpec:
|
||||
if getattr(model_cls, "handles_checkpoint_quantization", False):
|
||||
quant_config = None
|
||||
@@ -472,6 +590,7 @@ def resolve_transformer_quant_load_spec(
|
||||
nunchaku_config=nunchaku_config,
|
||||
model_cls=model_cls,
|
||||
safetensors_list=safetensors_list,
|
||||
component_name=component_name,
|
||||
)
|
||||
for adapter in adapters:
|
||||
adapter.prepare()
|
||||
@@ -496,6 +615,9 @@ def _needs_device_weight_postprocess(
|
||||
) -> bool:
|
||||
"""Return whether post-load weight processing needs CUDA/NPU tensors."""
|
||||
quant_name = _get_quant_config_name(quant_config)
|
||||
if quant_name == "modelopt_fp8":
|
||||
return True
|
||||
|
||||
serialized_flag_by_quant_name = {
|
||||
"fp8": "is_checkpoint_fp8_serialized",
|
||||
"mxfp8": "is_checkpoint_fp8_serialized",
|
||||
@@ -516,20 +638,24 @@ def _build_transformer_quant_adapters(
|
||||
nunchaku_config: Optional[NunchakuConfig],
|
||||
model_cls: type[nn.Module],
|
||||
safetensors_list: list[str],
|
||||
component_name: str | None,
|
||||
) -> list[_TransformerQuantAdapter]:
|
||||
adapters: list[_TransformerQuantAdapter] = [
|
||||
_Flux2Nvfp4FallbackAdapter(
|
||||
cls_name=cls_name,
|
||||
server_args=server_args,
|
||||
quant_config=quant_config,
|
||||
component_name=component_name,
|
||||
),
|
||||
_ModelOptFp8OffloadAdapter(
|
||||
server_args=server_args,
|
||||
quant_config=quant_config,
|
||||
component_name=component_name,
|
||||
),
|
||||
_BitsAndBytes4BitAdapter(
|
||||
server_args=server_args,
|
||||
quant_config=quant_config,
|
||||
component_name=component_name,
|
||||
),
|
||||
]
|
||||
if nunchaku_config is not None:
|
||||
|
||||
@@ -11,8 +11,8 @@ class WeightLoadPlan:
|
||||
checkpoint_load_device: torch.device
|
||||
# Device required while running process_weights_after_loading; None means unchanged.
|
||||
weight_postprocess_device: torch.device | None = None
|
||||
# Delay non-FSDP component CPU offload until after weight postprocessing.
|
||||
defer_component_cpu_offload: bool = False
|
||||
# Delay final CPU placement until after device-side weight postprocessing.
|
||||
defer_cpu_placement: bool = False
|
||||
# keep the complete mapped checkpoint state dict on the load device
|
||||
load_full_state_dict_on_device: bool = False
|
||||
|
||||
@@ -22,7 +22,7 @@ class WeightLoadPlan:
|
||||
*,
|
||||
checkpoint_load_device: torch.device,
|
||||
needs_device_weight_postprocess: bool,
|
||||
component_cpu_offload: bool,
|
||||
component_starts_on_cpu: bool,
|
||||
load_full_state_dict_on_device: bool = False,
|
||||
) -> "WeightLoadPlan":
|
||||
# if on-device weight postprocessing is required, load directly to device to speedup loading
|
||||
@@ -32,8 +32,8 @@ class WeightLoadPlan:
|
||||
return cls(
|
||||
checkpoint_load_device=checkpoint_load_device,
|
||||
weight_postprocess_device=weight_postprocess_device,
|
||||
defer_component_cpu_offload=(
|
||||
needs_device_weight_postprocess and component_cpu_offload
|
||||
defer_cpu_placement=(
|
||||
needs_device_weight_postprocess and component_starts_on_cpu
|
||||
),
|
||||
load_full_state_dict_on_device=load_full_state_dict_on_device,
|
||||
)
|
||||
|
||||
@@ -160,7 +160,7 @@ class BatchAdmissionController:
|
||||
self._mode = getattr(server_args, "batching_mode", "dynamic")
|
||||
self._user_max_batch_size = max(1, int(server_args.batching_max_size))
|
||||
self._model_path = server_args.model_path
|
||||
self._offload = bool(server_args.layerwise_offload_components)
|
||||
self._offload = server_args.has_layerwise_offload_components()
|
||||
self._device_memory_gb = self._get_device_memory_gb(gpu_id)
|
||||
self._rules = load_batching_config(server_args.batching_config)
|
||||
self._pipeline_config = server_args.pipeline_config
|
||||
|
||||
@@ -289,13 +289,18 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
|
||||
# apply layerwise offload after lora is applied while building LoRAPipeline
|
||||
# otherwise empty offloaded weights could fail lora converting
|
||||
if self.server_args.layerwise_offload_components:
|
||||
if self.server_args.has_layerwise_offload_components():
|
||||
configure_layerwise_offload_modules(
|
||||
self.pipeline.modules,
|
||||
self.server_args,
|
||||
component_names=self.server_args.layerwise_offload_components,
|
||||
component_names=(
|
||||
None
|
||||
if self.server_args.component_residency is not None
|
||||
else self.server_args.layerwise_offload_components
|
||||
),
|
||||
warn_missing=(
|
||||
self.server_args.is_arg_explicitly_set(
|
||||
self.server_args.component_residency is not None
|
||||
or self.server_args.is_arg_explicitly_set(
|
||||
"layerwise_offload_components"
|
||||
)
|
||||
or self.server_args.is_arg_explicitly_set("dit_layerwise_offload")
|
||||
|
||||
+156
-122
@@ -1,17 +1,21 @@
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Mapping, MutableMapping, Protocol, Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
ComponentResidencyError,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
ComponentOffloadStrategy,
|
||||
ComponentResidencyStrategy,
|
||||
LayerwiseOffloadStrategy,
|
||||
ResidentStrategy,
|
||||
VanillaD2HStrategy,
|
||||
is_fsdp_managed_module,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
@@ -28,45 +32,28 @@ logger = init_logger(__name__)
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ComponentUse:
|
||||
"""Describes one stage/use-site access to a pipeline component."""
|
||||
"""One ordered stage access to a pipeline component."""
|
||||
|
||||
stage_name: str
|
||||
# Pipeline module key: transformer / video_dit / text_encoder / ...
|
||||
component_name: str
|
||||
# Model-specific phase for sequential components, e.g. stage1 or stage2.
|
||||
# TODO: Replace this with ordered timeline identity. In an all-sequential
|
||||
# pipeline, use-site identity should come from the declared ComponentUse
|
||||
# order instead of a per-use `phase` field.
|
||||
phase: str | None = None
|
||||
# Whether the manager may prepare this component for the next request.
|
||||
preferred_ready_after_request: bool = False
|
||||
# Whether cross-stage prefetch may prepare this use before the use-site.
|
||||
allow_prefetch: bool = True
|
||||
# Whether this use is expensive enough that earlier timeline prefetch matters.
|
||||
# TODO: Replace this boolean hint with a budget-aware lookahead planner:
|
||||
# estimate memory/load cost and reuse distance, keep small and early-request
|
||||
# components resident within budget, prefetch as soon as VRAM slack appears,
|
||||
# and release completed components only when the budget requires it.
|
||||
memory_intensive: bool = False
|
||||
# Optional module dtype required by this use-site.
|
||||
target_dtype: torch.dtype | None = None
|
||||
# Some components are intentionally kept ready between warmup and the first
|
||||
# real request to avoid measuring a cold H2D in the user-visible request.
|
||||
keep_ready_after_warmup: bool = False
|
||||
start_at_stage_entry: bool = True
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ResidencyState:
|
||||
"""
|
||||
Necessary internal runtime info of ComponentResidencyManager
|
||||
"""
|
||||
"""Request-local state shared with component strategies."""
|
||||
|
||||
stages: Sequence["ComponentResidencyStage"] = ()
|
||||
stage_index: int = -1
|
||||
stage_name: str | None = None
|
||||
next_stage_name: str | None = None
|
||||
current_use: ComponentUse | None = None
|
||||
# the ComponentUses of the preceding stages
|
||||
future_uses: tuple[ComponentUse, ...] = ()
|
||||
batch_is_warmup: bool = False
|
||||
|
||||
@@ -92,31 +79,30 @@ def build_component_residency_strategy(
|
||||
module: nn.Module,
|
||||
server_args: ServerArgs,
|
||||
) -> ComponentResidencyStrategy:
|
||||
residency_mode = server_args.residency_mode(component_name)
|
||||
if is_layerwise_offloaded_module(module):
|
||||
return LayerwiseOffloadStrategy()
|
||||
if residency_mode == LAYERWISE_OFFLOAD:
|
||||
raise ComponentResidencyError(
|
||||
f"Component {component_name!r} resolved to layerwise-offload, but its "
|
||||
"loaded module did not enable layerwise offload"
|
||||
)
|
||||
if residency_mode == COMPONENT_OFFLOAD and is_fsdp_managed_module(module):
|
||||
raise ComponentResidencyError(
|
||||
f"Component {component_name!r} resolved to component-offload, but it "
|
||||
"was loaded as an FSDP-managed module"
|
||||
)
|
||||
if (
|
||||
not current_platform.is_mps()
|
||||
and not server_args.use_fsdp_inference
|
||||
and not is_fsdp_managed_module(module)
|
||||
and server_args.should_cpu_offload_component(component_name)
|
||||
and residency_mode == COMPONENT_OFFLOAD
|
||||
):
|
||||
return VanillaD2HStrategy()
|
||||
return ComponentOffloadStrategy()
|
||||
return ResidentStrategy()
|
||||
|
||||
|
||||
class ComponentResidencyManager:
|
||||
"""Executor-owned component lifecycle coordinator. Provide hooks for a PipelineExecutor
|
||||
|
||||
Hooks are called around executor progress:
|
||||
before request: collect a flat ordered ComponentUse timeline.
|
||||
before stage: update current/next stage context only.
|
||||
begin use: finish previous active use, prepare current use, wait until ready.
|
||||
end use: finish or keep current use, then prefetch the next heavy timeline use.
|
||||
finish request: finish active use and schedule preferred next-request prefetch.
|
||||
|
||||
The manager instance is global and rebound to the active pipeline before request execution.
|
||||
This manager is designed only for sequential execution order for now
|
||||
"""
|
||||
"""Coordinate component placement over a sequential request timeline."""
|
||||
|
||||
def __init__(
|
||||
self, pipeline: ComponentResidencyPipeline, server_args: ServerArgs
|
||||
@@ -138,19 +124,24 @@ class ComponentResidencyManager:
|
||||
self._custom_strategies: dict[str, ComponentResidencyStrategy] = dict(
|
||||
pipeline.component_residency_strategies
|
||||
)
|
||||
self._strategy_cache: dict[
|
||||
str, tuple[nn.Module, ComponentResidencyStrategy]
|
||||
] = {}
|
||||
self._uses_seen: dict[str, ComponentUse] = {}
|
||||
self._modules_seen: dict[str, nn.Module] = {}
|
||||
|
||||
def refresh_pipeline(self, pipeline: ComponentResidencyPipeline) -> None:
|
||||
custom_strategies = dict(pipeline.component_residency_strategies)
|
||||
if pipeline is not self.pipeline:
|
||||
self._remove_nvtx_hooks()
|
||||
self.strategy_for.cache_clear()
|
||||
self._strategy_cache.clear()
|
||||
self._active_use = None
|
||||
self._active_use_module = None
|
||||
self._uses_seen.clear()
|
||||
self._modules_seen.clear()
|
||||
self._prefetched_use_keys.clear()
|
||||
elif custom_strategies != self._custom_strategies:
|
||||
self.strategy_for.cache_clear()
|
||||
self._strategy_cache.clear()
|
||||
self.pipeline = pipeline
|
||||
self._custom_strategies = custom_strategies
|
||||
self._stage_names_by_id = {
|
||||
@@ -159,16 +150,15 @@ class ComponentResidencyManager:
|
||||
|
||||
def refresh_server_args(self, server_args: ServerArgs) -> None:
|
||||
if server_args is not self.server_args:
|
||||
self.strategy_for.cache_clear()
|
||||
self._strategy_cache.clear()
|
||||
self.server_args = server_args
|
||||
|
||||
def begin_request(
|
||||
self,
|
||||
stages: Sequence[ComponentResidencyStage],
|
||||
batch: ResidencyBatch,
|
||||
batch: ResidencyBatch | list[ResidencyBatch],
|
||||
server_args: ServerArgs,
|
||||
) -> None:
|
||||
"""A hook called before processing an actual request"""
|
||||
self.refresh_server_args(server_args)
|
||||
self.state = ResidencyState(
|
||||
stages=stages,
|
||||
@@ -180,6 +170,7 @@ class ComponentResidencyManager:
|
||||
self._current_use_index = -1
|
||||
self._prefetched_use_keys.clear()
|
||||
self._uses_seen.clear()
|
||||
self._modules_seen.clear()
|
||||
self._stage_uses_by_index = [
|
||||
tuple(stage.component_uses(server_args, self.stage_name(stage)))
|
||||
for stage in stages
|
||||
@@ -191,9 +182,7 @@ class ComponentResidencyManager:
|
||||
@staticmethod
|
||||
def _is_warmup_batch(batch: ResidencyBatch | list[ResidencyBatch]) -> bool:
|
||||
if isinstance(batch, list):
|
||||
return bool(batch) and all(
|
||||
getattr(item, "is_warmup", False) for item in batch
|
||||
)
|
||||
return bool(batch) and all(item.is_warmup for item in batch)
|
||||
return batch.is_warmup
|
||||
|
||||
def before_stage(
|
||||
@@ -203,34 +192,69 @@ class ComponentResidencyManager:
|
||||
batch: ResidencyBatch,
|
||||
server_args: ServerArgs,
|
||||
) -> None:
|
||||
"""called after stage starts"""
|
||||
# update state before entering the stage
|
||||
self.state.stage_index = stage_index
|
||||
self.state.stage_name = self.stage_name(stage)
|
||||
self.state.next_stage_name = self._next_stage_name(stage_index)
|
||||
|
||||
def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
|
||||
"""Begin one sequential component use interval. this is idempotent
|
||||
def begin_stage(self) -> None:
|
||||
"""Prepare a stage that declares one uninterrupted component use."""
|
||||
stage_uses = self._stage_uses_by_index[self.state.stage_index]
|
||||
if len(stage_uses) == 1 and stage_uses[0].start_at_stage_entry:
|
||||
self.begin_use(stage_uses[0])
|
||||
|
||||
1. Finish the previous active use if this is a different timeline use.
|
||||
2. Prepare the current component.
|
||||
3. Wait until the current component is ready, then prefetch the next heavy use.
|
||||
def end_stage(self) -> None:
|
||||
"""Close the component interval owned by the current stage."""
|
||||
if self._active_use is None:
|
||||
return
|
||||
if self._active_use.stage_name != self.state.stage_name:
|
||||
return
|
||||
if self.state.future_uses and self._same_use(
|
||||
self._active_use, self.state.future_uses[0]
|
||||
):
|
||||
return
|
||||
self.finish_active_use()
|
||||
|
||||
def begin_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
|
||||
"""Begin one sequential component use interval.
|
||||
|
||||
Repeated calls for the same component/phase extend the active interval.
|
||||
"""
|
||||
if self._active_use is not None and self._same_use(self._active_use, use):
|
||||
previous_use = self._active_use
|
||||
if self._use_key(self._active_use) != self._use_key(use):
|
||||
self._mark_current_use(use)
|
||||
self._active_use = use
|
||||
self.state.current_use = use
|
||||
self._enable_nvtx_for_use(
|
||||
use,
|
||||
module
|
||||
or self._active_use_module
|
||||
or self.get_module(use.component_name),
|
||||
active_module = module
|
||||
if active_module is None:
|
||||
active_module = self._active_use_module
|
||||
if active_module is None:
|
||||
active_module = self.get_module(use.component_name)
|
||||
module_changed = (
|
||||
self._active_use_module is not None
|
||||
and active_module is not self._active_use_module
|
||||
)
|
||||
if module_changed:
|
||||
self._disable_active_nvtx()
|
||||
self._finish_use(
|
||||
previous_use,
|
||||
module=self._active_use_module,
|
||||
keep_on_warmup=False,
|
||||
force=True,
|
||||
)
|
||||
if active_module is not None and (
|
||||
self._active_use_module is None
|
||||
or module_changed
|
||||
or use.target_dtype != previous_use.target_dtype
|
||||
):
|
||||
active_module = self._prepare_forward_use(use, module=active_module)
|
||||
self._active_use = use
|
||||
self._active_use_module = active_module
|
||||
self.state.current_use = use
|
||||
self._enable_nvtx_for_use(use, active_module)
|
||||
return
|
||||
if self._active_use is not None:
|
||||
self._disable_active_nvtx()
|
||||
# finish previous active use
|
||||
self._finish_use(
|
||||
self._active_use,
|
||||
module=self._active_use_module,
|
||||
@@ -247,18 +271,17 @@ class ComponentResidencyManager:
|
||||
self._prefetch_next_memory_intensive_use()
|
||||
|
||||
def end_use(self, use: ComponentUse, module: nn.Module | None = None) -> None:
|
||||
"""End one sequential component use interval.
|
||||
|
||||
1. Finish or keep the current component.
|
||||
2. Clear it as the active use.
|
||||
3. Prefetch the next memory-intensive use without waiting.
|
||||
"""
|
||||
"""End one sequential component use interval."""
|
||||
if self._active_use is None or not self._same_use(self._active_use, use):
|
||||
return
|
||||
self._disable_active_nvtx()
|
||||
self._finish_use(
|
||||
self._active_use,
|
||||
module=self._active_use_module or module,
|
||||
module=(
|
||||
self._active_use_module
|
||||
if self._active_use_module is not None
|
||||
else module
|
||||
),
|
||||
keep_on_warmup=self._active_use.keep_ready_after_warmup,
|
||||
)
|
||||
self._active_use = None
|
||||
@@ -294,6 +317,33 @@ class ComponentResidencyManager:
|
||||
hooks.remove_hooks()
|
||||
del self._nvtx_hooks_by_use_key[key]
|
||||
|
||||
def forget_module(self, module: nn.Module | None) -> None:
|
||||
"""Drop manager-owned references before a component is deleted."""
|
||||
if module is None:
|
||||
return
|
||||
self.remove_nvtx_hooks_for_module(module)
|
||||
forgotten_component_names: set[str] = set()
|
||||
if self._active_use_module is module:
|
||||
forgotten_component_names.add(self._active_use.component_name)
|
||||
self._active_use = None
|
||||
self._active_use_module = None
|
||||
self.state.current_use = None
|
||||
for component_name, (cached_module, _) in list(self._strategy_cache.items()):
|
||||
if cached_module is module:
|
||||
del self._strategy_cache[component_name]
|
||||
forgotten_component_names.add(component_name)
|
||||
for component_name, seen_module in list(self._modules_seen.items()):
|
||||
if seen_module is module:
|
||||
del self._modules_seen[component_name]
|
||||
forgotten_component_names.add(component_name)
|
||||
for component_name in forgotten_component_names:
|
||||
self._uses_seen.pop(component_name, None)
|
||||
self._prefetched_use_keys = {
|
||||
key
|
||||
for key in self._prefetched_use_keys
|
||||
if key[1] not in forgotten_component_names
|
||||
}
|
||||
|
||||
def finish_active_use(self, *, prefetch_next: bool = True) -> None:
|
||||
"""Finish the currently active sequential use, if any."""
|
||||
if self._active_use is None:
|
||||
@@ -315,11 +365,13 @@ class ComponentResidencyManager:
|
||||
self, use: ComponentUse, module: nn.Module | None = None
|
||||
) -> nn.Module | None:
|
||||
"""Prepare a component that is about to run and wait until it is ready."""
|
||||
module = module or self.get_module(use.component_name)
|
||||
if module is None:
|
||||
module = self.get_module(use.component_name)
|
||||
if module is None:
|
||||
return None
|
||||
strategy = self.strategy_for(use.component_name, module)
|
||||
self._uses_seen[use.component_name] = use
|
||||
self._modules_seen[use.component_name] = module
|
||||
self.state.current_use = use
|
||||
strategy.prepare_for_use(module, use, self.state)
|
||||
strategy.wait_for_use(module, use, self.state)
|
||||
@@ -385,28 +437,23 @@ class ComponentResidencyManager:
|
||||
return ".".join(parts)
|
||||
|
||||
def _prefetch_use(self, use: ComponentUse) -> None:
|
||||
"""Prepare a future component opportunistically without waiting.
|
||||
|
||||
This is called for memory-intensive future uses where H2D placement can
|
||||
overlap with the current stage.
|
||||
"""
|
||||
"""Prepare a future memory-intensive component without waiting."""
|
||||
if not use.allow_prefetch:
|
||||
return
|
||||
module = self.get_module(use.component_name)
|
||||
if module is None:
|
||||
return
|
||||
strategy = self.strategy_for(use.component_name, module)
|
||||
if isinstance(strategy, VanillaD2HStrategy) and self._active_use is not None:
|
||||
# Avoid making two vanilla-offloaded heavy components resident before
|
||||
# a budget-aware planner can prove the overlap is safe.
|
||||
if (
|
||||
isinstance(strategy, ComponentOffloadStrategy)
|
||||
and self._active_use is not None
|
||||
):
|
||||
return
|
||||
if is_resident_layerwise_module(module):
|
||||
# A layerwise DiT holding a large resident set must not be prefetched
|
||||
# during a prior peer stage (e.g. text encoding): co-residing can lead
|
||||
# to OOMs. Pin it lazily at the DiT's own use-site.
|
||||
return
|
||||
|
||||
self._uses_seen[use.component_name] = use
|
||||
self._modules_seen[use.component_name] = module
|
||||
if strategy.prefetch_for_use(module, use, self.state):
|
||||
self._prefetched_use_keys.add(self._use_key(use))
|
||||
|
||||
@@ -416,29 +463,32 @@ class ComponentResidencyManager:
|
||||
*,
|
||||
module: nn.Module | None = None,
|
||||
keep_on_warmup: bool,
|
||||
force: bool = False,
|
||||
) -> None:
|
||||
"""finish a specific use by keeping them resident or call finish_use hook"""
|
||||
module = module or self.get_module(use.component_name)
|
||||
if module is None:
|
||||
module = self._modules_seen.get(use.component_name)
|
||||
if module is None:
|
||||
module = self.get_module(use.component_name)
|
||||
if module is None:
|
||||
return
|
||||
should_keep = (
|
||||
keep_on_warmup and self.state.batch_is_warmup
|
||||
) or self._should_keep_after_use(use)
|
||||
if should_keep:
|
||||
return
|
||||
if not force:
|
||||
should_keep = (
|
||||
keep_on_warmup and self.state.batch_is_warmup
|
||||
) or self._should_keep_after_use(use)
|
||||
if should_keep:
|
||||
return
|
||||
strategy = self.strategy_for(use.component_name, module)
|
||||
was_on_cuda = self._module_on_cuda(module)
|
||||
strategy.finish_use(module, use, self.state)
|
||||
self._empty_cache_after_large_release(use, strategy, module, was_on_cuda)
|
||||
|
||||
def finish_request(self) -> None:
|
||||
# 1. Close the currently active sequential use.
|
||||
self.finish_active_use(prefetch_next=False)
|
||||
# 2. Pick components that should be ready for the next request.
|
||||
preferred_uses = self._preferred_request_end_uses()
|
||||
# 3. Finish everything else, or prepare preferred uses for request tail.
|
||||
for component_name, use in list(self._uses_seen.items()):
|
||||
module = self.get_module(component_name)
|
||||
module = self._modules_seen.get(component_name)
|
||||
if module is None:
|
||||
module = self.get_module(component_name)
|
||||
if module is None:
|
||||
continue
|
||||
if self.state.batch_is_warmup and use.keep_ready_after_warmup:
|
||||
@@ -449,21 +499,13 @@ class ComponentResidencyManager:
|
||||
keep_single_dit = self._should_keep_single_dit(component_name, module)
|
||||
if not preferred and keep_single_dit:
|
||||
continue
|
||||
# A preferred component is normally prefetched for the next request.
|
||||
# Do not let that performance hint override CPU/layerwise offload for
|
||||
# a single DiT, which must obey the selected memory policy.
|
||||
preferred = preferred and (
|
||||
not self._is_single_dit_component(component_name) or keep_single_dit
|
||||
)
|
||||
strategy = self.strategy_for(component_name, module)
|
||||
if preferred and not self.state.batch_is_warmup:
|
||||
strategy.prepare_after_request(module, use, self.state)
|
||||
else:
|
||||
was_on_cuda = self._module_on_cuda(module)
|
||||
strategy.finish_request(module, use, self.state, preferred=preferred)
|
||||
self._empty_cache_after_large_release(
|
||||
use, strategy, module, was_on_cuda
|
||||
)
|
||||
was_on_cuda = self._module_on_cuda(module)
|
||||
strategy.finish_request(module, use, self.state, preferred=preferred)
|
||||
self._empty_cache_after_large_release(use, strategy, module, was_on_cuda)
|
||||
|
||||
def stage_name(self, stage: ComponentResidencyStage) -> str:
|
||||
return self._stage_names_by_id.get(id(stage), stage.__class__.__name__)
|
||||
@@ -480,17 +522,23 @@ class ComponentResidencyManager:
|
||||
module = self.pipeline.modules.get(component_name)
|
||||
return module if isinstance(module, nn.Module) else None
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def strategy_for(
|
||||
self, component_name: str, module: nn.Module
|
||||
) -> ComponentResidencyStrategy:
|
||||
"""Return the pre-registered strategy for a specific component"""
|
||||
cached = self._strategy_cache.get(component_name)
|
||||
if cached is not None and cached[0] is module:
|
||||
return cached[1]
|
||||
custom_strategy = self._custom_strategies.get(component_name)
|
||||
if custom_strategy is not None:
|
||||
return custom_strategy
|
||||
return build_component_residency_strategy(
|
||||
component_name, module, self.server_args
|
||||
)
|
||||
if custom_strategy is None:
|
||||
strategy = build_component_residency_strategy(
|
||||
component_name,
|
||||
module,
|
||||
self.server_args,
|
||||
)
|
||||
else:
|
||||
strategy = custom_strategy
|
||||
self._strategy_cache[component_name] = (module, strategy)
|
||||
return strategy
|
||||
|
||||
def _next_stage_name(self, stage_index: int) -> str | None:
|
||||
next_index = stage_index + 1
|
||||
@@ -511,9 +559,6 @@ class ComponentResidencyManager:
|
||||
for index in range(self._current_use_index + 1, len(self._ordered_uses)):
|
||||
if self._same_use(self._ordered_uses[index], use):
|
||||
return index
|
||||
for index, candidate in enumerate(self._ordered_uses):
|
||||
if self._same_use(candidate, use):
|
||||
return index
|
||||
return None
|
||||
|
||||
def _prefetch_next_memory_intensive_use(self) -> None:
|
||||
@@ -526,10 +571,7 @@ class ComponentResidencyManager:
|
||||
return
|
||||
|
||||
def _should_keep_after_use(self, use: ComponentUse) -> bool:
|
||||
future_component_names = {
|
||||
future.component_name for future in self.state.future_uses
|
||||
}
|
||||
if use.component_name in future_component_names:
|
||||
if self.state.future_uses and self._same_use(use, self.state.future_uses[0]):
|
||||
return True
|
||||
module = self.get_module(use.component_name)
|
||||
if module is not None and self._should_keep_single_dit(
|
||||
@@ -539,12 +581,6 @@ class ComponentResidencyManager:
|
||||
return False
|
||||
|
||||
def _should_keep_single_dit(self, component_name: str, module: nn.Module) -> bool:
|
||||
"""Keep a single DiT resident only when its effective strategy is resident.
|
||||
|
||||
The single-DiT fast path is a performance optimization, not a memory
|
||||
policy. In particular, it must not override explicit or auto-selected
|
||||
CPU/layerwise offload.
|
||||
"""
|
||||
if not self._is_single_dit_component(component_name):
|
||||
return False
|
||||
return isinstance(self.strategy_for(component_name, module), ResidentStrategy)
|
||||
@@ -556,7 +592,6 @@ class ComponentResidencyManager:
|
||||
)
|
||||
|
||||
def _preferred_request_end_use(self) -> ComponentUse | None:
|
||||
"""Returns a ComponentUse preferred to be resident after a request finishes, to prepare for next request"""
|
||||
for uses in self._stage_uses_by_index:
|
||||
for use in uses:
|
||||
if use.preferred_ready_after_request:
|
||||
@@ -609,7 +644,6 @@ class ComponentResidencyManager:
|
||||
module: nn.Module,
|
||||
was_on_cuda: bool,
|
||||
) -> None:
|
||||
"""explicitly empty cache after potential release of large component"""
|
||||
if not use.memory_intensive:
|
||||
return
|
||||
released_cuda_storage = was_on_cuda and not self._module_on_cuda(module)
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Parsing and resolution for component residency CLI modes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS,
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_VAE_GROUP,
|
||||
component_base_name,
|
||||
is_dit_component_name,
|
||||
is_image_encoder_component_name,
|
||||
is_text_encoder_component_name,
|
||||
)
|
||||
|
||||
RESIDENT = "resident"
|
||||
COMPONENT_OFFLOAD = "component-offload"
|
||||
LAYERWISE_OFFLOAD = "layerwise-offload"
|
||||
COMPONENT_RESIDENCY_MODES = frozenset(
|
||||
(
|
||||
RESIDENT,
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class ComponentResidencyError(ValueError):
|
||||
"""Invalid or unsupported user-selected component residency."""
|
||||
|
||||
|
||||
COMPONENT_RESIDENCY_GROUPS = frozenset(
|
||||
(
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS,
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_VAE_GROUP,
|
||||
)
|
||||
)
|
||||
COMPONENT_RESIDENCY_GROUP_PRECEDENCE = (
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_VAE_GROUP,
|
||||
)
|
||||
|
||||
|
||||
def normalize_component_residency(
|
||||
assignments: str | Sequence[str] | Mapping[str, str] | None,
|
||||
) -> dict[str, str] | None:
|
||||
if assignments is None:
|
||||
return None
|
||||
|
||||
if isinstance(assignments, Mapping):
|
||||
entries = assignments.items()
|
||||
else:
|
||||
values = [assignments] if isinstance(assignments, str) else assignments
|
||||
parsed_entries: list[tuple[str, str]] = []
|
||||
for value in values:
|
||||
if not isinstance(value, str):
|
||||
raise ComponentResidencyError(
|
||||
f"Invalid component residency assignment: {value!r}"
|
||||
)
|
||||
for assignment in value.split(","):
|
||||
assignment = assignment.strip()
|
||||
if not assignment:
|
||||
continue
|
||||
if "=" not in assignment:
|
||||
raise ComponentResidencyError(
|
||||
"Component residency must use COMPONENT=MODE, got "
|
||||
f"{assignment!r}"
|
||||
)
|
||||
selector, mode = assignment.split("=", 1)
|
||||
parsed_entries.append((selector, mode))
|
||||
entries = parsed_entries
|
||||
|
||||
normalized: dict[str, str] = {}
|
||||
for raw_selector, raw_mode in entries:
|
||||
if not isinstance(raw_selector, str) or not isinstance(raw_mode, str):
|
||||
raise ComponentResidencyError(
|
||||
"Invalid component residency assignment: "
|
||||
f"{raw_selector!r}={raw_mode!r}"
|
||||
)
|
||||
selector = raw_selector.strip().replace("-", "_").lower()
|
||||
mode = raw_mode.strip().replace("_", "-").lower()
|
||||
if not selector:
|
||||
raise ComponentResidencyError(
|
||||
"Component residency selector cannot be empty"
|
||||
)
|
||||
if mode not in COMPONENT_RESIDENCY_MODES:
|
||||
expected = ", ".join(sorted(COMPONENT_RESIDENCY_MODES))
|
||||
raise ComponentResidencyError(
|
||||
f"Invalid component residency mode {raw_mode!r} for "
|
||||
f"{selector!r}; expected one of: {expected}"
|
||||
)
|
||||
normalized[selector] = mode
|
||||
|
||||
return normalized or None
|
||||
|
||||
|
||||
def component_residency_selector_matches(component_name: str, selector: str) -> bool:
|
||||
if selector == LAYERWISE_OFFLOAD_ALL_COMPONENTS:
|
||||
return True
|
||||
if selector == LAYERWISE_OFFLOAD_DIT_GROUP:
|
||||
return is_dit_component_name(component_name)
|
||||
if selector == LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP:
|
||||
return is_text_encoder_component_name(component_name)
|
||||
if selector == LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP:
|
||||
return is_image_encoder_component_name(component_name)
|
||||
if selector == LAYERWISE_OFFLOAD_VAE_GROUP:
|
||||
return component_base_name(component_name) in (
|
||||
"vae",
|
||||
"video_vae",
|
||||
"audio_vae",
|
||||
"delight_vae",
|
||||
"hy3dshape_vae",
|
||||
"paint_vae",
|
||||
)
|
||||
return component_name == selector
|
||||
|
||||
|
||||
def resolve_component_residency_mode(
|
||||
component_name: str, assignments: Mapping[str, str] | None
|
||||
) -> str | None:
|
||||
if not assignments:
|
||||
return None
|
||||
|
||||
exact_mode = assignments.get(component_name)
|
||||
if exact_mode is not None:
|
||||
return exact_mode
|
||||
|
||||
for selector in COMPONENT_RESIDENCY_GROUP_PRECEDENCE:
|
||||
mode = assignments.get(selector)
|
||||
if mode is not None and component_residency_selector_matches(
|
||||
component_name, selector
|
||||
):
|
||||
return mode
|
||||
return assignments.get(LAYERWISE_OFFLOAD_ALL_COMPONENTS)
|
||||
|
||||
|
||||
def resolve_diffusers_pipeline_offload(
|
||||
assignments: Mapping[str, str] | None,
|
||||
) -> bool | None:
|
||||
if assignments is None:
|
||||
return None
|
||||
if LAYERWISE_OFFLOAD in assignments.values():
|
||||
raise ComponentResidencyError(
|
||||
"--component-residency layerwise-offload requires the native SGLang "
|
||||
"backend"
|
||||
)
|
||||
|
||||
pipeline_mode = assignments.get(LAYERWISE_OFFLOAD_ALL_COMPONENTS)
|
||||
if len(assignments) == 1 and pipeline_mode is not None:
|
||||
return pipeline_mode == COMPONENT_OFFLOAD
|
||||
raise ComponentResidencyError(
|
||||
"The diffusers backend supports only pipeline-wide residency; use "
|
||||
"--component-residency all=resident or all=component-offload"
|
||||
)
|
||||
+229
@@ -0,0 +1,229 @@
|
||||
"""Runtime strategies used by the component residency manager."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributed.fsdp import FSDPModule
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
ResidencyState,
|
||||
)
|
||||
|
||||
|
||||
def _module_to_local_device(
|
||||
module: nn.Module, *, dtype: torch.dtype | None = None
|
||||
) -> None:
|
||||
device = get_local_torch_device()
|
||||
tensor = _module_reference_tensor(module)
|
||||
if tensor is not None and tensor.device == device:
|
||||
if dtype is None or tensor.dtype == dtype:
|
||||
return
|
||||
if dtype is None:
|
||||
module.to(device, non_blocking=True)
|
||||
else:
|
||||
module.to(device, dtype=dtype, non_blocking=True)
|
||||
|
||||
|
||||
def _module_reference_tensor(module: nn.Module) -> torch.Tensor | None:
|
||||
tensor = next(module.parameters(), None)
|
||||
if tensor is None:
|
||||
tensor = next(module.buffers(), None)
|
||||
return tensor
|
||||
|
||||
|
||||
def _module_ready_on_local_device(
|
||||
module: nn.Module, *, dtype: torch.dtype | None = None
|
||||
) -> bool:
|
||||
tensor = _module_reference_tensor(module)
|
||||
if tensor is None:
|
||||
return True
|
||||
if tensor.device != get_local_torch_device():
|
||||
return False
|
||||
return dtype is None or tensor.dtype == dtype
|
||||
|
||||
|
||||
def is_fsdp_managed_module(module: nn.Module) -> bool:
|
||||
return isinstance(module, FSDPModule)
|
||||
|
||||
|
||||
class ComponentResidencyStrategy:
|
||||
"""Controls one component's device placement around declared use intervals."""
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def finish_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def finish_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
*,
|
||||
preferred: bool,
|
||||
) -> None:
|
||||
if not preferred:
|
||||
self.finish_use(module, use, state)
|
||||
|
||||
def prefetch_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> bool:
|
||||
self.prepare_for_use(module, use, state)
|
||||
return True
|
||||
|
||||
|
||||
class ResidentStrategy(ComponentResidencyStrategy):
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
if is_fsdp_managed_module(module):
|
||||
return
|
||||
_module_to_local_device(module, dtype=use.target_dtype)
|
||||
|
||||
|
||||
class ComponentOffloadStrategy(ComponentResidencyStrategy):
|
||||
"""Move a complete component between CPU and device around each use."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._prefetch_stream: object | None = None
|
||||
self._ready_events: dict[str, object] = {}
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
_module_to_local_device(module, dtype=use.target_dtype)
|
||||
|
||||
def wait_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
ready_event = self._ready_events.get(use.component_name)
|
||||
if ready_event is None or not current_platform.is_cuda():
|
||||
return
|
||||
torch.get_device_module().current_stream().wait_event(ready_event)
|
||||
|
||||
def prefetch_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> bool:
|
||||
if not current_platform.is_cuda():
|
||||
self.prepare_for_use(module, use, state)
|
||||
return True
|
||||
if _module_ready_on_local_device(module, dtype=use.target_dtype):
|
||||
return True
|
||||
if self._prefetch_stream is None:
|
||||
self._prefetch_stream = torch.get_device_module().Stream(
|
||||
device=get_local_torch_device()
|
||||
)
|
||||
with torch.get_device_module().stream(self._prefetch_stream):
|
||||
_module_to_local_device(module, dtype=use.target_dtype)
|
||||
event = torch.get_device_module().Event()
|
||||
event.record(self._prefetch_stream)
|
||||
self._ready_events[use.component_name] = event
|
||||
return True
|
||||
|
||||
def finish_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.wait_for_use(module, use, state)
|
||||
tensor = _module_reference_tensor(module)
|
||||
if tensor is not None and tensor.device.type != "cpu":
|
||||
module.to("cpu", non_blocking=True)
|
||||
self._ready_events.pop(use.component_name, None)
|
||||
|
||||
def finish_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
*,
|
||||
preferred: bool,
|
||||
) -> None:
|
||||
if preferred and state.batch_is_warmup:
|
||||
self.prepare_for_use(module, use, state)
|
||||
self.wait_for_use(module, use, state)
|
||||
return
|
||||
self.finish_use(module, use, state)
|
||||
|
||||
|
||||
class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
|
||||
"""Run the lifecycle of an already configured layerwise component."""
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
if isinstance(module, LayerwiseOffloadableModuleMixin):
|
||||
module.prepare_for_next_req()
|
||||
|
||||
def finish_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
if not isinstance(module, LayerwiseOffloadableModuleMixin):
|
||||
return
|
||||
for manager in module.layerwise_offload_managers:
|
||||
manager.release_all()
|
||||
|
||||
def finish_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
*,
|
||||
preferred: bool,
|
||||
) -> None:
|
||||
if preferred:
|
||||
self.prepare_for_use(module, use, state)
|
||||
else:
|
||||
self.finish_use(module, use, state)
|
||||
-508
@@ -1,508 +0,0 @@
|
||||
"""
|
||||
Basic Component Resident Strategy Utilities for defining usage of components, to let ComponentResidencyManager to coordinate
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
ResidencyState,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _module_to_local_device(
|
||||
module: nn.Module, *, dtype: torch.dtype | None = None
|
||||
) -> None:
|
||||
device = get_local_torch_device()
|
||||
tensor = _module_reference_tensor(module)
|
||||
if tensor is not None and tensor.device == device:
|
||||
if dtype is None or tensor.dtype == dtype:
|
||||
return
|
||||
if dtype is None:
|
||||
module.to(device, non_blocking=True)
|
||||
else:
|
||||
module.to(device, dtype=dtype, non_blocking=True)
|
||||
|
||||
|
||||
def _module_reference_tensor(module: nn.Module) -> torch.Tensor | None:
|
||||
tensor = next(module.parameters(), None)
|
||||
if tensor is None:
|
||||
tensor = next(module.buffers(), None)
|
||||
return tensor
|
||||
|
||||
|
||||
def _module_ready_on_local_device(
|
||||
module: nn.Module, *, dtype: torch.dtype | None = None
|
||||
) -> bool:
|
||||
tensor = _module_reference_tensor(module)
|
||||
if tensor is None:
|
||||
return True
|
||||
if tensor.device != get_local_torch_device():
|
||||
return False
|
||||
return dtype is None or tensor.dtype == dtype
|
||||
|
||||
|
||||
def is_fsdp_managed_module(module: nn.Module) -> bool:
|
||||
return module.__class__.__name__.startswith("FSDP")
|
||||
|
||||
|
||||
class ComponentResidencyStrategy:
|
||||
"""Baseclass for describing how a component should be treated (regarding where its weights locates)
|
||||
|
||||
e.g., a LayerwiseOffloadStrategy would override:
|
||||
enter: to prefetch some layers before DiT is used, and
|
||||
exits: to release GPU weight snapshot after DiT is used
|
||||
to achieve desired behavior
|
||||
|
||||
"""
|
||||
|
||||
name = "resident"
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.enter(module)
|
||||
|
||||
def wait_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
"""Wait for the preparation to be ready, only applicable for async device syncs"""
|
||||
pass
|
||||
|
||||
def finish_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
"""Finish a specific component use"""
|
||||
self.exit(module)
|
||||
|
||||
def prepare_after_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
"""Called after a request is finished, to prepare for the upcoming request"""
|
||||
pass
|
||||
|
||||
def finish_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
*,
|
||||
preferred: bool,
|
||||
) -> None:
|
||||
if preferred:
|
||||
self.prepare_for_use(module, use, state)
|
||||
self.wait_for_use(module, use, state)
|
||||
else:
|
||||
self.finish_use(module, use, state)
|
||||
|
||||
def prefetch_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> bool:
|
||||
self.prepare_for_use(module, use, state)
|
||||
return True
|
||||
|
||||
def enter(self, module: nn.Module) -> None:
|
||||
pass
|
||||
|
||||
def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class ResidentStrategy(ComponentResidencyStrategy):
|
||||
name = "resident"
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
if is_fsdp_managed_module(module):
|
||||
return
|
||||
_module_to_local_device(module, dtype=use.target_dtype)
|
||||
|
||||
|
||||
class SnapshotModuleResidency:
|
||||
"""Reusable snapshot-based module residency primitive.
|
||||
|
||||
This helper only knows how to:
|
||||
- keep CPU parameter/buffer snapshots,
|
||||
- prefetch a module (H2D) to the local device on a CUDA side stream
|
||||
- release a module by rebinding tensors to those snapshots,
|
||||
- track and wait for readiness events.
|
||||
|
||||
It deliberately does not know about pipeline stages, phases, or model-specific
|
||||
ordering. Strategy subclasses decide when each primitive is called.
|
||||
"""
|
||||
|
||||
def __init__(self, *, pin_cpu_memory: bool, enable_async_prefetch: bool) -> None:
|
||||
self.pin_cpu_memory = pin_cpu_memory
|
||||
self.enable_async_prefetch = enable_async_prefetch
|
||||
self._cpu_param_snapshots: dict[str, dict[str, torch.Tensor]] = {}
|
||||
self._cpu_buffer_snapshots: dict[str, dict[str, torch.Tensor]] = {}
|
||||
self._prefetch_stream: object | None = None
|
||||
self._ready_events: dict[str, object] = {}
|
||||
|
||||
@staticmethod
|
||||
def is_on_gpu(module: nn.Module | None) -> bool:
|
||||
if module is None:
|
||||
return False
|
||||
param = next(module.parameters(), None)
|
||||
return param is not None and param.device.type == "cuda"
|
||||
|
||||
def is_ready(self, component_name: str) -> bool:
|
||||
return component_name in self._ready_events
|
||||
|
||||
def wait_ready(self, component_name: str) -> None:
|
||||
"""wait for the (H2D) stream to be ready"""
|
||||
ready_event = self._ready_events.get(component_name)
|
||||
if ready_event is None or not current_platform.is_cuda():
|
||||
return
|
||||
torch.get_device_module().current_stream().wait_event(ready_event)
|
||||
|
||||
def record_ready(self, component_name: str, module: nn.Module | None) -> None:
|
||||
if not current_platform.is_cuda():
|
||||
self._ready_events.pop(component_name, None)
|
||||
return
|
||||
if not self.is_on_gpu(module):
|
||||
self._ready_events.pop(component_name, None)
|
||||
return
|
||||
event = torch.get_device_module().Event()
|
||||
event.record(torch.get_device_module().current_stream())
|
||||
self._ready_events[component_name] = event
|
||||
|
||||
@staticmethod
|
||||
def _clone_cpu_tensor_snapshot(
|
||||
tensor: torch.Tensor, *, pin_memory: bool
|
||||
) -> torch.Tensor:
|
||||
snapshot = tensor.detach()
|
||||
if snapshot.device.type == "cpu":
|
||||
if pin_memory and not snapshot.is_pinned():
|
||||
return snapshot.pin_memory()
|
||||
return snapshot
|
||||
|
||||
cpu_tensor = snapshot.to("cpu")
|
||||
if pin_memory:
|
||||
return cpu_tensor.pin_memory()
|
||||
return cpu_tensor
|
||||
|
||||
def _should_pin_memory(self) -> bool:
|
||||
return bool(self.pin_cpu_memory and torch.get_device_module().is_available())
|
||||
|
||||
def capture(self, component_name: str, module: nn.Module) -> None:
|
||||
"""Capture a CPU snapshot for a component"""
|
||||
if component_name in self._cpu_param_snapshots:
|
||||
return
|
||||
|
||||
pin_memory = self._should_pin_memory()
|
||||
self._cpu_param_snapshots[component_name] = {
|
||||
name: self._clone_cpu_tensor_snapshot(param.data, pin_memory=pin_memory)
|
||||
for name, param in module.named_parameters()
|
||||
}
|
||||
self._cpu_buffer_snapshots[component_name] = {
|
||||
name: self._clone_cpu_tensor_snapshot(buffer.data, pin_memory=pin_memory)
|
||||
for name, buffer in module.named_buffers()
|
||||
}
|
||||
|
||||
def release_to_snapshot(
|
||||
self,
|
||||
component_name: str,
|
||||
module: nn.Module,
|
||||
*,
|
||||
copy_runtime_buffers: bool = False,
|
||||
) -> None:
|
||||
"""Release CUDA storages by rebinding tensors to cached CPU snapshots.
|
||||
|
||||
This does not call `module.to("cpu")`. Instead, parameter and buffer
|
||||
storages are rebound to pre-captured CPU tensors so CUDA storages can be
|
||||
released by the allocator without an explicit D2H transfer.
|
||||
"""
|
||||
param_snapshots = self._cpu_param_snapshots.get(component_name)
|
||||
buffer_snapshots = self._cpu_buffer_snapshots.get(component_name)
|
||||
if param_snapshots is None or buffer_snapshots is None:
|
||||
module.to("cpu")
|
||||
self._ready_events.pop(component_name, None)
|
||||
return
|
||||
|
||||
pin_memory = self._should_pin_memory()
|
||||
for name, param in module.named_parameters():
|
||||
snapshot = param_snapshots.get(name)
|
||||
if snapshot is None:
|
||||
snapshot = self._clone_cpu_tensor_snapshot(
|
||||
param.data, pin_memory=pin_memory
|
||||
)
|
||||
param_snapshots[name] = snapshot
|
||||
param.data = snapshot
|
||||
|
||||
for name, buffer in module.named_buffers():
|
||||
snapshot = buffer_snapshots.get(name)
|
||||
if snapshot is None:
|
||||
snapshot = self._clone_cpu_tensor_snapshot(
|
||||
buffer.data, pin_memory=pin_memory
|
||||
)
|
||||
buffer_snapshots[name] = snapshot
|
||||
if copy_runtime_buffers:
|
||||
# Preserve runtime-updated buffers (e.g., lazily built caches) when
|
||||
# releasing back to CPU snapshots.
|
||||
if buffer.device.type == "cuda":
|
||||
snapshot.copy_(
|
||||
buffer.detach().to(device="cpu", dtype=snapshot.dtype)
|
||||
)
|
||||
elif buffer.device.type == "cpu":
|
||||
snapshot.copy_(buffer.detach().to(dtype=snapshot.dtype))
|
||||
buffer.data = snapshot
|
||||
|
||||
self._ready_events.pop(component_name, None)
|
||||
|
||||
def _supports_async_prefetch(self) -> bool:
|
||||
return self.enable_async_prefetch and current_platform.is_cuda()
|
||||
|
||||
def _get_prefetch_stream(self):
|
||||
"""returns a stream is async-prefetch is enabled"""
|
||||
if not self._supports_async_prefetch():
|
||||
return None
|
||||
if self._prefetch_stream is None:
|
||||
self._prefetch_stream = torch.get_device_module().Stream(
|
||||
device=get_local_torch_device()
|
||||
)
|
||||
return self._prefetch_stream
|
||||
|
||||
def prefetch_to_device(self, component_name: str, module: nn.Module | None) -> None:
|
||||
if module is None:
|
||||
self._ready_events.pop(component_name, None)
|
||||
return
|
||||
prefetch_stream = self._get_prefetch_stream()
|
||||
if prefetch_stream is None:
|
||||
# if the async prefetching is disabled
|
||||
module.to(get_local_torch_device(), non_blocking=True)
|
||||
self.record_ready(component_name, module)
|
||||
return
|
||||
with torch.get_device_module().stream(prefetch_stream):
|
||||
module.to(get_local_torch_device(), non_blocking=True)
|
||||
event = torch.get_device_module().Event()
|
||||
event.record(prefetch_stream)
|
||||
self._ready_events[component_name] = event
|
||||
|
||||
|
||||
class SnapshotStrategy(ComponentResidencyStrategy):
|
||||
"""Snapshot residency: async H2D before use and light snapshot release after use."""
|
||||
|
||||
name = "snapshot"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
pin_cpu_memory: bool,
|
||||
enable_async_prefetch: bool,
|
||||
copy_runtime_buffers_on_release: bool = False,
|
||||
) -> None:
|
||||
self._snapshot_residency = SnapshotModuleResidency(
|
||||
pin_cpu_memory=pin_cpu_memory,
|
||||
enable_async_prefetch=enable_async_prefetch,
|
||||
)
|
||||
self._copy_runtime_buffers_on_release = copy_runtime_buffers_on_release
|
||||
|
||||
def capture(self, component_name: str, module: nn.Module) -> None:
|
||||
self._snapshot_residency.capture(component_name, module)
|
||||
|
||||
def is_ready(self, component_name: str) -> bool:
|
||||
return self._snapshot_residency.is_ready(component_name)
|
||||
|
||||
def record_ready(self, component_name: str, module: nn.Module | None) -> None:
|
||||
self._snapshot_residency.record_ready(component_name, module)
|
||||
|
||||
def prefetch_component(self, component_name: str, module: nn.Module | None) -> None:
|
||||
if SnapshotModuleResidency.is_on_gpu(module):
|
||||
self._snapshot_residency.record_ready(component_name, module)
|
||||
return
|
||||
self._snapshot_residency.prefetch_to_device(component_name, module)
|
||||
|
||||
def wait_component_ready(self, component_name: str) -> None:
|
||||
self._snapshot_residency.wait_ready(component_name)
|
||||
|
||||
def release_component(self, component_name: str, module: nn.Module) -> None:
|
||||
self._snapshot_residency.release_to_snapshot(
|
||||
component_name,
|
||||
module,
|
||||
copy_runtime_buffers=self._copy_runtime_buffers_on_release,
|
||||
)
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.prefetch_component(use.component_name, module)
|
||||
|
||||
def wait_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.wait_component_ready(use.component_name)
|
||||
|
||||
def finish_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.release_component(use.component_name, module)
|
||||
|
||||
def prepare_after_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.prepare_for_use(module, use, state)
|
||||
|
||||
|
||||
class VanillaD2HStrategy(ComponentResidencyStrategy):
|
||||
"""A strategy that performs native torch D2H and H2D for a component"""
|
||||
|
||||
name = "vanilla"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._prefetch_stream: object | None = None
|
||||
self._ready_events: dict[str, object] = {}
|
||||
|
||||
def prepare_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
_module_to_local_device(module, dtype=use.target_dtype)
|
||||
|
||||
def wait_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
ready_event = self._ready_events.get(use.component_name)
|
||||
if ready_event is None or not current_platform.is_cuda():
|
||||
return
|
||||
torch.get_device_module().current_stream().wait_event(ready_event)
|
||||
|
||||
def prefetch_for_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> bool:
|
||||
if not current_platform.is_cuda():
|
||||
self.prepare_for_use(module, use, state)
|
||||
return True
|
||||
if _module_ready_on_local_device(module, dtype=use.target_dtype):
|
||||
return True
|
||||
if self._prefetch_stream is None:
|
||||
self._prefetch_stream = torch.get_device_module().Stream(
|
||||
device=get_local_torch_device()
|
||||
)
|
||||
with torch.get_device_module().stream(self._prefetch_stream):
|
||||
_module_to_local_device(module, dtype=use.target_dtype)
|
||||
event = torch.get_device_module().Event()
|
||||
event.record(self._prefetch_stream)
|
||||
self._ready_events[use.component_name] = event
|
||||
return True
|
||||
|
||||
def enter(self, module: nn.Module) -> None:
|
||||
param = next(module.parameters(), None)
|
||||
if param is not None and param.device.type == "cpu":
|
||||
_module_to_local_device(module)
|
||||
|
||||
def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
|
||||
param = next(module.parameters(), None)
|
||||
if param is not None and param.device.type == "cuda":
|
||||
module.to("cpu", non_blocking=True)
|
||||
|
||||
def finish_use(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.wait_for_use(module, use, state)
|
||||
self.exit(module)
|
||||
self._ready_events.pop(use.component_name, None)
|
||||
|
||||
def prepare_after_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.prefetch_for_use(module, use, state)
|
||||
|
||||
def finish_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
*,
|
||||
preferred: bool,
|
||||
) -> None:
|
||||
if preferred and state.batch_is_warmup:
|
||||
self.prepare_for_use(module, use, state)
|
||||
self.wait_for_use(module, use, state)
|
||||
return
|
||||
if not preferred:
|
||||
self.finish_use(module, use, state)
|
||||
|
||||
|
||||
class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
|
||||
"""A wrapper around LayerwiseOffloadManager to fit in a ComponentResidencyStrategy"""
|
||||
|
||||
name = "layerwise"
|
||||
|
||||
def enter(self, module: nn.Module) -> None:
|
||||
if isinstance(module, LayerwiseOffloadableModuleMixin):
|
||||
module.prepare_for_next_req()
|
||||
|
||||
def exit(self, module: nn.Module, next_module: nn.Module | None = None) -> None:
|
||||
if not isinstance(module, LayerwiseOffloadableModuleMixin):
|
||||
return
|
||||
for manager in module.layerwise_offload_managers:
|
||||
manager.release_all()
|
||||
|
||||
def prepare_after_request(
|
||||
self,
|
||||
module: nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
) -> None:
|
||||
self.prepare_for_use(module, use, state)
|
||||
+91
-25
@@ -6,12 +6,19 @@ from typing import Any, Dict, List, Set, Tuple
|
||||
import torch
|
||||
from torch.distributed.tensor import DTensor
|
||||
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_RESIDENCY_GROUPS,
|
||||
LAYERWISE_OFFLOAD,
|
||||
ComponentResidencyError,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS,
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
RESIDENCY_POLICIES,
|
||||
RESIDENCY_POLICY_LEADING,
|
||||
RESIDENCY_POLICY_STRIDED,
|
||||
is_dit_component_name,
|
||||
layerwise_component_matches_any_selection,
|
||||
normalize_layerwise_offload_components,
|
||||
)
|
||||
@@ -328,14 +335,14 @@ class LayerwiseOffloadManager:
|
||||
|
||||
self.register_forward_hooks()
|
||||
self._configured = True
|
||||
logger.info(
|
||||
logger.debug(
|
||||
f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, num resident layers: {self.resident_layers}, total num layers: {self.num_layers}, residency policy: {self.residency_policy}"
|
||||
)
|
||||
if self.residency_policy == RESIDENCY_POLICY_STRIDED and self._streamed_order:
|
||||
# Printed because the layout is the whole point of the policy, and
|
||||
# "did it actually stride?" is otherwise only answerable from a
|
||||
# profile.
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"Strided residency streams layers %s (%d of %d)",
|
||||
list(self._streamed_order),
|
||||
len(self._streamed_order),
|
||||
@@ -800,13 +807,13 @@ class LayerwiseOffloadableModuleMixin:
|
||||
configured_layer_names.append(layer_name)
|
||||
|
||||
if configured_layer_names:
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"Enabled layerwise offload for %s on modules: %s",
|
||||
self.__class__.__name__,
|
||||
configured_layer_names,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"No layerwise-offloadable ModuleList found for %s. Candidates: %s",
|
||||
self.__class__.__name__,
|
||||
self.layer_names,
|
||||
@@ -920,7 +927,7 @@ def get_layerwise_offload_component_names_for_pipeline(
|
||||
return [
|
||||
component_name
|
||||
for component_name, module in modules.items()
|
||||
if isinstance(module, LayerwiseOffloadableModuleMixin)
|
||||
if isinstance(module, torch.nn.Module)
|
||||
]
|
||||
|
||||
explicit_component_names = selected_component_names - {LAYERWISE_OFFLOAD_DIT_GROUP}
|
||||
@@ -932,10 +939,12 @@ def get_layerwise_offload_component_names_for_pipeline(
|
||||
):
|
||||
selected_pipeline_component_names.append(component_name)
|
||||
continue
|
||||
if (
|
||||
select_dit_group
|
||||
and isinstance(module, LayerwiseOffloadableModuleMixin)
|
||||
and module.layerwise_offload_dit_group_enabled
|
||||
if select_dit_group and (
|
||||
is_dit_component_name(component_name)
|
||||
or (
|
||||
isinstance(module, LayerwiseOffloadableModuleMixin)
|
||||
and module.layerwise_offload_dit_group_enabled
|
||||
)
|
||||
):
|
||||
selected_pipeline_component_names.append(component_name)
|
||||
return selected_pipeline_component_names
|
||||
@@ -969,12 +978,42 @@ def configure_layerwise_offload_modules(
|
||||
selected_component_names is not None
|
||||
and LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_component_names
|
||||
)
|
||||
selected_pipeline_component_names = (
|
||||
get_layerwise_offload_component_names_for_pipeline(
|
||||
modules,
|
||||
normalized_component_names,
|
||||
exact_layerwise_selectors = {
|
||||
selector
|
||||
for selector, mode in (server_args.component_residency or {}).items()
|
||||
if mode == LAYERWISE_OFFLOAD and selector not in COMPONENT_RESIDENCY_GROUPS
|
||||
}
|
||||
if server_args.component_residency is not None:
|
||||
selected_pipeline_component_names = [
|
||||
component_name
|
||||
for component_name, module in modules.items()
|
||||
if server_args.residency_mode(component_name) == LAYERWISE_OFFLOAD
|
||||
and (
|
||||
isinstance(module, torch.nn.Module)
|
||||
or component_name in exact_layerwise_selectors
|
||||
)
|
||||
]
|
||||
else:
|
||||
selected_pipeline_component_names = (
|
||||
get_layerwise_offload_component_names_for_pipeline(
|
||||
modules,
|
||||
normalized_component_names,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
if (
|
||||
warn_missing
|
||||
and server_args.component_residency is not None
|
||||
and server_args.disagg_role == RoleType.MONOLITHIC
|
||||
):
|
||||
missing_component_names = sorted(exact_layerwise_selectors - modules.keys())
|
||||
if missing_component_names:
|
||||
logger.warning(
|
||||
"Layerwise offload components are not currently loaded: %s. "
|
||||
"Available pipeline components: %s",
|
||||
missing_component_names,
|
||||
sorted(modules),
|
||||
)
|
||||
|
||||
if warn_missing and selected_component_names is not None and not select_all:
|
||||
explicit_component_names = selected_component_names - {
|
||||
@@ -998,21 +1037,45 @@ def configure_layerwise_offload_modules(
|
||||
sorted(modules),
|
||||
)
|
||||
|
||||
unsupported_component_names = [
|
||||
unsupported_component_names = [
|
||||
component_name
|
||||
for component_name in selected_pipeline_component_names
|
||||
if not isinstance(modules[component_name], LayerwiseOffloadableModuleMixin)
|
||||
]
|
||||
explicit_unsupported_component_names = [
|
||||
component_name
|
||||
for component_name in unsupported_component_names
|
||||
if (warn_missing and server_args.component_residency is None)
|
||||
or server_args.is_explicit_layerwise_offload_component(component_name)
|
||||
]
|
||||
if explicit_unsupported_component_names:
|
||||
raise ComponentResidencyError(
|
||||
"Components selected for layerwise-offload do not support it: "
|
||||
f"{sorted(explicit_unsupported_component_names)}"
|
||||
)
|
||||
if unsupported_component_names:
|
||||
for component_name in unsupported_component_names:
|
||||
server_args.record_component_layerwise_capability(
|
||||
component_name, supported=False
|
||||
)
|
||||
selected_pipeline_component_names = [
|
||||
component_name
|
||||
for component_name in selected_pipeline_component_names
|
||||
if not isinstance(modules[component_name], LayerwiseOffloadableModuleMixin)
|
||||
if component_name not in unsupported_component_names
|
||||
]
|
||||
if unsupported_component_names:
|
||||
logger.warning(
|
||||
"Layerwise offload components do not support layerwise offload: %s",
|
||||
sorted(unsupported_component_names),
|
||||
)
|
||||
logger.warning(
|
||||
"Auto layerwise selection skipped unsupported components; their "
|
||||
"existing placement remains active: %s",
|
||||
sorted(unsupported_component_names),
|
||||
)
|
||||
|
||||
for component_name in selected_pipeline_component_names:
|
||||
module = modules[component_name]
|
||||
if not isinstance(module, LayerwiseOffloadableModuleMixin):
|
||||
continue
|
||||
server_args.record_component_layerwise_capability(
|
||||
component_name, supported=True
|
||||
)
|
||||
module_id = id(module)
|
||||
if module_id in configured_module_ids:
|
||||
# avoid duplicated configures on a same module
|
||||
@@ -1022,14 +1085,17 @@ def configure_layerwise_offload_modules(
|
||||
|
||||
if not is_layerwise_offloaded_module(module):
|
||||
module.configure_layerwise_offload(server_args)
|
||||
if is_layerwise_offloaded_module(module):
|
||||
configured_component_names.append(component_name)
|
||||
if not is_layerwise_offloaded_module(module):
|
||||
raise ComponentResidencyError(
|
||||
f"Component {component_name!r} did not enable layerwise offload"
|
||||
)
|
||||
configured_component_names.append(component_name)
|
||||
|
||||
if configured_component_names:
|
||||
logger.info(
|
||||
"Enabled layerwise offload for pipeline components: %s",
|
||||
configured_component_names,
|
||||
)
|
||||
else:
|
||||
logger.info("No pipeline component supports layerwise offload.")
|
||||
elif warn_missing:
|
||||
logger.debug("No selected pipeline component enabled layerwise offload")
|
||||
return configured_component_names
|
||||
|
||||
+43
-46
@@ -24,13 +24,19 @@ LAYERWISE_OFFLOAD_DEFAULT_GROUP_COMPONENTS = (
|
||||
DIT_COMPONENT_NAMES = frozenset(
|
||||
{
|
||||
"transformer",
|
||||
"transformer_2",
|
||||
"video_dit",
|
||||
"video_dit_2",
|
||||
"audio_dit",
|
||||
"dual_tower_bridge",
|
||||
"delight_transformer",
|
||||
"hy3dshape_model",
|
||||
"paint_transformer",
|
||||
"unconditional_transformer",
|
||||
}
|
||||
)
|
||||
LEGACY_DIT_OFFLOAD_COMPONENT_NAMES = frozenset(
|
||||
{
|
||||
"connectors",
|
||||
"dual_tower_bridge",
|
||||
"vision_language_encoder",
|
||||
}
|
||||
)
|
||||
VAE_COMPONENT_NAMES = frozenset(
|
||||
@@ -41,9 +47,11 @@ VAE_COMPONENT_NAMES = frozenset(
|
||||
"vocoder",
|
||||
"spatial_upsampler",
|
||||
"condition_image_encoder",
|
||||
"diffusion_decoder",
|
||||
"delight_vae",
|
||||
"diffusion_decoder",
|
||||
"hy3dshape_vae",
|
||||
"paint_vae",
|
||||
"sound_tokenizer",
|
||||
}
|
||||
)
|
||||
DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
|
||||
@@ -55,17 +63,25 @@ DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES = frozenset(
|
||||
"paint_vae",
|
||||
}
|
||||
)
|
||||
CPU_OFFLOAD_FLAG_NAMES = (
|
||||
"dit_cpu_offload",
|
||||
"text_encoder_cpu_offload",
|
||||
"image_encoder_cpu_offload",
|
||||
"vae_cpu_offload",
|
||||
)
|
||||
CPU_OFFLOAD_ALL_COMPONENTS = "all"
|
||||
|
||||
|
||||
def component_base_name(component_name: str) -> str:
|
||||
prefix, separator, suffix = component_name.rpartition("_")
|
||||
if separator and suffix.isdigit():
|
||||
return prefix
|
||||
return component_name
|
||||
|
||||
|
||||
def is_dit_component_name(component_name: str) -> bool:
|
||||
return component_name in DIT_COMPONENT_NAMES
|
||||
return component_base_name(component_name) in DIT_COMPONENT_NAMES
|
||||
|
||||
|
||||
def is_legacy_dit_offload_component_name(component_name: str) -> bool:
|
||||
return (
|
||||
is_dit_component_name(component_name)
|
||||
or component_base_name(component_name) in LEGACY_DIT_OFFLOAD_COMPONENT_NAMES
|
||||
)
|
||||
|
||||
|
||||
def normalize_cpu_offload_components(
|
||||
@@ -107,11 +123,11 @@ def cpu_offload_component_matches(
|
||||
if component_name in selected_component_names:
|
||||
return True
|
||||
if LAYERWISE_OFFLOAD_DIT_GROUP in selected_component_names:
|
||||
return is_dit_component_name(component_name)
|
||||
return is_legacy_dit_offload_component_name(component_name)
|
||||
if LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP in selected_component_names:
|
||||
return is_text_encoder_component_name(component_name)
|
||||
if LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP in selected_component_names:
|
||||
return component_name in ("image_encoder", "condition_image_encoder")
|
||||
return is_legacy_image_encoder_offload_component_name(component_name)
|
||||
if LAYERWISE_OFFLOAD_VAE_GROUP in selected_component_names:
|
||||
return is_vae_component_name(component_name)
|
||||
return False
|
||||
@@ -124,11 +140,21 @@ def is_text_encoder_component_name(component_name: str) -> bool:
|
||||
|
||||
|
||||
def is_image_encoder_component_name(component_name: str) -> bool:
|
||||
return component_name == "image_encoder"
|
||||
return component_base_name(component_name) in {
|
||||
"image_encoder",
|
||||
"hy3dshape_conditioner",
|
||||
}
|
||||
|
||||
|
||||
def is_legacy_image_encoder_offload_component_name(component_name: str) -> bool:
|
||||
return (
|
||||
is_image_encoder_component_name(component_name)
|
||||
or component_base_name(component_name) == "condition_image_encoder"
|
||||
)
|
||||
|
||||
|
||||
def is_vae_component_name(component_name: str) -> bool:
|
||||
return component_name in VAE_COMPONENT_NAMES
|
||||
return component_base_name(component_name) in VAE_COMPONENT_NAMES
|
||||
|
||||
|
||||
def layerwise_component_matches_selection(
|
||||
@@ -138,6 +164,8 @@ def layerwise_component_matches_selection(
|
||||
"""if the provided component_name (unnormalized, e.g., text_encoder_2) matches with the selected_component_name (normalized)"""
|
||||
if selected_component_name == LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP:
|
||||
return is_text_encoder_component_name(component_name)
|
||||
if selected_component_name == LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP:
|
||||
return is_image_encoder_component_name(component_name)
|
||||
if selected_component_name == LAYERWISE_OFFLOAD_VAE_GROUP:
|
||||
# `vae` is a default-policy selector; AV-side decoders remain explicit-only
|
||||
return component_name in DEFAULT_LAYERWISE_VAE_COMPONENT_NAMES
|
||||
@@ -154,37 +182,6 @@ def layerwise_component_matches_any_selection(
|
||||
)
|
||||
|
||||
|
||||
def cpu_offload_flags_for_layerwise_components(
|
||||
component_names: Sequence[str],
|
||||
) -> tuple[str, ...]:
|
||||
component_names = normalize_layerwise_offload_components(component_names) or []
|
||||
if LAYERWISE_OFFLOAD_ALL_COMPONENTS in component_names:
|
||||
return CPU_OFFLOAD_FLAG_NAMES
|
||||
|
||||
flag_names: list[str] = []
|
||||
if LAYERWISE_OFFLOAD_DIT_GROUP in component_names:
|
||||
flag_names.append("dit_cpu_offload")
|
||||
|
||||
for component_name in component_names:
|
||||
if component_name == LAYERWISE_OFFLOAD_DIT_GROUP:
|
||||
continue
|
||||
if is_dit_component_name(component_name):
|
||||
flag_name = "dit_cpu_offload"
|
||||
elif is_text_encoder_component_name(component_name):
|
||||
flag_name = "text_encoder_cpu_offload"
|
||||
elif is_image_encoder_component_name(component_name):
|
||||
flag_name = "image_encoder_cpu_offload"
|
||||
elif is_vae_component_name(component_name):
|
||||
flag_name = "vae_cpu_offload"
|
||||
else:
|
||||
continue
|
||||
|
||||
if flag_name not in flag_names:
|
||||
flag_names.append(flag_name)
|
||||
|
||||
return tuple(flag_names)
|
||||
|
||||
|
||||
def expand_layerwise_offload_component_group(component_name: str) -> tuple[str, ...]:
|
||||
if component_name == LAYERWISE_OFFLOAD_DEFAULT_GROUP:
|
||||
return LAYERWISE_OFFLOAD_DEFAULT_GROUP_COMPONENTS
|
||||
|
||||
@@ -12,6 +12,9 @@ from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import (
|
||||
LTX2ConnectorConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
|
||||
|
||||
@@ -528,12 +531,18 @@ class LTX2ConnectorTransformer1d(nn.Module):
|
||||
return hidden_states, attention_mask
|
||||
|
||||
|
||||
class LTX2TextConnectors(nn.Module):
|
||||
class LTX2TextConnectors(nn.Module, LayerwiseOffloadableModuleMixin):
|
||||
"""
|
||||
Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio
|
||||
streams.
|
||||
"""
|
||||
|
||||
layerwise_offload_dit_group_enabled = False
|
||||
layer_names = [
|
||||
"video_connector.transformer_blocks",
|
||||
"audio_connector.transformer_blocks",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: LTX2ConnectorConfig,
|
||||
|
||||
@@ -23,6 +23,9 @@ from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder imp
|
||||
from sglang.multimodal_gen.runtime.layers.visual_embedding import (
|
||||
timestep_embedding,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -782,18 +785,27 @@ def _tile_intervals(
|
||||
]
|
||||
|
||||
|
||||
class LTX2VideoDiffusionDecoderModel(nn.Module):
|
||||
class LTX2VideoDiffusionDecoderModel(nn.Module, LayerwiseOffloadableModuleMixin):
|
||||
"""Checkpoint-level wrapper: the decoder plus the latent statistics.
|
||||
|
||||
`diffusion_decoder/` stores `latents_mean` / `latents_std` alongside a
|
||||
`decoder.` submodule, so this mirrors that layout rather than flattening it.
|
||||
"""
|
||||
|
||||
layerwise_offload_dit_group_enabled = False
|
||||
|
||||
def __init__(self, config: LTX25DiffusionDecoderConfig) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
latent_channels = config.arch_config.latent_channels
|
||||
self.decoder = LTX2VideoDiffusionDecoder3d(config)
|
||||
self.layer_names = [
|
||||
*(
|
||||
f"decoder.det_stages.{index}"
|
||||
for index in range(len(self.decoder.det_stages))
|
||||
),
|
||||
"decoder.diff_blocks",
|
||||
]
|
||||
self.register_buffer(
|
||||
"latents_mean", torch.zeros(latent_channels), persistent=True
|
||||
)
|
||||
|
||||
@@ -26,6 +26,7 @@ class IdeogramQwen3VLTextEncoder(TextEncoder):
|
||||
"""Language-only Qwen3-VL text encoder stored inside Ideogram checkpoints."""
|
||||
|
||||
_activation_layers = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
|
||||
layer_names = ["language_model.layers"]
|
||||
|
||||
def __init__(self, config: Ideogram4TextEncoderConfig) -> None:
|
||||
super().__init__(config)
|
||||
|
||||
@@ -10,6 +10,10 @@ import torch
|
||||
from torch import nn
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
|
||||
|
||||
class Snake1d(nn.Module):
|
||||
def __init__(self, hidden_dim: int, logscale: bool = True) -> None:
|
||||
@@ -128,9 +132,12 @@ def _cfg(config: dict[str, Any], *keys: str, default: Any = None) -> Any:
|
||||
return default
|
||||
|
||||
|
||||
class Cosmos3AVAEAudioTokenizer(nn.Module):
|
||||
class Cosmos3AVAEAudioTokenizer(nn.Module, LayerwiseOffloadableModuleMixin):
|
||||
"""Cosmos3 audio tokenizer: latents → waveform via an Oobleck decoder stack."""
|
||||
|
||||
layerwise_offload_dit_group_enabled = False
|
||||
layer_names = ["decoder.block"]
|
||||
|
||||
def __init__(self, config: dict[str, Any]) -> None:
|
||||
super().__init__()
|
||||
self.sample_rate = int(
|
||||
|
||||
@@ -232,7 +232,10 @@ class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase):
|
||||
with set_default_torch_dtype(default_dtype), torch.device("meta"):
|
||||
model = model_cls(**{"config": dit_config, "hf_config": hf_config})
|
||||
|
||||
use_fsdp = server_args.use_fsdp_inference
|
||||
use_fsdp = server_args.should_use_fsdp_for_component("transformer")
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
"transformer"
|
||||
)
|
||||
if current_platform.is_mps():
|
||||
use_fsdp = False
|
||||
logger.info("Disabling FSDP for MPS platform as it's not compatible")
|
||||
@@ -248,7 +251,7 @@ class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase):
|
||||
)
|
||||
shard_model(
|
||||
model,
|
||||
cpu_offload=server_args.dit_cpu_offload,
|
||||
cpu_offload=False,
|
||||
reshard_after_forward=True,
|
||||
mp_policy=mp_policy,
|
||||
mesh=device_mesh,
|
||||
@@ -270,6 +273,10 @@ class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase):
|
||||
updated_mapping,
|
||||
server_args: ServerArgs,
|
||||
):
|
||||
use_fsdp = server_args.should_use_fsdp_for_component("transformer")
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
"transformer"
|
||||
)
|
||||
# Create weight iterator for loading
|
||||
weight_iterator = safetensors_weights_iterator(safetensors_list)
|
||||
|
||||
@@ -281,7 +288,7 @@ class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase):
|
||||
get_local_torch_device(),
|
||||
default_dtype,
|
||||
strict=True,
|
||||
cpu_offload=server_args.dit_cpu_offload,
|
||||
cpu_offload=component_starts_on_cpu and not use_fsdp,
|
||||
param_names_mapping=param_names_mapping_fn,
|
||||
)
|
||||
|
||||
|
||||
@@ -309,7 +309,10 @@ class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
model = model_cls(**{"config": dit_config, "hf_config": hf_config})
|
||||
|
||||
# Check if we should use FSDP
|
||||
use_fsdp = server_args.use_fsdp_inference
|
||||
use_fsdp = server_args.should_use_fsdp_for_component("transformer")
|
||||
component_starts_on_cpu = server_args.should_start_component_on_cpu(
|
||||
"transformer"
|
||||
)
|
||||
if current_platform.is_mps():
|
||||
use_fsdp = False
|
||||
logger.info("Disabling FSDP for MPS platform as it's not compatible")
|
||||
@@ -325,7 +328,7 @@ class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
)
|
||||
shard_model(
|
||||
model,
|
||||
cpu_offload=server_args.dit_cpu_offload,
|
||||
cpu_offload=False,
|
||||
reshard_after_forward=True,
|
||||
mp_policy=mp_policy,
|
||||
mesh=device_mesh,
|
||||
@@ -355,7 +358,7 @@ class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
get_local_torch_device(),
|
||||
default_dtype,
|
||||
strict=True,
|
||||
cpu_offload=server_args.dit_cpu_offload,
|
||||
cpu_offload=component_starts_on_cpu and not use_fsdp,
|
||||
param_names_mapping=param_names_mapping_fn,
|
||||
)
|
||||
|
||||
|
||||
@@ -24,6 +24,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im
|
||||
ComponentResidencyStrategy,
|
||||
get_global_component_residency_manager,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
resolve_diffusers_pipeline_offload,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
@@ -455,11 +458,22 @@ class DiffusersPipeline(ComposedPipelineBase):
|
||||
raise
|
||||
|
||||
# Use CPU offload (all-or-nothing in diffusers) if any component offload is requested.
|
||||
explicit_pipeline_offload = resolve_diffusers_pipeline_offload(
|
||||
server_args.component_residency
|
||||
)
|
||||
any_offload = (
|
||||
server_args.dit_cpu_offload
|
||||
or server_args.text_encoder_cpu_offload
|
||||
or server_args.image_encoder_cpu_offload
|
||||
or server_args.vae_cpu_offload
|
||||
explicit_pipeline_offload
|
||||
if explicit_pipeline_offload is not None
|
||||
else any(
|
||||
server_args.should_cpu_offload_component(component_name)
|
||||
for component_name in {
|
||||
*pipe.components,
|
||||
"transformer",
|
||||
"text_encoder",
|
||||
"image_encoder",
|
||||
"vae",
|
||||
}
|
||||
)
|
||||
)
|
||||
if any_offload:
|
||||
device = get_local_torch_device()
|
||||
|
||||
@@ -333,7 +333,7 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
|
||||
@staticmethod
|
||||
def _component_device(server_args: ServerArgs, component_name: str) -> torch.device:
|
||||
if server_args.should_cpu_offload_component(component_name):
|
||||
if server_args.should_start_component_on_cpu(component_name):
|
||||
return torch.device("cpu")
|
||||
return get_local_torch_device()
|
||||
|
||||
@@ -528,15 +528,36 @@ class Hunyuan3D2Pipeline(ComposedPipelineBase):
|
||||
components: dict[str, Any] = {}
|
||||
|
||||
components["hy3dshape_model"] = self._load_dit_model(
|
||||
model_config["model"], ckpt["model"], device, dtype
|
||||
model_config["model"],
|
||||
ckpt["model"],
|
||||
(
|
||||
torch.device("cpu")
|
||||
if server_args.should_start_component_on_cpu("hy3dshape_model")
|
||||
else device
|
||||
),
|
||||
dtype,
|
||||
)
|
||||
|
||||
components["hy3dshape_vae"] = self._load_simple_component(
|
||||
model_config["vae"], ckpt.get("vae"), device, dtype
|
||||
model_config["vae"],
|
||||
ckpt.get("vae"),
|
||||
(
|
||||
torch.device("cpu")
|
||||
if server_args.should_start_component_on_cpu("hy3dshape_vae")
|
||||
else device
|
||||
),
|
||||
dtype,
|
||||
)
|
||||
|
||||
components["hy3dshape_conditioner"] = self._load_simple_component(
|
||||
model_config["conditioner"], ckpt.get("conditioner"), device, dtype
|
||||
model_config["conditioner"],
|
||||
ckpt.get("conditioner"),
|
||||
(
|
||||
torch.device("cpu")
|
||||
if server_args.should_start_component_on_cpu("hy3dshape_conditioner")
|
||||
else device
|
||||
),
|
||||
dtype,
|
||||
)
|
||||
|
||||
components["hy3dshape_scheduler"] = self._instantiate_component(
|
||||
|
||||
@@ -483,12 +483,17 @@ class LTX2TwoStageResidencyStrategy(ComponentResidencyStrategy):
|
||||
) -> None:
|
||||
self.exit_phase(self._phase(use))
|
||||
|
||||
def prepare_after_request(
|
||||
def finish_request(
|
||||
self,
|
||||
module: torch.nn.Module,
|
||||
use: ComponentUse,
|
||||
state: ResidencyState,
|
||||
*,
|
||||
preferred: bool,
|
||||
) -> None:
|
||||
if not preferred:
|
||||
self.finish_use(module, use, state)
|
||||
return
|
||||
phase = self._phase(use)
|
||||
if phase != self.manager._active_phase:
|
||||
self.enter_phase(phase)
|
||||
@@ -539,7 +544,7 @@ class LTX2ResidentResidencyStrategy(LTX2TwoStageResidencyStrategy):
|
||||
class LTX2TwoStageResidencyController:
|
||||
"""
|
||||
LTX-2.3 two-stage residency controller.
|
||||
It builds the selected LTX2 ComponentResidencyStrategy and keeps the
|
||||
It builds the selected LTX2 component residency strategy and keeps the
|
||||
thin stage adapter methods that are specific to two-stage LoRA flow.
|
||||
|
||||
Modes:
|
||||
|
||||
@@ -105,7 +105,7 @@ class QwenImageLayeredPipeline(QwenImageEditPipeline):
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
def create_before_denoising_stage():
|
||||
return QwenImageLayeredBeforeDenoisingStage(
|
||||
stage = QwenImageLayeredBeforeDenoisingStage(
|
||||
vae=self.get_module("vae"),
|
||||
text_encoder=self.get_module("text_encoder"),
|
||||
tokenizer=self.get_module("tokenizer"),
|
||||
@@ -117,6 +117,7 @@ class QwenImageLayeredPipeline(QwenImageEditPipeline):
|
||||
server_args.pipeline_config.text_encoder_precisions[0]
|
||||
],
|
||||
)
|
||||
return stage
|
||||
|
||||
self.add_stage_factory(
|
||||
RoleType.ENCODER,
|
||||
|
||||
@@ -58,10 +58,6 @@ class PipelineExecutor(ABC):
|
||||
batch: Any,
|
||||
server_args: ServerArgs,
|
||||
) -> None:
|
||||
if isinstance(batch, list):
|
||||
if not batch:
|
||||
return
|
||||
batch = batch[0]
|
||||
self.component_residency_manager.begin_request(stages, batch, server_args)
|
||||
|
||||
def before_stage(
|
||||
@@ -187,31 +183,7 @@ class PipelineExecutor(ABC):
|
||||
|
||||
stage_name = stage._active_component_stage_name()
|
||||
for use in stage.component_uses(server_args, stage_name):
|
||||
component_name = use.component_name
|
||||
if server_args.dit_cpu_offload and component_name in (
|
||||
"transformer",
|
||||
"transformer_2",
|
||||
"video_dit",
|
||||
"audio_dit",
|
||||
):
|
||||
return True
|
||||
if server_args.text_encoder_cpu_offload and component_name.startswith(
|
||||
"text_encoder"
|
||||
):
|
||||
return True
|
||||
if server_args.image_encoder_cpu_offload and component_name in (
|
||||
"image_encoder",
|
||||
"condition_image_encoder",
|
||||
):
|
||||
return True
|
||||
if server_args.vae_cpu_offload and component_name in (
|
||||
"vae",
|
||||
"video_vae",
|
||||
"audio_vae",
|
||||
"vocoder",
|
||||
"spatial_upsampler",
|
||||
"condition_image_encoder",
|
||||
):
|
||||
if server_args.should_cpu_offload_component(use.component_name):
|
||||
return True
|
||||
return False
|
||||
|
||||
@@ -223,7 +195,11 @@ class PipelineExecutor(ABC):
|
||||
run_stage,
|
||||
):
|
||||
with self._stage_execution_context(stage, server_args):
|
||||
return run_stage(stage, payload)
|
||||
self.component_residency_manager.begin_stage()
|
||||
try:
|
||||
return run_stage(stage, payload)
|
||||
finally:
|
||||
self.component_residency_manager.end_stage()
|
||||
|
||||
@abstractmethod
|
||||
def execute(
|
||||
|
||||
@@ -68,6 +68,9 @@ class PipelineStage(StageDedupMixin, ABC):
|
||||
# calling super().__init__() still see a consistent explicit-range gate.
|
||||
_current_use_nvtx: bool = False
|
||||
_current_batch_is_warmup: bool = False
|
||||
_component_residency_manager = None
|
||||
_registered_stage_name: str | None = None
|
||||
_profile_stage_name: str | None = None
|
||||
|
||||
def __init__(self):
|
||||
self.server_args = get_global_server_args()
|
||||
@@ -161,28 +164,21 @@ class PipelineStage(StageDedupMixin, ABC):
|
||||
self._profile_stage_name = stage_name
|
||||
|
||||
def _component_stage_name(self, stage_name: str | None = None) -> str:
|
||||
return (
|
||||
stage_name
|
||||
or getattr(self, "_registered_stage_name", None)
|
||||
or self.__class__.__name__
|
||||
)
|
||||
return stage_name or self._registered_stage_name or self.__class__.__name__
|
||||
|
||||
def _active_component_stage_name(self) -> str:
|
||||
"""Stage name reported by the residency manager.
|
||||
|
||||
Only valid between ``before_stage`` and ``after_stage``; outside
|
||||
that window the manager state still holds the previous stage's
|
||||
name. Use :meth:`_component_stage_name` for the static identity.
|
||||
During execution this comes from the manager; otherwise the registered
|
||||
stage name provides the static identity.
|
||||
"""
|
||||
manager = getattr(self, "_component_residency_manager", None)
|
||||
manager_state = getattr(manager, "state", None)
|
||||
manager_stage_name = getattr(manager_state, "stage_name", None)
|
||||
if manager_stage_name is not None:
|
||||
return manager_stage_name
|
||||
manager = self._component_residency_manager
|
||||
if manager is not None and manager.state.stage_name is not None:
|
||||
return manager.state.stage_name
|
||||
return self._component_stage_name()
|
||||
|
||||
def _active_profile_stage_name(self) -> str:
|
||||
return getattr(self, "_profile_stage_name", None) or self.__class__.__name__
|
||||
return self._profile_stage_name or self.__class__.__name__
|
||||
|
||||
def _finish_active_component_use(self) -> None:
|
||||
if self._component_residency_manager is not None:
|
||||
@@ -241,6 +237,24 @@ class PipelineStage(StageDedupMixin, ABC):
|
||||
with self._use_component(use, module) as component:
|
||||
yield component
|
||||
|
||||
def begin_declared_component_use(
|
||||
self,
|
||||
*,
|
||||
component_name: str,
|
||||
module=None,
|
||||
phase: str | None = None,
|
||||
target_dtype: torch.dtype | None = None,
|
||||
) -> None:
|
||||
"""Keep a declared component active until the next use interval begins."""
|
||||
if self._component_residency_manager is None:
|
||||
return
|
||||
use = self._declared_component_use(
|
||||
component_name=component_name,
|
||||
phase=phase,
|
||||
target_dtype=target_dtype,
|
||||
)
|
||||
self._component_residency_manager.begin_use(use, module=module)
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
|
||||
@@ -93,7 +93,7 @@ from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_c
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
is_fsdp_managed_module,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
@@ -1447,10 +1447,8 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
||||
torch.mps.current_allocated_memory(),
|
||||
)
|
||||
if self._component_residency_manager is not None:
|
||||
self._component_residency_manager.remove_nvtx_hooks_for_module(
|
||||
self.transformer
|
||||
)
|
||||
self._component_residency_manager.strategy_for.cache_clear()
|
||||
self._component_residency_manager.finish_active_use(prefetch_next=False)
|
||||
self._component_residency_manager.forget_module(self.transformer)
|
||||
del self.transformer
|
||||
if pipeline is not None and "transformer" in pipeline.modules:
|
||||
del pipeline.modules["transformer"]
|
||||
|
||||
@@ -530,7 +530,11 @@ class LTX2ImageEncodingStage(PipelineStage):
|
||||
configure_layerwise_offload_modules(
|
||||
modules,
|
||||
server_args,
|
||||
component_names=server_args.layerwise_offload_components,
|
||||
component_names=(
|
||||
None
|
||||
if server_args.component_residency is not None
|
||||
else server_args.layerwise_offload_components
|
||||
),
|
||||
warn_missing=False,
|
||||
)
|
||||
return True
|
||||
|
||||
+50
-40
@@ -30,6 +30,9 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_sp_world_size,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
|
||||
PipelineStage,
|
||||
@@ -498,8 +501,9 @@ class Cosmos3LatentPreparationStage(PipelineStage):
|
||||
)
|
||||
cond_indexes = [0]
|
||||
|
||||
with torch.no_grad():
|
||||
cond_latent = self._vae_encode(pixel_input).to(dtype)
|
||||
with self.use_declared_component(component_name="vae", module=self.vae):
|
||||
with torch.no_grad():
|
||||
cond_latent = self._vae_encode(pixel_input).to(dtype)
|
||||
|
||||
max_idx = max(cond_indexes)
|
||||
if max_idx >= num_latent_frames:
|
||||
@@ -565,6 +569,11 @@ class Cosmos3LatentPreparationStage(PipelineStage):
|
||||
self._prepare_action_latents(batch, generator, device, dtype)
|
||||
return batch
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
return [ComponentUse(self._component_stage_name(stage_name), "vae")]
|
||||
|
||||
@staticmethod
|
||||
def _resolve_domain_id(batch: Req) -> int:
|
||||
"""Resolve action embodiment domain ID; required for action generation."""
|
||||
@@ -930,21 +939,6 @@ class Cosmos3DenoisingStage(PipelineStage):
|
||||
action_start_frame_offset=action_start_frame_offset,
|
||||
)
|
||||
|
||||
def _manage_device_placement(self, server_args: ServerArgs):
|
||||
"""Move transformer to GPU if CPU offload is enabled."""
|
||||
if not server_args.dit_cpu_offload:
|
||||
return
|
||||
|
||||
# FSDP manages offloading internally
|
||||
if server_args.use_fsdp_inference:
|
||||
return
|
||||
|
||||
device = get_local_torch_device()
|
||||
# Load the model to GPU if it's on CPU
|
||||
if next(self.transformer.parameters()).device.type == "cpu":
|
||||
self.log_info("Moving transformer to GPU for inference")
|
||||
self.transformer.to(device)
|
||||
|
||||
@staticmethod
|
||||
def _cfg_active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool:
|
||||
"""Return True iff CFG should be applied at timestep ``t``.
|
||||
@@ -961,8 +955,6 @@ class Cosmos3DenoisingStage(PipelineStage):
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
|
||||
"""Run the denoising loop with CFG and optional I2V conditioning."""
|
||||
self._manage_device_placement(server_args)
|
||||
|
||||
latents = batch.latents
|
||||
sound_latents = batch.audio_latents
|
||||
action_latents = getattr(batch, "action_latents", None)
|
||||
@@ -1449,6 +1441,18 @@ class Cosmos3DenoisingStage(PipelineStage):
|
||||
return tuple(cfg_model_parallel_all_reduce(coeff * p) for p in out)
|
||||
return cfg_model_parallel_all_reduce(coeff * out)
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
return [
|
||||
ComponentUse(
|
||||
self._component_stage_name(stage_name),
|
||||
"transformer",
|
||||
preferred_ready_after_request=True,
|
||||
memory_intensive=True,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
class Cosmos3DecodingStage(PipelineStage):
|
||||
"""
|
||||
@@ -1590,16 +1594,11 @@ class Cosmos3DecodingStage(PipelineStage):
|
||||
)
|
||||
|
||||
device = batch.latents.device
|
||||
if server_args.vae_cpu_offload:
|
||||
self.vae.to(device)
|
||||
with self.use_declared_component(component_name="vae", module=self.vae):
|
||||
with torch.no_grad():
|
||||
decoded = self._decode_latents(batch.latents)
|
||||
|
||||
with torch.no_grad():
|
||||
decoded = self._decode_latents(batch.latents)
|
||||
|
||||
if server_args.vae_cpu_offload and not getattr(batch, "is_warmup", False):
|
||||
self.vae.to("cpu", non_blocking=True)
|
||||
|
||||
self.log_info(f"Decoded tensor shape: {decoded.shape}")
|
||||
self.log_debug("Decoded tensor shape: %s", decoded.shape)
|
||||
output = self._postprocess_tensor(decoded)
|
||||
|
||||
if self._guardrails and batch.use_guardrails is not False:
|
||||
@@ -1614,23 +1613,25 @@ class Cosmos3DecodingStage(PipelineStage):
|
||||
else:
|
||||
output = check_video_safety(output)
|
||||
elif not is_image_gen:
|
||||
self.log_info(f"Postprocessed video tensor shape: {output.shape}")
|
||||
self.log_debug("Postprocessed video tensor shape: %s", output.shape)
|
||||
|
||||
audio = None
|
||||
audio_sample_rate = None
|
||||
if self.sound_tokenizer is not None and batch.audio_latents is not None:
|
||||
if server_args.vae_cpu_offload:
|
||||
self.sound_tokenizer.to(device)
|
||||
with torch.no_grad():
|
||||
decoded_audio = self.sound_tokenizer.decode(
|
||||
batch.audio_latents.to(device)
|
||||
)
|
||||
with self.use_declared_component(
|
||||
component_name="sound_tokenizer", module=self.sound_tokenizer
|
||||
) as sound_tokenizer:
|
||||
assert sound_tokenizer is not None
|
||||
with torch.no_grad():
|
||||
decoded_audio = sound_tokenizer.decode(
|
||||
batch.audio_latents.to(device)
|
||||
)
|
||||
audio = decoded_audio.float().cpu()
|
||||
audio_sample_rate = self.sound_tokenizer.sample_rate
|
||||
if server_args.vae_cpu_offload and not getattr(batch, "is_warmup", False):
|
||||
self.sound_tokenizer.to("cpu", non_blocking=True)
|
||||
self.log_info(
|
||||
f"Decoded audio tensor shape: {tuple(audio.shape)} @ {audio_sample_rate} Hz"
|
||||
audio_sample_rate = sound_tokenizer.sample_rate
|
||||
self.log_debug(
|
||||
"Decoded audio tensor shape: %s @ %s Hz",
|
||||
tuple(audio.shape),
|
||||
audio_sample_rate,
|
||||
)
|
||||
|
||||
return OutputBatch(
|
||||
@@ -1641,3 +1642,12 @@ class Cosmos3DecodingStage(PipelineStage):
|
||||
metrics=batch.metrics if hasattr(batch, "metrics") else None,
|
||||
**action_metadata,
|
||||
)
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
stage_name = self._component_stage_name(stage_name)
|
||||
uses = [ComponentUse(stage_name, "vae", keep_ready_after_warmup=True)]
|
||||
if self.sound_tokenizer is not None:
|
||||
uses.append(ComponentUse(stage_name, "sound_tokenizer"))
|
||||
return uses
|
||||
|
||||
+11
-1
@@ -7,6 +7,9 @@ import json
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
@@ -60,7 +63,7 @@ class PromptEnhancementStage(PipelineStage):
|
||||
else:
|
||||
batch.prompt = enhanced
|
||||
|
||||
logger.info("PE enhanced prompt: %s", batch.prompt)
|
||||
logger.debug("PE enhanced prompt: %s", batch.prompt)
|
||||
return batch
|
||||
|
||||
def _enhance_single_prompt(
|
||||
@@ -96,3 +99,10 @@ class PromptEnhancementStage(PipelineStage):
|
||||
)
|
||||
|
||||
return output["text"].strip()
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
if not isinstance(self.pe_model, torch.nn.Module):
|
||||
return []
|
||||
return [ComponentUse(self._component_stage_name(stage_name), "pe")]
|
||||
|
||||
+25
-3
@@ -235,6 +235,19 @@ class GlmImageAR(PipelineStage):
|
||||
self.processor = processor
|
||||
self.vision_language_encoder = vision_language_encoder
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
if not isinstance(self.vision_language_encoder, torch.nn.Module):
|
||||
return []
|
||||
return [
|
||||
ComponentUse(
|
||||
self._component_stage_name(stage_name),
|
||||
"vision_language_encoder",
|
||||
memory_intensive=True,
|
||||
)
|
||||
]
|
||||
|
||||
@property
|
||||
def parallelism_type(self) -> StageParallelismType:
|
||||
return StageParallelismType.MAIN_RANK_ONLY_AND_SEND_TO_OTHERS
|
||||
@@ -759,7 +772,7 @@ class GlmImageAR(PipelineStage):
|
||||
prior_token_id = torch.cat(prior_token_ids, dim=0)
|
||||
prior_token_id = prior_token_id.to(device=device)
|
||||
time_end = time.time()
|
||||
logger.info(f"generate_prior_tokens time: {time_end - time_start}")
|
||||
logger.debug("generate_prior_tokens time: %.3fs", time_end - time_start)
|
||||
|
||||
batch.prior_token_id = prior_token_id
|
||||
batch.prior_token_image_ids = prior_token_image_ids
|
||||
@@ -827,7 +840,9 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
stage_name = self._component_stage_name(stage_name)
|
||||
uses: list[ComponentUse] = []
|
||||
uses = [ComponentUse(stage_name, "text_encoder", memory_intensive=True)]
|
||||
if self.vae is not None:
|
||||
uses.append(ComponentUse(stage_name, "vae", phase="condition_image"))
|
||||
if self.transformer is not None:
|
||||
uses.append(
|
||||
ComponentUse(
|
||||
@@ -1120,6 +1135,9 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
|
||||
prior_token_id = _repeat_to_batch(prior_token_id, batch_size)
|
||||
|
||||
# 3. Encode input prompt
|
||||
self.begin_declared_component_use(
|
||||
component_name="text_encoder", module=self.text_encoder
|
||||
)
|
||||
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
|
||||
prompt,
|
||||
do_classifier_free_guidance,
|
||||
@@ -1175,6 +1193,11 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
|
||||
latents_mean = latents_mean.to(device=device, dtype=vae_dtype)
|
||||
latents_std = latents_std.to(device=device, dtype=vae_dtype)
|
||||
|
||||
self.begin_declared_component_use(
|
||||
component_name="vae",
|
||||
module=self.vae,
|
||||
phase="condition_image",
|
||||
)
|
||||
for condition_image, condition_image_prior_token_id in zip(
|
||||
ar_condition_images, prior_token_image_ids
|
||||
):
|
||||
@@ -1182,7 +1205,6 @@ class GlmImageBeforeDenoisingStage(PipelineStage):
|
||||
condition_image, self.vae, device=device
|
||||
)
|
||||
condition_image = _repeat_to_batch(condition_image, batch_size)
|
||||
|
||||
condition_latent = retrieve_latents(
|
||||
self.vae.encode(condition_image),
|
||||
generator=generator,
|
||||
|
||||
-10
@@ -511,16 +511,6 @@ class HeliosChunkedDenoisingStage(PipelineStage):
|
||||
is_amplify_first_chunk = pipeline_config.is_amplify_first_chunk
|
||||
gamma = pipeline_config.gamma
|
||||
|
||||
transformer_use = ComponentUse(
|
||||
self.__class__.__name__,
|
||||
"transformer",
|
||||
phase="transformer",
|
||||
preferred_ready_after_request=True,
|
||||
memory_intensive=True,
|
||||
)
|
||||
manager = self._component_residency_manager
|
||||
manager.begin_use(transformer_use, module=self.transformer)
|
||||
|
||||
# Get encoder outputs (prompt_embeds is a list of tensors, one per encoder)
|
||||
prompt_embeds = batch.prompt_embeds
|
||||
if isinstance(prompt_embeds, list):
|
||||
|
||||
+24
-3
@@ -17,6 +17,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
|
||||
Hunyuan3D2PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
|
||||
@@ -165,6 +168,15 @@ class Hunyuan3DShapeBeforeDenoisingStage(PipelineStage):
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
return latents * getattr(scheduler, "init_noise_sigma", 1.0)
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
return [
|
||||
ComponentUse(
|
||||
self._component_stage_name(stage_name), "hy3dshape_conditioner"
|
||||
)
|
||||
]
|
||||
|
||||
def _find_conditioner_dtype(self, items_fn_name: str) -> torch.dtype | None:
|
||||
items_fn = getattr(self.conditioner, items_fn_name, None)
|
||||
if not callable(items_fn):
|
||||
@@ -285,6 +297,11 @@ class Hunyuan3DShapeDenoisingStage(DenoisingStage):
|
||||
def __init__(self, transformer: Any, scheduler: Any, **kwargs) -> None:
|
||||
super().__init__(transformer=transformer, scheduler=scheduler, **kwargs)
|
||||
|
||||
def _component_name_for_stage_module(self, module, default_name: str) -> str:
|
||||
if module is self.transformer:
|
||||
return "hy3dshape_model"
|
||||
return super()._component_name_for_stage_module(module, default_name)
|
||||
|
||||
def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs):
|
||||
"""Prepare Hunyuan3D-specific variables for the base denoising loop."""
|
||||
assert self.transformer is not None
|
||||
@@ -450,6 +467,11 @@ class Hunyuan3DShapeExportStage(PipelineStage):
|
||||
self.vae = vae
|
||||
self.config = config
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
return [ComponentUse(self._component_stage_name(stage_name), "hy3dshape_vae")]
|
||||
|
||||
@property
|
||||
def role_affinity(self):
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
@@ -468,15 +490,14 @@ class Hunyuan3DShapeExportStage(PipelineStage):
|
||||
]()
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
f"Could not load SurfaceExtractors for mc_algo={self.config.shape_mc_algo}"
|
||||
"Could not load SurfaceExtractors for mc_algo=%s",
|
||||
self.config.shape_mc_algo,
|
||||
)
|
||||
|
||||
latents = batch.latents
|
||||
|
||||
if self.config.shape_output_type != "latent":
|
||||
latents = 1.0 / self.vae.scale_factor * latents
|
||||
latents = self.vae(latents)
|
||||
|
||||
outputs = self.vae.latents2mesh(
|
||||
latents,
|
||||
bounds=self.config.shape_box_v,
|
||||
|
||||
+1
-1
@@ -268,7 +268,7 @@ class LongCatPromptRewriteStage(PipelineStage):
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=False,
|
||||
)
|
||||
logger.info("Rewritten prompts: %s", rewritten)
|
||||
logger.debug("Rewritten prompts: %s", rewritten)
|
||||
return rewritten
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
+9
-3
@@ -1,6 +1,9 @@
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
@@ -16,6 +19,11 @@ class LTX2TextConnectorStage(PipelineStage):
|
||||
super().__init__()
|
||||
self.connectors = connectors
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
return [ComponentUse(self._component_stage_name(stage_name), "connectors")]
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
|
||||
# Input: batch.prompt_embeds (from Gemma, [B, S, D])
|
||||
# Output: batch.prompt_embeds (Video Context), batch.audio_prompt_embeds (Audio Context)
|
||||
@@ -58,8 +66,7 @@ class LTX2TextConnectorStage(PipelineStage):
|
||||
"and attention mask when classifier-free guidance is enabled."
|
||||
)
|
||||
|
||||
# Official LTX-2.3 processes positive and negative prompts through
|
||||
# the connector independently; batching shifts output numerics.
|
||||
# Official LTX-2.3 processes positive and negative prompts separately.
|
||||
dtype = prompt_embeds.dtype
|
||||
pos_additive_mask = (prompt_attention_mask.to(torch.int64) - 1).to(
|
||||
dtype
|
||||
@@ -83,7 +90,6 @@ class LTX2TextConnectorStage(PipelineStage):
|
||||
batch.negative_audio_prompt_embeds = [neg_audio_embeds]
|
||||
batch.negative_attention_mask = neg_mask
|
||||
else:
|
||||
# Prepare additive mask for connectors (as per diffusers implementation)
|
||||
dtype = prompt_embeds.dtype
|
||||
additive_attention_mask = (prompt_attention_mask.to(torch.int64) - 1).to(
|
||||
dtype
|
||||
|
||||
+7
-13
@@ -75,13 +75,7 @@ class LTX2UpsampleStage(PipelineStage):
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
stage_name = self._component_stage_name(stage_name)
|
||||
uses = [
|
||||
ComponentUse(stage_name, "spatial_upsampler"),
|
||||
ComponentUse(stage_name, "vae"),
|
||||
]
|
||||
if self.audio_vae is not None:
|
||||
uses.append(ComponentUse(stage_name, "audio_vae"))
|
||||
return uses
|
||||
return [ComponentUse(stage_name, "spatial_upsampler")]
|
||||
|
||||
def _upsample_video_latents(
|
||||
self, latents: torch.Tensor, server_args: ServerArgs, device: torch.device
|
||||
@@ -93,12 +87,12 @@ class LTX2UpsampleStage(PipelineStage):
|
||||
device=device, dtype=latents.dtype
|
||||
)
|
||||
latents = latents * vae_std + vae_mean
|
||||
self.spatial_upsampler = self.spatial_upsampler.to(
|
||||
device=device, dtype=latents.dtype
|
||||
)
|
||||
latents = self.spatial_upsampler(latents)
|
||||
# Keep the small spatial upsampler resident after warmup; moving it
|
||||
# every request dominates the measured two-stage upsample latency.
|
||||
with self.use_declared_component(
|
||||
component_name="spatial_upsampler",
|
||||
module=self.spatial_upsampler,
|
||||
target_dtype=latents.dtype,
|
||||
) as spatial_upsampler:
|
||||
latents = spatial_upsampler(latents)
|
||||
latents = (latents - vae_mean) / vae_std
|
||||
return latents
|
||||
|
||||
|
||||
+1
-1
@@ -16,7 +16,7 @@ from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
|
||||
CacheDitConfig,
|
||||
disable_cache_on_transformer,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
is_fsdp_managed_module,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
|
||||
-1
@@ -270,7 +270,6 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
|
||||
raise ValueError(
|
||||
"MiniMaxH3TextEncodingStage direct encode requires a tokenizer component"
|
||||
)
|
||||
self._manage_text_encoder_use(0)
|
||||
with set_forward_context(current_timestep=0, attn_metadata=None):
|
||||
if plan.task == "ref2va":
|
||||
embeddings = self._encode_ref2va(batch, plan, encode_ids)
|
||||
|
||||
@@ -11,10 +11,14 @@ closer to the actual denoising step.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
|
||||
TextEncodingStage,
|
||||
@@ -105,6 +109,15 @@ class RealtimeTextState(BaseRealtimeState):
|
||||
class RealtimeTextEncodingStage(TextEncodingStage):
|
||||
"""Cache text encoder outputs across realtime chunks by prompt identity."""
|
||||
|
||||
def component_uses(
|
||||
self, server_args: ServerArgs, stage_name: str | None = None
|
||||
) -> list[ComponentUse]:
|
||||
# Cache hits return before encode_text reaches the declared use site.
|
||||
return [
|
||||
replace(use, start_at_stage_entry=False)
|
||||
for use in super().component_uses(server_args, stage_name)
|
||||
]
|
||||
|
||||
def _make_cache_key(self, batch: Req) -> tuple[Any, ...]:
|
||||
return (
|
||||
_normalize_prompt_value(batch.prompt),
|
||||
|
||||
@@ -492,20 +492,16 @@ class TextEncodingStage(ConditionEncodingStage):
|
||||
|
||||
return tok_kwargs
|
||||
|
||||
def _manage_text_encoder_use(self, encoder_index: int) -> None:
|
||||
manager = self._component_residency_manager
|
||||
if manager is None:
|
||||
return
|
||||
def _begin_text_encoder_use(self, encoder_index: int) -> None:
|
||||
component_name = (
|
||||
"text_encoder"
|
||||
if encoder_index == 0
|
||||
else f"text_encoder_{encoder_index + 1}"
|
||||
)
|
||||
use = self._declared_component_use(component_name=component_name)
|
||||
# TODO: Keep this begin-only interval until manager supports explicit
|
||||
# declared-use interval grouping. Wrapping each encoder call separately
|
||||
# can offload between positive and negative prompt encoding.
|
||||
manager.begin_use(use, module=self.text_encoders[encoder_index])
|
||||
self.begin_declared_component_use(
|
||||
component_name=component_name,
|
||||
module=self.text_encoders[encoder_index],
|
||||
)
|
||||
|
||||
def _forward_text_encoder(self, text_encoder, encoder_forward_kwargs):
|
||||
if not getattr(text_encoder, "uses_sglang_forward_context", True):
|
||||
@@ -703,7 +699,7 @@ class TextEncodingStage(ConditionEncodingStage):
|
||||
encoder_forward_kwargs["attention_mask"] = attention_mask
|
||||
if "use_cache" in inspect.signature(text_encoder.forward).parameters:
|
||||
encoder_forward_kwargs["use_cache"] = False
|
||||
self._manage_text_encoder_use(i)
|
||||
self._begin_text_encoder_use(i)
|
||||
dp_group = self._text_encode_dp_group(
|
||||
server_args, encoder_config, input_ids.shape[0], text_encoder
|
||||
)
|
||||
|
||||
@@ -46,10 +46,7 @@ class ServerArgsAutoTuner:
|
||||
|
||||
def __init__(self, server_args: ServerArgs):
|
||||
self.server_args = server_args
|
||||
self._explicit_memory_policy = self._has_explicit_memory_policy()
|
||||
self._explicit_layerwise_replacement_policy = (
|
||||
self._has_explicit_layerwise_replacement_policy()
|
||||
)
|
||||
self._explicit_dit_residency = self._has_explicit_dit_residency()
|
||||
|
||||
def _deployment_config(self) -> ModelDeploymentConfig:
|
||||
return self.server_args.pipeline_config.get_model_deployment_config()
|
||||
@@ -127,12 +124,9 @@ class ServerArgsAutoTuner:
|
||||
if args.performance_mode != "auto" or current_platform.is_cpu():
|
||||
return
|
||||
|
||||
# Explicitness is component-scoped below. For example, explicitly
|
||||
# disabling DiT layerwise offload must not freeze an unrelated,
|
||||
# implicit ``dit_cpu_offload=True`` default on a high-memory GPU.
|
||||
# Each mutation below already preserves its own explicit CLI flag.
|
||||
# Explicit placement is component-scoped; unmatched components still
|
||||
# receive automatic defaults.
|
||||
|
||||
explicit_cpu_components = args.is_arg_explicitly_set("cpu_offload_components")
|
||||
explicit_layerwise_components = (
|
||||
normalize_layerwise_offload_components(args.layerwise_offload_components)
|
||||
if args.is_arg_explicitly_set("layerwise_offload_components")
|
||||
@@ -180,8 +174,7 @@ class ServerArgsAutoTuner:
|
||||
if (
|
||||
args.dit_cpu_offload
|
||||
and "dit" in components
|
||||
and not args.is_arg_explicitly_set("dit_cpu_offload")
|
||||
and not explicit_cpu_components
|
||||
and args.explicit_residency_mode("transformer") is None
|
||||
and not explicit_dit_layerwise
|
||||
):
|
||||
args.dit_cpu_offload = False
|
||||
@@ -189,24 +182,21 @@ class ServerArgsAutoTuner:
|
||||
if (
|
||||
args.text_encoder_cpu_offload
|
||||
and LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP in components
|
||||
and not args.is_arg_explicitly_set("text_encoder_cpu_offload")
|
||||
and not explicit_cpu_components
|
||||
and args.explicit_residency_mode("text_encoder") is None
|
||||
):
|
||||
args.text_encoder_cpu_offload = False
|
||||
changed.append("text_encoder_cpu_offload=False")
|
||||
if (
|
||||
args.image_encoder_cpu_offload
|
||||
and LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP in components
|
||||
and not args.is_arg_explicitly_set("image_encoder_cpu_offload")
|
||||
and not explicit_cpu_components
|
||||
and args.explicit_residency_mode("image_encoder") is None
|
||||
):
|
||||
args.image_encoder_cpu_offload = False
|
||||
changed.append("image_encoder_cpu_offload=False")
|
||||
if (
|
||||
args.vae_cpu_offload
|
||||
and LAYERWISE_OFFLOAD_VAE_GROUP in components
|
||||
and not args.is_arg_explicitly_set("vae_cpu_offload")
|
||||
and not explicit_cpu_components
|
||||
and args.explicit_residency_mode("vae") is None
|
||||
):
|
||||
args.vae_cpu_offload = False
|
||||
changed.append("vae_cpu_offload=False")
|
||||
@@ -235,14 +225,21 @@ class ServerArgsAutoTuner:
|
||||
|
||||
# high-memory resident mode keeps both DiTs on GPU; unset auxiliary
|
||||
# placement should stay resident instead of using default layerwise
|
||||
for arg_name in (
|
||||
"text_encoder_cpu_offload",
|
||||
"image_encoder_cpu_offload",
|
||||
"vae_cpu_offload",
|
||||
if (
|
||||
args.text_encoder_cpu_offload
|
||||
and args.explicit_residency_mode("text_encoder") is None
|
||||
):
|
||||
if getattr(args, arg_name) and not args.is_arg_explicitly_set(arg_name):
|
||||
setattr(args, arg_name, False)
|
||||
changed.append(f"{arg_name}=False")
|
||||
args.text_encoder_cpu_offload = False
|
||||
changed.append("text_encoder_cpu_offload=False")
|
||||
if (
|
||||
args.image_encoder_cpu_offload
|
||||
and args.explicit_residency_mode("image_encoder") is None
|
||||
):
|
||||
args.image_encoder_cpu_offload = False
|
||||
changed.append("image_encoder_cpu_offload=False")
|
||||
if args.vae_cpu_offload and args.explicit_residency_mode("vae") is None:
|
||||
args.vae_cpu_offload = False
|
||||
changed.append("vae_cpu_offload=False")
|
||||
|
||||
if changed:
|
||||
logger.info(
|
||||
@@ -255,7 +252,7 @@ class ServerArgsAutoTuner:
|
||||
if (
|
||||
args.performance_mode == "auto"
|
||||
and args.num_gpus >= 2
|
||||
and not self._explicit_memory_policy
|
||||
and not self._explicit_dit_residency
|
||||
and self._auto_uses_dit_offload()
|
||||
and self._can_apply_fsdp_policy(require_memory_headroom=True)
|
||||
):
|
||||
@@ -280,7 +277,6 @@ class ServerArgsAutoTuner:
|
||||
if (
|
||||
args.layerwise_offload_components is not None
|
||||
or args.dit_layerwise_offload is True
|
||||
or args.is_arg_explicitly_set("cpu_offload_components")
|
||||
):
|
||||
return
|
||||
if not current_platform.is_cuda():
|
||||
@@ -301,12 +297,12 @@ class ServerArgsAutoTuner:
|
||||
args = self.server_args
|
||||
if (
|
||||
not self.could_override_server_args()
|
||||
or self._explicit_layerwise_replacement_policy
|
||||
or current_platform.is_cpu()
|
||||
or not current_platform.is_cuda()
|
||||
or envs.SGLANG_CACHE_DIT_ENABLED
|
||||
or args.use_fsdp_inference
|
||||
or args.layerwise_offload_components is not None
|
||||
or args.dit_layerwise_offload is True
|
||||
):
|
||||
return
|
||||
|
||||
@@ -315,17 +311,19 @@ class ServerArgsAutoTuner:
|
||||
layerwise_components.append(LAYERWISE_OFFLOAD_DIT_GROUP)
|
||||
|
||||
changed: list[str] = []
|
||||
if args.text_encoder_cpu_offload and not args.is_arg_explicitly_set(
|
||||
"text_encoder_cpu_offload"
|
||||
if (
|
||||
args.text_encoder_cpu_offload
|
||||
and args.explicit_residency_mode("text_encoder") is None
|
||||
):
|
||||
layerwise_components.append(LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP)
|
||||
changed.append(LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP)
|
||||
if args.image_encoder_cpu_offload and not args.is_arg_explicitly_set(
|
||||
"image_encoder_cpu_offload"
|
||||
if (
|
||||
args.image_encoder_cpu_offload
|
||||
and args.explicit_residency_mode("image_encoder") is None
|
||||
):
|
||||
layerwise_components.append(LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP)
|
||||
changed.append(LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP)
|
||||
if args.vae_cpu_offload and not args.is_arg_explicitly_set("vae_cpu_offload"):
|
||||
if args.vae_cpu_offload and args.explicit_residency_mode("vae") is None:
|
||||
layerwise_components.append(LAYERWISE_OFFLOAD_VAE_GROUP)
|
||||
changed.append(LAYERWISE_OFFLOAD_VAE_GROUP)
|
||||
|
||||
@@ -438,7 +436,6 @@ class ServerArgsAutoTuner:
|
||||
if (
|
||||
args.is_arg_explicitly_set("layerwise_offload_components")
|
||||
or args.dit_layerwise_offload is True
|
||||
or args.is_arg_explicitly_set("cpu_offload_components")
|
||||
):
|
||||
# The legacy --dit-layerwise-offload flag is a DiT-only selector.
|
||||
# Do not merge implicit defaults into that explicit mode.
|
||||
@@ -450,7 +447,7 @@ class ServerArgsAutoTuner:
|
||||
components = [
|
||||
component_name
|
||||
for component_name, arg_name in DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES
|
||||
if not args.is_arg_explicitly_set(arg_name)
|
||||
if args.explicit_residency_mode(component_name) is None
|
||||
]
|
||||
components = self._filter_high_memory_resident_components(components)
|
||||
if self._should_auto_enable_dit_layerwise_offload():
|
||||
@@ -503,8 +500,7 @@ class ServerArgsAutoTuner:
|
||||
or not current_platform.enable_dit_layerwise_offload_by_default()
|
||||
or envs.SGLANG_CACHE_DIT_ENABLED
|
||||
or args.use_fsdp_inference
|
||||
or args.is_arg_explicitly_set("dit_cpu_offload")
|
||||
or args.is_arg_explicitly_set("cpu_offload_components")
|
||||
or args.explicit_residency_mode("transformer") is not None
|
||||
):
|
||||
return False
|
||||
|
||||
@@ -544,28 +540,11 @@ class ServerArgsAutoTuner:
|
||||
)
|
||||
)
|
||||
|
||||
def _has_explicit_memory_policy(self) -> bool:
|
||||
def _has_explicit_dit_residency(self) -> bool:
|
||||
args = self.server_args
|
||||
return any(
|
||||
args.is_arg_explicitly_set(arg_name)
|
||||
for arg_name in (
|
||||
"use_fsdp_inference",
|
||||
"dit_cpu_offload",
|
||||
"dit_layerwise_offload",
|
||||
"layerwise_offload_components",
|
||||
"cpu_offload_components",
|
||||
)
|
||||
)
|
||||
|
||||
def _has_explicit_layerwise_replacement_policy(self) -> bool:
|
||||
args = self.server_args
|
||||
return any(
|
||||
args.is_arg_explicitly_set(arg_name)
|
||||
for arg_name in (
|
||||
"dit_layerwise_offload",
|
||||
"layerwise_offload_components",
|
||||
"cpu_offload_components",
|
||||
)
|
||||
return bool(
|
||||
args.is_arg_explicitly_set("use_fsdp_inference")
|
||||
or args.explicit_residency_mode("transformer") is not None
|
||||
)
|
||||
|
||||
def _has_explicit_parallel_policy(self) -> bool:
|
||||
|
||||
@@ -32,15 +32,26 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
|
||||
NunchakuConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
RESIDENT,
|
||||
normalize_component_residency,
|
||||
resolve_component_residency_mode,
|
||||
resolve_diffusers_pipeline_offload,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS,
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP,
|
||||
LAYERWISE_OFFLOAD_VAE_GROUP,
|
||||
RESIDENCY_POLICIES,
|
||||
RESIDENCY_POLICY_LEADING,
|
||||
cpu_offload_component_matches,
|
||||
cpu_offload_flags_for_layerwise_components,
|
||||
is_dit_component_name,
|
||||
is_image_encoder_component_name,
|
||||
is_legacy_dit_offload_component_name,
|
||||
is_text_encoder_component_name,
|
||||
is_vae_component_name,
|
||||
layerwise_component_matches_any_selection,
|
||||
@@ -319,6 +330,8 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
lora_target_modules: list[str] | None = None
|
||||
|
||||
# CPU offload parameters
|
||||
# Exact component keys or component groups mapped to a residency mode.
|
||||
component_residency: dict[str, str] | list[str] | str | None = None
|
||||
# Exact component keys from model_index.json, or a legacy component group.
|
||||
cpu_offload_components: list[str] | None = None
|
||||
dit_cpu_offload: bool | None = None
|
||||
@@ -340,6 +353,15 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
pin_cpu_memory: bool = True
|
||||
ltx2_two_stage_device_mode: str | None = None
|
||||
_explicit_arg_names: set[str] = field(default_factory=set, repr=False)
|
||||
_required_resident_components: set[str] = field(
|
||||
default_factory=set, init=False, repr=False
|
||||
)
|
||||
_fsdp_disabled_components: set[str] = field(
|
||||
default_factory=set, init=False, repr=False
|
||||
)
|
||||
_component_layerwise_capabilities: dict[str, bool] = field(
|
||||
default_factory=dict, init=False, repr=False
|
||||
)
|
||||
|
||||
# ComfyUI integration
|
||||
comfyui_mode: bool = False
|
||||
@@ -506,9 +528,10 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
|
||||
def _adjust_parameters(self):
|
||||
"""set defaults and normalize values."""
|
||||
self._normalize_component_residency()
|
||||
self._adjust_cpu_offload_components()
|
||||
auto_tuner = ServerArgsAutoTuner(self)
|
||||
auto_tuner.adjust_based_on_performance_mode()
|
||||
self._adjust_cpu_offload_components()
|
||||
if auto_tuner.could_override_server_args():
|
||||
self._adjust_offload()
|
||||
auto_tuner.maybe_adjust_auto_default_layerwise_offload()
|
||||
@@ -765,48 +788,23 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
self.image_encoder_cpu_offload = True
|
||||
|
||||
def _adjust_cpu_offload_components(self) -> None:
|
||||
"""Apply the unified CPU offload component selector, when provided."""
|
||||
"""Normalize the legacy component offload selector, when provided."""
|
||||
if self.cpu_offload_components is None:
|
||||
return
|
||||
normalized = normalize_cpu_offload_components(self.cpu_offload_components)
|
||||
self.cpu_offload_components = normalized if normalized is not None else []
|
||||
|
||||
legacy_flags = (
|
||||
"dit_cpu_offload",
|
||||
"text_encoder_cpu_offload",
|
||||
"image_encoder_cpu_offload",
|
||||
"vae_cpu_offload",
|
||||
def _normalize_component_residency(self) -> None:
|
||||
self.component_residency = normalize_component_residency(
|
||||
self.component_residency
|
||||
)
|
||||
conflicting_flags = [
|
||||
flag_name
|
||||
for flag_name in legacy_flags
|
||||
if self.is_arg_explicitly_set(flag_name)
|
||||
]
|
||||
if conflicting_flags:
|
||||
formatted_flags = ", ".join(
|
||||
"--" + flag_name.replace("_", "-") for flag_name in conflicting_flags
|
||||
)
|
||||
raise ValueError(
|
||||
"--cpu-offload-components cannot be combined with the legacy "
|
||||
f"CPU offload flags: {formatted_flags}"
|
||||
)
|
||||
|
||||
selected_components = (
|
||||
normalize_cpu_offload_components(self.cpu_offload_components) or []
|
||||
)
|
||||
self.cpu_offload_components = selected_components
|
||||
self.dit_cpu_offload = self.should_cpu_offload_component("transformer")
|
||||
self.text_encoder_cpu_offload = self.should_cpu_offload_component(
|
||||
"text_encoder"
|
||||
)
|
||||
self.image_encoder_cpu_offload = self.should_cpu_offload_component(
|
||||
"image_encoder"
|
||||
)
|
||||
self.vae_cpu_offload = self.should_cpu_offload_component("vae")
|
||||
|
||||
def _adjust_ltx2_two_stage_device_mode(self):
|
||||
if not self._is_ltx23_two_stage_pipeline():
|
||||
return
|
||||
|
||||
mode = self.ltx2_two_stage_device_mode
|
||||
env_mode = None
|
||||
if mode is None:
|
||||
env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE")
|
||||
mode = (
|
||||
@@ -823,6 +821,29 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
f"Expected one of {LTX2_TWO_STAGE_DEVICE_MODE_CHOICES}."
|
||||
)
|
||||
|
||||
explicit_nonresident_dits = {
|
||||
component_name: residency_mode
|
||||
for component_name in ("transformer", "transformer_2")
|
||||
if (residency_mode := self.explicit_residency_mode(component_name))
|
||||
in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD)
|
||||
}
|
||||
if mode == "resident" and explicit_nonresident_dits:
|
||||
configured = ", ".join(
|
||||
f"{name}={residency_mode}"
|
||||
for name, residency_mode in explicit_nonresident_dits.items()
|
||||
)
|
||||
if self.is_arg_explicitly_set("ltx2_two_stage_device_mode") or env_mode:
|
||||
raise ValueError(
|
||||
"ltx2_two_stage_device_mode=resident conflicts with explicit "
|
||||
f"component residency: {configured}"
|
||||
)
|
||||
mode = "original"
|
||||
logger.info(
|
||||
"Using ltx2_two_stage_device_mode=original because DiT offload "
|
||||
"was explicitly configured: %s",
|
||||
configured,
|
||||
)
|
||||
|
||||
self.ltx2_two_stage_device_mode = mode
|
||||
|
||||
def _resolve_default_ltx2_two_stage_device_mode(self) -> str:
|
||||
@@ -1282,6 +1303,13 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
|
||||
def _adjust_platform_specific(self):
|
||||
if current_platform.is_mps():
|
||||
if self.component_residency is not None and any(
|
||||
mode != RESIDENT for mode in self.component_residency.values()
|
||||
):
|
||||
raise ValueError(
|
||||
"--component-residency offload modes require CUDA; "
|
||||
"MPS supports only resident components"
|
||||
)
|
||||
self.use_fsdp_inference = False
|
||||
self.dit_layerwise_offload = False
|
||||
self.layerwise_offload_components = None
|
||||
@@ -1299,114 +1327,256 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
self.text_encoder_cpu_offload = False
|
||||
self.image_encoder_cpu_offload = False
|
||||
self.vae_cpu_offload = False
|
||||
self.cpu_offload_components = None
|
||||
|
||||
def is_arg_explicitly_set(self, arg_name: str) -> bool:
|
||||
return arg_name in self._explicit_arg_names
|
||||
|
||||
def should_cpu_offload_component(self, component_name: str) -> bool:
|
||||
if self.cpu_offload_components is not None:
|
||||
return cpu_offload_component_matches(
|
||||
def canonical_residency_mode(self, component_name: str) -> str | None:
|
||||
"""Resolve the canonical selector for one component, if present."""
|
||||
return resolve_component_residency_mode(
|
||||
component_name, self.component_residency
|
||||
)
|
||||
|
||||
def explicit_residency_mode(self, component_name: str) -> str | None:
|
||||
"""Resolve explicit controls in canonical-to-compatibility priority."""
|
||||
mode = self.canonical_residency_mode(component_name)
|
||||
if mode is not None:
|
||||
return mode
|
||||
|
||||
if self.is_explicit_layerwise_offload_component(component_name):
|
||||
return LAYERWISE_OFFLOAD
|
||||
|
||||
if self.is_arg_explicitly_set("cpu_offload_components"):
|
||||
if not self.cpu_offload_components:
|
||||
return RESIDENT
|
||||
if cpu_offload_component_matches(
|
||||
component_name, self.cpu_offload_components
|
||||
)
|
||||
if is_dit_component_name(component_name) or component_name in (
|
||||
"connectors",
|
||||
"unconditional_transformer",
|
||||
"vision_language_encoder",
|
||||
):
|
||||
return COMPONENT_OFFLOAD
|
||||
|
||||
legacy_flag = self._legacy_component_offload_flag(component_name)
|
||||
if legacy_flag is not None and self.is_arg_explicitly_set(legacy_flag):
|
||||
legacy_values = {
|
||||
"dit_cpu_offload": self.dit_cpu_offload,
|
||||
"text_encoder_cpu_offload": self.text_encoder_cpu_offload,
|
||||
"image_encoder_cpu_offload": self.image_encoder_cpu_offload,
|
||||
"vae_cpu_offload": self.vae_cpu_offload,
|
||||
}
|
||||
return COMPONENT_OFFLOAD if legacy_values[legacy_flag] else RESIDENT
|
||||
|
||||
# ``--dit-layerwise-offload false`` historically has no matching CPU
|
||||
# flag, but explicitly requests that the DiT not be layerwise-offloaded.
|
||||
if is_legacy_dit_offload_component_name(component_name) and (
|
||||
self.is_arg_explicitly_set("dit_layerwise_offload")
|
||||
):
|
||||
return bool(self.dit_cpu_offload)
|
||||
return RESIDENT
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _legacy_component_offload_flag(component_name: str) -> str | None:
|
||||
if is_legacy_dit_offload_component_name(component_name):
|
||||
return "dit_cpu_offload"
|
||||
if is_text_encoder_component_name(component_name):
|
||||
return bool(self.text_encoder_cpu_offload)
|
||||
return "text_encoder_cpu_offload"
|
||||
if is_image_encoder_component_name(component_name):
|
||||
return bool(self.image_encoder_cpu_offload)
|
||||
if is_vae_component_name(component_name) or component_name == "sound_tokenizer":
|
||||
return bool(self.vae_cpu_offload)
|
||||
return False
|
||||
return "image_encoder_cpu_offload"
|
||||
if is_vae_component_name(component_name):
|
||||
return "vae_cpu_offload"
|
||||
return None
|
||||
|
||||
def residency_mode(self, component_name: str) -> str:
|
||||
"""Return the effective residency mode for a loaded component."""
|
||||
if current_platform.is_cpu():
|
||||
return RESIDENT
|
||||
if component_name in self._required_resident_components:
|
||||
return RESIDENT
|
||||
|
||||
explicit_mode = self.explicit_residency_mode(component_name)
|
||||
if explicit_mode is not None:
|
||||
return explicit_mode
|
||||
|
||||
component_names = normalize_layerwise_offload_components(
|
||||
self.layerwise_offload_components
|
||||
)
|
||||
if self._component_layerwise_capabilities.get(component_name, True):
|
||||
if component_names and (
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS in component_names
|
||||
or layerwise_component_matches_any_selection(
|
||||
component_name, component_names
|
||||
)
|
||||
or (
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP in component_names
|
||||
and is_dit_component_name(component_name)
|
||||
)
|
||||
):
|
||||
return LAYERWISE_OFFLOAD
|
||||
|
||||
if self.cpu_offload_components is not None:
|
||||
if cpu_offload_component_matches(
|
||||
component_name, self.cpu_offload_components
|
||||
):
|
||||
return COMPONENT_OFFLOAD
|
||||
if is_legacy_dit_offload_component_name(component_name):
|
||||
return COMPONENT_OFFLOAD if self.dit_cpu_offload else RESIDENT
|
||||
if is_text_encoder_component_name(component_name):
|
||||
return COMPONENT_OFFLOAD if self.text_encoder_cpu_offload else RESIDENT
|
||||
if is_image_encoder_component_name(component_name):
|
||||
return COMPONENT_OFFLOAD if self.image_encoder_cpu_offload else RESIDENT
|
||||
if is_vae_component_name(component_name):
|
||||
return COMPONENT_OFFLOAD if self.vae_cpu_offload else RESIDENT
|
||||
return RESIDENT
|
||||
|
||||
def should_cpu_offload_component(self, component_name: str) -> bool:
|
||||
return self.residency_mode(component_name) == COMPONENT_OFFLOAD
|
||||
|
||||
def should_start_component_on_cpu(self, component_name: str) -> bool:
|
||||
return self.residency_mode(component_name) in (
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
)
|
||||
|
||||
def require_component_resident(
|
||||
self, component_name: str, *, feature_name: str
|
||||
) -> None:
|
||||
configured_mode = self.canonical_residency_mode(component_name)
|
||||
if configured_mode is not None and configured_mode != RESIDENT:
|
||||
raise ValueError(
|
||||
f"{feature_name} requires {component_name!r} to be resident; "
|
||||
f"got {configured_mode!r} from --component-residency"
|
||||
)
|
||||
self._required_resident_components.add(component_name)
|
||||
|
||||
def should_use_fsdp_for_component(self, component_name: str) -> bool:
|
||||
return bool(
|
||||
self.use_fsdp_inference
|
||||
and component_name not in self._fsdp_disabled_components
|
||||
and self.residency_mode(component_name) == RESIDENT
|
||||
)
|
||||
|
||||
def disable_fsdp_for_component(self, component_name: str) -> None:
|
||||
self._fsdp_disabled_components.add(component_name)
|
||||
|
||||
def record_component_layerwise_capability(
|
||||
self, component_name: str, *, supported: bool
|
||||
) -> None:
|
||||
self._component_layerwise_capabilities[component_name] = supported
|
||||
|
||||
def has_layerwise_offload_components(self) -> bool:
|
||||
return bool(
|
||||
self.dit_layerwise_offload
|
||||
or self.layerwise_offload_components
|
||||
or (
|
||||
self.component_residency
|
||||
and LAYERWISE_OFFLOAD in self.component_residency.values()
|
||||
)
|
||||
)
|
||||
|
||||
def should_configure_layerwise_offload_for_lazy_component(
|
||||
self, component_name: str
|
||||
) -> bool:
|
||||
"""Return whether a lazy-loaded component should try layerwise offload.
|
||||
"""Return whether a lazy-loaded component needs layerwise setup."""
|
||||
return self.residency_mode(component_name) == LAYERWISE_OFFLOAD
|
||||
|
||||
Lazy components are loaded after the normal pipeline-wide configuration
|
||||
pass, so they should only attempt layerwise configuration when their
|
||||
component name is covered by the selected layerwise scope.
|
||||
"""
|
||||
component_names = normalize_layerwise_offload_components(
|
||||
self.layerwise_offload_components
|
||||
)
|
||||
if not component_names:
|
||||
return False
|
||||
if LAYERWISE_OFFLOAD_ALL_COMPONENTS in component_names:
|
||||
def is_explicit_layerwise_offload_component(self, component_name: str) -> bool:
|
||||
if self.canonical_residency_mode(component_name) == LAYERWISE_OFFLOAD:
|
||||
return True
|
||||
return layerwise_component_matches_any_selection(
|
||||
component_name, component_names
|
||||
|
||||
if self.is_arg_explicitly_set("layerwise_offload_components"):
|
||||
selected_components = normalize_layerwise_offload_components(
|
||||
self.layerwise_offload_components
|
||||
)
|
||||
if selected_components and (
|
||||
LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected_components
|
||||
or layerwise_component_matches_any_selection(
|
||||
component_name, selected_components
|
||||
)
|
||||
or (
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP in selected_components
|
||||
and is_dit_component_name(component_name)
|
||||
)
|
||||
):
|
||||
return True
|
||||
|
||||
return bool(
|
||||
self.is_arg_explicitly_set("dit_layerwise_offload")
|
||||
and self.dit_layerwise_offload
|
||||
and is_dit_component_name(component_name)
|
||||
)
|
||||
|
||||
@property
|
||||
def is_dit_layerwise_offload_selected(self) -> bool:
|
||||
"""returns if dit is selected to be layerwise-offload"""
|
||||
component_names = self.layerwise_offload_components
|
||||
return bool(
|
||||
component_names
|
||||
and "dit_cpu_offload"
|
||||
in cpu_offload_flags_for_layerwise_components(component_names)
|
||||
)
|
||||
"""Return whether the primary DiT resolves to layerwise offload."""
|
||||
return self.residency_mode("transformer") == LAYERWISE_OFFLOAD
|
||||
|
||||
def _adjust_layerwise_offload_components(self):
|
||||
explicitly_set_component_names = normalize_layerwise_offload_components(
|
||||
selected_component_names = normalize_layerwise_offload_components(
|
||||
self.layerwise_offload_components
|
||||
)
|
||||
if self.dit_layerwise_offload:
|
||||
if explicitly_set_component_names is None:
|
||||
explicitly_set_component_names = [LAYERWISE_OFFLOAD_DIT_GROUP]
|
||||
elif LAYERWISE_OFFLOAD_DIT_GROUP not in explicitly_set_component_names:
|
||||
explicitly_set_component_names = [
|
||||
if selected_component_names is None:
|
||||
selected_component_names = [LAYERWISE_OFFLOAD_DIT_GROUP]
|
||||
elif LAYERWISE_OFFLOAD_DIT_GROUP not in selected_component_names:
|
||||
selected_component_names = [
|
||||
LAYERWISE_OFFLOAD_DIT_GROUP,
|
||||
*explicitly_set_component_names,
|
||||
*selected_component_names,
|
||||
]
|
||||
|
||||
if explicitly_set_component_names is not None:
|
||||
self.layerwise_offload_components = explicitly_set_component_names
|
||||
self._disable_non_dit_cpu_offload_for_layerwise_components(
|
||||
explicitly_set_component_names
|
||||
self.layerwise_offload_components = selected_component_names
|
||||
self._clear_non_dit_component_offload_for_layerwise_groups(
|
||||
selected_component_names or ()
|
||||
)
|
||||
|
||||
has_explicit_dit_offload = bool(
|
||||
self.canonical_residency_mode("transformer")
|
||||
in (COMPONENT_OFFLOAD, LAYERWISE_OFFLOAD)
|
||||
or self.is_explicit_layerwise_offload_component("transformer")
|
||||
or (
|
||||
self.is_arg_explicitly_set("cpu_offload_components")
|
||||
and cpu_offload_component_matches(
|
||||
"transformer", self.cpu_offload_components
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
def _disable_non_dit_cpu_offload_for_layerwise_components(
|
||||
self, component_names: list[str]
|
||||
) -> None:
|
||||
# non-DiT layerwise offload replaces the corresponding component-level CPU offload
|
||||
flag_names = cpu_offload_flags_for_layerwise_components(component_names)
|
||||
disabled_flag_names: list[str] = []
|
||||
|
||||
or (self.is_arg_explicitly_set("dit_cpu_offload") and self.dit_cpu_offload)
|
||||
)
|
||||
if (
|
||||
"text_encoder_cpu_offload" in flag_names
|
||||
and self.text_encoder_cpu_offload is not False
|
||||
self.is_arg_explicitly_set("dit_layerwise_offload")
|
||||
and not self.dit_layerwise_offload
|
||||
and not has_explicit_dit_offload
|
||||
):
|
||||
self.dit_cpu_offload = False
|
||||
|
||||
def _clear_non_dit_component_offload_for_layerwise_groups(
|
||||
self, selected_component_names: tuple[str, ...] | list[str]
|
||||
) -> None:
|
||||
selected = set(selected_component_names)
|
||||
select_all = LAYERWISE_OFFLOAD_ALL_COMPONENTS in selected
|
||||
disabled_explicit_flags: list[str] = []
|
||||
|
||||
if (select_all or LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP in selected) and (
|
||||
self.text_encoder_cpu_offload is not False
|
||||
):
|
||||
self.text_encoder_cpu_offload = False
|
||||
disabled_flag_names.append("text_encoder_cpu_offload")
|
||||
if (
|
||||
"image_encoder_cpu_offload" in flag_names
|
||||
and self.image_encoder_cpu_offload is not False
|
||||
if self.is_arg_explicitly_set("text_encoder_cpu_offload"):
|
||||
disabled_explicit_flags.append("text_encoder_cpu_offload")
|
||||
if (select_all or LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP in selected) and (
|
||||
self.image_encoder_cpu_offload is not False
|
||||
):
|
||||
self.image_encoder_cpu_offload = False
|
||||
disabled_flag_names.append("image_encoder_cpu_offload")
|
||||
if "vae_cpu_offload" in flag_names and self.vae_cpu_offload is not False:
|
||||
if self.is_arg_explicitly_set("image_encoder_cpu_offload"):
|
||||
disabled_explicit_flags.append("image_encoder_cpu_offload")
|
||||
if (select_all or LAYERWISE_OFFLOAD_VAE_GROUP in selected) and (
|
||||
self.vae_cpu_offload is not False
|
||||
):
|
||||
self.vae_cpu_offload = False
|
||||
disabled_flag_names.append("vae_cpu_offload")
|
||||
if self.is_arg_explicitly_set("vae_cpu_offload"):
|
||||
disabled_explicit_flags.append("vae_cpu_offload")
|
||||
|
||||
explicit_disabled_flag_names = [
|
||||
flag_name
|
||||
for flag_name in disabled_flag_names
|
||||
if self.is_arg_explicitly_set(flag_name)
|
||||
]
|
||||
if explicit_disabled_flag_names:
|
||||
if disabled_explicit_flags:
|
||||
logger.info(
|
||||
"Ignoring explicit CPU-offload flags because layerwise offload "
|
||||
"manages the same component weights: %s",
|
||||
", ".join(
|
||||
f"{flag_name}=False" for flag_name in explicit_disabled_flag_names
|
||||
),
|
||||
"Ignoring component-offload flags because layerwise offload "
|
||||
"controls the same component groups: %s",
|
||||
", ".join(disabled_explicit_flags),
|
||||
)
|
||||
|
||||
def _adjust_autocast(self):
|
||||
@@ -1906,7 +2076,20 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
default=ServerArgs.warmup_steps,
|
||||
help="The number of warmup steps to perform for each resolution.",
|
||||
)
|
||||
# layerwise offload
|
||||
# component residency and legacy offload controls
|
||||
parser.add_argument(
|
||||
"--component-residency",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=ServerArgs.component_residency,
|
||||
metavar="COMPONENT=MODE",
|
||||
help=(
|
||||
"Select resident, component-offload, or layerwise-offload for "
|
||||
"pipeline components. Exact model_index.json component keys override "
|
||||
"the dit, text_encoder, image_encoder, vae, and all groups. "
|
||||
"Components without an assignment keep their automatic placement."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dit-cpu-offload",
|
||||
action=StoreBoolean,
|
||||
@@ -1930,8 +2113,9 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
"Select component keys from model_index.json for coarse CPU offload. "
|
||||
"Use dit, text_encoder, image_encoder, or vae as group aliases; "
|
||||
"all selects every loaded module and none disables component offload. "
|
||||
"This unified option cannot be combined with the legacy "
|
||||
"per-component CPU offload flags."
|
||||
"This compatibility option can be combined with per-component CPU "
|
||||
"offload flags; selected components take component-offload while "
|
||||
"unmatched components retain their existing settings."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -1940,9 +2124,8 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
default=ServerArgs.dit_layerwise_offload,
|
||||
help="Enable layerwise CPU offload with async H2D prefetch overlap for DiTs. "
|
||||
"It selects only the DiT layerwise group. Cannot be used together with cache-dit "
|
||||
"(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. May be combined with "
|
||||
"--dit-cpu-offload, in which case DiT weights stay on host memory and only the "
|
||||
"layers needed for the current step are brought on-device (lowest peak GPU memory).",
|
||||
"(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. If legacy DiT offload "
|
||||
"flags are also provided, layerwise offload is the effective DiT mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layerwise-offload-components",
|
||||
@@ -2660,7 +2843,11 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
@classmethod
|
||||
def from_kwargs(cls, **kwargs: Any) -> "ServerArgs":
|
||||
cls._reject_retired_args(kwargs)
|
||||
explicit_arg_names = set(kwargs)
|
||||
explicit_arg_names = kwargs.get("_explicit_arg_names")
|
||||
if explicit_arg_names is None:
|
||||
explicit_arg_names = set(kwargs)
|
||||
else:
|
||||
explicit_arg_names = set(explicit_arg_names)
|
||||
|
||||
# Convert backend string to enum if necessary
|
||||
if "backend" in kwargs and isinstance(kwargs["backend"], str):
|
||||
@@ -2733,6 +2920,17 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
)
|
||||
|
||||
def _validate_offload(self):
|
||||
if (
|
||||
self.component_residency is not None
|
||||
and self.pipeline_config.task_type.is_action_gen()
|
||||
):
|
||||
raise ValueError(
|
||||
"--component-residency is not supported by action-generation "
|
||||
"pipelines; use their existing model-specific offload controls"
|
||||
)
|
||||
if self.backend == Backend.DIFFUSERS and self.component_residency is not None:
|
||||
resolve_diffusers_pipeline_offload(self.component_residency)
|
||||
|
||||
# validate dit_offload_prefetch_size
|
||||
if self.dit_offload_prefetch_size > 1 and (
|
||||
isinstance(self.dit_offload_prefetch_size, float)
|
||||
@@ -2817,16 +3015,11 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
)
|
||||
self.use_fsdp_inference = False
|
||||
|
||||
if self.layerwise_offload_components:
|
||||
if self.has_layerwise_offload_components():
|
||||
if self.dit_offload_prefetch_size < 0.0:
|
||||
raise ValueError("dit_offload_prefetch_size must be non-negative")
|
||||
|
||||
is_dit_layerwise_offload_selected = self.is_dit_layerwise_offload_selected
|
||||
if self.use_fsdp_inference and is_dit_layerwise_offload_selected:
|
||||
logger.warning(
|
||||
"layerwise offload is selected for DiT components, automatically disabling use_fsdp_inference."
|
||||
)
|
||||
self.use_fsdp_inference = False
|
||||
|
||||
if envs.SGLANG_CACHE_DIT_ENABLED and is_dit_layerwise_offload_selected:
|
||||
raise ValueError(
|
||||
@@ -2841,10 +3034,22 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
self.performance_mode == "memory"
|
||||
or self.is_arg_explicitly_set("layerwise_offload_components")
|
||||
or self.dit_layerwise_offload
|
||||
or (
|
||||
self.component_residency
|
||||
and LAYERWISE_OFFLOAD in self.component_residency.values()
|
||||
)
|
||||
):
|
||||
selected_components = list(self.layerwise_offload_components or ())
|
||||
if self.component_residency:
|
||||
selected_components.extend(
|
||||
selector
|
||||
for selector, mode in self.component_residency.items()
|
||||
if mode == LAYERWISE_OFFLOAD
|
||||
)
|
||||
selected_components = list(dict.fromkeys(selected_components))
|
||||
logger.info_once(
|
||||
"Using layerwise offload components: "
|
||||
f"{', '.join(self.layerwise_offload_components)}. "
|
||||
f"{', '.join(selected_components or ())}. "
|
||||
"This reduces peak GPU memory and can increase latency; use "
|
||||
"--performance-mode speed for GPU-resident defaults when memory allows."
|
||||
)
|
||||
@@ -2854,7 +3059,10 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
return
|
||||
if not current_platform.is_cuda():
|
||||
raise ValueError("--direct-gpu-weight-loading requires CUDA")
|
||||
if self.dit_cpu_offload or self.is_dit_layerwise_offload_selected:
|
||||
if (
|
||||
self.should_cpu_offload_component("transformer")
|
||||
or self.residency_mode("transformer") == LAYERWISE_OFFLOAD
|
||||
):
|
||||
raise ValueError(
|
||||
"--direct-gpu-weight-loading requires a GPU-resident DiT; disable "
|
||||
"DiT CPU and layerwise offload"
|
||||
|
||||
@@ -181,7 +181,7 @@ def _load_sglang_component(
|
||||
sgl_args: ServerArgs,
|
||||
component: ComponentType,
|
||||
library: str,
|
||||
text_encoder_cpu_offload: bool | None = None,
|
||||
component_starts_on_cpu: bool | None = None,
|
||||
) -> nn.Module:
|
||||
loader = ComponentLoader.for_component_type(component.value, library)
|
||||
if component == ComponentType.TEXT_ENCODER:
|
||||
@@ -189,7 +189,7 @@ def _load_sglang_component(
|
||||
comp_path,
|
||||
sgl_args,
|
||||
component.value,
|
||||
cpu_offload_flag=text_encoder_cpu_offload,
|
||||
component_starts_on_cpu=component_starts_on_cpu,
|
||||
)
|
||||
else:
|
||||
component_model = loader.load_customized(comp_path, sgl_args, component.value)
|
||||
@@ -611,7 +611,7 @@ class AccuracyEngine:
|
||||
sgl_args,
|
||||
component,
|
||||
library,
|
||||
text_encoder_cpu_offload=(
|
||||
component_starts_on_cpu=(
|
||||
False
|
||||
if component != ComponentType.TEXT_ENCODER or materialize_sgl_on_device
|
||||
else True
|
||||
|
||||
@@ -55,6 +55,14 @@ class TestAdapterLoaderOffloadTarget(unittest.TestCase):
|
||||
server_args = self._server_args(cpu_offload_components=["vae"])
|
||||
self.assertTrue(server_args.should_cpu_offload_component("diffusion_decoder"))
|
||||
|
||||
def test_layerwise_diffusion_decoder_starts_on_cpu(self):
|
||||
server_args = self._server_args(
|
||||
component_residency={"diffusion_decoder": "layerwise-offload"}
|
||||
)
|
||||
|
||||
self.assertFalse(server_args.should_cpu_offload_component("diffusion_decoder"))
|
||||
self.assertTrue(server_args.should_start_component_on_cpu("diffusion_decoder"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,398 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
|
||||
ComponentResidencyManager,
|
||||
ComponentUse,
|
||||
ResidencyState,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
ComponentOffloadStrategy,
|
||||
ResidentStrategy,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.text_encoding import (
|
||||
RealtimeTextEncodingStage,
|
||||
)
|
||||
|
||||
|
||||
def test_component_offload_releases_preferred_component_after_request():
|
||||
strategy = ComponentOffloadStrategy()
|
||||
strategy.finish_use = Mock()
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse(
|
||||
stage_name="TextEncodingStage",
|
||||
component_name="text_encoder",
|
||||
preferred_ready_after_request=True,
|
||||
)
|
||||
state = ResidencyState(batch_is_warmup=False)
|
||||
|
||||
strategy.finish_request(module, use, state, preferred=True)
|
||||
|
||||
strategy.finish_use.assert_called_once_with(module, use, state)
|
||||
|
||||
|
||||
def test_component_offload_keeps_preferred_component_after_warmup():
|
||||
strategy = ComponentOffloadStrategy()
|
||||
strategy.prepare_for_use = Mock()
|
||||
strategy.wait_for_use = Mock()
|
||||
strategy.finish_use = Mock()
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse(
|
||||
stage_name="TextEncodingStage",
|
||||
component_name="text_encoder",
|
||||
preferred_ready_after_request=True,
|
||||
)
|
||||
state = ResidencyState(batch_is_warmup=True)
|
||||
|
||||
strategy.finish_request(module, use, state, preferred=True)
|
||||
|
||||
strategy.prepare_for_use.assert_called_once_with(module, use, state)
|
||||
strategy.wait_for_use.assert_called_once_with(module, use, state)
|
||||
strategy.finish_use.assert_not_called()
|
||||
|
||||
|
||||
def test_request_tail_uses_dynamic_component_instance():
|
||||
pipeline = SimpleNamespace(
|
||||
modules={},
|
||||
_stage_name_mapping={},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
manager = ComponentResidencyManager(
|
||||
pipeline,
|
||||
SimpleNamespace(enable_layerwise_nvtx_marker=False),
|
||||
)
|
||||
strategy = Mock()
|
||||
strategy.prefetch_for_use.return_value = False
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse("DynamicStage", "dynamic_encoder")
|
||||
|
||||
manager.ensure_ready(use, module=module)
|
||||
manager.finish_request()
|
||||
|
||||
strategy.finish_request.assert_called_once_with(
|
||||
module, use, manager.state, preferred=False
|
||||
)
|
||||
|
||||
|
||||
def test_strategy_cache_replaces_stale_component_instance():
|
||||
server_args = SimpleNamespace(
|
||||
enable_layerwise_nvtx_marker=False,
|
||||
residency_mode=lambda _component_name: "resident",
|
||||
)
|
||||
pipeline = SimpleNamespace(
|
||||
modules={},
|
||||
_stage_name_mapping={},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
manager = ComponentResidencyManager(pipeline, server_args)
|
||||
first_module = torch.nn.Linear(2, 2)
|
||||
second_module = torch.nn.Linear(2, 2)
|
||||
|
||||
first_strategy = manager.strategy_for("transformer", first_module)
|
||||
second_strategy = manager.strategy_for("transformer", second_module)
|
||||
|
||||
assert isinstance(first_strategy, ResidentStrategy)
|
||||
assert isinstance(second_strategy, ResidentStrategy)
|
||||
assert first_strategy is not second_strategy
|
||||
assert manager._strategy_cache["transformer"][0] is second_module
|
||||
|
||||
|
||||
def test_forget_module_clears_active_manager_references():
|
||||
pipeline = SimpleNamespace(
|
||||
modules={},
|
||||
_stage_name_mapping={},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
manager = ComponentResidencyManager(
|
||||
pipeline,
|
||||
SimpleNamespace(enable_layerwise_nvtx_marker=False),
|
||||
)
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse("stage", "transformer")
|
||||
manager._active_use = use
|
||||
manager._active_use_module = module
|
||||
manager.state.current_use = use
|
||||
manager._uses_seen["transformer"] = use
|
||||
manager._modules_seen["transformer"] = module
|
||||
manager._prefetched_use_keys.add(("stage", "transformer", None))
|
||||
|
||||
manager.forget_module(module)
|
||||
|
||||
assert manager._active_use is None
|
||||
assert manager._active_use_module is None
|
||||
assert manager.state.current_use is None
|
||||
assert "transformer" not in manager._uses_seen
|
||||
assert "transformer" not in manager._modules_seen
|
||||
assert manager._prefetched_use_keys == set()
|
||||
|
||||
|
||||
def test_group_warmup_state_requires_every_batch_to_be_warmup():
|
||||
pipeline = SimpleNamespace(
|
||||
modules={},
|
||||
_stage_name_mapping={},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False)
|
||||
manager = ComponentResidencyManager(pipeline, server_args)
|
||||
|
||||
manager.begin_request(
|
||||
[],
|
||||
[
|
||||
SimpleNamespace(is_warmup=True),
|
||||
SimpleNamespace(is_warmup=False),
|
||||
],
|
||||
server_args,
|
||||
)
|
||||
|
||||
assert manager.state.batch_is_warmup is False
|
||||
|
||||
|
||||
class _Stage:
|
||||
def __init__(self, *uses: ComponentUse):
|
||||
self.uses = list(uses)
|
||||
|
||||
def component_uses(self, server_args, stage_name=None):
|
||||
return self.uses
|
||||
|
||||
|
||||
def _manager_for_stage(stage, modules):
|
||||
pipeline = SimpleNamespace(
|
||||
modules=modules,
|
||||
_stage_name_mapping={"stage": stage},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False)
|
||||
manager = ComponentResidencyManager(pipeline, server_args)
|
||||
manager.refresh_pipeline(pipeline)
|
||||
manager.begin_request([stage], SimpleNamespace(is_warmup=False), server_args)
|
||||
return manager, server_args
|
||||
|
||||
|
||||
def test_single_component_stage_is_prepared_at_stage_entry():
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse("stage", "text_encoder")
|
||||
stage = _Stage(use)
|
||||
manager, server_args = _manager_for_stage(stage, {"text_encoder": module})
|
||||
strategy = Mock()
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
|
||||
manager.before_stage(stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
|
||||
strategy.prepare_for_use.assert_called_once_with(module, use, manager.state)
|
||||
strategy.wait_for_use.assert_called_once_with(module, use, manager.state)
|
||||
|
||||
|
||||
def test_explicit_component_use_is_prepared_only_at_call_site():
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse("stage", "text_encoder", start_at_stage_entry=False)
|
||||
stage = _Stage(use)
|
||||
manager, server_args = _manager_for_stage(stage, {"text_encoder": module})
|
||||
strategy = Mock()
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
|
||||
manager.before_stage(stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
|
||||
strategy.prepare_for_use.assert_not_called()
|
||||
|
||||
manager.begin_use(use)
|
||||
|
||||
strategy.prepare_for_use.assert_called_once_with(module, use, manager.state)
|
||||
strategy.wait_for_use.assert_called_once_with(module, use, manager.state)
|
||||
|
||||
|
||||
def test_realtime_text_encoder_use_starts_at_call_site():
|
||||
stage = RealtimeTextEncodingStage.__new__(RealtimeTextEncodingStage)
|
||||
stage.text_encoders = [None]
|
||||
stage._registered_stage_name = None
|
||||
|
||||
uses = stage.component_uses(SimpleNamespace(), "RealtimeTextEncodingStage")
|
||||
|
||||
assert len(uses) == 1
|
||||
assert uses[0].component_name == "text_encoder"
|
||||
assert uses[0].start_at_stage_entry is False
|
||||
|
||||
|
||||
def test_qwen_layered_uses_loaded_text_encoder(monkeypatch):
|
||||
from sglang.multimodal_gen.runtime.pipelines import qwen_image
|
||||
|
||||
text_encoder = object()
|
||||
stage = SimpleNamespace()
|
||||
stage_kwargs = {}
|
||||
pipeline = qwen_image.QwenImageLayeredPipeline.__new__(
|
||||
qwen_image.QwenImageLayeredPipeline
|
||||
)
|
||||
pipeline.model_path = "model"
|
||||
pipeline.modules = {
|
||||
name: object()
|
||||
for name in (
|
||||
"text_encoder",
|
||||
"vae",
|
||||
"tokenizer",
|
||||
"processor",
|
||||
"transformer",
|
||||
"scheduler",
|
||||
)
|
||||
}
|
||||
pipeline.modules["text_encoder"] = text_encoder
|
||||
pipeline.add_stage_factory = lambda _role, factory, _name: factory()
|
||||
pipeline.add_standard_timestep_preparation_stage = lambda **_kwargs: None
|
||||
pipeline.add_standard_denoising_stage = lambda: None
|
||||
pipeline.add_standard_decoding_stage = lambda: None
|
||||
|
||||
def create_stage(**kwargs):
|
||||
stage_kwargs.update(kwargs)
|
||||
return stage
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image, "QwenImageLayeredBeforeDenoisingStage", create_stage
|
||||
)
|
||||
server_args = SimpleNamespace(
|
||||
pipeline_config=SimpleNamespace(
|
||||
vae_precision="bf16",
|
||||
text_encoder_precisions=("bf16",),
|
||||
)
|
||||
)
|
||||
|
||||
pipeline.create_pipeline_stages(server_args)
|
||||
|
||||
assert stage_kwargs["text_encoder"] is text_encoder
|
||||
|
||||
|
||||
def test_single_component_stage_is_finished_at_stage_exit():
|
||||
module = torch.nn.Linear(2, 2)
|
||||
use = ComponentUse("stage", "text_encoder")
|
||||
stage = _Stage(use)
|
||||
manager, server_args = _manager_for_stage(stage, {"text_encoder": module})
|
||||
strategy = Mock()
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
|
||||
manager.before_stage(stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
manager.end_stage()
|
||||
|
||||
strategy.finish_use.assert_called_once_with(module, use, manager.state)
|
||||
|
||||
|
||||
def test_adjacent_stages_reuse_the_same_component_interval():
|
||||
module = torch.nn.Linear(2, 2)
|
||||
first_use = ComponentUse("first", "text_encoder")
|
||||
second_use = ComponentUse("second", "text_encoder")
|
||||
first_stage = _Stage(first_use)
|
||||
second_stage = _Stage(second_use)
|
||||
pipeline = SimpleNamespace(
|
||||
modules={"text_encoder": module},
|
||||
_stage_name_mapping={"first": first_stage, "second": second_stage},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False)
|
||||
manager = ComponentResidencyManager(pipeline, server_args)
|
||||
manager.refresh_pipeline(pipeline)
|
||||
manager.begin_request(
|
||||
[first_stage, second_stage], SimpleNamespace(is_warmup=False), server_args
|
||||
)
|
||||
strategy = Mock()
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
|
||||
manager.before_stage(first_stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
manager.end_stage()
|
||||
manager.before_stage(second_stage, 1, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
manager.end_stage()
|
||||
|
||||
strategy.prepare_for_use.assert_called_once_with(module, first_use, manager.state)
|
||||
strategy.finish_use.assert_called_once_with(module, second_use, manager.state)
|
||||
|
||||
|
||||
def test_adjacent_same_component_replacement_finishes_old_instance():
|
||||
first_module = torch.nn.Linear(2, 2)
|
||||
second_module = torch.nn.Linear(2, 2)
|
||||
first_use = ComponentUse("first", "text_encoder")
|
||||
second_use = ComponentUse("second", "text_encoder")
|
||||
first_stage = _Stage(first_use)
|
||||
second_stage = _Stage(second_use)
|
||||
pipeline = SimpleNamespace(
|
||||
modules={"text_encoder": first_module},
|
||||
_stage_name_mapping={"first": first_stage, "second": second_stage},
|
||||
component_residency_strategies={},
|
||||
)
|
||||
server_args = SimpleNamespace(enable_layerwise_nvtx_marker=False)
|
||||
manager = ComponentResidencyManager(pipeline, server_args)
|
||||
manager.refresh_pipeline(pipeline)
|
||||
manager.begin_request(
|
||||
[first_stage, second_stage],
|
||||
SimpleNamespace(is_warmup=False),
|
||||
server_args,
|
||||
)
|
||||
first_strategy = Mock()
|
||||
second_strategy = Mock()
|
||||
manager.strategy_for = Mock(
|
||||
side_effect=lambda _component_name, module: (
|
||||
first_strategy if module is first_module else second_strategy
|
||||
)
|
||||
)
|
||||
|
||||
manager.begin_use(first_use, module=first_module)
|
||||
manager.begin_use(second_use, module=second_module)
|
||||
|
||||
first_strategy.finish_use.assert_called_once_with(
|
||||
first_module, first_use, manager.state
|
||||
)
|
||||
second_strategy.prepare_for_use.assert_called_once_with(
|
||||
second_module, second_use, manager.state
|
||||
)
|
||||
|
||||
|
||||
def test_multi_component_stage_controls_its_use_intervals():
|
||||
uses = (
|
||||
ComponentUse("stage", "text_encoder"),
|
||||
ComponentUse("stage", "vae"),
|
||||
)
|
||||
stage = _Stage(*uses)
|
||||
manager, server_args = _manager_for_stage(stage, {})
|
||||
manager.begin_use = Mock()
|
||||
|
||||
manager.before_stage(stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
|
||||
manager.begin_use.assert_not_called()
|
||||
|
||||
|
||||
def test_dynamic_component_is_prepared_when_stage_supplies_its_module():
|
||||
use = ComponentUse("stage", "dynamic_encoder")
|
||||
stage = _Stage(use)
|
||||
manager, server_args = _manager_for_stage(stage, {})
|
||||
strategy = Mock()
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
|
||||
manager.before_stage(stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_stage()
|
||||
module = torch.nn.Linear(2, 2)
|
||||
manager.begin_use(use, module=module)
|
||||
|
||||
strategy.prepare_for_use.assert_called_once_with(module, use, manager.state)
|
||||
strategy.wait_for_use.assert_called_once_with(module, use, manager.state)
|
||||
|
||||
|
||||
def test_component_is_not_kept_across_another_component_use():
|
||||
text_use = ComponentUse("stage", "text_encoder")
|
||||
stage = _Stage(
|
||||
text_use,
|
||||
ComponentUse("stage", "transformer"),
|
||||
ComponentUse("stage", "text_encoder", phase="second"),
|
||||
)
|
||||
module = torch.nn.Linear(2, 2)
|
||||
manager, server_args = _manager_for_stage(stage, {"text_encoder": module})
|
||||
strategy = Mock()
|
||||
manager.strategy_for = Mock(return_value=strategy)
|
||||
|
||||
manager.before_stage(stage, 0, SimpleNamespace(is_warmup=False), server_args)
|
||||
manager.begin_use(text_use)
|
||||
manager.end_use(text_use)
|
||||
|
||||
strategy.finish_use.assert_called_once_with(module, text_use, manager.state)
|
||||
@@ -50,6 +50,7 @@ class _DummyScheduler:
|
||||
|
||||
class _RecordingBeforeDenoisingStage(GlmImageBeforeDenoisingStage):
|
||||
def __init__(self):
|
||||
self.text_encoder = None
|
||||
self.transformer = SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
in_channels=4,
|
||||
|
||||
@@ -21,6 +21,9 @@ from sglang.multimodal_gen.runtime.models.dits.stable_diffusion import (
|
||||
StableDiffusionUNetConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
|
||||
from sglang.multimodal_gen.runtime.pipelines.hunyuan3d_pipeline import (
|
||||
Hunyuan3D2Pipeline,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.paint import (
|
||||
Hunyuan3DPaintPostprocessStage,
|
||||
Hunyuan3DPaintTexGenStage,
|
||||
@@ -188,6 +191,20 @@ class TestHunyuan3DWarmupOutput(unittest.TestCase):
|
||||
self.assertEqual(output.output_file_paths, [])
|
||||
|
||||
|
||||
class TestHunyuan3DComponentResidency(unittest.TestCase):
|
||||
def test_layerwise_texture_component_starts_on_cpu(self):
|
||||
server_args = SimpleNamespace(
|
||||
should_start_component_on_cpu=lambda component_name: (
|
||||
component_name == "paint_transformer"
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
Hunyuan3D2Pipeline._component_device(server_args, "paint_transformer"),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
class TestHunyuan3DPaintTurboSchedule(unittest.TestCase):
|
||||
def test_uses_standard_lcm_schedule_without_custom_timesteps(self):
|
||||
stage = Hunyuan3DPaintTexGenStage.__new__(Hunyuan3DPaintTexGenStage)
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp8Config,
|
||||
@@ -11,7 +13,7 @@ from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
|
||||
_ModelOptFp8OffloadAdapter,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers import (
|
||||
component_resident_strategies as component_resident_strategies_mod,
|
||||
component_residency_strategies as component_residency_strategies_mod,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers import (
|
||||
layerwise_offload as layerwise_offload_mod,
|
||||
@@ -20,10 +22,16 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im
|
||||
ComponentUse,
|
||||
build_component_residency_strategy,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
RESIDENT,
|
||||
ComponentResidencyError,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
ComponentOffloadStrategy,
|
||||
LayerwiseOffloadStrategy,
|
||||
ResidentStrategy,
|
||||
VanillaD2HStrategy,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
@@ -168,14 +176,34 @@ class _LayerwiseComponent(torch.nn.Module, LayerwiseOffloadableModuleMixin):
|
||||
|
||||
|
||||
class _TestServerArgs(SimpleNamespace):
|
||||
canonical_residency_mode = ServerArgs.canonical_residency_mode
|
||||
explicit_residency_mode = ServerArgs.explicit_residency_mode
|
||||
_legacy_component_offload_flag = staticmethod(
|
||||
ServerArgs._legacy_component_offload_flag
|
||||
)
|
||||
residency_mode = ServerArgs.residency_mode
|
||||
is_arg_explicitly_set = ServerArgs.is_arg_explicitly_set
|
||||
is_explicit_layerwise_offload_component = (
|
||||
ServerArgs.is_explicit_layerwise_offload_component
|
||||
)
|
||||
should_cpu_offload_component = ServerArgs.should_cpu_offload_component
|
||||
record_component_layerwise_capability = (
|
||||
ServerArgs.record_component_layerwise_capability
|
||||
)
|
||||
|
||||
|
||||
def _server_args(**kwargs):
|
||||
defaults = dict(
|
||||
component_residency=None,
|
||||
disagg_role=RoleType.MONOLITHIC,
|
||||
_required_resident_components=set(),
|
||||
_component_layerwise_capabilities={},
|
||||
_explicit_arg_names=set(),
|
||||
cpu_offload_components=None,
|
||||
use_fsdp_inference=False,
|
||||
dit_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
layerwise_offload_components=None,
|
||||
text_encoder_cpu_offload=False,
|
||||
image_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
@@ -501,11 +529,81 @@ def test_layerwise_configuration_all_selects_every_capable_component(monkeypatch
|
||||
assert is_layerwise_offloaded_module(transformer)
|
||||
|
||||
|
||||
def test_component_cpu_offload_strategy_remains_flag_driven():
|
||||
def test_explicit_layerwise_all_rejects_unsupported_modules():
|
||||
modules = {
|
||||
"text_encoder": _NestedEncoderDummyModel(),
|
||||
"unsupported_adapter": torch.nn.Linear(2, 2),
|
||||
"scheduler": object(),
|
||||
}
|
||||
|
||||
with pytest.raises(ComponentResidencyError, match="unsupported_adapter"):
|
||||
configure_layerwise_offload_modules(
|
||||
modules, _server_args(), component_names=["all"]
|
||||
)
|
||||
|
||||
|
||||
def test_explicit_layerwise_dit_rejects_unsupported_dit():
|
||||
modules = {"transformer": torch.nn.Linear(2, 2)}
|
||||
|
||||
with pytest.raises(ComponentResidencyError, match="transformer"):
|
||||
configure_layerwise_offload_modules(
|
||||
modules, _server_args(), component_names=["dit"]
|
||||
)
|
||||
|
||||
|
||||
def test_explicit_layerwise_exact_selector_rejects_non_module(monkeypatch):
|
||||
monkeypatch.setattr(layerwise_offload_mod.current_platform, "is_cpu", lambda: False)
|
||||
server_args = _server_args(component_residency={"scheduler": LAYERWISE_OFFLOAD})
|
||||
|
||||
with pytest.raises(ComponentResidencyError, match="scheduler"):
|
||||
configure_layerwise_offload_modules(
|
||||
{"scheduler": object()}, server_args, warn_missing=False
|
||||
)
|
||||
|
||||
|
||||
def test_auto_layerwise_skips_unsupported_component(monkeypatch):
|
||||
monkeypatch.setattr(layerwise_offload_mod.current_platform, "is_cpu", lambda: False)
|
||||
server_args = _server_args(
|
||||
layerwise_offload_components=["text_encoder"],
|
||||
text_encoder_cpu_offload=True,
|
||||
)
|
||||
|
||||
configured = configure_layerwise_offload_modules(
|
||||
{"text_encoder": torch.nn.Linear(2, 2)},
|
||||
server_args,
|
||||
component_names=["text_encoder"],
|
||||
warn_missing=False,
|
||||
)
|
||||
|
||||
assert configured == []
|
||||
assert server_args.residency_mode("text_encoder") == COMPONENT_OFFLOAD
|
||||
|
||||
|
||||
def test_canonical_selector_does_not_make_auto_layerwise_selection_strict(
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setattr(layerwise_offload_mod.current_platform, "is_cpu", lambda: False)
|
||||
server_args = _server_args(
|
||||
component_residency={"transformer": RESIDENT},
|
||||
layerwise_offload_components=["text_encoder"],
|
||||
text_encoder_cpu_offload=True,
|
||||
)
|
||||
|
||||
configured = configure_layerwise_offload_modules(
|
||||
{"text_encoder": torch.nn.Linear(2, 2)},
|
||||
server_args,
|
||||
warn_missing=True,
|
||||
)
|
||||
|
||||
assert configured == []
|
||||
assert server_args.residency_mode("text_encoder") == COMPONENT_OFFLOAD
|
||||
|
||||
|
||||
def test_legacy_cpu_offload_flag_selects_component_offload_strategy():
|
||||
strategy = build_component_residency_strategy(
|
||||
"text_encoder", _DummyModel(), _server_args(text_encoder_cpu_offload=True)
|
||||
)
|
||||
assert isinstance(strategy, VanillaD2HStrategy)
|
||||
assert isinstance(strategy, ComponentOffloadStrategy)
|
||||
|
||||
strategy = build_component_residency_strategy(
|
||||
"unknown_component", _DummyModel(), _server_args(text_encoder_cpu_offload=True)
|
||||
@@ -513,6 +611,26 @@ def test_component_cpu_offload_strategy_remains_flag_driven():
|
||||
assert isinstance(strategy, ResidentStrategy)
|
||||
|
||||
|
||||
def test_component_residency_strategy_selection_is_direct():
|
||||
for mode, strategy_type in (
|
||||
(RESIDENT, ResidentStrategy),
|
||||
(COMPONENT_OFFLOAD, ComponentOffloadStrategy),
|
||||
):
|
||||
strategy = build_component_residency_strategy(
|
||||
"text_encoder",
|
||||
_DummyModel(),
|
||||
_server_args(component_residency={"text_encoder": mode}),
|
||||
)
|
||||
assert isinstance(strategy, strategy_type)
|
||||
|
||||
|
||||
def test_explicit_layerwise_requires_component_support():
|
||||
server_args = _server_args(component_residency={"text_encoder": LAYERWISE_OFFLOAD})
|
||||
|
||||
with pytest.raises(ValueError, match="did not enable layerwise offload"):
|
||||
build_component_residency_strategy("text_encoder", _DummyModel(), server_args)
|
||||
|
||||
|
||||
def test_resident_strategy_prepares_local_device_without_dtype(monkeypatch):
|
||||
calls = []
|
||||
|
||||
@@ -520,7 +638,7 @@ def test_resident_strategy_prepares_local_device_without_dtype(monkeypatch):
|
||||
calls.append((module, dtype))
|
||||
|
||||
monkeypatch.setattr(
|
||||
component_resident_strategies_mod,
|
||||
component_residency_strategies_mod,
|
||||
"_module_to_local_device",
|
||||
fake_module_to_local_device,
|
||||
)
|
||||
@@ -542,11 +660,16 @@ def test_resident_strategy_keeps_fsdp_managed_module_owned_by_fsdp(monkeypatch):
|
||||
calls.append((module, dtype))
|
||||
|
||||
monkeypatch.setattr(
|
||||
component_resident_strategies_mod,
|
||||
component_residency_strategies_mod,
|
||||
"_module_to_local_device",
|
||||
fake_module_to_local_device,
|
||||
)
|
||||
module = type("FSDPDummyModel", (_DummyModel,), {})()
|
||||
monkeypatch.setattr(
|
||||
component_residency_strategies_mod,
|
||||
"is_fsdp_managed_module",
|
||||
lambda _module: True,
|
||||
)
|
||||
module = _DummyModel()
|
||||
|
||||
ResidentStrategy().prepare_for_use(
|
||||
module,
|
||||
@@ -886,7 +1009,11 @@ def test_disable_offload_short_circuits_residency_release(monkeypatch):
|
||||
assert tuple(param.shape) != (1,), name
|
||||
|
||||
# The exact call path the residency strategy takes on use-site switches.
|
||||
LayerwiseOffloadStrategy().exit(model)
|
||||
LayerwiseOffloadStrategy().finish_use(
|
||||
model,
|
||||
ComponentUse(stage_name="test", component_name="transformer"),
|
||||
SimpleNamespace(),
|
||||
)
|
||||
model.prepare_for_next_req()
|
||||
for name, param in model.named_parameters():
|
||||
assert tuple(param.shape) != (1,), name
|
||||
|
||||
@@ -420,6 +420,36 @@ class TestLTX25DiffusionDecoder(unittest.TestCase):
|
||||
# trailing frames that stage 4 crops.
|
||||
self.assertEqual(model.decoder.trailing_pad_latent_frames, 2)
|
||||
|
||||
def test_exposes_each_executable_block_group_for_layerwise_offload(self):
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
|
||||
LayerwiseOffloadableModuleMixin,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.decoders.ltx_2_5_diffusion_decoder import (
|
||||
LTX2VideoDiffusionDecoderModel,
|
||||
)
|
||||
|
||||
with torch.device("meta"):
|
||||
model = LTX2VideoDiffusionDecoderModel(self._config())
|
||||
|
||||
self.assertIsInstance(model, LayerwiseOffloadableModuleMixin)
|
||||
self.assertEqual(
|
||||
model.layer_names,
|
||||
[
|
||||
"decoder.det_stages.0",
|
||||
"decoder.det_stages.1",
|
||||
"decoder.det_stages.2",
|
||||
"decoder.det_stages.3",
|
||||
"decoder.diff_blocks",
|
||||
],
|
||||
)
|
||||
named_modules = dict(model.named_modules())
|
||||
for layer_name in model.layer_names:
|
||||
with self.subTest(layer_name=layer_name):
|
||||
self.assertIsInstance(named_modules[layer_name], nn.ModuleList)
|
||||
|
||||
def test_timestep_embedder_is_replicated_and_checkpoint_compatible(self):
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
@@ -12,12 +12,9 @@ from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
|
||||
class TestMiniMaxH3EncoderDevice(unittest.TestCase):
|
||||
"""`device` must name the compute side, not the parameter storage side.
|
||||
|
||||
`--text-encoder-cpu-offload` loads this encoder under an FSDP CPU offload
|
||||
policy: the sharded parameters sit on CPU and are all-gathered to the
|
||||
accelerator for the forward. Reporting the parameter device there sent
|
||||
`encode_ids` to build `input_ids`/`attention_mask`/`position_ids` on CPU
|
||||
while the forward ran on the accelerator, and the rope matmul died with
|
||||
"Expected all tensors to be on the same device ... mat2 is on cpu".
|
||||
Component offload stores parameters on CPU between uses. Input construction
|
||||
must still target the accelerator selected for the forward rather than infer
|
||||
the compute device from an offloaded parameter.
|
||||
"""
|
||||
|
||||
def _encoder_with_param_on(self, device: torch.device) -> MiniMaxH3Qwen3VLEncoder:
|
||||
|
||||
@@ -16,9 +16,9 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager im
|
||||
ComponentResidencyManager,
|
||||
ComponentUse,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_resident_strategies import (
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
|
||||
ComponentOffloadStrategy,
|
||||
ResidentStrategy,
|
||||
VanillaD2HStrategy,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils import nvtx_pytorch_hooks
|
||||
from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import (
|
||||
@@ -209,9 +209,6 @@ class _NoOpResidencyStrategy:
|
||||
def prefetch_for_use(self, module, use, state) -> bool:
|
||||
return False
|
||||
|
||||
def prepare_after_request(self, module, use, state) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _test_manager(
|
||||
modules: dict[str, torch.nn.Module],
|
||||
@@ -239,7 +236,9 @@ class TestComponentResidencyNvtxHooks(unittest.TestCase):
|
||||
manager.strategy_for = lambda _component_name, _module: ResidentStrategy()
|
||||
self.assertTrue(manager._should_keep_single_dit("transformer", module))
|
||||
|
||||
manager.strategy_for = lambda _component_name, _module: VanillaD2HStrategy()
|
||||
manager.strategy_for = (
|
||||
lambda _component_name, _module: ComponentOffloadStrategy()
|
||||
)
|
||||
self.assertFalse(manager._should_keep_single_dit("transformer", module))
|
||||
|
||||
def test_disabled_flag_is_noop(self) -> None:
|
||||
|
||||
@@ -1,9 +1,14 @@
|
||||
import contextlib
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.executors import pipeline_executor
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
|
||||
PipelineExecutor,
|
||||
@@ -34,18 +39,18 @@ def _batch():
|
||||
return SimpleNamespace(profile=False, is_warmup=False)
|
||||
|
||||
|
||||
class _TestServerArgs(SimpleNamespace):
|
||||
def should_cpu_offload_component(self, component_name):
|
||||
return self.component_modes.get(component_name) == COMPONENT_OFFLOAD
|
||||
|
||||
|
||||
def _server_args(**overrides):
|
||||
values = {
|
||||
"use_fsdp_inference": False,
|
||||
"dit_cpu_offload": False,
|
||||
"text_encoder_cpu_offload": False,
|
||||
"image_encoder_cpu_offload": False,
|
||||
"vae_cpu_offload": False,
|
||||
"dit_layerwise_offload": False,
|
||||
"layerwise_offload_components": (),
|
||||
"component_modes": {},
|
||||
}
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
return _TestServerArgs(**values)
|
||||
|
||||
|
||||
class _NoGradPlatform:
|
||||
@@ -110,14 +115,36 @@ def test_execute_group_with_profiling_uses_platform_inference_mode(monkeypatch):
|
||||
assert executor.group_grad_enabled is False
|
||||
|
||||
|
||||
def test_group_payload_is_forwarded_to_component_residency_manager():
|
||||
executor = _RecordingExecutor()
|
||||
executor.component_residency_manager = Mock()
|
||||
batches = [_batch(), _batch()]
|
||||
server_args = _server_args()
|
||||
|
||||
executor.begin_component_residency_request([], batches, server_args)
|
||||
|
||||
executor.component_residency_manager.begin_request.assert_called_once_with(
|
||||
[], batches, server_args
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("server_args", "component_names"),
|
||||
[
|
||||
(_server_args(use_fsdp_inference=True), ("transformer",)),
|
||||
(_server_args(dit_cpu_offload=True), ("transformer",)),
|
||||
(_server_args(text_encoder_cpu_offload=True), ("text_encoder",)),
|
||||
(_server_args(image_encoder_cpu_offload=True), ("image_encoder",)),
|
||||
(_server_args(vae_cpu_offload=True), ("vae",)),
|
||||
(
|
||||
_server_args(component_modes={"transformer": COMPONENT_OFFLOAD}),
|
||||
("transformer",),
|
||||
),
|
||||
(
|
||||
_server_args(component_modes={"text_encoder": COMPONENT_OFFLOAD}),
|
||||
("text_encoder",),
|
||||
),
|
||||
(
|
||||
_server_args(component_modes={"image_encoder": COMPONENT_OFFLOAD}),
|
||||
("image_encoder",),
|
||||
),
|
||||
(_server_args(component_modes={"vae": COMPONENT_OFFLOAD}), ("vae",)),
|
||||
],
|
||||
)
|
||||
def test_stage_context_preserves_version_counters_when_needed(
|
||||
@@ -137,16 +164,12 @@ def test_stage_context_preserves_version_counters_when_needed(
|
||||
@pytest.mark.parametrize(
|
||||
("server_args", "component_names"),
|
||||
[
|
||||
(_server_args(dit_layerwise_offload=True), ("transformer",)),
|
||||
(
|
||||
_server_args(
|
||||
text_encoder_cpu_offload=True,
|
||||
layerwise_offload_components=("transformer",),
|
||||
),
|
||||
_server_args(component_modes={"transformer": LAYERWISE_OFFLOAD}),
|
||||
("transformer",),
|
||||
),
|
||||
(
|
||||
_server_args(layerwise_offload_components=("text_encoder",)),
|
||||
_server_args(component_modes={"text_encoder": LAYERWISE_OFFLOAD}),
|
||||
("text_encoder",),
|
||||
),
|
||||
],
|
||||
|
||||
@@ -58,6 +58,14 @@ from sglang.multimodal_gen.registry import (
|
||||
get_non_diffusers_pipeline_name,
|
||||
is_known_non_diffusers_multimodal_model,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
|
||||
COMPONENT_OFFLOAD,
|
||||
LAYERWISE_OFFLOAD,
|
||||
RESIDENT,
|
||||
normalize_component_residency,
|
||||
resolve_component_residency_mode,
|
||||
resolve_diffusers_pipeline_offload,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
|
||||
QwenImageTransformer2DModel,
|
||||
)
|
||||
@@ -413,10 +421,8 @@ class TestServerArgsPathExpansion(unittest.TestCase):
|
||||
server_args = ServerArgs.from_cli_args(args, unknown_args)
|
||||
|
||||
self.assertEqual(server_args.cpu_offload_components, ["transformer", "vae"])
|
||||
self.assertTrue(server_args.dit_cpu_offload)
|
||||
self.assertTrue(server_args.vae_cpu_offload)
|
||||
self.assertFalse(server_args.text_encoder_cpu_offload)
|
||||
self.assertFalse(server_args.image_encoder_cpu_offload)
|
||||
self.assertEqual(server_args.residency_mode("transformer"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(server_args.residency_mode("vae"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_serve_cli_preserves_config_and_dynamic_unknown_args(self):
|
||||
from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
|
||||
@@ -867,8 +873,8 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
self.assertTrue(args.should_cpu_offload_component("transformer_2"))
|
||||
self.assertTrue(args.should_cpu_offload_component("audio_vae"))
|
||||
self.assertTrue(args.should_cpu_offload_component("connectors"))
|
||||
self.assertFalse(args.dit_cpu_offload)
|
||||
self.assertFalse(args.vae_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("vae"), RESIDENT)
|
||||
|
||||
def test_cpu_offload_components_all_matches_dynamic_components(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
@@ -893,10 +899,10 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
)
|
||||
|
||||
self.assertEqual(args.cpu_offload_components, [])
|
||||
self.assertFalse(args.dit_cpu_offload)
|
||||
self.assertFalse(args.text_encoder_cpu_offload)
|
||||
self.assertFalse(args.image_encoder_cpu_offload)
|
||||
self.assertFalse(args.vae_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("text_encoder"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("image_encoder"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("vae"), RESIDENT)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "cannot be combined"):
|
||||
self._from_dict_with_task_type(
|
||||
@@ -904,16 +910,249 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
kwargs={"cpu_offload_components": ["none", "vae"]},
|
||||
)
|
||||
|
||||
def test_cpu_offload_components_rejects_legacy_flag_conflict(self):
|
||||
with self.assertRaisesRegex(ValueError, "cannot be combined"):
|
||||
self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={
|
||||
"performance_mode": "manual",
|
||||
"cpu_offload_components": ["dit"],
|
||||
def test_cpu_offload_components_can_mix_with_legacy_flags(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={
|
||||
"performance_mode": "manual",
|
||||
"cpu_offload_components": ["vae"],
|
||||
"dit_cpu_offload": True,
|
||||
"text_encoder_cpu_offload": False,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.residency_mode("vae"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("text_encoder"), RESIDENT)
|
||||
|
||||
def test_legacy_component_cpu_offload_flags_remain_supported(self):
|
||||
cases = (
|
||||
("dit_cpu_offload", "transformer_2"),
|
||||
("text_encoder_cpu_offload", "text_encoder_3"),
|
||||
("image_encoder_cpu_offload", "image_encoder"),
|
||||
("vae_cpu_offload", "audio_vae"),
|
||||
)
|
||||
for flag_name, component_name in cases:
|
||||
with self.subTest(flag_name=flag_name, enabled=True):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={"performance_mode": "manual", flag_name: True},
|
||||
)
|
||||
self.assertEqual(args.residency_mode(component_name), COMPONENT_OFFLOAD)
|
||||
with self.subTest(flag_name=flag_name, enabled=False):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={"performance_mode": "manual", flag_name: False},
|
||||
)
|
||||
self.assertEqual(args.residency_mode(component_name), RESIDENT)
|
||||
|
||||
def test_component_residency_normalizes_assignments(self):
|
||||
self.assertEqual(
|
||||
normalize_component_residency(
|
||||
[
|
||||
"all=resident,dit=layerwise_offload",
|
||||
"transformer_2=component-offload",
|
||||
]
|
||||
),
|
||||
{
|
||||
"all": RESIDENT,
|
||||
"dit": LAYERWISE_OFFLOAD,
|
||||
"transformer_2": COMPONENT_OFFLOAD,
|
||||
},
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "COMPONENT=MODE"):
|
||||
normalize_component_residency(["dit"])
|
||||
with self.assertRaisesRegex(ValueError, "Invalid component residency mode"):
|
||||
normalize_component_residency(["dit=cpu"])
|
||||
|
||||
def test_component_residency_resolves_exact_group_and_all_precedence(self):
|
||||
assignments = normalize_component_residency(
|
||||
[
|
||||
"all=component-offload",
|
||||
"dit=layerwise-offload",
|
||||
"transformer_2=resident",
|
||||
"connectors=resident",
|
||||
]
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
resolve_component_residency_mode("transformer", assignments),
|
||||
LAYERWISE_OFFLOAD,
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_component_residency_mode("transformer_2", assignments),
|
||||
RESIDENT,
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_component_residency_mode("text_encoder_2", assignments),
|
||||
COMPONENT_OFFLOAD,
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_component_residency_mode("connectors", assignments), RESIDENT
|
||||
)
|
||||
|
||||
def test_component_residency_groups_exclude_legacy_helpers(self):
|
||||
assignments = normalize_component_residency(
|
||||
["dit=component-offload", "vae=component-offload"]
|
||||
)
|
||||
|
||||
for component_name in (
|
||||
"connectors",
|
||||
"dual_tower_bridge",
|
||||
"vision_language_encoder",
|
||||
"condition_image_encoder",
|
||||
"sound_tokenizer",
|
||||
"spatial_upsampler",
|
||||
"vocoder",
|
||||
):
|
||||
with self.subTest(component_name=component_name):
|
||||
self.assertIsNone(
|
||||
resolve_component_residency_mode(component_name, assignments)
|
||||
)
|
||||
|
||||
def test_component_residency_groups_include_native_texture_models(self):
|
||||
assignments = normalize_component_residency(
|
||||
["dit=layerwise-offload", "vae=component-offload"]
|
||||
)
|
||||
|
||||
for component_name in ("paint_transformer", "delight_transformer"):
|
||||
with self.subTest(component_name=component_name):
|
||||
self.assertEqual(
|
||||
resolve_component_residency_mode(component_name, assignments),
|
||||
LAYERWISE_OFFLOAD,
|
||||
)
|
||||
for component_name in ("paint_vae", "delight_vae"):
|
||||
with self.subTest(component_name=component_name):
|
||||
self.assertEqual(
|
||||
resolve_component_residency_mode(component_name, assignments),
|
||||
COMPONENT_OFFLOAD,
|
||||
)
|
||||
|
||||
def test_component_residency_overrides_matching_legacy_flags_only(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={
|
||||
"performance_mode": "manual",
|
||||
"component_residency": ["dit=resident"],
|
||||
"dit_cpu_offload": True,
|
||||
"text_encoder_cpu_offload": True,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("text_encoder"), COMPONENT_OFFLOAD)
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertTrue(args.text_encoder_cpu_offload)
|
||||
|
||||
def test_explicit_false_layerwise_keeps_dit_resident(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "Qwen/Qwen-Image",
|
||||
"dit_layerwise_offload": False,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
|
||||
def test_explicit_false_layerwise_preserves_other_explicit_dit_mode(self):
|
||||
for kwargs, expected_mode in (
|
||||
(
|
||||
{
|
||||
"dit_layerwise_offload": False,
|
||||
"dit_cpu_offload": True,
|
||||
},
|
||||
)
|
||||
COMPONENT_OFFLOAD,
|
||||
),
|
||||
(
|
||||
{
|
||||
"dit_layerwise_offload": False,
|
||||
"layerwise_offload_components": ["dit"],
|
||||
},
|
||||
LAYERWISE_OFFLOAD,
|
||||
),
|
||||
):
|
||||
with self.subTest(expected_mode=expected_mode):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={"model_path": "Qwen/Qwen-Image", **kwargs},
|
||||
)
|
||||
self.assertEqual(args.residency_mode("transformer"), expected_mode)
|
||||
|
||||
def test_exact_canonical_residency_preserves_unmatched_legacy_dit_scope(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "Qwen/Qwen-Image",
|
||||
"component_residency": ["transformer=resident"],
|
||||
"dit_layerwise_offload": False,
|
||||
"dit_cpu_offload": True,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("transformer_2"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("connectors"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_explicit_layerwise_takes_precedence_over_legacy_cpu_offload(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "Qwen/Qwen-Image",
|
||||
"performance_mode": "manual",
|
||||
"dit_cpu_offload": True,
|
||||
"dit_layerwise_offload": True,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), LAYERWISE_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("connectors"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_component_residency_overrides_only_matching_legacy_components(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={
|
||||
"performance_mode": "manual",
|
||||
"component_residency": ["text_encoder=resident"],
|
||||
"dit_cpu_offload": True,
|
||||
"text_encoder_cpu_offload": True,
|
||||
"image_encoder_cpu_offload": True,
|
||||
"vae_cpu_offload": True,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.residency_mode("text_encoder"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("image_encoder"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("vae"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_component_residency_fsdp_decision_is_component_scoped(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
kwargs={
|
||||
"performance_mode": "manual",
|
||||
"use_fsdp_inference": True,
|
||||
"component_residency": [
|
||||
"transformer=layerwise-offload",
|
||||
"text_encoder=resident",
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
self.assertFalse(args.should_use_fsdp_for_component("transformer"))
|
||||
self.assertTrue(args.should_use_fsdp_for_component("text_encoder"))
|
||||
|
||||
args.disable_fsdp_for_component("text_encoder")
|
||||
self.assertFalse(args.should_use_fsdp_for_component("text_encoder"))
|
||||
|
||||
def test_diffusers_component_residency_is_pipeline_wide(self):
|
||||
self.assertFalse(resolve_diffusers_pipeline_offload({"all": RESIDENT}))
|
||||
self.assertTrue(resolve_diffusers_pipeline_offload({"all": COMPONENT_OFFLOAD}))
|
||||
with self.assertRaisesRegex(ValueError, "pipeline-wide"):
|
||||
resolve_diffusers_pipeline_offload({"dit": COMPONENT_OFFLOAD})
|
||||
with self.assertRaisesRegex(ValueError, "native SGLang backend"):
|
||||
resolve_diffusers_pipeline_offload({"all": LAYERWISE_OFFLOAD})
|
||||
|
||||
def test_vae_cpu_offload_defaults_false_on_low_memory_gpu(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
@@ -953,7 +1192,7 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
self.assertFalse(args.text_encoder_cpu_offload)
|
||||
self.assertEqual(args.layerwise_offload_components, ["image_encoder", "vae"])
|
||||
|
||||
def test_layerwise_components_disable_matching_non_dit_cpu_offloads(self):
|
||||
def test_layerwise_components_override_matching_cpu_offload_modes(self):
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2V,
|
||||
memory_gb=16,
|
||||
@@ -963,32 +1202,30 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
"text_encoder_cpu_offload": True,
|
||||
"image_encoder_cpu_offload": True,
|
||||
"vae_cpu_offload": True,
|
||||
"layerwise_offload_components": [
|
||||
"text_encoder",
|
||||
"image_encoder",
|
||||
"video_dit",
|
||||
"vae",
|
||||
],
|
||||
},
|
||||
)
|
||||
args.layerwise_offload_components = [
|
||||
"text_encoder",
|
||||
"image_encoder",
|
||||
"video_dit",
|
||||
"vae",
|
||||
]
|
||||
args._adjust_layerwise_offload_components()
|
||||
|
||||
self.assertTrue(args.layerwise_offload_components)
|
||||
# dit_cpu_offload is complementary to DiT layerwise offload (keeps
|
||||
# weights off-device during load), so it must be preserved here.
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertFalse(args.text_encoder_cpu_offload)
|
||||
self.assertFalse(args.image_encoder_cpu_offload)
|
||||
self.assertFalse(args.vae_cpu_offload)
|
||||
for component_name in (
|
||||
"text_encoder",
|
||||
"image_encoder",
|
||||
"video_dit",
|
||||
"vae",
|
||||
):
|
||||
self.assertEqual(args.residency_mode(component_name), LAYERWISE_OFFLOAD)
|
||||
|
||||
def test_dit_layerwise_offload_preserves_dit_cpu_offload(self):
|
||||
"""Combining --dit-cpu-offload with --dit-layerwise-offload must keep both on.
|
||||
|
||||
dit_cpu_offload controls initial residency (host memory), while
|
||||
dit_layerwise_offload only swaps layers on/off device at inference.
|
||||
Force-disabling dit_cpu_offload here would push the full DiT to GPU at
|
||||
load time and OOM low-VRAM cards.
|
||||
"""
|
||||
def test_legacy_layerwise_wins_over_component_offload(self):
|
||||
"""Each component resolves to one mode under mixed legacy flags."""
|
||||
args = self._from_dict_with_task_type(
|
||||
ModelTaskType.T2I,
|
||||
memory_gb=32,
|
||||
@@ -1001,6 +1238,8 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertTrue(args.dit_layerwise_offload)
|
||||
self.assertEqual(args.layerwise_offload_components, ["dit"])
|
||||
self.assertEqual(args.residency_mode("transformer"), LAYERWISE_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("connectors"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_explicit_layerwise_false_keeps_independent_auto_residency(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
@@ -1013,6 +1252,7 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
|
||||
self.assertFalse(args.dit_layerwise_offload)
|
||||
self.assertFalse(args.dit_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
|
||||
def test_explicit_dit_cpu_offload_is_preserved_by_auto_residency(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
@@ -1026,8 +1266,9 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
|
||||
self.assertFalse(args.dit_layerwise_offload)
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_explicit_layerwise_true_preserves_initial_dit_residency(self):
|
||||
def test_explicit_layerwise_true_wins_over_auto_component_offload(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
@@ -1038,6 +1279,7 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
|
||||
self.assertTrue(args.dit_layerwise_offload)
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), LAYERWISE_OFFLOAD)
|
||||
|
||||
def test_explicit_vae_cpu_offload_is_preserved_by_auto_residency(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
@@ -1061,10 +1303,23 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertTrue(args.vae_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("vae"), COMPONENT_OFFLOAD)
|
||||
|
||||
def test_explicit_dit_layerwise_component_preserves_initial_residency(self):
|
||||
def test_cpu_offload_component_selector_keeps_unmatched_auto_defaults(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"performance_mode": "auto",
|
||||
"cpu_offload_components": ["vae"],
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.residency_mode("vae"), COMPONENT_OFFLOAD)
|
||||
self.assertEqual(args.residency_mode("transformer"), RESIDENT)
|
||||
self.assertEqual(args.residency_mode("text_encoder"), LAYERWISE_OFFLOAD)
|
||||
|
||||
def test_explicit_dit_layerwise_component_wins_over_auto_residency(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
@@ -1075,6 +1330,7 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertEqual(args.layerwise_offload_components, ["dit"])
|
||||
self.assertEqual(args.residency_mode("transformer"), LAYERWISE_OFFLOAD)
|
||||
|
||||
def test_pipeline_configs_declare_auto_tune_hints(self):
|
||||
qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config()
|
||||
@@ -1412,10 +1668,8 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
|
||||
self.assertTrue(args.layerwise_offload_components)
|
||||
self.assertFalse(args.use_fsdp_inference)
|
||||
# dit_cpu_offload is complementary to DiT layerwise offload:
|
||||
# layerwise only moves layers on/off device at runtime, while
|
||||
# dit_cpu_offload keeps the initial weights on host memory.
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertEqual(args.residency_mode("transformer"), LAYERWISE_OFFLOAD)
|
||||
self.assertFalse(args.text_encoder_cpu_offload)
|
||||
self.assertFalse(args.image_encoder_cpu_offload)
|
||||
self.assertEqual(args.dit_offload_prefetch_size, 2)
|
||||
@@ -1630,8 +1884,6 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
|
||||
# dit_cpu_offload defaults to True from _adjust_offload and is now
|
||||
# preserved alongside DiT layerwise offload (the two are complementary).
|
||||
self.assertTrue(args.dit_cpu_offload)
|
||||
self.assertEqual(args.layerwise_offload_components, ["dit"])
|
||||
|
||||
@@ -2710,10 +2962,19 @@ class TestDirectGpuWeightLoading(unittest.TestCase):
|
||||
def _args(self) -> ServerArgs:
|
||||
args = ServerArgs.__new__(ServerArgs)
|
||||
args.direct_gpu_weight_loading = True
|
||||
args.component_residency = None
|
||||
args.cpu_offload_components = None
|
||||
args.dit_cpu_offload = False
|
||||
args.dit_layerwise_offload = False
|
||||
args.layerwise_offload_components = []
|
||||
args.text_encoder_cpu_offload = False
|
||||
args.image_encoder_cpu_offload = False
|
||||
args.vae_cpu_offload = False
|
||||
args.use_fsdp_inference = False
|
||||
args.tp_size = 1
|
||||
args._explicit_arg_names = set()
|
||||
args._required_resident_components = set()
|
||||
args._component_layerwise_capabilities = {}
|
||||
return args
|
||||
|
||||
def test_cli_defaults_off_and_parses_explicit_enable(self):
|
||||
|
||||
@@ -249,19 +249,19 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
plan = WeightLoadPlan.for_component(
|
||||
checkpoint_load_device=device,
|
||||
needs_device_weight_postprocess=True,
|
||||
component_cpu_offload=True,
|
||||
component_starts_on_cpu=True,
|
||||
)
|
||||
|
||||
self.assertEqual(plan.checkpoint_load_device, device)
|
||||
self.assertEqual(plan.weight_postprocess_device, device)
|
||||
self.assertTrue(plan.defer_component_cpu_offload)
|
||||
self.assertTrue(plan.defer_cpu_placement)
|
||||
self.assertFalse(plan.load_full_state_dict_on_device)
|
||||
|
||||
def test_weight_load_plan_can_keep_full_state_dict_on_device(self):
|
||||
plan = WeightLoadPlan.for_component(
|
||||
checkpoint_load_device=torch.device("cuda:0"),
|
||||
needs_device_weight_postprocess=False,
|
||||
component_cpu_offload=False,
|
||||
component_starts_on_cpu=False,
|
||||
load_full_state_dict_on_device=True,
|
||||
)
|
||||
|
||||
@@ -270,7 +270,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
def test_unquantized_cpu_offload_loads_checkpoint_on_cpu(self):
|
||||
device = _resolve_checkpoint_load_device(
|
||||
torch.device("cuda:0"),
|
||||
component_cpu_offload=True,
|
||||
component_starts_on_cpu=True,
|
||||
runtime_quant_config=None,
|
||||
)
|
||||
|
||||
@@ -280,7 +280,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
runtime_device = torch.device("cuda:0")
|
||||
device = _resolve_checkpoint_load_device(
|
||||
runtime_device,
|
||||
component_cpu_offload=True,
|
||||
component_starts_on_cpu=True,
|
||||
runtime_quant_config=object(),
|
||||
)
|
||||
|
||||
@@ -290,7 +290,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
runtime_device = torch.device("cuda:0")
|
||||
device = _resolve_checkpoint_load_device(
|
||||
runtime_device,
|
||||
component_cpu_offload=False,
|
||||
component_starts_on_cpu=False,
|
||||
runtime_quant_config=None,
|
||||
)
|
||||
|
||||
@@ -314,6 +314,13 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
|
||||
warning.assert_called_once()
|
||||
|
||||
def test_modelopt_fp8_serialized_checkpoint_needs_device_postprocess(self):
|
||||
self.assertTrue(
|
||||
_needs_device_weight_postprocess(
|
||||
ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
|
||||
)
|
||||
)
|
||||
|
||||
def test_online_fp8_needs_device_weight_postprocess(self):
|
||||
self.assertTrue(_needs_device_weight_postprocess(Fp8Config()))
|
||||
self.assertFalse(
|
||||
|
||||
Reference in New Issue
Block a user