diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py index 4ecd974e0..e7e4cacb5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 3aa3490f9..9890b8bae 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -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.. 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.. 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: """ diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py index d4e32fb48..97d1b6ef1 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 42af6581d..590c2c1f0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 420837f95..81b721e0f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -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]: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py index 4f266d411..f6fbbcc12 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 226b195d3..a07a33e8a 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py index 38fa68338..7fc01c04b 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -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", diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py b/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py index 95b3a9f94..31fc1ec1d 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py @@ -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"] diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index c94286221..bd0bc7f73 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py index 56faa2654..9d5904a4b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py index a7299708c..9f1077188 100644 --- a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py +++ b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py @@ -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"} ) diff --git a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py index 675f8224f..ec629ee37 100644 --- a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py @@ -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( diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 2c1f29036..ac41dceed 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -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( diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index c4c4c2138..e349daaa7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -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, diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py b/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py index 6b7d8c8db..613949cf6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_loader_fallback.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index 47fc32558..6dfae598f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -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",