minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] refactor: unify component residency controls (#34736)

Browse Source
This commit is contained in:
Mick
2026-08-16 11:24:48 +08:00
committed by GitHub
parent f68517f644
commit e9fe58139f
70 changed files with 2629 additions and 1423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/README.md
+18 -1
View File
@@ -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`:
python/sglang/multimodal_gen/runtime/launch_server.py
+4 -1
View File
@@ -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)
python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py
+1 -1
View File
@@ -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"
python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py
+14 -6
View File
@@ -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())
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+18 -95
View File
@@ -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
python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py
+1 -1
View File
@@ -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"
python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
-27
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py
+5 -1
View File
@@ -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(
python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py
+1 -1
View File
@@ -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():
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+17 -77
View File
@@ -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:
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+20 -14
View File
@@ -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,
python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py
+4 -2
View File
@@ -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)
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+4 -2
View File
@@ -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", ()
python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py
+1 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py
+4 -2
View File
@@ -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)
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+11 -13
View File
@@ -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
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+143 -17
View File
@@ -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:
python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py
+5 -5
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py
+1 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+8 -3
View File
@@ -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")
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py
+156 -122
View File
@@ -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)
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency.py
+162
View File
@@ -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"
)
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py
+229
View File
@@ -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)
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_resident_strategies.py
-508
View File
@@ -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)
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
+91 -25
View File
@@ -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
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload_components.py
+43 -46
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py
+10 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/models/decoders/ltx_2_5_diffusion_decoder.py
+13 -1
View File
@@ -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
)
python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py
+1
View File
@@ -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)
python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py
+8 -1
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py
+10 -3
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py
+6 -3
View File
@@ -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,
)
python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
+18 -4
View File
@@ -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()
python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py
+25 -4
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+7 -2
View File
@@ -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:
python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py
+2 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py
+6 -30
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
+28 -14
View File
@@ -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]:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+3 -5
View File
@@ -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"]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
+5 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
+50 -40
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ernie_image_pe.py
+11 -1
View File
@@ -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")]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/glm_image.py
+25 -3
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py
-10
View File
@@ -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):
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py
+24 -3
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/longcat_image.py
+1 -1
View File
@@ -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()
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py
+9 -3
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py
+7 -13
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
+1 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py
-1
View File
@@ -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)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/text_encoding.py
+13
View File
@@ -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),
python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py
+6 -10
View File
@@ -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
)
python/sglang/multimodal_gen/runtime/server_args/auto_tune.py
+36 -57
View File
@@ -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:
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+330 -122
View File
@@ -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"
python/sglang/multimodal_gen/test/single_test_file/component_accuracy/engine.py
+3 -3
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_adapter_loader_offload_target.py
+8
View File
@@ -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()
python/sglang/multimodal_gen/test/unit/test_component_residency.py
+398
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_glm_image_multi_output.py
+1
View File
@@ -50,6 +50,7 @@ class _DummyScheduler:
class _RecordingBeforeDenoisingStage(GlmImageBeforeDenoisingStage):
def __init__(self):
self.text_encoder = None
self.transformer = SimpleNamespace(
config=SimpleNamespace(
in_channels=4,
python/sglang/multimodal_gen/test/unit/test_hunyuan3d_native_texture_models.py
+17
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
+136 -9
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_ltx2_5_config.py
+30
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_minimax_h3_encoder_device.py
+3 -6
View File
@@ -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:
python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py
+5 -6
View File
@@ -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:
python/sglang/multimodal_gen/test/unit/test_pipeline_executor.py
+40 -17
View File
@@ -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",),
),
],
python/sglang/multimodal_gen/test/unit/test_server_args.py
+306 -45
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+13 -6
View File
@@ -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(