[diffusion] refactor: consolidate plain state-dict component loaders (#38128)

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):

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

python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py

@@ -1,127 +1,11 @@
-from copy import deepcopy
+from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
-
+    MOVADualTowerConfig,
-import torch
+)
-
-from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 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"
+    config_classes = {"dual_tower_bridge": MOVADualTowerConfig}
-
-    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

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

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"
+    config_classes = {"diffusion_decoder": LTX25DiffusionDecoderConfig}
-
+    default_precision_attr = "vae_precision"
-    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

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(
+    def build_model_config(
-        self,
+        self, component_model_path, model_config, server_args, component_name
-        component_model_path: str,
-        server_args: ServerArgs,
-        component_name: str = "image_encoder",
     ):
-        """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(
+        return encoder_config
-            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,
-        )

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(
+    def build_model_config(self, config, component_name):
-        self, component_model_path: str, server_args: ServerArgs, component_name: str
+        return {"config": config}
-    ):
-        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

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 collections.abc import Generator
-from typing import cast
 
 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,
+    checkpoint_weights_iterator,
-    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,
 )
 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="",
+            to_cpu=to_cpu,
-            fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
+            key_filter=include_checkpoint_weight,
-            allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
+            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
         )
-
+        encoder_config = self.build_model_config(
-        # TODO(mick): had to throw an exception for different text-encoder arch
+            component_model_path, model_config, server_args, component_name
-        encoder_index = self._extract_encoder_index(
-            self.structural_component_name(component_name)
         )
-        assert encoder_index < len(
+        encoder_config.post_diffusers_config_update()
-            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()
         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():
+            model_cls, _ = ModelRegistry.resolve_model_cls(
-                architectures = getattr(model_config, "architectures", [])
+                model_config.arch_config.architectures
-                model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+            )
-                enable_image_understanding = isinstance(
+            model_config.enable_image_understanding = isinstance(
-                    server_args.pipeline_config,
+                server_args.pipeline_config,
-                    (QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig),
+                (QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig),
-                )
+            )
-                model_config.enable_image_understanding = enable_image_understanding
+            # longcat consumes the padded body without an attention cache
-                # LongCat feeds its padded body to the DiT, so it must mask
+            model_config.honor_cache_free_padding_mask = isinstance(
-                # padding on the cache-free path; scoped so others are unchanged.
+                server_args.pipeline_config, LongCatImagePipelineConfig
-                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
-                model = model_cls(model_config)
+            )
 
             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

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,
+                weight_files=safetensors_list,
-                device=local_torch_device,
+                server_args=component_server_args,
-                hsdp_replicate_dim=server_args.hsdp_replicate_dim,
+                component_name=component_name,
-                hsdp_shard_dim=server_args.hsdp_shard_dim,
                 component_starts_on_cpu=component_starts_on_cpu,
-                pin_cpu_memory=component_server_args.pin_cpu_memory,
+                dtype=quant_spec.param_dtype,
-                fsdp_inference=use_fsdp,
-                param_dtype=quant_spec.param_dtype,
-                reduce_dtype=torch.float32,
-                output_dtype=None,
-                strict=False,
                 weight_load_plan=weight_load_plan,
                 checkpoint_key_filter=checkpoint_key_filter,
                 weights_iterator=(

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:
+        vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-            with (
+        vae = initialize_model(
-                set_default_torch_dtype(vae_dtype),
+            vae_cls,
-                skip_init_modules(),
+            {"config": vae_config},
-                torch.device("meta"),
+            vae_dtype,
-            ):
+            torch.device("meta") if direct_gpu_weight_loading else None,
-                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+        )
-                vae = vae_cls(vae_config)
+        if not direct_gpu_weight_loading:
-        else:
+            vae = vae.to(target_device)
-            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)
 
         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:
+        load_model_state_dict(
-            _match_checkpoint_dtypes(loaded, target_state)
+            vae,
-        vae.load_state_dict(
             loaded,
             strict=strict_load,
             assign=keep_mapping,

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"
+    config_classes = {"vocoder": LTXVocoderConfig}
-
+    default_precision_attr = "audio_vae_precision"
-    def load_customized(
+    default_dtype = torch.float32
-        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

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 = initialize_model(
-        model = model_cls(**init_params)
+        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():
+    process_model_weights_after_loading(model)
-        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()
     model.post_load_weights()
 
-    for n, p in chain(model.named_parameters(), model.named_buffers()):
+    finalize_loaded_model(model)
-        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
 
     # 4. deferred cpu offload
     if defer_cpu_placement:

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
 
 

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

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,
 )

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,

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),
         ):

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,
         )
 

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,
+            )

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()
+            weight = torch.ones(2, 2)
-            variant.touch()
+            save_file({"weight": weight}, canonical)
+            save_file({"weight": weight.half()}, variant)
 
-            (
+            state = dict(checkpoint_weights_iterator(str(model_dir)))
-                hf_folder,
+            self.assertEqual(state["weight"].dtype, torch.float32)
-                weight_files,
+            torch.testing.assert_close(state["weight"], weight)
-                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)])
 
 
 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"))

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):
+    def test_assignment_preserves_mixed_dtypes_and_matching_storage(self):
-        loaded = {"w": torch.zeros(4, dtype=torch.float32)}
+        model = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16)
-        before = loaded["w"]
+        model.register_buffer("scale", torch.zeros(4, dtype=torch.float32))
-        _match_checkpoint_dtypes(loaded, {"w": torch.zeros(4, dtype=torch.float32)})
+        weights = {
-        self.assertIs(loaded["w"], before)
+            "weight": torch.ones(4, 4, dtype=torch.float32),
-
+            "scale": torch.ones(4, dtype=torch.float32),
-    def test_a_mismatched_tensor_is_converted(self):
+        }
-        loaded = {"w": torch.zeros(4, dtype=torch.float32)}
+        checkpoint_weight = weights["weight"]
-        _match_checkpoint_dtypes(loaded, {"w": torch.zeros(4, dtype=torch.bfloat16)})
+        load_model_state_dict(model, weights, assign=True)
-        self.assertEqual(loaded["w"].dtype, torch.bfloat16)
+        self.assertEqual(model.weight.dtype, torch.bfloat16)
-
+        self.assertEqual(model.scale.dtype, torch.float32)
-    def test_a_tensor_the_module_does_not_want_is_left_alone(self):
+        self.assertEqual(model.scale.data_ptr(), weights["scale"].data_ptr())
-        loaded = {"extra": torch.zeros(4, dtype=torch.float32)}
+        self.assertNotEqual(model.weight.data_ptr(), checkpoint_weight.data_ptr())
-        before = loaded["extra"]
+        self.assertTrue(torch.equal(model.weight.float(), checkpoint_weight))
-        _match_checkpoint_dtypes(loaded, {})
-        self.assertIs(loaded["extra"], before)
 
 
-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)
+        load_model_state_dict(module, loaded)
-        module.load_state_dict(loaded, strict=True)
 
         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)
+        load_model_state_dict(module, loaded)
-        module.load_state_dict(loaded, strict=True)
 
         self.assertIn("0.parametrizations.weight.original0", module.state_dict())