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

[diffusion] chore: preserve exact component identity during loading (#36875)

Browse Source
This commit is contained in:
Mick
2026-08-31 14:17:57 +08:00
committed by GitHub
parent 5e679b0cad
commit 2cb3f32b03
10 changed files with 498 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/base.py
+2 -1
View File
@@ -7,7 +7,7 @@ import os
from collections.abc import Callable
from dataclasses import asdict, dataclass, field, fields
from enum import Enum, auto
from typing import Any
from typing import Any, ClassVar
import numpy as np
import PIL
@@ -198,6 +198,7 @@ def maybe_unpad_latents(latents, batch):
class PipelineConfig:
"""The base configuration class for a generation pipeline."""
native_only_components: ClassVar[tuple[str, ...]] = ()
task_type: ModelTaskType = ModelTaskType.I2I
skip_input_image_preprocess: bool = False
# Components that cannot fall back to a native Transformers/Diffusers
python/sglang/multimodal_gen/runtime/disaggregation/roles.py
+20 -2
View File
@@ -1,11 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
"""Role definitions for diffusion pipeline disaggregation."""
from collections.abc import Mapping
from enum import Enum
_ROLE_ALIASES = {"denoising": "denoiser"}
def _matches_component_type(component_name: str, component_type: str) -> bool:
prefix = f"{component_type}_"
return component_name == component_type or (
component_name.startswith(prefix) and component_name[len(prefix) :].isdigit()
)
class RoleType(str, Enum):
MONOLITHIC = "monolithic"
ENCODER = "encoder"
@@ -86,21 +94,31 @@ def filter_modules_for_role(
role: "RoleType",
*,
extra_allowed_modules: set[str] | None = None,
structural_component_names: Mapping[str, str] | None = None,
) -> list[str]:
"""Filter module names to only those needed by the given role."""
if role in (RoleType.MONOLITHIC, RoleType.SERVER):
return module_names
extra_allowed_modules = extra_allowed_modules or set()
structural_component_names = structural_component_names or {}
filtered = []
for name in module_names:
module_role = get_module_role(name)
structural_name = structural_component_names.get(name, name)
module_role = get_module_role(structural_name)
if module_role is None:
filtered.append(name)
elif module_role == role:
filtered.append(name)
elif name in extra_allowed_modules:
elif (
name in extra_allowed_modules
or structural_name in extra_allowed_modules
or any(
_matches_component_type(structural_name, component_type)
for component_type in extra_allowed_modules
)
):
filtered.append(name)
return filtered
python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py
+2 -1
View File
@@ -71,7 +71,8 @@ class AdapterLoader(PlainStateDictComponentLoader):
server_args, component_name, precision_attr="dit_precision"
)
config_cls = self._CONFIG_CLASSES[component_name]
component_type = self.structural_component_type(component_name)
config_cls = self._CONFIG_CLASSES[component_type]
with set_default_torch_dtype(default_dtype), skip_init_modules():
adapter_cfg = config_cls()
adapter_cfg.update_model_arch(config)
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+56 -16
View File
@@ -156,8 +156,17 @@ class ComponentLoader(ABC):
def __init__(self, device=None) -> None:
self.device = device
self.component_architecture: str | None = None
self.component_type: str | None = None
self._native_load_manages_placement = False
def structural_component_name(self, component_name: str) -> str:
"""Return the config slot without changing the exact policy key."""
return self.component_type or component_name
def structural_component_type(self, component_name: str) -> str:
"""Return the normalized loader role for an exact component key."""
return _normalize_component_type(self.structural_component_name(component_name))
@staticmethod
def target_device(component_starts_on_cpu: bool) -> torch.device:
if component_starts_on_cpu:
@@ -190,10 +199,23 @@ class ComponentLoader(ABC):
) -> bool:
return self.supports_direct_gpu_weight_loading
def is_native_only_component(
self, server_args: ServerArgs, component_name: str
) -> bool:
native_only_components = server_args.pipeline_config.native_only_components
return any(
name in native_only_components
for name in (
component_name,
self.structural_component_name(component_name),
self.structural_component_type(component_name),
)
)
def should_raise_customized_load_error(
self, server_args: ServerArgs, component_name: str
) -> bool:
return component_name in server_args.pipeline_config.native_only_components
return self.is_native_only_component(server_args, component_name)
def validate_native_fallback(
self, _server_args: ServerArgs, _component_name: str
@@ -429,11 +451,19 @@ class ComponentLoader(ABC):
"""
Load the component using the native library (transformers/diffusers).
"""
precision = (
resolve_component_precision(server_args, component_name)
if component_name is not None
else None
)
precision = None
if component_name is not None:
precision_names = dict.fromkeys(
(
component_name,
self.structural_component_name(component_name),
self.structural_component_type(component_name),
)
)
for precision_name in precision_names:
precision = resolve_component_precision(server_args, precision_name)
if precision is not None:
break
load_kwargs = {}
if precision is not None:
load_kwargs["torch_dtype"] = precision
@@ -571,7 +601,7 @@ class ComponentLoader(ABC):
@classmethod
def for_component_type(
cls,
component_name: str,
component_type: str,
transformers_or_diffusers: str,
component_architecture: str | None = None,
) -> "ComponentLoader":
@@ -579,37 +609,43 @@ class ComponentLoader(ABC):
Factory method to create a component loader for a specific component type.
Args:
component_name: Type of component (e.g., "vae", "text_encoder", "transformer", "scheduler")
component_type: Structural role (e.g. "vae" or "text_encoder")
transformers_or_diffusers: Whether the component is from transformers or diffusers
"""
cls._ensure_loaders_registered()
# Map of component types to their loader classes and expected library
component_name = _normalize_component_type(component_name)
structural_component_name = component_type
loader_type = _normalize_component_type(component_type)
transformers_or_diffusers = cls.resolve_transformers_or_diffusers(
transformers_or_diffusers, component_name
transformers_or_diffusers, loader_type
)
if component_name in component_name_to_loader_cls:
if loader_type in component_name_to_loader_cls:
loader_cls: Type[ComponentLoader] = component_name_to_loader_cls[
component_name
loader_type
]
expected_library = loader_cls.expected_library
# Assert that the library matches what's expected for this component type
assert (
transformers_or_diffusers == expected_library
), f"{component_name} must be loaded from {expected_library}, got {transformers_or_diffusers}"
), f"{loader_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
loader = loader_cls()
loader.component_type = structural_component_name
loader.component_architecture = component_architecture
return loader
# For unknown component types, use a generic loader
logger.warning(
"No specific loader found for component type: %s. Using generic loader.",
component_name,
loader_type,
)
return GenericComponentLoader(transformers_or_diffusers, component_architecture)
loader = GenericComponentLoader(
transformers_or_diffusers, component_architecture
)
loader.component_type = structural_component_name
return loader
class PlainStateDictComponentLoader(ComponentLoader):
@@ -759,6 +795,7 @@ class PipelineComponentLoader:
component_architecture: str | None = None,
component_attn_backend: Any = None,
component_attn_name: str | None = None,
component_type: str | None = None,
):
"""
Load a pipeline component.
@@ -768,11 +805,14 @@ class PipelineComponentLoader:
component_model_path: Path to the component model
transformers_or_diffusers: Whether the component is from transformers or diffusers
component_architecture: the class name of the module
component_type: structural config slot when it differs from the exact key
"""
# Get the appropriate loader for this component type
loader = ComponentLoader.for_component_type(
component_name, transformers_or_diffusers, component_architecture
component_type or component_name,
transformers_or_diffusers,
component_architecture,
)
try:
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+3 -1
View File
@@ -744,7 +744,9 @@ class TextEncoderLoader(ComponentLoader):
)
# TODO(mick): had to throw an exception for different text-encoder arch
encoder_index = self._extract_encoder_index(component_name)
encoder_index = self._extract_encoder_index(
self.structural_component_name(component_name)
)
assert encoder_index < len(
server_args.pipeline_config.text_encoder_configs
) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions)
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+48 -41
View File
@@ -29,7 +29,6 @@ from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
resolve_transformer_gguf_to_load,
resolve_transformer_quant_load_spec,
)
from sglang.multimodal_gen.runtime.loader.utils import _normalize_component_type
from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
from sglang.multimodal_gen.runtime.platforms import (
@@ -96,56 +95,60 @@ def _warn_if_expected_param_dtype_missing(
def _server_args_for_transformer_component(
server_args: ServerArgs, component_name: str
server_args: ServerArgs,
component_name: str,
structural_component_name: str | None = None,
) -> ServerArgs:
"""Mask global quantized override flags for secondary transformer components."""
structural_component_name = structural_component_name or component_name
_, separator, suffix = structural_component_name.rpartition("_")
is_secondary = structural_component_name == "unconditional_transformer" or (
bool(separator) and suffix.isdigit() and int(suffix) >= 2
)
component_weights_path = server_args.component_weights_paths.get(component_name)
component_quantization = server_args.component_quantizations.get(component_name)
component_ignored_layers = server_args.component_quantization_ignored_layers.get(
component_name
)
if (
has_exact_override = (
component_weights_path is not None
or component_quantization is not None
or component_ignored_layers is not None
):
component_server_args = copy.copy(server_args)
if component_weights_path is not None:
component_server_args.transformer_weights_path = component_weights_path
component_server_args.nunchaku_config = None
logger.info(
"Using transformer_weights_path override for %s: %s",
component_name,
component_weights_path,
)
if component_quantization is not None:
component_server_args.quantization = component_quantization
logger.info(
"Using quantization override %s for %s",
component_quantization,
component_name,
)
if component_ignored_layers is not None:
component_server_args.quantization_ignored_layers = component_ignored_layers
return component_server_args
if component_name not in ("transformer_2", "unconditional_transformer"):
return server_args
if (
server_args.transformer_weights_path is None
and server_args.nunchaku_config is None
):
)
has_global_weights = (
server_args.transformer_weights_path is not None
or server_args.nunchaku_config is not None
)
if not has_exact_override and not (is_secondary and has_global_weights):
return server_args
component_server_args = copy.copy(server_args)
component_server_args.transformer_weights_path = None
component_server_args.nunchaku_config = None
logger.info(
"Ignoring global transformer_weights_path for %s; keep it on the base "
"checkpoint unless a per-component override path is provided.",
component_name,
)
if is_secondary:
component_server_args.transformer_weights_path = None
component_server_args.nunchaku_config = None
if has_global_weights:
logger.info(
"Ignoring global transformer weight overrides for %s; keep them "
"on the primary component unless an exact override is provided.",
component_name,
)
if component_weights_path is not None:
component_server_args.transformer_weights_path = component_weights_path
component_server_args.nunchaku_config = None
logger.info(
"Using transformer_weights_path override for %s: %s",
component_name,
component_weights_path,
)
if component_quantization is not None:
component_server_args.quantization = component_quantization
logger.info(
"Using quantization override %s for %s",
component_quantization,
component_name,
)
if component_ignored_layers is not None:
component_server_args.quantization_ignored_layers = component_ignored_layers
return component_server_args
@@ -183,7 +186,9 @@ class TransformerLoader(ComponentLoader):
self, server_args: ServerArgs, component_name: str
) -> bool:
component_server_args = _server_args_for_transformer_component(
server_args, component_name
server_args,
component_name,
self.structural_component_name(component_name),
)
# Don't let a quantized load quietly fall back to the unquantized native
# model. That would drop the requested precision and bury the real error.
@@ -238,7 +243,9 @@ class TransformerLoader(ComponentLoader):
):
"""Load the transformer based on the model path, and inference args."""
component_server_args = _server_args_for_transformer_component(
server_args, component_name
server_args,
component_name,
self.structural_component_name(component_name),
)
# 1. hf config
@@ -261,7 +268,7 @@ class TransformerLoader(ComponentLoader):
# 2. dit config
# Config from Diffusers supersedes sgl_diffusion's model config
component_type = _normalize_component_type(component_name)
component_type = self.structural_component_type(component_name)
server_args.model_paths[component_name] = component_model_path
if component_type in (
"transformer",
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+51 -21
View File
@@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
)
from sglang.multimodal_gen.runtime.loader.utils import (
_list_safetensors_files,
_normalize_component_type,
checkpoint_bytes,
keep_checkpoint_mapped,
set_default_torch_dtype,
@@ -96,9 +97,9 @@ def _require_native_loader_for_quantized_vae(
def _backfill_ltx2_audio_vae_latent_stats(
loaded: dict[str, torch.Tensor], component_name: str
loaded: dict[str, torch.Tensor], component_type: str
) -> None:
if component_name != "audio_vae":
if component_type != "audio_vae":
return
mean_key = "per_channel_statistics.mean-of-means"
std_key = "per_channel_statistics.std-of-means"
@@ -128,9 +129,9 @@ def _convert_conv3d_weights_to_channels_last_3d(module: nn.Module) -> int:
def _should_use_channels_last_3d(
server_args: ServerArgs | None, component_name: str
server_args: ServerArgs | None, component_type: str
) -> bool:
if component_name not in (
if component_type not in (
"vae",
"video_vae",
) or not (current_platform.is_cuda() or current_platform.is_rocm()):
@@ -232,7 +233,12 @@ def _rehome_cast_weights_to_file(
def _hold_decoder_weights_in_decode_dtype(
vae, server_args: ServerArgs, component_name: str, component_model_path: str = ""
vae,
server_args: ServerArgs,
component_name: str,
component_model_path: str = "",
*,
component_type: str | None = None,
) -> None:
"""Round decoder weights to their decode compute dtype at load.
@@ -244,7 +250,8 @@ def _hold_decoder_weights_in_decode_dtype(
restreaming a third of it per tile and holding all 36 blocks on a 12 GiB
card for the decode.
"""
if component_name not in ("vae", "video_vae"):
component_type = component_type or _normalize_component_type(component_name)
if component_type not in ("vae", "video_vae"):
return
if envs.SGLANG_DIFFUSION_DISABLE_EARLY_VAE_DECODER_CAST:
return
@@ -398,6 +405,21 @@ class VAELoader(ComponentLoader):
) -> bool:
return component_name in ("vae", "video_vae")
def select_weight_files(
self,
safetensors_list: list[str],
component_model_path: str,
server_args: ServerArgs,
component_name: str,
vae_precision: str,
) -> list[str]:
return server_args.pipeline_config.select_vae_weight_files(
safetensors_list=safetensors_list,
component_model_path=component_model_path,
component_name=self.structural_component_type(component_name),
vae_precision=vae_precision,
)
def component_load_precision(
self, server_args: ServerArgs, component_name: str
) -> str | None:
@@ -459,9 +481,7 @@ class VAELoader(ComponentLoader):
)
config = get_diffusers_component_config(component_path=component_model_path)
server_args.model_paths[component_name] = component_model_path
native_only = component_name in getattr(
server_args.pipeline_config, "native_only_components", ()
)
native_only = self.is_native_only_component(server_args, component_name)
_require_native_loader_for_quantized_vae(
config,
component_name,
@@ -474,10 +494,11 @@ class VAELoader(ComponentLoader):
class_name is not None
), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
if component_name in ("vae", "video_vae"):
component_type = self.structural_component_type(component_name)
if component_type in ("vae", "video_vae"):
pipeline_vae_config_attr = "vae_config"
pipeline_vae_precision = "vae_precision"
elif component_name in ("audio_vae",):
elif component_type == "audio_vae":
pipeline_vae_config_attr = "audio_vae_config"
pipeline_vae_precision = "audio_vae_precision"
else:
@@ -529,14 +550,18 @@ class VAELoader(ComponentLoader):
trust_remote_code=server_args.trust_remote_code,
)
vae = vae.to(device=target_device, dtype=vae_dtype)
if _should_use_channels_last_3d(server_args, component_name):
if _should_use_channels_last_3d(server_args, component_type):
n = _convert_conv3d_weights_to_channels_last_3d(vae)
if n > 0:
logger.info(
"VAE: converted %d Conv3d weights to channels_last_3d", n
)
_hold_decoder_weights_in_decode_dtype(
vae, server_args, component_name, component_model_path
vae,
server_args,
component_name,
component_model_path,
component_type=component_type,
)
vae = current_platform.optimize_vae(vae)
return vae
@@ -567,11 +592,12 @@ class VAELoader(ComponentLoader):
safetensors_list = [component_weights_path]
else:
safetensors_list = _list_safetensors_files(component_weights_path)
safetensors_list = server_args.pipeline_config.select_vae_weight_files(
safetensors_list=safetensors_list,
component_model_path=component_weights_path,
component_name=component_name,
vae_precision=vae_precision,
safetensors_list = self.select_weight_files(
safetensors_list,
component_weights_path,
server_args,
component_name,
vae_precision,
)
assert (
@@ -596,7 +622,7 @@ class VAELoader(ComponentLoader):
loaded = {}
for sf_path in safetensors_list:
loaded.update(safetensors_load_file(sf_path))
_backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
_backfill_ltx2_audio_vae_latent_stats(loaded, component_type)
strict_load = native_only
# `loaded` holds views into the safetensors mapping. When the component
# starts on the CPU and the host cannot afford copies of the whole
@@ -639,13 +665,17 @@ class VAELoader(ComponentLoader):
if unexpected_keys:
logger.warning("VAE unexpected keys: %s", unexpected_keys)
if _should_use_channels_last_3d(server_args, component_name):
if _should_use_channels_last_3d(server_args, component_type):
n = _convert_conv3d_weights_to_channels_last_3d(vae)
if n > 0:
logger.info("VAE: converted %d Conv3d weights to channels_last_3d", n)
_hold_decoder_weights_in_decode_dtype(
vae, server_args, component_name, component_weights_path
vae,
server_args,
component_name,
component_weights_path,
component_type=component_type,
)
vae = current_platform.optimize_vae(vae)
return vae
python/sglang/multimodal_gen/runtime/loader/component_loaders/vl_encoder_loader.py
+5 -4
View File
@@ -22,9 +22,9 @@ class VisionLanguageEncoderLoader(ComponentLoader):
self,
component_model_path: str,
server_args: ServerArgs,
transformers_or_diffusers: str = "vision_language_encoder",
component_name: str = "vision_language_encoder",
) -> Any:
if transformers_or_diffusers == "vision_language_encoder":
if self.structural_component_type(component_name) == "vision_language_encoder":
if server_args.srt_encoder_url is not None:
health_url = server_args.srt_encoder_url.rstrip("/") + "/health"
@@ -59,7 +59,7 @@ class VisionLanguageEncoderLoader(ComponentLoader):
revision=server_args.revision,
)
target_device = self.target_device(
server_args.should_start_component_on_cpu("vision_language_encoder")
server_args.should_start_component_on_cpu(component_name)
)
model = GlmImageForConditionalGeneration.from_pretrained(
component_model_path,
@@ -70,5 +70,6 @@ class VisionLanguageEncoderLoader(ComponentLoader):
return model
else:
raise ValueError(
f"Unsupported library for VisionLanguageEncoder: {transformers_or_diffusers}"
f"Unsupported component type for VisionLanguageEncoder: "
f"{self.structural_component_type(component_name)}"
)
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
+48 -14
View File
@@ -21,6 +21,7 @@ from sglang.multimodal_gen.runtime.disaggregation.roles import (
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PipelineComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.utils import _normalize_component_type
from sglang.multimodal_gen.runtime.managers.memory_managers.component_loading_order import (
ComponentLoadSpec,
order_component_load_specs,
@@ -73,6 +74,7 @@ class ComposedPipelineBase(ABC):
is_video_pipeline: bool = False # To be overridden by video pipelines
# should contains only the modules to be loaded
_required_config_modules: list[str] = []
_unfiltered_required_config_modules: tuple[str, ...] = ()
_extra_config_module_map: dict[str, str] = {}
server_args: ServerArgs | None = None
modules: dict[str, Any] = {}
@@ -118,7 +120,8 @@ class ComposedPipelineBase(ABC):
)
if base_required_config_modules is None:
raise NotImplementedError("Subclass must set _required_config_modules")
self._required_config_modules = list(base_required_config_modules)
self._unfiltered_required_config_modules = tuple(base_required_config_modules)
self._required_config_modules = list(self._unfiltered_required_config_modules)
self._extra_config_module_map = dict(self._extra_config_module_map)
# Filter modules based on disaggregation role
@@ -131,6 +134,7 @@ class ComposedPipelineBase(ABC):
extra_allowed_modules=self._get_extra_allowed_modules_for_role(
self._disagg_role, task_name
),
structural_component_names=self._extra_config_module_map,
)
skipped = set(original_modules) - set(self._required_config_modules)
if skipped:
@@ -308,11 +312,36 @@ class ComposedPipelineBase(ABC):
get_diffusers_component_config,
)
required = set(self.required_config_modules)
for module_name in full_model_index:
if module_name in required:
continue # will be loaded normally
cfg_attr = self._CONFIG_ATTR_MAP.get(module_name)
loaded_components = set(self.required_config_modules)
loaded_structural_components = {
self._extra_config_module_map.get(name, name)
for name in self.required_config_modules
}
skipped_components: list[tuple[str, str]] = []
seen_component_keys = loaded_components | loaded_structural_components
for component_name in self._unfiltered_required_config_modules:
structural_name = self._extra_config_module_map.get(
component_name, component_name
)
if (
component_name in seen_component_keys
or structural_name in seen_component_keys
or (
component_name not in full_model_index
and structural_name not in full_model_index
)
):
continue
skipped_components.append((component_name, structural_name))
seen_component_keys.update((component_name, structural_name))
for structural_name in full_model_index:
if structural_name not in seen_component_keys:
skipped_components.append((structural_name, structural_name))
for component_name, structural_name in skipped_components:
cfg_attr = self._CONFIG_ATTR_MAP.get(
_normalize_component_type(structural_name)
)
if cfg_attr is None:
continue # not a config we need to patch
@@ -322,7 +351,7 @@ class ComposedPipelineBase(ABC):
try:
component_path = self._resolve_component_path(
server_args, module_name, module_name
server_args, component_name, structural_name
)
hf_config = get_diffusers_component_config(
component_path=component_path
@@ -336,7 +365,7 @@ class ComposedPipelineBase(ABC):
"Disagg role=%s: initialized %s config from HF JSON "
"(spatial_compression_ratio=%s)",
self._disagg_role.value,
module_name,
component_name,
getattr(
getattr(pipeline_cfg, "arch_config", None),
"spatial_compression_ratio",
@@ -348,7 +377,7 @@ class ComposedPipelineBase(ABC):
"Disagg role=%s: failed to read HF config for skipped "
"component %s: %s",
self._disagg_role.value,
module_name,
component_name,
e,
)
@@ -460,9 +489,11 @@ class ComposedPipelineBase(ABC):
if self._disagg_role != RoleType.MONOLITHIC:
self._init_skipped_component_configs(model_index, server_args)
declared_modules = model_index
model_index = {
required_module: model_index[required_module]
for required_module in self.required_config_modules
if required_module in model_index
}
for module_name in self.required_config_modules:
@@ -477,15 +508,17 @@ class ComposedPipelineBase(ABC):
module_name,
extra_module_value,
)
if extra_module_value in model_index:
if extra_module_value in declared_modules:
logger.info(
"Using module %s for %s", extra_module_value, module_name
)
model_index[module_name] = model_index[extra_module_value]
model_index[module_name] = declared_modules[extra_module_value]
continue
else:
raise ValueError(
f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}"
f"Required module key: {module_name} value: "
f"{declared_modules.get(module_name)} was not found in "
f"declared modules {declared_modules.keys()}"
)
# all the component models used by the pipeline
@@ -581,7 +614,8 @@ class ComposedPipelineBase(ABC):
matched_backend_key,
)
module, memory_usage = PipelineComponentLoader.load_component(
component_name=load_module_name,
component_name=module_name,
component_type=load_module_name,
component_model_path=component_model_path,
transformers_or_diffusers=transformers_or_diffusers,
server_args=server_args,
@@ -590,7 +624,7 @@ class ComposedPipelineBase(ABC):
component_attn_name=matched_backend_key or module_name,
)
self.memory_usages[load_module_name] = memory_usage
self.memory_usages[module_name] = memory_usage
if module_name in loaded_components:
logger.warning("Overwriting module %s", module_name)
python/sglang/multimodal_gen/test/unit/test_component_loader_identity.py
+263
View File
@@ -0,0 +1,263 @@
# SPDX-License-Identifier: Apache-2.0
import unittest
from types import SimpleNamespace
from unittest.mock import Mock, patch
import torch
import torch.nn as nn
from sglang.multimodal_gen.runtime.disaggregation.roles import (
RoleType,
filter_modules_for_role,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.adapter_loader import (
AdapterLoader,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentLoader,
PipelineComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
_server_args_for_transformer_component,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import VAELoader
from sglang.multimodal_gen.runtime.loader.component_loaders.vl_encoder_loader import (
VisionLanguageEncoderLoader,
)
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase,
)
class _AliasPipeline(ComposedPipelineBase):
def initialize_pipeline(self, _server_args):
pass
def create_pipeline_stages(self, _server_args):
pass
class TestComponentLoaderIdentity(unittest.TestCase):
def test_structural_identity_selects_roles_and_preserves_exact_policy_keys(self):
aliases = {
"conditioning": "text_encoder_2",
"refiner": "transformer_3",
"decoder": "video_vae",
}
self.assertEqual(
filter_modules_for_role(
aliases, RoleType.ENCODER, structural_component_names=aliases
),
["conditioning"],
)
self.assertEqual(
filter_modules_for_role(
aliases, RoleType.DENOISER, structural_component_names=aliases
),
["refiner"],
)
self.assertEqual(
filter_modules_for_role(
aliases, RoleType.DECODER, structural_component_names=aliases
),
["decoder"],
)
loader = ComponentLoader.for_component_type(
"duration_head_2", "ltx2", "LTX2DurationHeadModel"
)
self.assertIsInstance(loader, AdapterLoader)
self.assertEqual(
loader.structural_component_name("auxiliary_head"), "duration_head_2"
)
self.assertEqual(
loader.structural_component_type("auxiliary_head"), "duration_head"
)
def test_declared_alias_loads_by_exact_key_and_structural_source(self):
pipeline = object.__new__(_AliasPipeline)
pipeline.model_path = "/model"
pipeline.memory_usages = {}
pipeline._disagg_role = RoleType.MONOLITHIC
pipeline._required_config_modules = ["auxiliary_head"]
pipeline._extra_config_module_map = {"auxiliary_head": "duration_head_2"}
pipeline._load_config = lambda: {
"_class_name": "TestPipeline",
"_diffusers_version": "0",
"duration_head_2": ["ltx2", "LTX2DurationHeadModel"],
"scheduler": ["diffusers", "Scheduler"],
}
server_args = SimpleNamespace(
component_paths={},
component_direct_gpu_weight_loading=set(),
resolve_component_attention_backend=lambda *_names: (None, None),
)
with patch.object(
PipelineComponentLoader,
"load_component",
return_value=(nn.Linear(1, 1), 0.25),
) as load_component:
loaded = pipeline.load_modules(server_args)
self.assertIn("auxiliary_head", loaded)
self.assertEqual(pipeline.memory_usages, {"auxiliary_head": 0.25})
load_component.assert_called_once_with(
component_name="auxiliary_head",
component_type="duration_head_2",
component_model_path="/model/duration_head_2",
transformers_or_diffusers="ltx2",
server_args=server_args,
component_architecture="LTX2DurationHeadModel",
component_attn_backend=None,
component_attn_name="auxiliary_head",
)
def test_skipped_alias_keeps_exact_override_and_structural_config(self):
pipeline = object.__new__(_AliasPipeline)
pipeline.model_path = "/model"
pipeline._disagg_role = RoleType.ENCODER
pipeline._required_config_modules = []
pipeline._unfiltered_required_config_modules = ("vae_2",)
pipeline._extra_config_module_map = {"vae_2": "video_vae_2"}
vae_config = Mock()
server_args = SimpleNamespace(
component_paths={"vae_2": "/exact/decoder"},
pipeline_config=SimpleNamespace(vae_config=vae_config),
)
with (
patch(
"sglang.multimodal_gen.runtime.pipelines_core."
"composed_pipeline_base.prepare_diffusers_component_path_for_loading",
return_value="/resolved/decoder",
) as prepare_path,
patch(
"sglang.multimodal_gen.runtime.utils.hf_diffusers_utils."
"get_diffusers_component_config",
return_value={"sample_size": 32},
) as get_config,
):
pipeline._init_skipped_component_configs(
{
"vae_2": ["diffusers", "VideoVAE"],
"video_vae_2": ["diffusers", "VideoVAE"],
},
server_args,
)
prepare_path.assert_called_once_with("/exact/decoder")
get_config.assert_called_once_with(component_path="/resolved/decoder")
vae_config.update_model_arch.assert_called_once_with({"sample_size": 32})
def test_structural_aliases_select_loader_config_and_weight_behavior(self):
selected = ["selected.safetensors"]
select_weight_files = Mock(return_value=selected)
vae_loader = VAELoader()
vae_loader.component_type = "video_vae"
server_args = SimpleNamespace(
pipeline_config=SimpleNamespace(select_vae_weight_files=select_weight_files)
)
self.assertIs(
vae_loader.select_weight_files(
["candidate.safetensors"],
"/decoder",
server_args,
"decoder",
"fp32",
),
selected,
)
select_weight_files.assert_called_once_with(
safetensors_list=["candidate.safetensors"],
component_model_path="/decoder",
component_name="video_vae",
vae_precision="fp32",
)
def test_transformer_exact_override_wins_without_leaking_global_flags(self):
server_args = SimpleNamespace(
component_weights_paths={},
component_quantizations={},
component_quantization_ignored_layers={},
transformer_weights_path="global.safetensors",
nunchaku_config="global-nunchaku",
)
secondary = _server_args_for_transformer_component(
server_args, "refiner", "transformer_3"
)
self.assertIsNone(secondary.transformer_weights_path)
self.assertIsNone(secondary.nunchaku_config)
server_args.component_weights_paths["refiner"] = "refiner.safetensors"
exact = _server_args_for_transformer_component(
server_args, "refiner", "transformer_3"
)
self.assertEqual(exact.transformer_weights_path, "refiner.safetensors")
self.assertIs(
_server_args_for_transformer_component(
server_args, "denoiser", "transformer"
),
server_args,
)
def test_native_fallback_prioritizes_exact_component_precision(self):
loader = ComponentLoader()
loader.component_type = "video_vae_2"
server_args = SimpleNamespace(
pipeline_config=SimpleNamespace(),
revision="test-revision",
trust_remote_code=False,
)
native_model = nn.Linear(1, 1)
with (
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.resolve_component_precision",
side_effect=[torch.float32],
) as resolve_precision,
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.prepare_diffusers_component_path_for_loading",
return_value="/resolved/decoder",
),
patch("diffusers.AutoModel.from_pretrained", return_value=native_model),
):
loaded = loader.load_native("/decoder", server_args, "diffusers", "decoder")
self.assertIs(loaded, native_model)
self.assertEqual(
[call.args[1] for call in resolve_precision.call_args_list], ["decoder"]
)
def test_vision_language_loader_uses_exact_residency_key(self):
requested = []
server_args = SimpleNamespace(
srt_encoder_url=None,
trust_remote_code=False,
revision=None,
should_start_component_on_cpu=lambda name: requested.append(name) or True,
)
loader = VisionLanguageEncoderLoader()
loader.component_type = "vision_language_encoder"
with (
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"vl_encoder_loader.get_hf_config",
return_value=object(),
),
patch(
"transformers.GlmImageForConditionalGeneration.from_pretrained",
return_value=nn.Linear(1, 1),
),
):
loader.load_customized("unused", server_args, "prompt_conditioner")
self.assertEqual(requested, ["prompt_conditioner"])
if __name__ == "__main__":
unittest.main()