diff --git a/python/sglang/multimodal_gen/configs/models/base.py b/python/sglang/multimodal_gen/configs/models/base.py index a789b25af..df483a631 100644 --- a/python/sglang/multimodal_gen/configs/models/base.py +++ b/python/sglang/multimodal_gen/configs/models/base.py @@ -21,6 +21,9 @@ class ArchConfig: stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=list ) # mapping from huggingface weight names to custom names + param_names_mapping: dict[str, str | tuple[str, int, int]] = field( + default_factory=dict + ) extra_attrs: Dict[str, Any] = field(default_factory=dict) def __getattr__(self, name: str): @@ -53,6 +56,9 @@ class ModelConfig: # sglang-diffusion-specific parameters here # i.e. STA, quantization, teacache + def post_diffusers_config_update(self) -> None: + """Normalize external configuration before constructing the runtime model.""" + def __getattr__(self, name): # Only called if 'name' is not found in ModelConfig directly if hasattr(self.arch_config, name): diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py index bbbf37bb9..589411032 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py @@ -1,5 +1,3 @@ -import re - from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import ( LTX2ConnectorConfig, ) @@ -9,97 +7,12 @@ from sglang.multimodal_gen.configs.models.adapter.ltx_2_duration_head import ( from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) -from sglang.multimodal_gen.runtime.loader.utils import ( - load_safetensors_state_dict, - set_default_torch_dtype, - skip_init_modules, -) -from sglang.multimodal_gen.runtime.models.registry import ModelRegistry -from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.precision import resolve_precision class AdapterLoader(PlainStateDictComponentLoader): - """Loader for small adapter-style modules (e.g., LTX-2 connectors). - - This loader intentionally avoids FSDP sharding and just: - 1) Instantiates the module from `config.json`. - 2) Loads the safetensors state_dict (single-file or sharded). - """ - component_names = ["connectors", "duration_head"] - expected_library = "diffusers" - # `update_model_arch` fills each from the component's `config.json`. - _CONFIG_CLASSES = { + config_classes = { "connectors": LTX2ConnectorConfig, "duration_head": LTX2DurationHeadConfig, } - - def load_customized( - self, - component_model_path: str, - server_args: ServerArgs, - component_name: str = "connectors", - *args, - ): - 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 - ) - - cls_name = config.pop("_class_name", None) - if cls_name is None: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported." - ) - - config.pop("_diffusers_version", None) - config.pop("_name_or_path", None) - - server_args.model_paths[component_name] = component_model_path - - model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) - - # Not a fixed name: connectors follow DiT offload, while the duration - # head stays resident unless selected explicitly. - target_device = self.target_device( - server_args.should_start_component_on_cpu(component_name) - ) - default_dtype = resolve_precision( - server_args, component_name, precision_attr="dit_precision" - ) - - 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) - model = model_cls(adapter_cfg).to(device=target_device, dtype=default_dtype) - - loaded = load_safetensors_state_dict(component_weights_path) - mapping = adapter_cfg.arch_config.param_names_mapping - loaded = {_remap_connector_key(k, mapping): v for k, v in loaded.items()} - - missing, unexpected = model.load_state_dict(loaded, strict=False) - # `strict=False` because a checkpoint carries either the shared - # `text_proj_in` or the per-modality projections, never both. Anything - # else uninitialized would surface later as garbage embeddings. - if missing or unexpected: - raise ValueError( - f"Adapter weights at '{component_weights_path}' do not match the " - f"instantiated {cls_name}. Missing: {sorted(missing)}. " - f"Unexpected: {sorted(unexpected)}. This usually means the " - "adapter config or its weight-name mapping is wrong." - ) - - return model - - -def _remap_connector_key(key: str, param_names_mapping: dict[str, str]) -> str: - for pattern, replacement in param_names_mapping.items(): - key, replaced = re.subn(pattern, replacement, key) - if replaced: - break - return key 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 e7e4cacb5..db9e76d7e 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 @@ -1,127 +1,11 @@ -from copy import deepcopy - -import torch - -from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( + MOVADualTowerConfig, +) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) -from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model -from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files -from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan -from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( - RESIDENT, -) -from sglang.multimodal_gen.runtime.models.registry import ModelRegistry -from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.runtime.utils.precision import resolve_precision - -logger = init_logger(__name__) class BridgeLoader(PlainStateDictComponentLoader): - """Loader for MOVA dual tower bridge with FSDP support.""" - - pipeline_bridge_config_attr: str = "bridge_config" - component_names = ["dual_tower_bridge"] - expected_library = "diffusers" - - 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: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported." - ) - server_args.model_paths[component_name] = component_model_path - - # Try to get bridge config from pipeline config, fallback to creating one - bridge_config = getattr( - server_args.pipeline_config, self.pipeline_bridge_config_attr, None - ) - if bridge_config is not None: - bridge_config.update_model_arch(config) - else: - # Create a minimal config from hf_config - from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( - MOVADualTowerConfig, - ) - - bridge_config = MOVADualTowerConfig() - bridge_config.update_model_arch(config) - - model_cls, _ = ModelRegistry.resolve_model_cls(class_name) - - # Find all safetensors files - safetensors_list = _list_safetensors_files(component_weights_path) - if not safetensors_list: - raise ValueError(f"No safetensors files found in {component_weights_path}") - - default_dtype = resolve_precision( - server_args, component_name, precision_attr="dit_precision" - ) - - logger.info( - "Loading %s from %s safetensors files, default_dtype: %s", - class_name, - len(safetensors_list), - default_dtype, - ) - - use_fsdp = server_args.should_use_fsdp_for_component(component_name) - component_starts_on_cpu = server_args.should_start_component_on_cpu( - component_name - ) - - # Use the FSDP loader when FSDP is requested or shard rules are declared. - fsdp_shard_conditions = getattr(model_cls, "_fsdp_shard_conditions", None) - if ( - 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 - model = maybe_load_fsdp_model( - model_cls=model_cls, - init_params={"config": bridge_config, "hf_config": hf_config}, - weight_dir_list=safetensors_list, - device=local_torch_device, - hsdp_replicate_dim=server_args.hsdp_replicate_dim, - hsdp_shard_dim=server_args.hsdp_shard_dim, - component_starts_on_cpu=component_starts_on_cpu, - pin_cpu_memory=server_args.pin_cpu_memory, - fsdp_inference=use_fsdp, - param_dtype=default_dtype, - reduce_dtype=torch.float32, - output_dtype=None, - strict=False, - weight_load_plan=WeightLoadPlan( - checkpoint_load_device=local_torch_device - ), - ) - else: - # Fallback to simple loading (for non-FSDP or legacy models) - model = model_cls.from_pretrained( - component_model_path, torch_dtype=default_dtype - ) - target_device = self.target_device(component_starts_on_cpu) - model = model.to(device=target_device, dtype=default_dtype) - - total_params = sum(p.numel() for p in model.parameters()) - logger.info("Loaded bridge model with %.2fM parameters", total_params / 1e6) - - return model + config_classes = {"dual_tower_bridge": MOVADualTowerConfig} 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 e1116d4d5..646e6e6ca 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 @@ -7,6 +7,7 @@ import os import pkgutil import traceback from abc import ABC +from collections.abc import Callable, Iterator from typing import Any, Type import torch @@ -21,17 +22,29 @@ from transformers import ( ) from transformers.quantizers import AutoHfQuantizer +from sglang.multimodal_gen.configs.models.base import ModelConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.attention.selector import ( ComponentAttentionBackendNotAppliedError, component_attn_backend_context_manager, get_component_attn_backend_context, ) +from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model from sglang.multimodal_gen.runtime.loader.utils import ( + _list_safetensors_files, _normalize_component_type, component_name_to_loader_cls, + finalize_loaded_model, format_component_residency, get_memory_usage_of_component, + get_param_names_mapping, + hf_to_custom_state_dict, + initialize_model, + load_model_state_dict, +) +from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan +from sglang.multimodal_gen.runtime.loader.weight_utils import ( + checkpoint_weights_iterator, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( RESIDENT, @@ -40,6 +53,8 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import ( is_fsdp_managed_module, ) +from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT +from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( @@ -48,7 +63,10 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( prepare_diffusers_component_path_for_loading, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision +from sglang.multimodal_gen.runtime.utils.precision import ( + resolve_component_precision, + resolve_precision, +) from sglang.multimodal_gen.runtime.weights.source import ( materialize_weight, resolve_weight, @@ -450,7 +468,7 @@ class ComponentLoader(ABC): # a parallel capability declaration for FSDP support. server_args.disable_fsdp_for_component(component_name) if isinstance(component, nn.Module): - component = component.eval() + component = finalize_loaded_model(component) if ( not is_fsdp_managed_module(component) and not self._native_load_manages_placement @@ -678,6 +696,55 @@ class ComponentLoader(ABC): class WeightOverrideComponentLoader(ComponentLoader): """Base for loaders that consume an exact weights-only override.""" + ignored_checkpoint_prefixes: tuple[str, ...] = () + + def load_state_dict_model( + self, + model_cls: type[nn.Module], + init_params: dict[str, Any], + weight_files: list[str], + server_args: ServerArgs, + component_name: str, + dtype: torch.dtype, + *, + component_starts_on_cpu: bool, + weight_load_plan: WeightLoadPlan | None = None, + checkpoint_key_filter: Callable[[str], bool] | None = None, + weights_iterator: Iterator[tuple[str, torch.Tensor]] | None = None, + ) -> nn.Module: + """Restore mapped model state with optional TP/FSDP materialization.""" + return maybe_load_fsdp_model( + model_cls=model_cls, + init_params=init_params, + weight_dir_list=weight_files, + device=get_local_torch_device(), + hsdp_replicate_dim=server_args.hsdp_replicate_dim, + hsdp_shard_dim=server_args.hsdp_shard_dim, + component_starts_on_cpu=component_starts_on_cpu, + pin_cpu_memory=server_args.pin_cpu_memory, + fsdp_inference=server_args.should_use_fsdp_for_component(component_name), + param_dtype=dtype, + reduce_dtype=torch.float32, + strict=False, + weight_load_plan=weight_load_plan, + checkpoint_key_filter=checkpoint_key_filter, + weights_iterator=weights_iterator, + ) + + def validate_checkpoint_keys( + self, missing: list[str] | set[str], unexpected: list[str], component_name: str + ) -> None: + unexpected = [ + name + for name in unexpected + if not name.startswith(self.ignored_checkpoint_prefixes) + ] + if missing or unexpected: + raise ComponentCheckpointUnsupportedError( + f"Checkpoint weights do not match {component_name!r}. " + f"Missing: {sorted(missing)}. Unexpected: {sorted(unexpected)}." + ) + def resolve_component_weight_override( self, server_args: ServerArgs, component_name: str ) -> str | None: @@ -711,7 +778,97 @@ class OnlineQuantizationComponentLoader(WeightOverrideComponentLoader): class PlainStateDictComponentLoader(WeightOverrideComponentLoader): - """Base for native loaders whose current materializer expects plain weights.""" + """Construct registered modules and restore a complete plain state dict.""" + + expected_library = "diffusers" + config_classes: dict[str, type[ModelConfig]] = {} + default_precision_attr = "dit_precision" + default_dtype = torch.bfloat16 + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, component_name: str + ) -> nn.Module: + config = self.load_component_config(component_model_path, component_name) + class_name = config.pop("_class_name", None) or self.component_architecture + if class_name is None: + raise ComponentCheckpointUnsupportedError( + f"{component_name!r} must declare _class_name in config.json " + "or its architecture in model_index.json" + ) + weights_path = self.resolve_component_weights_path( + component_model_path, server_args, component_name + ) + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + model_config = self.build_model_config(config, component_name) + dtype = self.resolve_dtype(server_args, component_name) + component_starts_on_cpu = server_args.should_start_component_on_cpu( + component_name + ) + server_args.model_paths[component_name] = component_model_path + if issubclass(model_cls, BaseDiT): + weight_files = _list_safetensors_files(weights_path) + return self.load_state_dict_model( + model_cls, + {"config": model_config, "hf_config": config}, + weight_files, + server_args, + component_name, + dtype, + component_starts_on_cpu=component_starts_on_cpu, + weights_iterator=( + None if weight_files else checkpoint_weights_iterator(weights_path) + ), + ) + + target_device = self.target_device(component_starts_on_cpu) + model = initialize_model( + model_cls, + model_config + if isinstance(model_config, dict) + else {"config": model_config}, + dtype, + ).to(target_device) + + try: + state_dict, _ = hf_to_custom_state_dict( + checkpoint_weights_iterator(weights_path), + get_param_names_mapping( + model_config.arch_config.param_names_mapping + if isinstance(model_config, ModelConfig) + else {} + ), + valid_target_names=set(model.state_dict()), + strict=True, + ) + missing, unexpected = load_model_state_dict(model, state_dict, strict=False) + except (RuntimeError, ValueError) as error: + raise ComponentCheckpointUnsupportedError( + f"Cannot restore checkpoint for {component_name!r}: {error}" + ) from error + self.validate_checkpoint_keys(missing, unexpected, component_name) + return model + + def build_model_config( + self, config: dict[str, Any], component_name: str + ) -> ModelConfig | dict[str, Any]: + config_cls = self.config_classes.get( + self.structural_component_type(component_name) + ) + if config_cls is not None: + model_config = config_cls() + model_config.update_model_arch(config) + return model_config + return {key: value for key, value in config.items() if not key.startswith("_")} + + def resolve_dtype( + self, server_args: ServerArgs, component_name: str + ) -> torch.dtype: + try: + return resolve_precision( + server_args, component_name, precision_attr=self.default_precision_attr + ) + except AttributeError: + return self.default_dtype def component_load_precision( self, server_args: ServerArgs, component_name: str diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py index b4b7926b6..425324ea5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py @@ -6,57 +6,11 @@ from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder imp from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) -from sglang.multimodal_gen.runtime.loader.utils import ( - load_safetensors_state_dict, - set_default_torch_dtype, - skip_init_modules, -) -from sglang.multimodal_gen.runtime.models.registry import ModelRegistry -from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.precision import resolve_precision class DiffusionDecoderLoader(PlainStateDictComponentLoader): - """Loader for the standalone, replicated LTX-2.5 diffusion decoder.""" + """Standalone, replicated LTX-2.5 diffusion decoder.""" component_names = ["diffusion_decoder"] - expected_library = "diffusers" - - def load_customized( - self, - component_model_path: str, - server_args: ServerArgs, - component_name: str = "diffusion_decoder", - *args, - ): - 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 - ) - class_name = config.pop("_class_name", None) - if class_name is None: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported." - ) - config.pop("_diffusers_version", None) - config.pop("_name_or_path", None) - - server_args.model_paths[component_name] = component_model_path - model_cls, _ = ModelRegistry.resolve_model_cls(class_name) - target_device = self.target_device( - server_args.should_start_component_on_cpu(component_name) - ) - dtype = resolve_precision( - server_args, component_name, precision_attr="vae_precision" - ) - - decoder_config = LTX25DiffusionDecoderConfig() - decoder_config.update_model_arch(config) - with set_default_torch_dtype(dtype), skip_init_modules(): - model = model_cls(decoder_config).to(device=target_device, dtype=dtype) - - model.load_state_dict( - load_safetensors_state_dict(component_weights_path), strict=True - ) - return model + config_classes = {"diffusion_decoder": LTX25DiffusionDecoderConfig} + default_precision_attr = "vae_precision" 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 97d1b6ef1..bdaf74683 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 @@ -1,15 +1,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import ( TextEncoderLoader, - _resolve_and_configure_encoder_quantization, ) -from sglang.multimodal_gen.runtime.models.encoders.base import finalize_encoder_folding from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - get_diffusers_component_config, -) -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) class ImageEncoderLoader(TextEncoderLoader): @@ -23,52 +15,9 @@ class ImageEncoderLoader(TextEncoderLoader): component_name, server_args.pipeline_config.image_encoder_precision ) - def load_customized( - self, - component_model_path: str, - server_args: ServerArgs, - component_name: str = "image_encoder", + def build_model_config( + self, component_model_path, model_config, server_args, component_name ): - """Load the text encoders based on the model path, and inference args.""" - component_weights_path = self.resolve_component_weights_path( - component_model_path, - server_args, - component_name, - ) - # model_config: PretrainedConfig = get_hf_config( - # model=model_path, - # trust_remote_code=server_args.trust_remote_code, - # revision=server_args.revision, - # model_override_args=None, - # ) - model_config = get_diffusers_component_config( - component_path=component_model_path - ) - encoder_config = server_args.pipeline_config.image_encoder_config encoder_config.update_model_arch(model_config) - _resolve_and_configure_encoder_quantization( - encoder_config, - model_config, - component_model_path, - component_weights_path, - component_name, - server_args.component_quantizations.get(component_name), - ) - # real dims are populated now; resolve fold vs replicate - finalize_encoder_folding( - encoder_config, - server_args.encoder_parallel, - ) - - # Always start with local device; load_model will adjust for offload if needed - # TODO(will): add support for other dtypes - image_encoder_dtype = self.component_load_precision(server_args, component_name) - assert image_encoder_dtype is not None - return self.load_model( - component_weights_path, - encoder_config, - server_args, - image_encoder_dtype, - component_name=component_name, - ) + return encoder_config diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py index 511066434..8148612ab 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py @@ -1,71 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -import torch - from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) -from sglang.multimodal_gen.runtime.loader.utils import ( - load_safetensors_state_dict, - set_default_torch_dtype, - skip_init_modules, -) -from sglang.multimodal_gen.runtime.models.registry import ModelRegistry -from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision -from sglang.multimodal_gen.utils import PRECISION_TO_TYPE - -logger = init_logger(__name__) class SoundTokenizerLoader(PlainStateDictComponentLoader): component_names = ["sound_tokenizer"] - expected_library = "diffusers" + default_precision_attr = "vae_precision" + # the native tokenizer is decoder-only; encoder weights are unused + ignored_checkpoint_prefixes = ("encoder.",) - 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 - ) - class_name = config.pop("_class_name", None) or self.component_architecture - assert class_name is not None, ( - "Sound tokenizer class name must be available from component config." - ) - - server_args.model_paths[component_name] = component_model_path - - dtype = resolve_component_precision(server_args, component_name) - if dtype is None: - try: - dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - except AttributeError: - dtype = torch.bfloat16 - target_device = self.target_device( - server_args.should_start_component_on_cpu(component_name) - ) - - with set_default_torch_dtype(dtype), skip_init_modules(): - model_cls, _ = ModelRegistry.resolve_model_cls(class_name) - model = model_cls(config).to(device=target_device, dtype=dtype) - - loaded = load_safetensors_state_dict(component_weights_path) - incompatible = model.load_state_dict(loaded, strict=False) - missing = getattr(incompatible, "missing_keys", []) - # The tokenizer is decoder-only; the checkpoint's encoder weights are - # expected leftovers, so they're excluded from the load warning. - unexpected = [ - k - for k in getattr(incompatible, "unexpected_keys", []) - if not k.startswith("encoder.") - ] - if missing or unexpected: - logger.warning( - "Loaded sound_tokenizer with missing_keys=%d unexpected_keys=%d", - len(missing), - len(unexpected), - ) - model.eval() - return model + def build_model_config(self, config, component_name): + return {"config": config} 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 c7f9ac656..147b0e9fe 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 @@ -1,9 +1,6 @@ -import dataclasses -import glob import os import re -from collections.abc import Callable, Generator, Iterable -from typing import cast +from collections.abc import Generator import torch import transformers @@ -67,20 +64,14 @@ from sglang.multimodal_gen.runtime.loader.gguf_weights import ( remap_gguf_tensor_meta, ) from sglang.multimodal_gen.runtime.loader.utils import ( - _list_safetensors_files, checkpoint_bytes, get_param_names_mapping, + initialize_model, + keep_checkpoint_mapped, set_default_torch_dtype, - skip_init_modules, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( - filter_files_not_needed_for_inference, - pt_weights_iterator, - safetensors_weights_iterator, -) -from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( - host_copies_would_not_fit, - host_memory_available_bytes, + checkpoint_weights_iterator, ) from sglang.multimodal_gen.runtime.models.encoders.base import ( EncoderTensorParallelMixin, @@ -101,10 +92,10 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import ( get_quant_config, get_quant_config_from_safetensors_metadata, inspect_comfy_quant_markers, + process_model_weights_after_loading, resolve_comfy_checkpoint_quantization, ) from sglang.multimodal_gen.utils import PRECISION_TO_TYPE -from sglang.srt.environ import envs from sglang.srt.layers.linear import LinearBase as SrtLinearBase from sglang.srt.layers.quantization.fp8 import Fp8Config as SrtFp8Config from sglang.srt.layers.quantization.unquant import ( @@ -113,7 +104,6 @@ from sglang.srt.layers.quantization.unquant import ( from sglang.srt.model_loader.checkpoint_quantization import ( resolve_checkpoint_quant_spec, ) -from sglang.srt.model_loader.post_load import stage_module_for_post_load logger = init_logger(__name__) @@ -340,7 +330,7 @@ def _resolve_and_configure_encoder_quantization( explicit_quantization: str | None = None, ignored_layers: list[str] | None = None, ) -> type[nn.Module]: - architectures = getattr(model_config, "architectures", []) + architectures = model_config.arch_config.architectures try: model_cls, _ = ModelRegistry.resolve_model_cls(architectures) except Exception as resolution_error: @@ -384,35 +374,6 @@ def _resolve_and_configure_encoder_quantization( return model_cls -def _process_quantized_encoder_weights( - model: nn.Module, - process_device: torch.device | None, - component_name: str, -) -> int: - processed_layers = 0 - for module in model.modules(): - if not isinstance(module, (LinearBase, SrtLinearBase)): - continue - quant_method = module.quant_method - if quant_method is None or isinstance( - quant_method, - (UnquantizedLinearMethod, SrtUnquantizedLinearMethod), - ): - continue - if process_device is None: - quant_method.process_weights_after_loading(module) - else: - with stage_module_for_post_load(module, process_device): - quant_method.process_weights_after_loading(module) - processed_layers += 1 - if processed_layers == 0: - raise ValueError( - f"The {component_name!r} checkpoint declares quantization, but the " - "model did not construct any quantized linear layers" - ) - return processed_layers - - def _require_quantized_encoder_layers( model: nn.Module, component_name: str, @@ -463,21 +424,6 @@ def _require_quantized_encoder_layers( ) -def _keep_this_checkpoint_mapped(model_path: str) -> bool: - """Whether this encoder's weights should stay on their file mapping.""" - weight_bytes = checkpoint_bytes(model_path) - if not host_copies_would_not_fit(weight_bytes): - return False - logger.info( - "Text encoder checkpoint is %.2f GiB against %.2f GiB of host memory, " - "so its compatible weights stay on the checkpoint mapping instead of " - "being copied in.", - weight_bytes / 1024**3, - host_memory_available_bytes() / 1024**3, - ) - return True - - class TextEncoderLoader(OnlineQuantizationComponentLoader): """Loader for text encoders.""" @@ -491,7 +437,7 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): if override is not None: return override return server_args.pipeline_config.text_encoder_precisions[ - self._extract_encoder_index(component_name) + self._extract_encoder_index(self.structural_component_name(component_name)) ] def should_raise_customized_load_error( @@ -510,22 +456,6 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): f"no {current_platform.device_type} implementation" ) - @dataclasses.dataclass - class Source: - """A source for weights.""" - - model_or_path: str - """The model ID or path.""" - - prefix: str = "" - """A prefix to prepend to all weights.""" - - fall_back_to_pt: bool = True - """Whether .pt weights can be used.""" - - allow_patterns_overrides: list[str] | None = None - """If defined, weights will load exclusively using these patterns.""" - def resolve_native_transformers_model_class(self, config: PretrainedConfig) -> type: """Resolve the concrete transformers class for a text encoder. @@ -543,149 +473,23 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): return transformers_model_class return transformers.AutoModel - def _prepare_weights( - self, - model_name_or_path: str, - fall_back_to_pt: bool, - allow_patterns_overrides: list[str] | None, - key_filter: Callable[[str], bool] | None = None, - ) -> tuple[str, list[str], bool]: - """Prepare weights for the model. - - If the model is not local, it will be downloaded.""" - # model_name_or_path = (self._maybe_download_from_modelscope( - # model_name_or_path, revision) or model_name_or_path) - - if os.path.isfile(model_name_or_path): - if model_name_or_path.endswith(".safetensors"): - return os.path.dirname(model_name_or_path), [model_name_or_path], True - if fall_back_to_pt and model_name_or_path.endswith((".bin", ".pt")): - return os.path.dirname(model_name_or_path), [model_name_or_path], False - raise ValueError( - "Native encoder weight overrides currently support one " - f"safetensors, bin, or pt file, got {model_name_or_path!r}" - ) - if not os.path.isdir(model_name_or_path): - raise ValueError( - f"Model path must be a local file or directory: {model_name_or_path!r}" - ) - - use_safetensors = False - index_file = SAFE_WEIGHTS_INDEX_NAME - allow_patterns = ["*.safetensors", "*.bin"] - - if fall_back_to_pt: - allow_patterns += ["*.pt"] - - if allow_patterns_overrides is not None: - allow_patterns = allow_patterns_overrides - - hf_folder = model_name_or_path - - hf_weights_files: list[str] = [] - for pattern in allow_patterns: - if pattern == "*.safetensors": - hf_weights_files = _list_safetensors_files( - hf_folder, - index_file=index_file, - key_filter=key_filter, - ) - else: - hf_weights_files = glob.glob(os.path.join(hf_folder, pattern)) - if len(hf_weights_files) > 0: - if pattern == "*.safetensors": - use_safetensors = True - break - - if not use_safetensors: - hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) - - if len(hf_weights_files) == 0: - raise RuntimeError( - f"Cannot find any model weights with `{model_name_or_path}`" - ) - - # Sort weight files when SGLANG_SORT_WEIGHT_FILES >= 0 (default). - # Staggering is not applicable to text-encoder loading (no TP split). - if envs.SGLANG_SORT_WEIGHT_FILES.get() >= 0: - hf_weights_files.sort() - - return hf_folder, hf_weights_files, use_safetensors - - def _get_weights_iterator( - self, - source: "Source", - to_cpu: bool, - key_filter: Callable[[str], bool] | None = None, - ) -> Generator[tuple[str, torch.Tensor], None, None]: - """get an iterator for the model weights based on the load format.""" - source_key_filter: Callable[[str], bool] | None - if key_filter is None: - source_key_filter = None - else: - - def include_source_weight(name: str) -> bool: - return key_filter(source.prefix + name) - - source_key_filter = include_source_weight - - hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( - source.model_or_path, - source.fall_back_to_pt, - source.allow_patterns_overrides, - key_filter=source_key_filter, - ) - if use_safetensors: - weights_iterator = safetensors_weights_iterator( - hf_weights_files, - to_cpu=to_cpu, - key_filter=source_key_filter, - ) - else: - weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu) - if source_key_filter is not None: - weights_iterator = ( - (name, tensor) - for name, tensor in weights_iterator - if source_key_filter(name) - ) - - # apply the prefix. - return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) - def _get_all_weights( self, model: EncoderTensorParallelMixin, model_path: str, to_cpu: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: - key_filter = model.should_materialize_checkpoint_weight - def include_checkpoint_weight(name: str) -> bool: - return not name.endswith(".comfy_quant") and key_filter(name) + return not name.endswith( + ".comfy_quant" + ) and model.should_materialize_checkpoint_weight(name) - primary_weights = TextEncoderLoader.Source( + yield from checkpoint_weights_iterator( model_path, - prefix="", - fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True), - allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None), + to_cpu=to_cpu, + key_filter=include_checkpoint_weight, + index_file=SAFE_WEIGHTS_INDEX_NAME, ) - yield from self._get_weights_iterator( - primary_weights, - to_cpu, - include_checkpoint_weight, - ) - - secondary_weights = cast( - Iterable[TextEncoderLoader.Source], - getattr(model, "secondary_weights", ()), - ) - for source in secondary_weights: - yield from self._get_weights_iterator( - source, - to_cpu, - include_checkpoint_weight, - ) def load_customized( self, @@ -700,35 +504,13 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): server_args, component_name, ) - diffusers_pretrained_config = get_config( - component_model_path, trust_remote_code=True - ) model_config = get_diffusers_component_config( component_path=component_model_path ) - - # TODO(mick): had to throw an exception for different text-encoder arch - encoder_index = self._extract_encoder_index( - self.structural_component_name(component_name) + encoder_config = self.build_model_config( + component_model_path, model_config, server_args, component_name ) - assert encoder_index < len( - server_args.pipeline_config.text_encoder_configs - ) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions) - - encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index] - encoder_config.update_model_arch(model_config) - encoder_config.generation_config = load_dict( - os.path.join(component_model_path, "generation_config.json") - ) - - if encoder_index == 0: - for key, value in diffusers_pretrained_config.__dict__.items(): - setattr(encoder_config.arch_config, key, value) - post_diffusers_config_update = getattr( - encoder_config, "post_diffusers_config_update", None - ) - if post_diffusers_config_update is not None: - post_diffusers_config_update() + encoder_config.post_diffusers_config_update() model_cls = _resolve_and_configure_encoder_quantization( encoder_config, model_config, @@ -778,6 +560,34 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): f"Failed to load quantized native {component_name!r}: {error}" ) from error + def build_model_config( + self, + component_model_path: str, + model_config: dict, + server_args: ServerArgs, + component_name: str, + ) -> EncoderConfig: + diffusers_pretrained_config = get_config( + component_model_path, trust_remote_code=True + ) + 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) + + encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index] + encoder_config.update_model_arch(model_config) + encoder_config.generation_config = load_dict( + os.path.join(component_model_path, "generation_config.json") + ) + + if encoder_index == 0: + for key, value in diffusers_pretrained_config.__dict__.items(): + setattr(encoder_config.arch_config, key, value) + return encoder_config + @staticmethod def _extract_encoder_index(component_name: str) -> int: """ @@ -872,20 +682,20 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): use_tensor_parallel_group(encoder_tp_group), set_default_torch_dtype(PRECISION_TO_TYPE[dtype]), ): - with model_device, skip_init_modules(): - architectures = getattr(model_config, "architectures", []) - model_cls, _ = ModelRegistry.resolve_model_cls(architectures) - enable_image_understanding = isinstance( - server_args.pipeline_config, - (QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig), - ) - model_config.enable_image_understanding = enable_image_understanding - # LongCat feeds its padded body to the DiT, so it must mask - # padding on the cache-free path; scoped so others are unchanged. - model_config.honor_cache_free_padding_mask = isinstance( - server_args.pipeline_config, LongCatImagePipelineConfig - ) - model = model_cls(model_config) + model_cls, _ = ModelRegistry.resolve_model_cls( + model_config.arch_config.architectures + ) + model_config.enable_image_understanding = isinstance( + server_args.pipeline_config, + (QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig), + ) + # longcat consumes the padded body without an attention cache + model_config.honor_cache_free_padding_mask = isinstance( + server_args.pipeline_config, LongCatImagePipelineConfig + ) + model = initialize_model( + model_cls, {"config": model_config}, param_dtype, model_device + ) if not isinstance(model, EncoderTensorParallelMixin): raise TypeError( @@ -904,15 +714,11 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): ) if component_starts_on_cpu and ( - current_platform.is_mps() or _keep_this_checkpoint_mapped(model_path) + current_platform.is_mps() + or keep_checkpoint_mapped( + weight_bytes=checkpoint_bytes(model_path), component=component_name + ) ): - # The encoder is layered immediately after this loader returns, - # so compatible CPU safetensors can stay mapped instead of being - # copied. On MPS that is always the right call -- the memory is - # unified. On any host it becomes the only call once the - # checkpoint is larger than host memory, because the copy is - # what does not fit: H3's encoder is 62.13 GiB against a 32 GiB - # target. model._keep_checkpoint_mapping = True weights_to_load = {name for name, _ in model.named_parameters()} @@ -931,6 +737,9 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): if isinstance(quant_config, QuantoInt8Config): checkpoint_weights = normalize_quanto_int8_weights(checkpoint_weights) loaded_weights = model.load_weights(checkpoint_weights) + self.validate_checkpoint_keys( + weights_to_load - loaded_weights, [], component_name + ) if quant_config is not None and not isinstance(quant_config, GGUFConfig): postprocess_device: torch.device | None = local_torch_device @@ -939,10 +748,10 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): and quant_config.is_checkpoint_int8_serialized ): postprocess_device = None - processed_layers = _process_quantized_encoder_weights( + processed_layers = process_model_weights_after_loading( model, postprocess_device, - component_name, + quantized_only=True, ) logger.info( "Processed %d %s linear layers for %s", @@ -961,39 +770,5 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader): model = model.to("cpu") else: model = model.to(local_torch_device) - # We only enable strict check for non-quantized models - # that have loaded weights tracking currently. - # if loaded_weights is not None: - weights_not_loaded = weights_to_load - loaded_weights - if weights_not_loaded: - # NOTE: - # If we silently continue with uninitialized weights, the text encoder can - # produce NaNs/garbage embeddings that later fail stage verification in a - # hard-to-debug way (e.g., `prompt_embeds` fails the NaN check). - # - # We allow a small set of known-optional parameters to be missing, but - # default to strict behavior for the rest. - allowed_missing_patterns = ( - getattr(model, "_allowed_missing_weights_patterns", []) or [] - ) - unexpected_missing = { - n - for n in weights_not_loaded - if not any(pat in n for pat in allowed_missing_patterns) - } - if unexpected_missing: - raise ValueError( - "Following text encoder weights were not initialized from checkpoint: " - f"{sorted(unexpected_missing)}. " - "This usually indicates a checkpoint/model-arch mismatch or a broken " - "weight-name mapping. If these are truly optional, set " - "`model._allowed_missing_weights_patterns` to whitelist patterns." - ) - logger.warning( - "Following (allowed) text encoder weights were not initialized from " - "checkpoint: %s (allowed patterns: %s)", - sorted(weights_not_loaded), - allowed_missing_patterns, - ) return model 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 ea8b4d940..a9fa3993b 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 @@ -16,7 +16,6 @@ from sglang.multimodal_gen.runtime.layers.attention.selector import ( from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( 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 from sglang.multimodal_gen.runtime.loader.minimax_h3_weights import ( comfy_quant_key_filter, @@ -525,20 +524,14 @@ class TransformerLoader(OnlineQuantizationComponentLoader): # Model construction resolves attention implementations, so apply the # quantization-specific default around FSDP initialization and loading. with attn_backend_context: - model = maybe_load_fsdp_model( + model = self.load_state_dict_model( model_cls=model_cls, init_params=init_params, - weight_dir_list=safetensors_list, - device=local_torch_device, - hsdp_replicate_dim=server_args.hsdp_replicate_dim, - hsdp_shard_dim=server_args.hsdp_shard_dim, + weight_files=safetensors_list, + server_args=component_server_args, + component_name=component_name, component_starts_on_cpu=component_starts_on_cpu, - pin_cpu_memory=component_server_args.pin_cpu_memory, - fsdp_inference=use_fsdp, - param_dtype=quant_spec.param_dtype, - reduce_dtype=torch.float32, - output_dtype=None, - strict=False, + dtype=quant_spec.param_dtype, weight_load_plan=weight_load_plan, checkpoint_key_filter=checkpoint_key_filter, weights_iterator=( 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 c74805aaf..6d2b4c0b8 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 @@ -1,14 +1,12 @@ import hashlib import importlib.util import os -from collections.abc import Iterable import torch import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file from safetensors.torch import safe_open from safetensors.torch import save_file as safetensors_save_file -from torch.nn.utils import parametrize from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig @@ -25,10 +23,12 @@ 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, + adopt_plain_weight_norm_state, checkpoint_bytes, + initialize_model, keep_checkpoint_mapped, + load_model_state_dict, set_default_torch_dtype, - skip_init_modules, ) from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, @@ -283,56 +283,6 @@ def _hold_decoder_weights_in_decode_dtype( ) -def _match_checkpoint_dtypes(loaded: dict, target_state: dict) -> dict: - """Convert checkpoint tensors whose dtype differs from their parameter's. - - Assignment replaces the parameter rather than writing through it, so a - mismatched dtype would silently change the module's. Converting makes a - copy, which is the point: only the tensors that already match can stay on - the mapping. - """ - for name, tensor in list(loaded.items()): - param = target_state.get(name) - if param is not None and param.dtype != tensor.dtype: - loaded[name] = tensor.to(dtype=param.dtype) - return loaded - - -def _adopt_plain_weight_norm_state( - module: nn.Module, loaded_names: Iterable[str] -) -> int: - """Make deparameterized checkpoint weights native module state. - - PyTorch's weight-norm load hook accepts legacy ``weight_g``/``weight_v`` - tensors, while inference exports commonly fold those tensors into one - plain ``weight``. Removing only the matching parametrizations preserves - that already-computed weight exactly and leaves every other parameterized - module untouched. - """ - state_names = set(module.state_dict()) - module_by_name = dict(module.named_modules()) - owners: set[str] = set() - for name in loaded_names: - if name == "weight": - owner_name = "" - elif name.endswith(".weight"): - owner_name = name.removesuffix(".weight") - else: - continue - state_prefix = f"{owner_name}." if owner_name else "" - if { - f"{state_prefix}parametrizations.weight.original0", - f"{state_prefix}parametrizations.weight.original1", - }.issubset(state_names): - owners.add(owner_name) - - for owner_name in sorted(owners): - parametrize.remove_parametrizations( - module_by_name[owner_name], "weight", leave_parametrized=True - ) - return len(owners) - - def _vae_checkpoint_arch_metadata_names( vae_config: VAEConfig, target_state: dict[str, torch.Tensor], @@ -444,7 +394,7 @@ def _assign_direct_gpu_vae_state( vae_config: VAEConfig, ) -> tuple[int, tuple[str, ...]]: """Stream a complete standard VAE state directly onto its target device.""" - num_deparameterized = _adopt_plain_weight_norm_state( + num_deparameterized = adopt_plain_weight_norm_state( vae, _vae_checkpoint_tensor_names(weight_files) ) target_state, slots = _direct_gpu_vae_state_slots(vae, component_name) @@ -664,21 +614,15 @@ class VAELoader(WeightOverrideComponentLoader): return vae # Load from ModelRegistry (standard VAE classes) - if direct_gpu_weight_loading: - with ( - set_default_torch_dtype(vae_dtype), - skip_init_modules(), - torch.device("meta"), - ): - vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) - vae = vae_cls(vae_config) - else: - with ( - set_default_torch_dtype(vae_dtype), - skip_init_modules(), - ): - vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) - vae = vae_cls(vae_config).to(target_device) + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vae = initialize_model( + vae_cls, + {"config": vae_config}, + vae_dtype, + torch.device("meta") if direct_gpu_weight_loading else None, + ) + if not direct_gpu_weight_loading: + vae = vae.to(target_device) if os.path.isfile(component_weights_path): if not component_weights_path.endswith(".safetensors"): @@ -729,7 +673,7 @@ class VAELoader(WeightOverrideComponentLoader): for sf_path in safetensors_list: loaded.update(safetensors_load_file(sf_path)) _backfill_ltx2_audio_vae_latent_stats(loaded, component_type) - num_deparameterized = _adopt_plain_weight_norm_state(vae, loaded) + num_deparameterized = adopt_plain_weight_norm_state(vae, loaded) target_state = vae.state_dict() consumed_metadata = _consume_vae_checkpoint_arch_metadata( loaded, vae_config, target_state @@ -759,9 +703,8 @@ class VAELoader(WeightOverrideComponentLoader): component=f"{component_name or 'vae'} (VAE)", ) ) - if keep_mapping: - _match_checkpoint_dtypes(loaded, target_state) - vae.load_state_dict( + load_model_state_dict( + vae, loaded, strict=strict_load, assign=keep_mapping, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py index e616ea4f8..c93f5f25f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py @@ -1,80 +1,13 @@ -import re +import torch +from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) -from sglang.multimodal_gen.runtime.loader.utils import ( - load_safetensors_state_dict, - set_default_torch_dtype, - skip_init_modules, -) -from sglang.multimodal_gen.runtime.models.registry import ModelRegistry -from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision -from sglang.multimodal_gen.utils import PRECISION_TO_TYPE - -logger = init_logger(__name__) class VocoderLoader(PlainStateDictComponentLoader): component_names = ["vocoder"] - expected_library = "diffusers" - - 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 - ) - class_name = config.pop("_class_name", None) or self.component_architecture - assert class_name is not None, ( - "Vocoder class name must be available from component config or pipeline config." - ) - - server_args.model_paths[component_name] = component_model_path - - from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import ( - LTXVocoderConfig, - ) - - vocoder_config = LTXVocoderConfig() - vocoder_config.update_model_arch(config) - - resolved_vocoder_dtype = resolve_component_precision(server_args, "vocoder") - vocoder_dtype = ( - resolved_vocoder_dtype - if resolved_vocoder_dtype is not None - else PRECISION_TO_TYPE["fp32"] - ) - - component_starts_on_cpu = server_args.should_start_component_on_cpu( - component_name - ) - target_device = self.target_device(component_starts_on_cpu) - - with set_default_torch_dtype(vocoder_dtype), skip_init_modules(): - vocoder_cls, _ = ModelRegistry.resolve_model_cls(class_name) - vocoder = vocoder_cls(vocoder_config).to(target_device) - - loaded = load_safetensors_state_dict(component_weights_path) - mapping = vocoder_config.arch_config.param_names_mapping - loaded = {_remap_vocoder_key(k, mapping): v for k, v in loaded.items()} - - missing_keys, unexpected_keys = vocoder.load_state_dict(loaded, strict=False) - # A half-loaded vocoder produces plausible but wrong audio. - if missing_keys or unexpected_keys: - raise ValueError( - f"Vocoder weights at '{component_weights_path}' do not match the " - f"instantiated {class_name}. Missing: {sorted(missing_keys)}. " - f"Unexpected: {sorted(unexpected_keys)}." - ) - return vocoder - - -def _remap_vocoder_key(key: str, param_names_mapping: dict[str, str]) -> str: - # Applied in order, not first-match: one key can need several rules. - for pattern, replacement in param_names_mapping.items(): - key = re.sub(pattern, replacement, key) - return key + config_classes = {"vocoder": LTXVocoderConfig} + default_precision_attr = "audio_vae_precision" + default_dtype = torch.float32 diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index ca8ec495b..c8b01d1d6 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -8,7 +8,6 @@ from collections import Counter, defaultdict from collections.abc import Callable, Generator -from itertools import chain from types import MethodType from typing import Any @@ -40,9 +39,10 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import ( ) from sglang.multimodal_gen.runtime.loader import rank_local_checkpoint from sglang.multimodal_gen.runtime.loader.utils import ( + finalize_loaded_model, get_param_names_mapping, hf_to_custom_state_dict, - set_default_torch_dtype, + initialize_model, ) from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan from sglang.multimodal_gen.runtime.loader.weight_utils import ( @@ -50,10 +50,10 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import ( ) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.quantization_utils import ( + process_model_weights_after_loading, +) from sglang.multimodal_gen.utils import set_mixed_precision_policy -from sglang.srt.utils import is_npu - -_is_npu = is_npu() logger = init_logger(__name__) @@ -293,8 +293,9 @@ def maybe_load_fsdp_model( mp_policy=mp_policy, ) - with set_default_torch_dtype(default_torch_dtype), torch.device("meta"): - model = model_cls(**init_params) + model = initialize_model( + model_cls, init_params, default_torch_dtype, torch.device("meta") + ) # Check if we should use FSDP use_fsdp = fsdp_inference @@ -447,26 +448,10 @@ def maybe_load_fsdp_model( # move to device to perform postprocessing _move_to_device_preserving_meta(model, weight_postprocess_device) - for _, module in model.named_modules(): - quant_method = getattr(module, "quant_method", None) - if quant_method is not None and hasattr( - quant_method, "process_weights_after_loading" - ): - if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod): - # Activate the NZ format for storing weights, - # which is a specific optimization for Ascend NPU - torch.npu.config.allow_internal_format = True - quant_method.process_weights_after_loading(module) - if _is_npu: - torch.npu.empty_cache() + process_model_weights_after_loading(model) model.post_load_weights() - for n, p in chain(model.named_parameters(), model.named_buffers()): - if p.is_meta: - raise RuntimeError(f"Unexpected param or buffer {n} on meta device.") - # Avoid unintended computation graph accumulation during inference - if isinstance(p, torch.nn.Parameter): - p.requires_grad = False + finalize_loaded_model(model) # 4. deferred cpu offload if defer_cpu_placement: diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index 7b59cd02a..6a601d6d0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -10,12 +10,14 @@ import json import os import re from collections import defaultdict -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterable, Iterator +from itertools import chain from typing import Any, Dict, Type import torch from safetensors.torch import load_file as safetensors_load_file from torch import nn +from torch.nn.utils import parametrize from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.weights.source import ( @@ -45,6 +47,78 @@ def set_default_torch_dtype(dtype: torch.dtype): torch.set_default_dtype(old_dtype) +def initialize_model( + model_cls: type[nn.Module], + init_params: dict[str, Any], + dtype: torch.dtype, + device: torch.device | None = None, +) -> nn.Module: + """Construct a checkpoint-backed module without initializing replaceable weights.""" + with ( + set_default_torch_dtype(dtype), + skip_init_modules(), + device if device is not None else contextlib.nullcontext(), + ): + return model_cls(**init_params) + + +def finalize_loaded_model(model: nn.Module) -> nn.Module: + """Reject unmaterialized state and freeze parameters before inference.""" + for name, tensor in chain(model.named_parameters(), model.named_buffers()): + if tensor.is_meta: + raise RuntimeError(f"Unexpected param or buffer {name} on meta device.") + if isinstance(tensor, nn.Parameter): + tensor.requires_grad = False + return model.eval() + + +def adopt_plain_weight_norm_state( + module: nn.Module, loaded_names: Iterable[str] +) -> int: + """Restore folded weights without recomputing their checkpoint values.""" + state_names = set(module.state_dict()) + module_by_name = dict(module.named_modules()) + owners: set[str] = set() + for name in loaded_names: + if name == "weight": + owner_name = "" + elif name.endswith(".weight"): + owner_name = name.removesuffix(".weight") + else: + continue + state_prefix = f"{owner_name}." if owner_name else "" + if { + f"{state_prefix}parametrizations.weight.original0", + f"{state_prefix}parametrizations.weight.original1", + }.issubset(state_names): + owners.add(owner_name) + + for owner_name in sorted(owners): + parametrize.remove_parametrizations( + module_by_name[owner_name], "weight", leave_parametrized=True + ) + return len(owners) + + +def load_model_state_dict( + model: nn.Module, + state_dict: dict[str, torch.Tensor], + *, + strict: bool = True, + assign: bool = False, +): + """Restore plain module state, preserving constructor-declared mixed dtypes.""" + adopt_plain_weight_norm_state(model, state_dict) + if assign: + target_state = model.state_dict() + # assignment replaces storage; unlike copy loading it does not cast + for name, tensor in state_dict.items(): + target = target_state.get(name) + if target is not None and tensor.dtype != target.dtype: + state_dict[name] = tensor.to(dtype=target.dtype) + return model.load_state_dict(state_dict, strict=strict, assign=assign) + + def get_param_names_mapping( mapping_dict: dict[str, str | tuple[str, int, int]], ) -> Callable[[str], tuple[str, Any, Any]]: @@ -111,6 +185,8 @@ def hf_to_custom_state_dict( hf_param_sd: dict[str, torch.Tensor] | Iterator[tuple[str, torch.Tensor]], param_names_mapping: Callable[[str], tuple[str, Any, Any]], valid_target_names: set[str] | None = None, + *, + strict: bool = False, ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]: """ Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary. @@ -149,6 +225,10 @@ def hf_to_custom_state_dict( num_params_to_merge, ) if merge_index is not None: + if strict and merge_index in to_merge_params[target_param_name]: + raise ValueError( + f"Duplicate checkpoint slice for {target_param_name!r}" + ) to_merge_params[target_param_name][merge_index] = full_tensor if len(to_merge_params[target_param_name]) == num_params_to_merge: # cat at output dim according to the merge_index order @@ -161,6 +241,8 @@ def hf_to_custom_state_dict( else: continue existing_tensor = custom_param_sd.get(target_param_name) + if strict and existing_tensor is not None: + raise ValueError(f"Duplicate checkpoint mapping for {target_param_name!r}") if existing_tensor is not None and existing_tensor.dtype != full_tensor.dtype: existing_is_quantized = existing_tensor.dtype in _QUANTIZED_DTYPES current_is_quantized = full_tensor.dtype in _QUANTIZED_DTYPES @@ -180,6 +262,8 @@ def hf_to_custom_state_dict( full_tensor.dtype, ) custom_param_sd[target_param_name] = full_tensor + if strict and to_merge_params: + raise ValueError(f"Incomplete checkpoint slices for {sorted(to_merge_params)}") return custom_param_sd, reverse_param_names_mapping diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index 44f6309a5..5cfa32016 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -19,6 +19,10 @@ from torch.distributed.tensor import DTensor from tqdm.auto import tqdm from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.loader.utils import ( + _DEFAULT_SAFETENSORS_INDEX, + _list_safetensors_files, +) from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan from sglang.multimodal_gen.runtime.loader.weight_readers import ( FALLBACK_READER, @@ -33,6 +37,40 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +def checkpoint_weights_iterator( + model_path: str, + *, + to_cpu: bool = True, + key_filter: Callable[[str], bool] | None = None, + index_file: str = _DEFAULT_SAFETENSORS_INDEX, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Read a materialized component checkpoint, preferring indexed safetensors.""" + files = _list_safetensors_files( + model_path, index_file=index_file, key_filter=key_filter + ) + if files: + yield from safetensors_weights_iterator( + files, to_cpu=to_cpu, key_filter=key_filter + ) + return + if os.path.isfile(model_path): + files = [model_path] if model_path.endswith((".bin", ".pt")) else [] + else: + for suffix in ("*.bin", "*.pt"): + files = filter_files_not_needed_for_inference( + sorted(str(path) for path in Path(model_path).glob(suffix)) + ) + if files: + break + if not files: + raise ValueError( + f"No safetensors, bin, or pt checkpoint found at {model_path!r}" + ) + for name, tensor in pt_weights_iterator(files, to_cpu=to_cpu): + if key_filter is None or key_filter(name): + yield name, tensor + + def _disable_runai_streamer_rank_discovery_collective() -> None: """RunAI Model Streamer's ``find_local_ranks()`` fires a full-world collective on the first ``stream_files()`` of every streamer instance even diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py index a27690857..1f8d87bd5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py @@ -35,9 +35,11 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader ) from sglang.multimodal_gen.runtime.loader.fsdp_load import ( load_model_from_full_model_state_dict, +) +from sglang.multimodal_gen.runtime.loader.utils import ( + get_param_names_mapping, set_default_torch_dtype, ) -from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.hunyuan3d_paint import ( Hunyuan3DPaintUNet, ) diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 1b8ff0519..d0813d8ea 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -6,8 +6,14 @@ import struct from pathlib import Path from typing import Any, Callable, Dict, List, Optional +import torch from safetensors import safe_open +from torch import nn +from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) from sglang.multimodal_gen.runtime.layers.quantization import ( QuantizationConfig, get_quantization_config, @@ -26,15 +32,77 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_w4a8_conf KitchenW4A8Config, ) from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config +from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( + UnquantizedEmbeddingMethod, + VocabParallelEmbedding, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.srt.layers.linear import LinearBase as SrtLinearBase from sglang.srt.layers.modelopt_utils import canonicalize_modelopt_quant_algo +from sglang.srt.layers.quantization.unquant import ( + UnquantizedEmbeddingMethod as SrtUnquantizedEmbeddingMethod, +) +from sglang.srt.layers.quantization.unquant import ( + UnquantizedLinearMethod as SrtUnquantizedLinearMethod, +) +from sglang.srt.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding as SrtVocabParallelEmbedding, +) from sglang.srt.model_loader.checkpoint_quantization import ( resolve_checkpoint_quant_spec, ) +from sglang.srt.model_loader.post_load import stage_module_for_post_load +from sglang.srt.utils import is_npu logger = init_logger(__name__) +def process_model_weights_after_loading( + model: nn.Module, + process_device: torch.device | None = None, + *, + quantized_only: bool = False, +) -> int: + """Process native and SRT layers once, optionally staging one layer at a time.""" + processed_layers = 0 + for module in model.modules(): + if not isinstance( + module, + ( + LinearBase, + SrtLinearBase, + VocabParallelEmbedding, + SrtVocabParallelEmbedding, + ), + ): + continue + method = module.quant_method + if method is None: + continue + unquantized = isinstance( + method, + ( + UnquantizedLinearMethod, + SrtUnquantizedLinearMethod, + UnquantizedEmbeddingMethod, + SrtUnquantizedEmbeddingMethod, + ), + ) + if quantized_only and unquantized: + continue + if is_npu() and not unquantized: + torch.npu.config.allow_internal_format = True + if process_device is None: + method.process_weights_after_loading(module) + else: + with stage_module_for_post_load(module, process_device): + method.process_weights_after_loading(module) + if is_npu(): + torch.npu.empty_cache() + processed_layers += 1 + return processed_layers + + def inspect_comfy_quant_markers( safetensors_list: list[str], param_name_mapper: Callable[[str], str] | None = None, 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 df61480f6..d5da2eff9 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,7 +5,6 @@ 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 ( @@ -229,76 +228,6 @@ 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": { @@ -358,7 +287,7 @@ class TestComponentQuantizationAdmission(unittest.TestCase): ), patch( "sglang.multimodal_gen.runtime.loader.component_loaders." - "adapter_loader.ModelRegistry.resolve_model_cls" + "component_loader.ModelRegistry.resolve_model_cls" ) as resolve_model, self.assertRaises(ComponentCheckpointUnsupportedError), ): 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 4f0861a81..b87244665 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 @@ -41,6 +41,9 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): ), component_weights_paths={}, component_quantizations={}, + component_quantization_ignored_layers={}, + component_paths={}, + batching_max_size=1, component_precisions={}, encoder_parallel="replicate", resolve_component_attention_backend=lambda _name: (None, None), @@ -61,7 +64,7 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): def _config_patch(self, config): return mock.patch( "sglang.multimodal_gen.runtime.loader.component_loaders." - "image_encoder_loader.get_diffusers_component_config", + "text_encoder_loader.get_diffusers_component_config", return_value=config, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_plain_state_dict_loaders.py b/python/sglang/multimodal_gen/test/unit/test_plain_state_dict_loaders.py new file mode 100644 index 000000000..ae02219a3 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_plain_state_dict_loaders.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +from dataclasses import dataclass, field +from unittest.mock import patch + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import ( + LTX2ConnectorConfig, +) +from sglang.multimodal_gen.configs.models.adapter.ltx_2_duration_head import ( + LTX2DurationHeadConfig, +) +from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig +from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import ( + MOVADualTowerConfig, +) +from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder import ( + LTX25DiffusionDecoderConfig, +) +from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentCheckpointUnsupportedError, + PipelineComponentLoader, + PlainStateDictComponentLoader, +) +from sglang.multimodal_gen.runtime.loader.utils import set_default_torch_dtype +from sglang.multimodal_gen.runtime.models.registry import ModelRegistry +from sglang.multimodal_gen.runtime.server_args import ServerArgs + +# real architectures with reduced widths; checkpoint names use the external layout +CASES = [ + ( + "dual_tower_bridge", + "DualTowerConditionalBridge", + MOVADualTowerConfig, + { + "visual_layers": 1, + "audio_layers": 1, + "visual_hidden_dim": 16, + "audio_hidden_dim": 16, + "head_dim": 8, + }, + ), + ( + "duration_head", + "LTX2DurationHead", + LTX2DurationHeadConfig, + { + "video_cross_attention_dim": 8, + "audio_cross_attention_dim": 8, + "pooler_hidden_dim": 8, + "num_pooler_heads": 2, + "mlp_hidden_dim": 8, + }, + ), + ( + "connectors", + "LTX2TextConnectors", + LTX2ConnectorConfig, + { + "caption_channels": 8, + "text_proj_in_factor": 2, + "per_modality_projections": True, + "video_hidden_dim": 8, + "audio_hidden_dim": 8, + "video_connector_num_attention_heads": 2, + "video_connector_attention_head_dim": 4, + "video_connector_num_layers": 1, + "video_connector_num_learnable_registers": 2, + "audio_connector_num_attention_heads": 2, + "audio_connector_attention_head_dim": 4, + "audio_connector_num_layers": 1, + "audio_connector_num_learnable_registers": 2, + }, + ), + ( + "diffusion_decoder", + "LTX2VideoDiffusionDecoderModel", + LTX25DiffusionDecoderConfig, + { + "latent_channels": 4, + "decoder_head_dim": 8, + "decoder_t_emb_dim": 8, + "decoder_stage_channels": [8, 8, 8, 8, 8], + "decoder_stage_depths": [1, 1, 1, 1, 1], + "decoder_upsample_channel_reductions": [1, 1, 1, 1], + }, + ), + ( + "vocoder", + "LTX2VocoderWithBWE", + LTXVocoderConfig, + { + "hidden_channels": 32, + "upsample_factors": [2], + "upsample_kernel_sizes": [4], + "resnet_kernel_sizes": [3], + "resnet_dilations": [[1, 3, 5]], + "bwe_hidden_channels": 32, + "bwe_upsample_factors": [2], + "bwe_upsample_kernel_sizes": [4], + "bwe_resnet_kernel_sizes": [3], + "bwe_resnet_dilations": [[1, 3, 5]], + "input_sampling_rate": 16000, + "output_sampling_rate": 32000, + }, + ), +] + + +def _write_checkpoint(path, config, weights): + path.mkdir(exist_ok=True) + (path / "config.json").write_text(json.dumps(config)) + save_file( + {name: tensor.contiguous() for name, tensor in weights.items()}, + path / "model.safetensors", + ) + + +@pytest.mark.parametrize("role,class_name,config_cls,raw_config", CASES) +@pytest.mark.parametrize("residency", ["component-offload", "resident"]) +@pytest.mark.parametrize("precision", ["fp32", "bf16"]) +def test_real_components_restore_weights_and_exact_policy( + tmp_path, role, class_name, config_cls, raw_config, residency, precision +): + if residency == "resident" and not torch.cuda.is_available(): + pytest.skip("CUDA is required") + config = config_cls() + config.update_model_arch(raw_config) + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + dtype = torch.float32 if precision == "fp32" else torch.bfloat16 + with set_default_torch_dtype(dtype): + reference = model_cls(config).eval() + if role != "vocoder": + reference = reference.to(dtype=dtype) + # custom linear constructors allocate empty storage, not checkpoint values + generator = torch.Generator().manual_seed(0) + with torch.no_grad(): + for parameter in reference.parameters(): + parameter.uniform_(-0.1, 0.1, generator=generator) + weights = {} + for name, tensor in reference.state_dict().items(): + assert torch.isfinite(tensor).all(), name + name = name.replace("video_aggregate_embed.", "video_text_proj_in.").replace( + "audio_aggregate_embed.", "audio_text_proj_in." + ) + if role == "vocoder": + name = name.replace(".conv_pre.", ".conv_in.").replace( + ".conv_post.", ".conv_out." + ) + name = name.replace(".act_post.", ".act_out.").replace( + ".ups.", ".upsamplers." + ) + name = name.replace(".resblocks.", ".resnets.").replace( + ".downsample.lowpass.filter", ".downsample.filter" + ) + weights[name] = tensor + component = tmp_path / role + _write_checkpoint(component, {"_class_name": class_name, **raw_config}, weights) + # the exact key differs from the structural role, including for vocoders + name = role + "_2" + args = ServerArgs( + model_path="x", + component_precisions={name: precision}, + component_residency={name: residency}, + component_weights_paths={name: str(component / "model.safetensors")}, + ) + model, _ = PipelineComponentLoader.load_component( + name, str(component), "diffusers", args, component_type=role + ) + device = "cuda" if residency == "resident" else "cpu" + assert not model.training + assert args.model_paths[name] == str(component) + for key, expected in reference.state_dict().items(): + actual = model.state_dict()[key] + assert actual.device.type == device + assert actual.dtype == expected.dtype + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + # constructor-owned nonpersistent filters must survive common placement too + for key, expected in reference.named_buffers(): + torch.testing.assert_close( + model.get_buffer(key).cpu(), expected, rtol=0, atol=0 + ) + if role == "duration_head": + reference = reference.to(device) + inputs = torch.ones(1, 2, 8, device=device, dtype=dtype) + with torch.inference_mode(): + torch.testing.assert_close(model(inputs), reference(inputs), rtol=0, atol=0) + + +class _DecoderOnly(nn.Module): + def __init__(self, config): + super().__init__() + self.decoder = nn.Linear(config["width"], 2) + + +class _WeightNormModule(nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.utils.parametrizations.weight_norm( + nn.Linear(2, 4, bias=False) + ) + + +def test_plain_component_restores_folded_weight_norm(tmp_path): + expected = torch.arange(8, dtype=torch.float32).reshape(4, 2) + _write_checkpoint( + tmp_path, {"_class_name": "TestWeightNormModule"}, {"proj.weight": expected} + ) + args = ServerArgs( + model_path="x", + component_precisions={"auxiliary": "fp32"}, + component_residency={"auxiliary": "component-offload"}, + ) + with patch.dict(ModelRegistry.registered_models): + ModelRegistry.register_model("TestWeightNormModule", _WeightNormModule) + model, _ = PlainStateDictComponentLoader().load( + str(tmp_path), args, "auxiliary", "diffusers" + ) + torch.testing.assert_close(model.proj.weight, expected, rtol=0, atol=0) + assert set(model.state_dict()) == {"proj.weight"} + + +@dataclass +class _MappedArch(ArchConfig): + param_names_mapping: dict = field( + default_factory=lambda: { + r"^q.weight$": ("proj.weight", 0, 2), + r"^k.weight$": ("proj.weight", 1, 2), + } + ) + + +@dataclass +class _MappedConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=_MappedArch) + + +class _MappedModule(nn.Module): + def __init__(self, config): + super().__init__() + self.proj = nn.Linear(2, 4, bias=False) + + +@pytest.mark.parametrize("invalid", [None, "collision", "incomplete"]) +def test_plain_components_use_shared_fused_weight_mapping(tmp_path, invalid): + weights = {"q.weight": torch.ones(2, 2), "k.weight": torch.zeros(2, 2)} + if invalid == "collision": + weights["proj.weight"] = torch.ones(4, 2) + elif invalid == "incomplete": + weights.pop("k.weight") + _write_checkpoint(tmp_path, {"_class_name": "TestMappedModule"}, weights) + loader = PlainStateDictComponentLoader() + loader.config_classes = {"auxiliary": _MappedConfig} + args = ServerArgs( + model_path="x", component_residency={"auxiliary": "component-offload"} + ) + with patch.dict(ModelRegistry.registered_models): + ModelRegistry.register_model("TestMappedModule", _MappedModule) + if invalid: + with pytest.raises(ComponentCheckpointUnsupportedError): + loader.load(str(tmp_path), args, "auxiliary", "diffusers") + else: + model, _ = loader.load(str(tmp_path), args, "auxiliary", "diffusers") + expected = torch.cat([weights["q.weight"], weights["k.weight"]]).bfloat16() + torch.testing.assert_close(model.proj.weight, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("invalid", [None, "missing", "unexpected", "shape"]) +def test_sound_tokenizer_only_ignores_encoder_weights(tmp_path, invalid): + reference = _DecoderOnly({"width": 4}) + weights = dict(reference.state_dict(), **{"encoder.weight": torch.ones(2, 4)}) + if invalid == "missing": + weights.pop("decoder.bias") + elif invalid == "unexpected": + weights["other.weight"] = torch.ones(2, 4) + elif invalid == "shape": + weights["decoder.weight"] = torch.ones(3, 4) + _write_checkpoint(tmp_path, {"_class_name": "TestDecoderOnly", "width": 4}, weights) + args = ServerArgs( + model_path="x", component_residency={"sound_tokenizer": "component-offload"} + ) + with patch.dict(ModelRegistry.registered_models): + ModelRegistry.register_model("TestDecoderOnly", _DecoderOnly) + if invalid: + with pytest.raises(ComponentCheckpointUnsupportedError): + PipelineComponentLoader.load_component( + "sound_tokenizer", str(tmp_path), "diffusers", args + ) + else: + model, _ = PipelineComponentLoader.load_component( + "sound_tokenizer", str(tmp_path), "diffusers", args + ) + torch.testing.assert_close( + model.decoder.weight, + reference.decoder.weight, + rtol=0, + atol=0, + ) 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 b1cd198fb..36ea9f522 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 @@ -5,11 +5,13 @@ from pathlib import Path from types import SimpleNamespace from unittest import mock +import pytest import torch import transformers from safetensors.torch import save_file from torch import nn +from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.layers.quantization.comfy_nvfp4 import ( ComfyFullPrecisionNvfp4LinearMethod, @@ -35,11 +37,13 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader TextEncoderLoader, _configure_encoder_quantization, _get_encoder_quant_config, - _process_quantized_encoder_weights, _require_quantized_encoder_layers, _resolve_and_configure_encoder_quantization, ) from sglang.multimodal_gen.runtime.loader.gguf_weights import GGUFTensorMeta +from sglang.multimodal_gen.runtime.loader.weight_utils import ( + checkpoint_weights_iterator, +) from sglang.multimodal_gen.runtime.models.encoders.base import ( EncoderTensorParallelMixin, TextEncoder, @@ -49,9 +53,67 @@ from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import ( MiniMaxH3Qwen3VLEncoder, ) from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextModel +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.quantization_utils import ( + process_model_weights_after_loading, +) from sglang.srt.layers.linear import LinearBase as SrtLinearBase +@pytest.mark.parametrize("missing", [False, True]) +@pytest.mark.parametrize("competing_index", [False, True]) +def test_native_encoder_restoration_checks_checkpoint_before_fallback( + tmp_path, missing, competing_index +): + config = transformers.T5Config( + vocab_size=32, + d_model=8, + d_kv=4, + d_ff=16, + num_layers=1, + num_heads=2, + architectures=["T5EncoderModel"], + ) + reference = transformers.T5EncoderModel(config) + config.save_pretrained(tmp_path) + weights = {name: tensor.clone() for name, tensor in reference.state_dict().items()} + required = "encoder.final_layer_norm.weight" + if missing: + weights.pop(required) + save_file(weights, tmp_path / "model.safetensors") + if competing_index: + # LTX-2 ships both indexes in text_encoder; only the HF one owns this model + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {name: "model.safetensors" for name in weights}}) + ) + alternate = "diffusion_pytorch_model.safetensors" + save_file({"diffusion.weight": torch.ones(2, 2)}, tmp_path / alternate) + (tmp_path / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"diffusion.weight": alternate}}) + ) + alternate_weights = dict(checkpoint_weights_iterator(str(tmp_path))) + assert set(alternate_weights) == {"diffusion.weight"} + args = ServerArgs( + model_path="x", + component_precisions={"text_encoder": "fp32"}, + component_residency={"text_encoder": "component-offload"}, + ) + args.pipeline_config.text_encoder_configs = (T5Config(),) + loader = TextEncoderLoader() + with mock.patch.object( + loader, "load_native", side_effect=AssertionError("fallback") + ): + if missing: + with pytest.raises(ComponentCheckpointUnsupportedError, match=required): + loader.load(str(tmp_path), args, "text_encoder", "transformers") + else: + model, _ = loader.load(str(tmp_path), args, "text_encoder", "transformers") + torch.testing.assert_close( + model.state_dict()[required], weights[required], rtol=0, atol=0 + ) + assert all(not parameter.requires_grad for parameter in model.parameters()) + + class TestTextEncoderWeightDiscovery(unittest.TestCase): def test_prepare_weights_prefers_canonical_over_fp16_variant(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -59,22 +121,13 @@ class TestTextEncoderWeightDiscovery(unittest.TestCase): canonical = model_dir / "model.safetensors" variant = model_dir / "model.fp16.safetensors" - canonical.touch() - variant.touch() + weight = torch.ones(2, 2) + save_file({"weight": weight}, canonical) + save_file({"weight": weight.half()}, variant) - ( - hf_folder, - weight_files, - use_safetensors, - ) = TextEncoderLoader()._prepare_weights( - str(model_dir), - fall_back_to_pt=True, - allow_patterns_overrides=None, - ) - - self.assertEqual(hf_folder, str(model_dir)) - self.assertTrue(use_safetensors) - self.assertEqual(weight_files, [str(canonical)]) + state = dict(checkpoint_weights_iterator(str(model_dir))) + self.assertEqual(state["weight"].dtype, torch.float32) + torch.testing.assert_close(state["weight"], weight) class TestTextEncoderClassResolution(unittest.TestCase): @@ -757,7 +810,10 @@ class TestTextEncoderQuantization(unittest.TestCase): ), ): _resolve_and_configure_encoder_quantization( - SimpleNamespace(architectures=[architecture], quant_config=None), + SimpleNamespace( + arch_config=SimpleNamespace(architectures=[architecture]), + quant_config=None, + ), component_config, "/model/text_encoder", "/model/text_encoder", @@ -791,7 +847,10 @@ class TestTextEncoderQuantization(unittest.TestCase): ): _resolve_and_configure_encoder_quantization( SimpleNamespace( - architectures=["ThirdPartyTextEncoder"], quant_config=None + arch_config=SimpleNamespace( + architectures=["ThirdPartyTextEncoder"] + ), + quant_config=None, ), { "quantization_config": { @@ -812,7 +871,10 @@ class TestTextEncoderQuantization(unittest.TestCase): ): _resolve_and_configure_encoder_quantization( SimpleNamespace( - architectures=["ThirdPartyTextEncoder"], quant_config=None + arch_config=SimpleNamespace( + architectures=["ThirdPartyTextEncoder"] + ), + quant_config=None, ), { "quantization_config": { @@ -887,10 +949,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): quant_method = _RecordingQuantMethod() model = _SRTQuantizedLinear(quant_method) - processed = _process_quantized_encoder_weights( + processed = process_model_weights_after_loading( model, torch.device("cpu"), - "image_encoder", + quantized_only=True, ) self.assertEqual(processed, 1) @@ -926,10 +988,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): quant_method = _RecordingQuantMethod() model = _QuantizedEncoder(quant_method) - processed = _process_quantized_encoder_weights( + processed = process_model_weights_after_loading( model, torch.device("cpu"), - "text_encoder", + quantized_only=True, ) self.assertEqual(processed, 1) @@ -941,10 +1003,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): quant_method = _RecordingQuantMethod() model = _QuantizedEncoder(quant_method) - processed = _process_quantized_encoder_weights( + processed = process_model_weights_after_loading( model, torch.device("cuda", torch.cuda.current_device()), - "text_encoder", + quantized_only=True, ) self.assertEqual(processed, 1) @@ -957,10 +1019,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): model = _QuantizedEncoder(_RecordingQuantMethod(error=RuntimeError("boom"))) with self.assertRaisesRegex(RuntimeError, "boom"): - _process_quantized_encoder_weights( + process_model_weights_after_loading( model, torch.device("cuda", torch.cuda.current_device()), - "text_encoder", + quantized_only=True, ) self.assertEqual(model.quantized.weight.device, torch.device("cpu")) 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 4bd8cbe9a..163ac4752 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -27,18 +27,17 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp ComponentCheckpointUnsupportedError, ) from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import ( - _adopt_plain_weight_norm_state, _assign_direct_gpu_vae_state, _backfill_ltx2_audio_vae_latent_stats, _consume_vae_checkpoint_arch_metadata, _direct_gpu_vae_state_slots, - _match_checkpoint_dtypes, _require_native_loader_for_quantized_vae, _should_use_channels_last_3d, ) from sglang.multimodal_gen.runtime.loader.utils import ( checkpoint_bytes, keep_checkpoint_mapped, + load_model_state_dict, ) from sglang.multimodal_gen.runtime.managers.memory_managers import ( host_memory_budget, @@ -47,6 +46,7 @@ from sglang.multimodal_gen.runtime.models.vaes import wanvae from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.decoding_av import ( LTX2AVDecodingStage, ) +from sglang.test.test_utils import CustomTestCase class _FakeServerArgs: @@ -131,28 +131,26 @@ class TestKeepCheckpointMapped(unittest.TestCase): ) -class TestMatchCheckpointDtypes(unittest.TestCase): +class TestMatchCheckpointDtypes(CustomTestCase): """Assignment replaces a parameter, so only matching dtypes may stay mapped.""" - def test_a_matching_tensor_is_left_alone(self): - loaded = {"w": torch.zeros(4, dtype=torch.float32)} - before = loaded["w"] - _match_checkpoint_dtypes(loaded, {"w": torch.zeros(4, dtype=torch.float32)}) - self.assertIs(loaded["w"], before) - - def test_a_mismatched_tensor_is_converted(self): - loaded = {"w": torch.zeros(4, dtype=torch.float32)} - _match_checkpoint_dtypes(loaded, {"w": torch.zeros(4, dtype=torch.bfloat16)}) - self.assertEqual(loaded["w"].dtype, torch.bfloat16) - - def test_a_tensor_the_module_does_not_want_is_left_alone(self): - loaded = {"extra": torch.zeros(4, dtype=torch.float32)} - before = loaded["extra"] - _match_checkpoint_dtypes(loaded, {}) - self.assertIs(loaded["extra"], before) + def test_assignment_preserves_mixed_dtypes_and_matching_storage(self): + model = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16) + model.register_buffer("scale", torch.zeros(4, dtype=torch.float32)) + weights = { + "weight": torch.ones(4, 4, dtype=torch.float32), + "scale": torch.ones(4, dtype=torch.float32), + } + checkpoint_weight = weights["weight"] + load_model_state_dict(model, weights, assign=True) + self.assertEqual(model.weight.dtype, torch.bfloat16) + self.assertEqual(model.scale.dtype, torch.float32) + self.assertEqual(model.scale.data_ptr(), weights["scale"].data_ptr()) + self.assertNotEqual(model.weight.data_ptr(), checkpoint_weight.data_ptr()) + self.assertTrue(torch.equal(model.weight.float(), checkpoint_weight)) -class TestPlainWeightNormCheckpoint(unittest.TestCase): +class TestPlainWeightNormCheckpoint(CustomTestCase): def test_adopts_a_folded_weight_without_reconstructing_it(self): module = nn.Sequential( torch.nn.utils.parametrizations.weight_norm( @@ -162,8 +160,7 @@ class TestPlainWeightNormCheckpoint(unittest.TestCase): expected = torch.arange(18, dtype=torch.float32).reshape(3, 2, 3) / 19 loaded = {"0.weight": expected} - self.assertEqual(_adopt_plain_weight_norm_state(module, loaded), 1) - module.load_state_dict(loaded, strict=True) + load_model_state_dict(module, loaded) self.assertEqual(set(module.state_dict()), {"0.weight"}) self.assertTrue(torch.equal(module[0].weight, expected)) @@ -180,8 +177,7 @@ class TestPlainWeightNormCheckpoint(unittest.TestCase): "0.weight_v": original_state["0.parametrizations.weight.original1"].clone(), } - self.assertEqual(_adopt_plain_weight_norm_state(module, loaded), 0) - module.load_state_dict(loaded, strict=True) + load_model_state_dict(module, loaded) self.assertIn("0.parametrizations.weight.original0", module.state_dict())