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

[diffusion] refactor: remove component loader capability switches (#36824)

Browse Source
This commit is contained in:
Mick
2026-09-02 09:48:50 +08:00
committed by GitHub
parent 66de38f30c
commit 22195709d1
17 changed files with 421 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py
+13 -7
View File
@@ -27,12 +27,14 @@ class BridgeLoader(PlainStateDictComponentLoader):
component_names = ["dual_tower_bridge"]
expected_library = "diffusers"
supports_fsdp_inference = True
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
):
config = self.load_component_config(component_model_path, component_name)
component_weights_path = self.resolve_component_weights_path(
component_model_path, server_args, component_name
)
hf_config = deepcopy(config)
class_name = config.pop("_class_name", None)
if class_name is None:
@@ -60,9 +62,9 @@ class BridgeLoader(PlainStateDictComponentLoader):
model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
# Find all safetensors files
safetensors_list = _list_safetensors_files(component_model_path)
safetensors_list = _list_safetensors_files(component_weights_path)
if not safetensors_list:
raise ValueError(f"No safetensors files found in {component_model_path}")
raise ValueError(f"No safetensors files found in {component_weights_path}")
default_dtype = resolve_precision(
server_args, component_name, precision_attr="dit_precision"
@@ -82,10 +84,14 @@ class BridgeLoader(PlainStateDictComponentLoader):
# Use the FSDP loader when FSDP is requested or shard rules are declared.
fsdp_shard_conditions = getattr(model_cls, "_fsdp_shard_conditions", None)
if use_fsdp or (
server_args.residency_mode(component_name) == RESIDENT
and server_args.hsdp_shard_dim is not None
and fsdp_shard_conditions
if (
component_weights_path != component_model_path
or 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
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+137 -77
View File
@@ -131,18 +131,6 @@ class ComponentLoader(ABC):
# diffusers or transformers
expected_library: str = ""
# --attention-backend primarily selects the DiT backend. Auxiliary
# components may fall back when that global choice is incompatible; an
# explicit --component-attention-backends entry remains strict.
allow_global_attention_backend_fallback = True
# Gates only --component-quantizations.<name>. Quantization declared by a
# checkpoint is discovered and admitted by the component's normal loader.
supports_online_quantization_override = False
# Gates only --component-direct-gpu-weight-loading.<name>. The checkpoint
# source stays component-specific because its streaming ABI is loader-owned.
supports_direct_gpu_weight_loading = False
supports_fsdp_inference = False
_loaders_registered = False
def __init_subclass__(cls, **kwargs):
@@ -194,10 +182,56 @@ class ComponentLoader(ABC):
)
return None
def supports_direct_gpu_weight_loading_for_component(
self, _component_name: str
def resolve_component_weight_override(
self, server_args: ServerArgs, component_name: str
) -> str | None:
"""Return the consumed weights-only override or reject it."""
override = server_args.component_weights_paths.get(component_name)
if override is not None:
raise ComponentCheckpointUnsupportedError(
f"{component_name!r} does not support a weights-only override; "
f"use --component-paths.{component_name} to replace its config "
"and weights together"
)
return None
def resolve_component_quantization_override(
self, server_args: ServerArgs, component_name: str
) -> str | None:
"""Return the consumed online quantization override or reject it."""
quantization = server_args.component_quantizations.get(component_name)
if quantization is not None:
raise ComponentCheckpointUnsupportedError(
f"{component_name!r} does not support an explicit quantization "
"override; use a self-describing quantized component checkpoint "
"when supported"
)
return None
def resolve_component_direct_gpu_loading(
self, server_args: ServerArgs, component_name: str
) -> bool:
return self.supports_direct_gpu_weight_loading
"""Return whether this load consumes the exact direct-GPU request."""
requested = server_args.should_direct_gpu_weight_load_component(component_name)
if requested:
raise ComponentCheckpointUnsupportedError(
f"{component_name!r} does not support direct GPU weight loading"
)
return False
def component_attention_backend_context(
self,
attn_backend: Any,
component_attn_name: str | None,
require_backend_selection: bool,
):
"""Build the attention-selection context used by this loader."""
return component_attn_backend_context_manager(
attn_backend,
component_name=component_attn_name,
allow_global_backend_fallback=True,
require_backend_selection=require_backend_selection,
)
def is_native_only_component(
self, server_args: ServerArgs, component_name: str
@@ -223,15 +257,6 @@ class ComponentLoader(ABC):
"""Validate that fallback preserves the exact component's runtime contract."""
pass
def disable_unsupported_component_fsdp(
self, server_args: ServerArgs, component_name: str
) -> None:
if (
not self.supports_fsdp_inference
and server_args.should_use_fsdp_for_component(component_name)
):
server_args.disable_fsdp_for_component(component_name)
def _load_customized_with_context(
self,
component_model_path: str,
@@ -239,14 +264,12 @@ class ComponentLoader(ABC):
component_name: str,
attn_backend: Any,
component_attn_name: str | None,
allow_global_backend_fallback: bool,
require_backend_selection: bool,
) -> AutoModel:
with component_attn_backend_context_manager(
with self.component_attention_backend_context(
attn_backend,
component_name=component_attn_name,
allow_global_backend_fallback=allow_global_backend_fallback,
require_backend_selection=require_backend_selection,
component_attn_name,
require_backend_selection,
):
load_kwargs = self.customized_load_kwargs_for_component(
server_args, component_name
@@ -263,14 +286,12 @@ class ComponentLoader(ABC):
transformers_or_diffusers: str,
attn_backend: Any,
component_attn_name: str | None,
allow_global_backend_fallback: bool,
require_backend_selection: bool,
) -> AutoModel:
with component_attn_backend_context_manager(
with self.component_attention_backend_context(
attn_backend,
component_name=component_attn_name,
allow_global_backend_fallback=allow_global_backend_fallback,
require_backend_selection=require_backend_selection,
component_attn_name,
require_backend_selection,
):
component = self.load_native(
component_model_path,
@@ -300,23 +321,12 @@ class ComponentLoader(ABC):
"""
self._native_load_manages_placement = False
self.component_load_precision(server_args, component_name)
if server_args.should_direct_gpu_weight_load_component(
component_name
) and not self.supports_direct_gpu_weight_loading_for_component(component_name):
raise ComponentCheckpointUnsupportedError(
f"{component_name!r} does not support direct GPU weight loading"
)
self.disable_unsupported_component_fsdp(server_args, component_name)
component_quantization = server_args.component_quantizations.get(component_name)
if (
component_quantization is not None
and not self.supports_online_quantization_override
):
raise ValueError(
f"{component_name!r} does not support an explicit quantization "
"override; "
"use a self-describing quantized component checkpoint when supported"
)
component_weight_override = self.resolve_component_weight_override(
server_args, component_name
)
self.resolve_component_quantization_override(server_args, component_name)
self.resolve_component_direct_gpu_loading(server_args, component_name)
fsdp_requested = server_args.should_use_fsdp_for_component(component_name)
gpu_mem_before_loading = current_platform.get_available_gpu_memory()
logger.info(
@@ -361,7 +371,6 @@ class ComponentLoader(ABC):
component_name,
component_attn_backend,
component_attn_name,
self.allow_global_attention_backend_fallback,
require_backend_selection,
)
source = "sgl-diffusion"
@@ -375,9 +384,26 @@ class ComponentLoader(ABC):
if require_backend_selection:
raise
native_loader_required = isinstance(e, NativeComponentLoaderRequired)
if self.should_raise_customized_load_error(server_args, component_name):
if native_loader_required and component_weight_override is not None:
raise ComponentCheckpointUnsupportedError(
f"{component_name!r} requires its library loader, which cannot "
"consume a weights-only override; use "
f"--component-paths.{component_name} to replace its config "
"and weights together"
) from e
if (
component_weight_override is not None
or self.should_raise_customized_load_error(server_args, component_name)
):
if native_loader_required:
raise
if component_weight_override is not None:
raise RuntimeError(
f"Failed to load the weights-only override for "
f"{component_name!r}; fallback would ignore it. Use "
f"--component-paths.{component_name} when the checkpoint "
"also requires a different config or library loader."
) from e
traceback.print_exc()
raise RuntimeError(
f"Failed to load customized {component_name}; native fallback "
@@ -403,7 +429,6 @@ class ComponentLoader(ABC):
transformers_or_diffusers,
component_attn_backend,
component_attn_name,
self.allow_global_attention_backend_fallback,
require_backend_selection,
)
source = "native"
@@ -417,6 +442,13 @@ class ComponentLoader(ABC):
logger.error("Load %s failed", component_name)
consumed = 0.0
else:
if fsdp_requested and (
not isinstance(component, nn.Module)
or not is_fsdp_managed_module(component)
):
# The returned module is the source of truth. Loaders do not need
# a parallel capability declaration for FSDP support.
server_args.disable_fsdp_for_component(component_name)
if isinstance(component, nn.Module):
component = component.eval()
if (
@@ -492,11 +524,6 @@ class ComponentLoader(ABC):
resolved_component_name,
feature_name="Transformers quantized component",
)
if server_args.should_use_fsdp_for_component(resolved_component_name):
raise ComponentCheckpointUnsupportedError(
"Transformers-managed quantized components do not support "
"SGLang FSDP loading"
)
load_kwargs["device_map"] = {
"": self.target_device(component_starts_on_cpu=False)
}
@@ -648,7 +675,42 @@ class ComponentLoader(ABC):
return loader
class PlainStateDictComponentLoader(ComponentLoader):
class WeightOverrideComponentLoader(ComponentLoader):
"""Base for loaders that consume an exact weights-only override."""
def resolve_component_weight_override(
self, server_args: ServerArgs, component_name: str
) -> str | None:
return server_args.component_weights_paths.get(component_name)
def validate_component_weight_override(self, _override: str) -> None:
pass
def resolve_component_weights_path(
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str,
) -> str:
override = self.resolve_component_weight_override(server_args, component_name)
if override is None:
return component_model_path
self.validate_component_weight_override(override)
weights_path = materialize_weight(resolve_weight(override))
logger.info("Using weight override for %s: %s", component_name, weights_path)
return weights_path
class OnlineQuantizationComponentLoader(WeightOverrideComponentLoader):
"""Base for loaders that also consume an online quantization override."""
def resolve_component_quantization_override(
self, server_args: ServerArgs, component_name: str
) -> str | None:
return server_args.component_quantizations.get(component_name)
class PlainStateDictComponentLoader(WeightOverrideComponentLoader):
"""Base for native loaders whose current materializer expects plain weights."""
def component_load_precision(
@@ -682,19 +744,6 @@ class PlainStateDictComponentLoader(ComponentLoader):
self.ensure_plain_state_dict_checkpoint(config, component_name)
return config
def resolve_component_weights_path(
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str,
) -> str:
override = server_args.component_weights_paths.get(component_name)
if override is None:
return component_model_path
weights_path = materialize_weight(resolve_weight(override))
logger.info("Using weight override for %s: %s", component_name, weights_path)
return weights_path
class ImageProcessorLoader(ComponentLoader):
"""Loader for image processor."""
@@ -769,10 +818,6 @@ class TokenizerLoader(ComponentLoader):
class GenericComponentLoader(ComponentLoader):
"""Generic loader for components that don't have a specific loader."""
# An unknown out-of-tree component may itself be the primary transformer.
# Require it to opt into fallback through a registered component loader.
allow_global_attention_backend_fallback = False
def __init__(
self, library="transformers", component_architecture: str | None = None
) -> None:
@@ -780,6 +825,21 @@ class GenericComponentLoader(ComponentLoader):
self.library = library
self.component_architecture = component_architecture
def component_attention_backend_context(
self,
attn_backend: Any,
component_attn_name: str | None,
require_backend_selection: bool,
):
# An unknown out-of-tree component may itself be the primary transformer.
# Require it to opt into fallback through a registered component loader.
return component_attn_backend_context_manager(
attn_backend,
component_name=component_attn_name,
allow_global_backend_fallback=False,
require_backend_selection=require_backend_selection,
)
class PipelineComponentLoader:
"""
python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
+1 -1
View File
@@ -30,7 +30,7 @@ class ImageEncoderLoader(TextEncoderLoader):
component_name: str = "image_encoder",
):
"""Load the text encoders based on the model path, and inference args."""
component_weights_path = self.resolve_model_weights_path(
component_weights_path = self.resolve_component_weights_path(
component_model_path,
server_args,
component_name,
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+9 -31
View File
@@ -56,8 +56,8 @@ from sglang.multimodal_gen.runtime.layers.quantization.quanto_int8 import (
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
ComponentLoader,
NativeComponentLoaderRequired,
OnlineQuantizationComponentLoader,
uses_native_transformers_quantization,
)
from sglang.multimodal_gen.runtime.loader.gguf_weights import (
@@ -102,10 +102,6 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import (
inspect_comfy_quant_markers,
resolve_comfy_checkpoint_quantization,
)
from sglang.multimodal_gen.runtime.weights.source import (
materialize_weight,
resolve_weight,
)
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
from sglang.srt.environ import envs
from sglang.srt.layers.linear import LinearBase as SrtLinearBase
@@ -265,7 +261,10 @@ def _configure_encoder_quantization(
explicit_quantization: str | None = None,
ignored_layers: list[str] | None = None,
) -> None:
if getattr(model_cls, "manages_checkpoint_quantization", False):
if (
issubclass(model_cls, EncoderTensorParallelMixin)
and model_cls.checkpoint_quantization_backend == "model"
):
if explicit_quantization is not None:
raise ComponentCheckpointUnsupportedError(
f"{component_name!r} manages its own checkpoint quantization and "
@@ -493,12 +492,11 @@ def _keep_this_checkpoint_mapped(model_path: str) -> bool:
return True
class TextEncoderLoader(ComponentLoader):
class TextEncoderLoader(OnlineQuantizationComponentLoader):
"""Loader for text encoders."""
component_names = ["text_encoder"]
expected_library = "transformers"
supports_online_quantization_override = True
def component_load_precision(
self, server_args: ServerArgs, component_name: str
@@ -518,33 +516,13 @@ class TextEncoderLoader(ComponentLoader):
or component_name in server_args.component_quantizations
)
@staticmethod
def resolve_model_weights_path(
component_model_path: str,
server_args: ServerArgs,
component_name: str,
) -> str:
weights_override = server_args.component_weights_paths.get(component_name)
if weights_override is None:
return component_model_path
if names_gguf_checkpoint(weights_override):
def validate_component_weight_override(self, override: str) -> None:
if names_gguf_checkpoint(override):
if not current_platform.is_cuda():
raise ValueError(
"GGUF encoder checkpoints require CUDA; the GGML kernels have "
f"no {current_platform.device_type} implementation"
)
if server_args.should_use_fsdp_for_component(component_name):
raise ValueError(
f"GGUF encoder checkpoint {component_name!r} is incompatible "
"with FSDP; select resident or layerwise placement"
)
model_weights_path = materialize_weight(resolve_weight(weights_override))
logger.info(
"Using weight-file override for %s: %s",
component_name,
model_weights_path,
)
return model_weights_path
@dataclasses.dataclass
class Source:
@@ -731,7 +709,7 @@ class TextEncoderLoader(ComponentLoader):
component_starts_on_cpu: bool | None = None,
):
"""Load the text encoders based on the model path, and inference args."""
component_weights_path = self.resolve_model_weights_path(
component_weights_path = self.resolve_component_weights_path(
component_model_path,
server_args,
component_name,
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+15 -6
View File
@@ -13,7 +13,7 @@ from sglang.multimodal_gen.runtime.layers.attention.selector import (
get_global_forced_attn_backend,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentLoader,
OnlineQuantizationComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
from sglang.multimodal_gen.runtime.loader.gguf_weights import gguf_weights_iterator
@@ -152,13 +152,9 @@ def _server_args_for_transformer_component(
return component_server_args
class TransformerLoader(ComponentLoader):
class TransformerLoader(OnlineQuantizationComponentLoader):
"""Shared loader for (video/audio) DiT transformers."""
allow_global_attention_backend_fallback = False
supports_online_quantization_override = True
supports_fsdp_inference = True
component_names = [
"transformer",
"unconditional_transformer",
@@ -167,6 +163,19 @@ class TransformerLoader(ComponentLoader):
]
expected_library = "diffusers"
def component_attention_backend_context(
self,
attn_backend,
component_attn_name: str | None,
require_backend_selection: bool,
):
return component_attn_backend_context_manager(
attn_backend,
component_name=component_attn_name,
allow_global_backend_fallback=False,
require_backend_selection=require_backend_selection,
)
def customized_load_kwargs_for_component(
self, server_args: ServerArgs, component_name: str
) -> dict[str, bool]:
python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py
-1
View File
@@ -198,7 +198,6 @@ def _load_explicit_config(
class UpsamplerLoader(PlainStateDictComponentLoader):
component_names = ["spatial_upsampler"]
expected_library = "diffusers"
supports_component_weight_override = True
def load_customized(
self,
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+13 -36
View File
@@ -15,8 +15,8 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
ComponentLoader,
NativeComponentLoaderRequired,
WeightOverrideComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.utils import (
_list_safetensors_files,
@@ -42,10 +42,6 @@ from sglang.multimodal_gen.runtime.utils.precision import (
resolve_component_precision,
resolve_decode_precision,
)
from sglang.multimodal_gen.runtime.weights.source import (
materialize_weight,
resolve_weight,
)
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
from sglang.srt.model_loader.checkpoint_quantization import (
resolve_checkpoint_quant_spec,
@@ -393,17 +389,21 @@ def _assign_direct_gpu_vae_state(
)
class VAELoader(ComponentLoader):
class VAELoader(WeightOverrideComponentLoader):
"""Shared loader for (video/audio) VAE modules."""
component_names = ["vae", "audio_vae", "video_vae"]
expected_library = "diffusers"
supports_direct_gpu_weight_loading = True
def supports_direct_gpu_weight_loading_for_component(
self, component_name: str
def resolve_component_direct_gpu_loading(
self, server_args: ServerArgs, component_name: str
) -> bool:
return component_name in ("vae", "video_vae")
requested = server_args.should_direct_gpu_weight_load_component(component_name)
if requested and component_name not in ("vae", "video_vae"):
raise ComponentCheckpointUnsupportedError(
f"Direct GPU loading is not implemented for {component_name!r}"
)
return requested
def select_weight_files(
self,
@@ -425,25 +425,6 @@ class VAELoader(ComponentLoader):
) -> str | None:
return server_args.component_precisions.get(component_name)
@staticmethod
def resolve_model_weights_path(
component_model_path: str,
server_args: ServerArgs,
component_name: str,
) -> str:
weights_override = getattr(server_args, "component_weights_paths", {}).get(
component_name
)
if weights_override is None:
return component_model_path
model_weights_path = materialize_weight(resolve_weight(weights_override))
logger.info(
"Using weight-file override for %s: %s",
component_name,
model_weights_path,
)
return model_weights_path
def customized_load_kwargs_for_component(
self, server_args: ServerArgs, component_name: str
) -> dict[str, bool]:
@@ -467,14 +448,10 @@ class VAELoader(ComponentLoader):
cpu_offload_flag: bool = False,
):
"""Load the VAE based on the model path, and inference args."""
direct_gpu_weight_loading = server_args.should_direct_gpu_weight_load_component(
component_name
direct_gpu_weight_loading = self.resolve_component_direct_gpu_loading(
server_args, component_name
)
if direct_gpu_weight_loading and component_name not in ("vae", "video_vae"):
raise ComponentCheckpointUnsupportedError(
f"Direct GPU loading is not implemented for {component_name!r}"
)
component_weights_path = self.resolve_model_weights_path(
component_weights_path = self.resolve_component_weights_path(
component_model_path,
server_args,
component_name,
python/sglang/multimodal_gen/runtime/models/encoders/base.py
-6
View File
@@ -160,9 +160,6 @@ class EncoderTensorParallelMixin:
_encoder_tp_group: GroupCoordinator | None = None
checkpoint_quantization_backend = "diffusion"
packed_modules_mapping: dict[str, list[str]] = {}
# Some encoders own checkpoint quantization end to end because their weight
# states or sharding contract cannot use the generic loader lifecycle.
manages_checkpoint_quantization = False
@staticmethod
def should_materialize_checkpoint_weight(name: str) -> bool:
@@ -196,9 +193,6 @@ class TextEncoder(
# Qwen2_5_VLCausalLMOutputWithPast). Off by default so a new encoder is
# replicated rather than silently broken; flip it once dp is verified there.
supports_dp_encode = False
# Some encoders own checkpoint quantization end to end because their weight
# states or sharding contract cannot use the generic loader lifecycle.
manages_checkpoint_quantization = False
layerwise_offload_dit_group_enabled = False
layer_names = [
"layers",
python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py
+1 -1
View File
@@ -25,7 +25,7 @@ from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextMod
class IdeogramQwen3VLTextEncoder(TextEncoder):
"""Language-only Qwen3-VL text encoder stored inside Ideogram checkpoints."""
manages_checkpoint_quantization = True
checkpoint_quantization_backend = "model"
_activation_layers = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)
layer_names = ["language_model.layers"]
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+1 -10
View File
@@ -1788,16 +1788,7 @@ class ServerArgs(DisaggServerArgsMixin):
component_paths: dict[str, str] = {}
component_weights_paths = dict(self.component_weights_paths)
for component, path in self.component_paths.items():
supports_weight_file_override = (
is_dit_component_name(component)
or is_text_encoder_component_name(component)
or is_image_encoder_component_name(component)
or is_vae_component_name(component)
)
if (
not supports_weight_file_override
or not is_explicit_weight_file_reference(path)
):
if not is_explicit_weight_file_reference(path):
component_paths[component] = path
continue
existing = component_weights_paths.get(component)
python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py
+35 -5
View File
@@ -332,9 +332,24 @@ class TestComponentAttentionBackendScope(unittest.TestCase):
captured_context = get_component_attn_backend_context()
return object()
def component_attention_backend_context(
self,
attn_backend,
component_attn_name: str | None,
require_backend_selection: bool,
):
return component_attn_backend_context_manager(
attn_backend,
component_name=component_attn_name,
allow_global_backend_fallback=allow_global_backend_fallback,
require_backend_selection=require_backend_selection,
)
class _Args:
component_precisions = {}
component_quantizations = {}
component_weights_paths = {}
pipeline_config = SimpleNamespace(native_only_components=())
@staticmethod
def requested_component_attention_backend(_component_name):
@@ -348,7 +363,6 @@ class TestComponentAttentionBackendScope(unittest.TestCase):
def should_use_fsdp_for_component(_component_name):
return False
_Loader.allow_global_attention_backend_fallback = allow_global_backend_fallback
with (
patch.object(ComponentLoader, "for_component_type", return_value=_Loader()),
patch(
@@ -378,10 +392,22 @@ class TestComponentAttentionBackendScope(unittest.TestCase):
self.assertFalse(context.allow_global_backend_fallback)
def test_builtin_loader_scopes(self):
self.assertFalse(TransformerLoader.allow_global_attention_backend_fallback)
self.assertFalse(GenericComponentLoader.allow_global_attention_backend_fallback)
self.assertTrue(TextEncoderLoader.allow_global_attention_backend_fallback)
self.assertTrue(VAELoader.allow_global_attention_backend_fallback)
cases = (
(TransformerLoader(), "transformer", False),
(GenericComponentLoader(), "custom", False),
(TextEncoderLoader(), "text_encoder", True),
(VAELoader(), "vae", True),
)
for loader, component_name, expected_fallback in cases:
with (
self.subTest(loader=loader.__class__.__name__),
loader.component_attention_backend_context(None, component_name, False),
):
context = get_component_attn_backend_context()
self.assertIsNotNone(context)
self.assertEqual(
context.allow_global_backend_fallback, expected_fallback
)
def test_explicit_backend_must_be_consumed(self):
with self.assertRaisesRegex(
@@ -414,6 +440,8 @@ class TestComponentAttentionBackendScope(unittest.TestCase):
class _Args:
component_precisions = {}
component_quantizations = {}
component_weights_paths = {}
pipeline_config = SimpleNamespace(native_only_components=())
@staticmethod
def requested_component_attention_backend(_component_name):
@@ -479,6 +507,7 @@ class TestComponentAttentionBackendScope(unittest.TestCase):
class _Args:
component_precisions = {}
component_quantizations = {}
component_weights_paths = {}
pipeline_config = SimpleNamespace(native_only_components=())
@staticmethod
@@ -530,6 +559,7 @@ class TestComponentAttentionBackendScope(unittest.TestCase):
class _Args:
component_precisions = {}
component_quantizations = {}
component_weights_paths = {}
pipeline_config = SimpleNamespace(native_only_components=())
@staticmethod
python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py
+179 -2
View File
@@ -5,6 +5,9 @@ import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
from torch import nn
from sglang.multimodal_gen.runtime.loader.component_loaders.adapter_loader import (
AdapterLoader,
)
@@ -14,6 +17,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.bridge_loader import
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
ComponentLoader,
NativeComponentLoaderRequired,
PlainStateDictComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.diffusion_decoder_loader import (
@@ -38,6 +42,42 @@ class _TestLoader(PlainStateDictComponentLoader):
pass
class _WeightOverrideLoader(PlainStateDictComponentLoader):
def __init__(self, error: Exception):
super().__init__()
self.error = error
self.native_called = False
def load_customized(self, *args, **kwargs):
raise self.error
def load_native(self, *args, **kwargs):
self.native_called = True
return object()
class _ModuleLoader(ComponentLoader):
def load_customized(self, *args, **kwargs):
return nn.Linear(1, 1)
def _loader_server_args(component_weights_paths, *, fsdp_requested=False):
disabled_components = set()
return SimpleNamespace(
component_precisions={},
component_quantizations={},
component_weights_paths=component_weights_paths,
pipeline_config=SimpleNamespace(native_only_components=()),
resolve_component_attention_backend=lambda _name: (None, None),
requested_component_attention_backend=lambda _name: None,
should_direct_gpu_weight_load_component=lambda _name: False,
should_start_component_on_cpu=lambda _name: True,
should_use_fsdp_for_component=lambda _name: fsdp_requested,
disable_fsdp_for_component=lambda name: disabled_components.add(name),
disabled_components=disabled_components,
)
class TestComponentQuantizationAdmission(unittest.TestCase):
def test_plain_loader_admits_its_exact_precision(self):
server_args = SimpleNamespace(component_precisions={"vocoder": "fp16"})
@@ -60,6 +100,7 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
server_args = SimpleNamespace(
component_precisions={},
component_quantizations={},
component_weights_paths={},
should_direct_gpu_weight_load_component=lambda component: component
== "vocoder",
)
@@ -75,12 +116,14 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
server_args = SimpleNamespace(
component_precisions={},
component_quantizations={},
component_weights_paths={},
should_direct_gpu_weight_load_component=lambda component: component
== "audio_vae",
)
with self.assertRaisesRegex(
ComponentCheckpointUnsupportedError, "does not support direct GPU"
ComponentCheckpointUnsupportedError,
"Direct GPU loading is not implemented",
):
VAELoader().load("/model/audio_vae", server_args, "audio_vae", "diffusers")
@@ -107,6 +150,71 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
"/cache/vocoder.safetensors",
)
def test_unsupported_loader_rejects_weight_override_before_loading(self):
server_args = _loader_server_args({"pe": "/weights/model.safetensors"})
with self.assertRaisesRegex(
ComponentCheckpointUnsupportedError,
r"--component-paths\.pe.*config and weights",
):
ComponentLoader().load("/base/pe", server_args, "pe", "transformers")
def test_weight_override_failure_never_falls_back_to_base_component(self):
loader = _WeightOverrideLoader(ValueError("incompatible checkpoint"))
server_args = _loader_server_args(
{"text_encoder": "/weights/model.safetensors"}
)
with (
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.current_platform.get_available_gpu_memory",
return_value=10.0,
),
self.assertRaisesRegex(RuntimeError, "fallback would ignore it"),
):
loader.load(
"/base/text_encoder", server_args, "text_encoder", "transformers"
)
self.assertFalse(loader.native_called)
def test_library_fallback_requires_a_complete_component_override(self):
loader = _WeightOverrideLoader(
NativeComponentLoaderRequired("delegate to Transformers")
)
server_args = _loader_server_args(
{"text_encoder": "/weights/model.safetensors"}
)
with (
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.current_platform.get_available_gpu_memory",
return_value=10.0,
),
self.assertRaisesRegex(
ComponentCheckpointUnsupportedError,
r"--component-paths\.text_encoder.*config and weights",
),
):
loader.load(
"/base/text_encoder", server_args, "text_encoder", "transformers"
)
self.assertFalse(loader.native_called)
def test_materialized_module_decides_fsdp_support(self):
server_args = _loader_server_args({}, fsdp_requested=True)
with patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.current_platform.get_available_gpu_memory",
side_effect=(10.0, 9.0),
):
_ModuleLoader().load("/base/module", server_args, "module", "diffusers")
self.assertEqual(server_args.disabled_components, {"module"})
def test_plain_checkpoint_config_is_accepted(self):
config = {"_class_name": "TestModel"}
@@ -119,6 +227,76 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
self.assertIs(loaded, config)
def test_bridge_consumes_exact_component_weight_override(self):
loader = BridgeLoader()
bridge_config = SimpleNamespace(update_model_arch=lambda _config: None)
server_args = SimpleNamespace(
component_weights_paths={
"dual_tower_bridge": "owner/repo/bridge.safetensors"
},
model_paths={},
pipeline_config=SimpleNamespace(bridge_config=bridge_config),
should_use_fsdp_for_component=lambda _name: False,
should_start_component_on_cpu=lambda _name: False,
hsdp_replicate_dim=1,
hsdp_shard_dim=1,
pin_cpu_memory=False,
)
with (
patch.object(
loader,
"load_component_config",
return_value={"_class_name": "MOVADualTowerModel"},
),
patch.object(
loader,
"resolve_component_weights_path",
return_value="/cache/bridge.safetensors",
) as resolve_weights,
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"bridge_loader._list_safetensors_files",
return_value=["/cache/bridge.safetensors"],
) as list_weights,
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"bridge_loader.ModelRegistry.resolve_model_cls",
return_value=(nn.Linear, None),
),
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"bridge_loader.resolve_precision",
return_value=torch.bfloat16,
),
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"bridge_loader.get_local_torch_device",
return_value=torch.device("cpu"),
),
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"bridge_loader.maybe_load_fsdp_model",
return_value=nn.Linear(1, 1),
) as load_weights,
):
loaded = loader.load_customized(
"/base/dual_tower_bridge",
server_args,
"dual_tower_bridge",
)
self.assertIsInstance(loaded, nn.Linear)
resolve_weights.assert_called_once_with(
"/base/dual_tower_bridge", server_args, "dual_tower_bridge"
)
list_weights.assert_called_once_with("/cache/bridge.safetensors")
self.assertEqual(
load_weights.call_args.kwargs["weight_dir_list"],
["/cache/bridge.safetensors"],
)
self.assertFalse(load_weights.call_args.kwargs["fsdp_inference"])
def test_all_quantization_metadata_layouts_fail_closed(self):
configs = {
"quantization_config": {
@@ -217,7 +395,6 @@ class TestComponentQuantizationAdmission(unittest.TestCase):
load_weights.assert_not_called()
def test_upsampler_uses_exact_component_weight_override(self):
self.assertTrue(UpsamplerLoader.supports_component_weight_override)
server_args = SimpleNamespace(
component_weights_paths={"spatial_upsampler": "owner/repo/upsampler"}
)
python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py
+1
View File
@@ -246,6 +246,7 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
from_pretrained=mock.Mock(return_value=loaded_encoder)
)
server_args = SimpleNamespace(
component_weights_paths={},
component_quantizations={},
component_precisions={},
pipeline_config=SimpleNamespace(
python/sglang/multimodal_gen/test/unit/test_server_args.py
+3 -3
View File
@@ -188,7 +188,7 @@ class TestServerArgsPathExpansion(unittest.TestCase):
},
)
def test_supplemental_weight_file_remains_a_component_path(self):
def test_any_explicit_component_weight_file_keeps_base_config(self):
args = self._from_dict_without_model_resolution(
{
"model_path": "/data/my-model",
@@ -198,11 +198,11 @@ class TestServerArgsPathExpansion(unittest.TestCase):
}
)
self.assertEqual(args.component_paths, {})
self.assertEqual(
args.component_paths,
args.component_weights_paths,
{"conditioning_projection": "owner/repo/projection.safetensors"},
)
self.assertEqual(args.component_weights_paths, {})
def test_component_attention_backends_are_normalized(self):
args = self._from_dict_without_model_resolution(
python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
+3 -3
View File
@@ -792,12 +792,12 @@ class TestTextEncoderQuantization(unittest.TestCase):
"text_encoder",
)
def test_model_managed_quantization_bypasses_generic_lifecycle(self):
def test_model_quantization_backend_bypasses_generic_lifecycle(self):
model_config = SimpleNamespace(quant_config=None)
with mock.patch.object(
TextEncoder,
"manages_checkpoint_quantization",
True,
"checkpoint_quantization_backend",
"model",
):
_configure_encoder_quantization(
model_config,
python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py
+2 -32
View File
@@ -3,20 +3,12 @@ import unittest
from types import SimpleNamespace
from unittest import mock
from sglang.multimodal_gen.runtime.loader.component_loaders.bridge_loader import (
BridgeLoader,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentLoader,
NativeComponentLoaderRequired,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
TransformerLoader,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
RESIDENT,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
class TestTransformerLoaderFallbackAdmission(unittest.TestCase):
@@ -94,7 +86,7 @@ class TestTransformerLoaderFallbackAdmission(unittest.TestCase):
)
native.assert_not_called()
server_args.should_use_fsdp_for_component.assert_called_once_with(
server_args.should_use_fsdp_for_component.assert_called_with(
"transformer_2"
)
@@ -118,9 +110,7 @@ class TestTransformerLoaderFallbackAdmission(unittest.TestCase):
self.assertIsNotNone(component)
self.assertEqual(consumed, 0.0)
native.assert_called_once()
server_args.should_use_fsdp_for_component.assert_called_once_with(
"transformer_2"
)
server_args.should_use_fsdp_for_component.assert_called_with("transformer_2")
def test_parallel_execution_rejects_native_fallback(self):
cases = (
@@ -146,26 +136,6 @@ class TestTransformerLoaderFallbackAdmission(unittest.TestCase):
)
)
def test_only_fsdp_materializers_keep_the_component_request(self):
server_args = ServerArgs.__new__(ServerArgs)
server_args.use_fsdp_inference = True
server_args._fsdp_disabled_components = set()
server_args.residency_mode = lambda _component: RESIDENT
ComponentLoader().disable_unsupported_component_fsdp(
server_args, "text_encoder"
)
self.assertFalse(server_args.should_use_fsdp_for_component("text_encoder"))
TransformerLoader().disable_unsupported_component_fsdp(
server_args, "transformer"
)
BridgeLoader().disable_unsupported_component_fsdp(
server_args, "dual_tower_bridge"
)
self.assertTrue(server_args.should_use_fsdp_for_component("transformer"))
self.assertTrue(server_args.should_use_fsdp_for_component("dual_tower_bridge"))
if __name__ == "__main__":
unittest.main()
python/sglang/multimodal_gen/test/unit/test_vae_loader.py
+8 -4
View File
@@ -16,7 +16,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
Wan2_2_I2V_A14B_Config,
WanT2V480PConfig,
)
from sglang.multimodal_gen.runtime.loader.component_loaders import vae_loader
from sglang.multimodal_gen.runtime.loader.component_loaders import (
component_loader,
vae_loader,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
)
@@ -49,6 +52,7 @@ class _FakeServerArgs:
self.revision = "test-revision"
self.trust_remote_code = True
self.layerwise_components = set()
self.component_weights_paths = {}
self.component_quantizations = {}
self.component_precisions = {}
self.component_direct_gpu_weight_loading = {}
@@ -327,15 +331,15 @@ class TestVAELoader(unittest.TestCase):
}
with (
patch.object(vae_loader, "resolve_weight", return_value="resolved"),
patch.object(component_loader, "resolve_weight", return_value="resolved"),
patch.object(
vae_loader,
component_loader,
"materialize_weight",
return_value="/cache/audio.safetensors",
),
):
self.assertEqual(
loader.resolve_model_weights_path(
loader.resolve_component_weights_path(
"/base/audio_vae", server_args, "audio_vae"
),
"/cache/audio.safetensors",