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

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

Browse Source
This commit is contained in:
Mick
2026-09-06 13:52:54 +08:00
committed by GitHub
parent 67e3ccda97
commit e3f7097591
21 changed files with 898 additions and 976 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/base.py
+6
View File
@@ -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 -88
View File
@@ -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
+4 -120
View File
@@ -1,127 +1,11 @@
from copy import deepcopy
import torch
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
MOVADualTowerConfig,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PlainStateDictComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
RESIDENT,
)
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.precision import resolve_precision
logger = init_logger(__name__)
class BridgeLoader(PlainStateDictComponentLoader):
"""Loader for MOVA dual tower bridge with FSDP support."""
pipeline_bridge_config_attr: str = "bridge_config"
component_names = ["dual_tower_bridge"]
expected_library = "diffusers"
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
):
config = self.load_component_config(component_model_path, component_name)
component_weights_path = self.resolve_component_weights_path(
component_model_path, server_args, component_name
)
hf_config = deepcopy(config)
class_name = config.pop("_class_name", None)
if class_name is None:
raise ValueError(
"Model config does not contain a _class_name attribute. "
"Only diffusers format is supported."
)
server_args.model_paths[component_name] = component_model_path
# Try to get bridge config from pipeline config, fallback to creating one
bridge_config = getattr(
server_args.pipeline_config, self.pipeline_bridge_config_attr, None
)
if bridge_config is not None:
bridge_config.update_model_arch(config)
else:
# Create a minimal config from hf_config
from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
MOVADualTowerConfig,
)
bridge_config = MOVADualTowerConfig()
bridge_config.update_model_arch(config)
model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
# Find all safetensors files
safetensors_list = _list_safetensors_files(component_weights_path)
if not safetensors_list:
raise ValueError(f"No safetensors files found in {component_weights_path}")
default_dtype = resolve_precision(
server_args, component_name, precision_attr="dit_precision"
)
logger.info(
"Loading %s from %s safetensors files, default_dtype: %s",
class_name,
len(safetensors_list),
default_dtype,
)
use_fsdp = server_args.should_use_fsdp_for_component(component_name)
component_starts_on_cpu = server_args.should_start_component_on_cpu(
component_name
)
# Use the FSDP loader when FSDP is requested or shard rules are declared.
fsdp_shard_conditions = getattr(model_cls, "_fsdp_shard_conditions", None)
if (
component_weights_path != component_model_path
or use_fsdp
or (
server_args.residency_mode(component_name) == RESIDENT
and server_args.hsdp_shard_dim is not None
and fsdp_shard_conditions
)
):
local_torch_device = get_local_torch_device()
# Load with FSDP support
model = maybe_load_fsdp_model(
model_cls=model_cls,
init_params={"config": bridge_config, "hf_config": hf_config},
weight_dir_list=safetensors_list,
device=local_torch_device,
hsdp_replicate_dim=server_args.hsdp_replicate_dim,
hsdp_shard_dim=server_args.hsdp_shard_dim,
component_starts_on_cpu=component_starts_on_cpu,
pin_cpu_memory=server_args.pin_cpu_memory,
fsdp_inference=use_fsdp,
param_dtype=default_dtype,
reduce_dtype=torch.float32,
output_dtype=None,
strict=False,
weight_load_plan=WeightLoadPlan(
checkpoint_load_device=local_torch_device
),
)
else:
# Fallback to simple loading (for non-FSDP or legacy models)
model = model_cls.from_pretrained(
component_model_path, torch_dtype=default_dtype
)
target_device = self.target_device(component_starts_on_cpu)
model = model.to(device=target_device, dtype=default_dtype)
total_params = sum(p.numel() for p in model.parameters())
logger.info("Loaded bridge model with %.2fM parameters", total_params / 1e6)
return model
config_classes = {"dual_tower_bridge": MOVADualTowerConfig}
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+160 -3
View File
@@ -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
+3 -49
View File
@@ -6,57 +6,11 @@ from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder imp
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PlainStateDictComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.utils import (
load_safetensors_state_dict,
set_default_torch_dtype,
skip_init_modules,
)
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.precision import resolve_precision
class DiffusionDecoderLoader(PlainStateDictComponentLoader):
"""Loader for the standalone, replicated LTX-2.5 diffusion decoder."""
"""Standalone, replicated LTX-2.5 diffusion decoder."""
component_names = ["diffusion_decoder"]
expected_library = "diffusers"
def load_customized(
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str = "diffusion_decoder",
*args,
):
config = self.load_component_config(component_model_path, component_name)
component_weights_path = self.resolve_component_weights_path(
component_model_path, server_args, component_name
)
class_name = config.pop("_class_name", None)
if class_name is None:
raise ValueError(
"Model config does not contain a _class_name attribute. "
"Only diffusers format is supported."
)
config.pop("_diffusers_version", None)
config.pop("_name_or_path", None)
server_args.model_paths[component_name] = component_model_path
model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
target_device = self.target_device(
server_args.should_start_component_on_cpu(component_name)
)
dtype = resolve_precision(
server_args, component_name, precision_attr="vae_precision"
)
decoder_config = LTX25DiffusionDecoderConfig()
decoder_config.update_model_arch(config)
with set_default_torch_dtype(dtype), skip_init_modules():
model = model_cls(decoder_config).to(device=target_device, dtype=dtype)
model.load_state_dict(
load_safetensors_state_dict(component_weights_path), strict=True
)
return model
config_classes = {"diffusion_decoder": LTX25DiffusionDecoderConfig}
default_precision_attr = "vae_precision"
python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
+3 -54
View File
@@ -1,15 +1,7 @@
from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
TextEncoderLoader,
_resolve_and_configure_encoder_quantization,
)
from sglang.multimodal_gen.runtime.models.encoders.base import finalize_encoder_folding
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
get_diffusers_component_config,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
class ImageEncoderLoader(TextEncoderLoader):
@@ -23,52 +15,9 @@ class ImageEncoderLoader(TextEncoderLoader):
component_name, server_args.pipeline_config.image_encoder_precision
)
def load_customized(
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str = "image_encoder",
def build_model_config(
self, component_model_path, model_config, server_args, component_name
):
"""Load the text encoders based on the model path, and inference args."""
component_weights_path = self.resolve_component_weights_path(
component_model_path,
server_args,
component_name,
)
# model_config: PretrainedConfig = get_hf_config(
# model=model_path,
# trust_remote_code=server_args.trust_remote_code,
# revision=server_args.revision,
# model_override_args=None,
# )
model_config = get_diffusers_component_config(
component_path=component_model_path
)
encoder_config = server_args.pipeline_config.image_encoder_config
encoder_config.update_model_arch(model_config)
_resolve_and_configure_encoder_quantization(
encoder_config,
model_config,
component_model_path,
component_weights_path,
component_name,
server_args.component_quantizations.get(component_name),
)
# real dims are populated now; resolve fold vs replicate
finalize_encoder_folding(
encoder_config,
server_args.encoder_parallel,
)
# Always start with local device; load_model will adjust for offload if needed
# TODO(will): add support for other dtypes
image_encoder_dtype = self.component_load_precision(server_args, component_name)
assert image_encoder_dtype is not None
return self.load_model(
component_weights_path,
encoder_config,
server_args,
image_encoder_dtype,
component_name=component_name,
)
return encoder_config
python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py
+5 -61
View File
@@ -1,71 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
import torch
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PlainStateDictComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.utils import (
load_safetensors_state_dict,
set_default_torch_dtype,
skip_init_modules,
)
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
logger = init_logger(__name__)
class SoundTokenizerLoader(PlainStateDictComponentLoader):
component_names = ["sound_tokenizer"]
expected_library = "diffusers"
default_precision_attr = "vae_precision"
# the native tokenizer is decoder-only; encoder weights are unused
ignored_checkpoint_prefixes = ("encoder.",)
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
):
config = self.load_component_config(component_model_path, component_name)
component_weights_path = self.resolve_component_weights_path(
component_model_path, server_args, component_name
)
class_name = config.pop("_class_name", None) or self.component_architecture
assert class_name is not None, (
"Sound tokenizer class name must be available from component config."
)
server_args.model_paths[component_name] = component_model_path
dtype = resolve_component_precision(server_args, component_name)
if dtype is None:
try:
dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
except AttributeError:
dtype = torch.bfloat16
target_device = self.target_device(
server_args.should_start_component_on_cpu(component_name)
)
with set_default_torch_dtype(dtype), skip_init_modules():
model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
model = model_cls(config).to(device=target_device, dtype=dtype)
loaded = load_safetensors_state_dict(component_weights_path)
incompatible = model.load_state_dict(loaded, strict=False)
missing = getattr(incompatible, "missing_keys", [])
# The tokenizer is decoder-only; the checkpoint's encoder weights are
# expected leftovers, so they're excluded from the load warning.
unexpected = [
k
for k in getattr(incompatible, "unexpected_keys", [])
if not k.startswith("encoder.")
]
if missing or unexpected:
logger.warning(
"Loaded sound_tokenizer with missing_keys=%d unexpected_keys=%d",
len(missing),
len(unexpected),
)
model.eval()
return model
def build_model_config(self, config, component_name):
return {"config": config}
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+68 -293
View File
@@ -1,9 +1,6 @@
import dataclasses
import glob
import os
import re
from collections.abc import Callable, Generator, Iterable
from typing import cast
from collections.abc import Generator
import torch
import transformers
@@ -67,20 +64,14 @@ from sglang.multimodal_gen.runtime.loader.gguf_weights import (
remap_gguf_tensor_meta,
)
from sglang.multimodal_gen.runtime.loader.utils import (
_list_safetensors_files,
checkpoint_bytes,
get_param_names_mapping,
initialize_model,
keep_checkpoint_mapped,
set_default_torch_dtype,
skip_init_modules,
)
from sglang.multimodal_gen.runtime.loader.weight_utils import (
filter_files_not_needed_for_inference,
pt_weights_iterator,
safetensors_weights_iterator,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import (
host_copies_would_not_fit,
host_memory_available_bytes,
checkpoint_weights_iterator,
)
from sglang.multimodal_gen.runtime.models.encoders.base import (
EncoderTensorParallelMixin,
@@ -101,10 +92,10 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import (
get_quant_config,
get_quant_config_from_safetensors_metadata,
inspect_comfy_quant_markers,
process_model_weights_after_loading,
resolve_comfy_checkpoint_quantization,
)
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
from sglang.srt.environ import envs
from sglang.srt.layers.linear import LinearBase as SrtLinearBase
from sglang.srt.layers.quantization.fp8 import Fp8Config as SrtFp8Config
from sglang.srt.layers.quantization.unquant import (
@@ -113,7 +104,6 @@ from sglang.srt.layers.quantization.unquant import (
from sglang.srt.model_loader.checkpoint_quantization import (
resolve_checkpoint_quant_spec,
)
from sglang.srt.model_loader.post_load import stage_module_for_post_load
logger = init_logger(__name__)
@@ -340,7 +330,7 @@ def _resolve_and_configure_encoder_quantization(
explicit_quantization: str | None = None,
ignored_layers: list[str] | None = None,
) -> type[nn.Module]:
architectures = getattr(model_config, "architectures", [])
architectures = model_config.arch_config.architectures
try:
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
except Exception as resolution_error:
@@ -384,35 +374,6 @@ def _resolve_and_configure_encoder_quantization(
return model_cls
def _process_quantized_encoder_weights(
model: nn.Module,
process_device: torch.device | None,
component_name: str,
) -> int:
processed_layers = 0
for module in model.modules():
if not isinstance(module, (LinearBase, SrtLinearBase)):
continue
quant_method = module.quant_method
if quant_method is None or isinstance(
quant_method,
(UnquantizedLinearMethod, SrtUnquantizedLinearMethod),
):
continue
if process_device is None:
quant_method.process_weights_after_loading(module)
else:
with stage_module_for_post_load(module, process_device):
quant_method.process_weights_after_loading(module)
processed_layers += 1
if processed_layers == 0:
raise ValueError(
f"The {component_name!r} checkpoint declares quantization, but the "
"model did not construct any quantized linear layers"
)
return processed_layers
def _require_quantized_encoder_layers(
model: nn.Module,
component_name: str,
@@ -463,21 +424,6 @@ def _require_quantized_encoder_layers(
)
def _keep_this_checkpoint_mapped(model_path: str) -> bool:
"""Whether this encoder's weights should stay on their file mapping."""
weight_bytes = checkpoint_bytes(model_path)
if not host_copies_would_not_fit(weight_bytes):
return False
logger.info(
"Text encoder checkpoint is %.2f GiB against %.2f GiB of host memory, "
"so its compatible weights stay on the checkpoint mapping instead of "
"being copied in.",
weight_bytes / 1024**3,
host_memory_available_bytes() / 1024**3,
)
return True
class TextEncoderLoader(OnlineQuantizationComponentLoader):
"""Loader for text encoders."""
@@ -491,7 +437,7 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
if override is not None:
return override
return server_args.pipeline_config.text_encoder_precisions[
self._extract_encoder_index(component_name)
self._extract_encoder_index(self.structural_component_name(component_name))
]
def should_raise_customized_load_error(
@@ -510,22 +456,6 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
f"no {current_platform.device_type} implementation"
)
@dataclasses.dataclass
class Source:
"""A source for weights."""
model_or_path: str
"""The model ID or path."""
prefix: str = ""
"""A prefix to prepend to all weights."""
fall_back_to_pt: bool = True
"""Whether .pt weights can be used."""
allow_patterns_overrides: list[str] | None = None
"""If defined, weights will load exclusively using these patterns."""
def resolve_native_transformers_model_class(self, config: PretrainedConfig) -> type:
"""Resolve the concrete transformers class for a text encoder.
@@ -543,149 +473,23 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
return transformers_model_class
return transformers.AutoModel
def _prepare_weights(
self,
model_name_or_path: str,
fall_back_to_pt: bool,
allow_patterns_overrides: list[str] | None,
key_filter: Callable[[str], bool] | None = None,
) -> tuple[str, list[str], bool]:
"""Prepare weights for the model.
If the model is not local, it will be downloaded."""
# model_name_or_path = (self._maybe_download_from_modelscope(
# model_name_or_path, revision) or model_name_or_path)
if os.path.isfile(model_name_or_path):
if model_name_or_path.endswith(".safetensors"):
return os.path.dirname(model_name_or_path), [model_name_or_path], True
if fall_back_to_pt and model_name_or_path.endswith((".bin", ".pt")):
return os.path.dirname(model_name_or_path), [model_name_or_path], False
raise ValueError(
"Native encoder weight overrides currently support one "
f"safetensors, bin, or pt file, got {model_name_or_path!r}"
)
if not os.path.isdir(model_name_or_path):
raise ValueError(
f"Model path must be a local file or directory: {model_name_or_path!r}"
)
use_safetensors = False
index_file = SAFE_WEIGHTS_INDEX_NAME
allow_patterns = ["*.safetensors", "*.bin"]
if fall_back_to_pt:
allow_patterns += ["*.pt"]
if allow_patterns_overrides is not None:
allow_patterns = allow_patterns_overrides
hf_folder = model_name_or_path
hf_weights_files: list[str] = []
for pattern in allow_patterns:
if pattern == "*.safetensors":
hf_weights_files = _list_safetensors_files(
hf_folder,
index_file=index_file,
key_filter=key_filter,
)
else:
hf_weights_files = glob.glob(os.path.join(hf_folder, pattern))
if len(hf_weights_files) > 0:
if pattern == "*.safetensors":
use_safetensors = True
break
if not use_safetensors:
hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
if len(hf_weights_files) == 0:
raise RuntimeError(
f"Cannot find any model weights with `{model_name_or_path}`"
)
# Sort weight files when SGLANG_SORT_WEIGHT_FILES >= 0 (default).
# Staggering is not applicable to text-encoder loading (no TP split).
if envs.SGLANG_SORT_WEIGHT_FILES.get() >= 0:
hf_weights_files.sort()
return hf_folder, hf_weights_files, use_safetensors
def _get_weights_iterator(
self,
source: "Source",
to_cpu: bool,
key_filter: Callable[[str], bool] | None = None,
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""get an iterator for the model weights based on the load format."""
source_key_filter: Callable[[str], bool] | None
if key_filter is None:
source_key_filter = None
else:
def include_source_weight(name: str) -> bool:
return key_filter(source.prefix + name)
source_key_filter = include_source_weight
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
source.model_or_path,
source.fall_back_to_pt,
source.allow_patterns_overrides,
key_filter=source_key_filter,
)
if use_safetensors:
weights_iterator = safetensors_weights_iterator(
hf_weights_files,
to_cpu=to_cpu,
key_filter=source_key_filter,
)
else:
weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
if source_key_filter is not None:
weights_iterator = (
(name, tensor)
for name, tensor in weights_iterator
if source_key_filter(name)
)
# apply the prefix.
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
def _get_all_weights(
self,
model: EncoderTensorParallelMixin,
model_path: str,
to_cpu: bool,
) -> Generator[tuple[str, torch.Tensor], None, None]:
key_filter = model.should_materialize_checkpoint_weight
def include_checkpoint_weight(name: str) -> bool:
return not name.endswith(".comfy_quant") and key_filter(name)
return not name.endswith(
".comfy_quant"
) and model.should_materialize_checkpoint_weight(name)
primary_weights = TextEncoderLoader.Source(
yield from checkpoint_weights_iterator(
model_path,
prefix="",
fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
to_cpu=to_cpu,
key_filter=include_checkpoint_weight,
index_file=SAFE_WEIGHTS_INDEX_NAME,
)
yield from self._get_weights_iterator(
primary_weights,
to_cpu,
include_checkpoint_weight,
)
secondary_weights = cast(
Iterable[TextEncoderLoader.Source],
getattr(model, "secondary_weights", ()),
)
for source in secondary_weights:
yield from self._get_weights_iterator(
source,
to_cpu,
include_checkpoint_weight,
)
def load_customized(
self,
@@ -700,35 +504,13 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
server_args,
component_name,
)
diffusers_pretrained_config = get_config(
component_model_path, trust_remote_code=True
)
model_config = get_diffusers_component_config(
component_path=component_model_path
)
# TODO(mick): had to throw an exception for different text-encoder arch
encoder_index = self._extract_encoder_index(
self.structural_component_name(component_name)
encoder_config = self.build_model_config(
component_model_path, model_config, server_args, component_name
)
assert encoder_index < len(
server_args.pipeline_config.text_encoder_configs
) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions)
encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index]
encoder_config.update_model_arch(model_config)
encoder_config.generation_config = load_dict(
os.path.join(component_model_path, "generation_config.json")
)
if encoder_index == 0:
for key, value in diffusers_pretrained_config.__dict__.items():
setattr(encoder_config.arch_config, key, value)
post_diffusers_config_update = getattr(
encoder_config, "post_diffusers_config_update", None
)
if post_diffusers_config_update is not None:
post_diffusers_config_update()
encoder_config.post_diffusers_config_update()
model_cls = _resolve_and_configure_encoder_quantization(
encoder_config,
model_config,
@@ -778,6 +560,34 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
f"Failed to load quantized native {component_name!r}: {error}"
) from error
def build_model_config(
self,
component_model_path: str,
model_config: dict,
server_args: ServerArgs,
component_name: str,
) -> EncoderConfig:
diffusers_pretrained_config = get_config(
component_model_path, trust_remote_code=True
)
encoder_index = self._extract_encoder_index(
self.structural_component_name(component_name)
)
assert encoder_index < len(
server_args.pipeline_config.text_encoder_configs
) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions)
encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index]
encoder_config.update_model_arch(model_config)
encoder_config.generation_config = load_dict(
os.path.join(component_model_path, "generation_config.json")
)
if encoder_index == 0:
for key, value in diffusers_pretrained_config.__dict__.items():
setattr(encoder_config.arch_config, key, value)
return encoder_config
@staticmethod
def _extract_encoder_index(component_name: str) -> int:
"""
@@ -872,20 +682,20 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
use_tensor_parallel_group(encoder_tp_group),
set_default_torch_dtype(PRECISION_TO_TYPE[dtype]),
):
with model_device, skip_init_modules():
architectures = getattr(model_config, "architectures", [])
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
enable_image_understanding = isinstance(
server_args.pipeline_config,
(QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig),
)
model_config.enable_image_understanding = enable_image_understanding
# LongCat feeds its padded body to the DiT, so it must mask
# padding on the cache-free path; scoped so others are unchanged.
model_config.honor_cache_free_padding_mask = isinstance(
server_args.pipeline_config, LongCatImagePipelineConfig
)
model = model_cls(model_config)
model_cls, _ = ModelRegistry.resolve_model_cls(
model_config.arch_config.architectures
)
model_config.enable_image_understanding = isinstance(
server_args.pipeline_config,
(QwenImageEditPipelineConfig, LongCatImageEditPipelineConfig),
)
# longcat consumes the padded body without an attention cache
model_config.honor_cache_free_padding_mask = isinstance(
server_args.pipeline_config, LongCatImagePipelineConfig
)
model = initialize_model(
model_cls, {"config": model_config}, param_dtype, model_device
)
if not isinstance(model, EncoderTensorParallelMixin):
raise TypeError(
@@ -904,15 +714,11 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
)
if component_starts_on_cpu and (
current_platform.is_mps() or _keep_this_checkpoint_mapped(model_path)
current_platform.is_mps()
or keep_checkpoint_mapped(
weight_bytes=checkpoint_bytes(model_path), component=component_name
)
):
# The encoder is layered immediately after this loader returns,
# so compatible CPU safetensors can stay mapped instead of being
# copied. On MPS that is always the right call -- the memory is
# unified. On any host it becomes the only call once the
# checkpoint is larger than host memory, because the copy is
# what does not fit: H3's encoder is 62.13 GiB against a 32 GiB
# target.
model._keep_checkpoint_mapping = True
weights_to_load = {name for name, _ in model.named_parameters()}
@@ -931,6 +737,9 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
if isinstance(quant_config, QuantoInt8Config):
checkpoint_weights = normalize_quanto_int8_weights(checkpoint_weights)
loaded_weights = model.load_weights(checkpoint_weights)
self.validate_checkpoint_keys(
weights_to_load - loaded_weights, [], component_name
)
if quant_config is not None and not isinstance(quant_config, GGUFConfig):
postprocess_device: torch.device | None = local_torch_device
@@ -939,10 +748,10 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
and quant_config.is_checkpoint_int8_serialized
):
postprocess_device = None
processed_layers = _process_quantized_encoder_weights(
processed_layers = process_model_weights_after_loading(
model,
postprocess_device,
component_name,
quantized_only=True,
)
logger.info(
"Processed %d %s linear layers for %s",
@@ -961,39 +770,5 @@ class TextEncoderLoader(OnlineQuantizationComponentLoader):
model = model.to("cpu")
else:
model = model.to(local_torch_device)
# We only enable strict check for non-quantized models
# that have loaded weights tracking currently.
# if loaded_weights is not None:
weights_not_loaded = weights_to_load - loaded_weights
if weights_not_loaded:
# NOTE:
# If we silently continue with uninitialized weights, the text encoder can
# produce NaNs/garbage embeddings that later fail stage verification in a
# hard-to-debug way (e.g., `prompt_embeds` fails the NaN check).
#
# We allow a small set of known-optional parameters to be missing, but
# default to strict behavior for the rest.
allowed_missing_patterns = (
getattr(model, "_allowed_missing_weights_patterns", []) or []
)
unexpected_missing = {
n
for n in weights_not_loaded
if not any(pat in n for pat in allowed_missing_patterns)
}
if unexpected_missing:
raise ValueError(
"Following text encoder weights were not initialized from checkpoint: "
f"{sorted(unexpected_missing)}. "
"This usually indicates a checkpoint/model-arch mismatch or a broken "
"weight-name mapping. If these are truly optional, set "
"`model._allowed_missing_weights_patterns` to whitelist patterns."
)
logger.warning(
"Following (allowed) text encoder weights were not initialized from "
"checkpoint: %s (allowed patterns: %s)",
sorted(weights_not_loaded),
allowed_missing_patterns,
)
return model
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+5 -12
View File
@@ -16,7 +16,6 @@ from sglang.multimodal_gen.runtime.layers.attention.selector import (
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
OnlineQuantizationComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
from sglang.multimodal_gen.runtime.loader.gguf_weights import gguf_weights_iterator
from sglang.multimodal_gen.runtime.loader.minimax_h3_weights import (
comfy_quant_key_filter,
@@ -525,20 +524,14 @@ class TransformerLoader(OnlineQuantizationComponentLoader):
# Model construction resolves attention implementations, so apply the
# quantization-specific default around FSDP initialization and loading.
with attn_backend_context:
model = maybe_load_fsdp_model(
model = self.load_state_dict_model(
model_cls=model_cls,
init_params=init_params,
weight_dir_list=safetensors_list,
device=local_torch_device,
hsdp_replicate_dim=server_args.hsdp_replicate_dim,
hsdp_shard_dim=server_args.hsdp_shard_dim,
weight_files=safetensors_list,
server_args=component_server_args,
component_name=component_name,
component_starts_on_cpu=component_starts_on_cpu,
pin_cpu_memory=component_server_args.pin_cpu_memory,
fsdp_inference=use_fsdp,
param_dtype=quant_spec.param_dtype,
reduce_dtype=torch.float32,
output_dtype=None,
strict=False,
dtype=quant_spec.param_dtype,
weight_load_plan=weight_load_plan,
checkpoint_key_filter=checkpoint_key_filter,
weights_iterator=(
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+16 -73
View File
@@ -1,14 +1,12 @@
import hashlib
import importlib.util
import os
from collections.abc import Iterable
import torch
import torch.nn as nn
from safetensors.torch import load_file as safetensors_load_file
from safetensors.torch import safe_open
from safetensors.torch import save_file as safetensors_save_file
from torch.nn.utils import parametrize
from sglang.multimodal_gen import envs
from sglang.multimodal_gen.configs.models.vaes.base import VAEConfig
@@ -25,10 +23,12 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
from sglang.multimodal_gen.runtime.loader.utils import (
_list_safetensors_files,
_normalize_component_type,
adopt_plain_weight_norm_state,
checkpoint_bytes,
initialize_model,
keep_checkpoint_mapped,
load_model_state_dict,
set_default_torch_dtype,
skip_init_modules,
)
from sglang.multimodal_gen.runtime.loader.weight_utils import (
safetensors_weights_iterator,
@@ -283,56 +283,6 @@ def _hold_decoder_weights_in_decode_dtype(
)
def _match_checkpoint_dtypes(loaded: dict, target_state: dict) -> dict:
"""Convert checkpoint tensors whose dtype differs from their parameter's.
Assignment replaces the parameter rather than writing through it, so a
mismatched dtype would silently change the module's. Converting makes a
copy, which is the point: only the tensors that already match can stay on
the mapping.
"""
for name, tensor in list(loaded.items()):
param = target_state.get(name)
if param is not None and param.dtype != tensor.dtype:
loaded[name] = tensor.to(dtype=param.dtype)
return loaded
def _adopt_plain_weight_norm_state(
module: nn.Module, loaded_names: Iterable[str]
) -> int:
"""Make deparameterized checkpoint weights native module state.
PyTorch's weight-norm load hook accepts legacy ``weight_g``/``weight_v``
tensors, while inference exports commonly fold those tensors into one
plain ``weight``. Removing only the matching parametrizations preserves
that already-computed weight exactly and leaves every other parameterized
module untouched.
"""
state_names = set(module.state_dict())
module_by_name = dict(module.named_modules())
owners: set[str] = set()
for name in loaded_names:
if name == "weight":
owner_name = ""
elif name.endswith(".weight"):
owner_name = name.removesuffix(".weight")
else:
continue
state_prefix = f"{owner_name}." if owner_name else ""
if {
f"{state_prefix}parametrizations.weight.original0",
f"{state_prefix}parametrizations.weight.original1",
}.issubset(state_names):
owners.add(owner_name)
for owner_name in sorted(owners):
parametrize.remove_parametrizations(
module_by_name[owner_name], "weight", leave_parametrized=True
)
return len(owners)
def _vae_checkpoint_arch_metadata_names(
vae_config: VAEConfig,
target_state: dict[str, torch.Tensor],
@@ -444,7 +394,7 @@ def _assign_direct_gpu_vae_state(
vae_config: VAEConfig,
) -> tuple[int, tuple[str, ...]]:
"""Stream a complete standard VAE state directly onto its target device."""
num_deparameterized = _adopt_plain_weight_norm_state(
num_deparameterized = adopt_plain_weight_norm_state(
vae, _vae_checkpoint_tensor_names(weight_files)
)
target_state, slots = _direct_gpu_vae_state_slots(vae, component_name)
@@ -664,21 +614,15 @@ class VAELoader(WeightOverrideComponentLoader):
return vae
# Load from ModelRegistry (standard VAE classes)
if direct_gpu_weight_loading:
with (
set_default_torch_dtype(vae_dtype),
skip_init_modules(),
torch.device("meta"),
):
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
vae = vae_cls(vae_config)
else:
with (
set_default_torch_dtype(vae_dtype),
skip_init_modules(),
):
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
vae = vae_cls(vae_config).to(target_device)
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
vae = initialize_model(
vae_cls,
{"config": vae_config},
vae_dtype,
torch.device("meta") if direct_gpu_weight_loading else None,
)
if not direct_gpu_weight_loading:
vae = vae.to(target_device)
if os.path.isfile(component_weights_path):
if not component_weights_path.endswith(".safetensors"):
@@ -729,7 +673,7 @@ class VAELoader(WeightOverrideComponentLoader):
for sf_path in safetensors_list:
loaded.update(safetensors_load_file(sf_path))
_backfill_ltx2_audio_vae_latent_stats(loaded, component_type)
num_deparameterized = _adopt_plain_weight_norm_state(vae, loaded)
num_deparameterized = adopt_plain_weight_norm_state(vae, loaded)
target_state = vae.state_dict()
consumed_metadata = _consume_vae_checkpoint_arch_metadata(
loaded, vae_config, target_state
@@ -759,9 +703,8 @@ class VAELoader(WeightOverrideComponentLoader):
component=f"{component_name or 'vae'} (VAE)",
)
)
if keep_mapping:
_match_checkpoint_dtypes(loaded, target_state)
vae.load_state_dict(
load_model_state_dict(
vae,
loaded,
strict=strict_load,
assign=keep_mapping,
python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py
+5 -72
View File
@@ -1,80 +1,13 @@
import re
import torch
from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
PlainStateDictComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.utils import (
load_safetensors_state_dict,
set_default_torch_dtype,
skip_init_modules,
)
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
logger = init_logger(__name__)
class VocoderLoader(PlainStateDictComponentLoader):
component_names = ["vocoder"]
expected_library = "diffusers"
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
):
config = self.load_component_config(component_model_path, component_name)
component_weights_path = self.resolve_component_weights_path(
component_model_path, server_args, component_name
)
class_name = config.pop("_class_name", None) or self.component_architecture
assert class_name is not None, (
"Vocoder class name must be available from component config or pipeline config."
)
server_args.model_paths[component_name] = component_model_path
from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import (
LTXVocoderConfig,
)
vocoder_config = LTXVocoderConfig()
vocoder_config.update_model_arch(config)
resolved_vocoder_dtype = resolve_component_precision(server_args, "vocoder")
vocoder_dtype = (
resolved_vocoder_dtype
if resolved_vocoder_dtype is not None
else PRECISION_TO_TYPE["fp32"]
)
component_starts_on_cpu = server_args.should_start_component_on_cpu(
component_name
)
target_device = self.target_device(component_starts_on_cpu)
with set_default_torch_dtype(vocoder_dtype), skip_init_modules():
vocoder_cls, _ = ModelRegistry.resolve_model_cls(class_name)
vocoder = vocoder_cls(vocoder_config).to(target_device)
loaded = load_safetensors_state_dict(component_weights_path)
mapping = vocoder_config.arch_config.param_names_mapping
loaded = {_remap_vocoder_key(k, mapping): v for k, v in loaded.items()}
missing_keys, unexpected_keys = vocoder.load_state_dict(loaded, strict=False)
# A half-loaded vocoder produces plausible but wrong audio.
if missing_keys or unexpected_keys:
raise ValueError(
f"Vocoder weights at '{component_weights_path}' do not match the "
f"instantiated {class_name}. Missing: {sorted(missing_keys)}. "
f"Unexpected: {sorted(unexpected_keys)}."
)
return vocoder
def _remap_vocoder_key(key: str, param_names_mapping: dict[str, str]) -> str:
# Applied in order, not first-match: one key can need several rules.
for pattern, replacement in param_names_mapping.items():
key = re.sub(pattern, replacement, key)
return key
config_classes = {"vocoder": LTXVocoderConfig}
default_precision_attr = "audio_vae_precision"
default_dtype = torch.float32
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+10 -25
View File
@@ -8,7 +8,6 @@
from collections import Counter, defaultdict
from collections.abc import Callable, Generator
from itertools import chain
from types import MethodType
from typing import Any
@@ -40,9 +39,10 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
)
from sglang.multimodal_gen.runtime.loader import rank_local_checkpoint
from sglang.multimodal_gen.runtime.loader.utils import (
finalize_loaded_model,
get_param_names_mapping,
hf_to_custom_state_dict,
set_default_torch_dtype,
initialize_model,
)
from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
from sglang.multimodal_gen.runtime.loader.weight_utils import (
@@ -50,10 +50,10 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import (
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.quantization_utils import (
process_model_weights_after_loading,
)
from sglang.multimodal_gen.utils import set_mixed_precision_policy
from sglang.srt.utils import is_npu
_is_npu = is_npu()
logger = init_logger(__name__)
@@ -293,8 +293,9 @@ def maybe_load_fsdp_model(
mp_policy=mp_policy,
)
with set_default_torch_dtype(default_torch_dtype), torch.device("meta"):
model = model_cls(**init_params)
model = initialize_model(
model_cls, init_params, default_torch_dtype, torch.device("meta")
)
# Check if we should use FSDP
use_fsdp = fsdp_inference
@@ -447,26 +448,10 @@ def maybe_load_fsdp_model(
# move to device to perform postprocessing
_move_to_device_preserving_meta(model, weight_postprocess_device)
for _, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if quant_method is not None and hasattr(
quant_method, "process_weights_after_loading"
):
if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod):
# Activate the NZ format for storing weights,
# which is a specific optimization for Ascend NPU
torch.npu.config.allow_internal_format = True
quant_method.process_weights_after_loading(module)
if _is_npu:
torch.npu.empty_cache()
process_model_weights_after_loading(model)
model.post_load_weights()
for n, p in chain(model.named_parameters(), model.named_buffers()):
if p.is_meta:
raise RuntimeError(f"Unexpected param or buffer {n} on meta device.")
# Avoid unintended computation graph accumulation during inference
if isinstance(p, torch.nn.Parameter):
p.requires_grad = False
finalize_loaded_model(model)
# 4. deferred cpu offload
if defer_cpu_placement:
python/sglang/multimodal_gen/runtime/loader/utils.py
+85 -1
View File
@@ -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
+38
View File
@@ -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
+3 -1
View File
@@ -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
+68
View File
@@ -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
+1 -72
View File
@@ -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
+4 -1
View File
@@ -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
+304
View File
@@ -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
+89 -27
View File
@@ -5,11 +5,13 @@ from pathlib import Path
from types import SimpleNamespace
from unittest import mock
import pytest
import torch
import transformers
from safetensors.torch import save_file
from torch import nn
from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
from sglang.multimodal_gen.runtime.layers.linear import LinearBase
from sglang.multimodal_gen.runtime.layers.quantization.comfy_nvfp4 import (
ComfyFullPrecisionNvfp4LinearMethod,
@@ -35,11 +37,13 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader
TextEncoderLoader,
_configure_encoder_quantization,
_get_encoder_quant_config,
_process_quantized_encoder_weights,
_require_quantized_encoder_layers,
_resolve_and_configure_encoder_quantization,
)
from sglang.multimodal_gen.runtime.loader.gguf_weights import GGUFTensorMeta
from sglang.multimodal_gen.runtime.loader.weight_utils import (
checkpoint_weights_iterator,
)
from sglang.multimodal_gen.runtime.models.encoders.base import (
EncoderTensorParallelMixin,
TextEncoder,
@@ -49,9 +53,67 @@ from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
MiniMaxH3Qwen3VLEncoder,
)
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextModel
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.quantization_utils import (
process_model_weights_after_loading,
)
from sglang.srt.layers.linear import LinearBase as SrtLinearBase
@pytest.mark.parametrize("missing", [False, True])
@pytest.mark.parametrize("competing_index", [False, True])
def test_native_encoder_restoration_checks_checkpoint_before_fallback(
tmp_path, missing, competing_index
):
config = transformers.T5Config(
vocab_size=32,
d_model=8,
d_kv=4,
d_ff=16,
num_layers=1,
num_heads=2,
architectures=["T5EncoderModel"],
)
reference = transformers.T5EncoderModel(config)
config.save_pretrained(tmp_path)
weights = {name: tensor.clone() for name, tensor in reference.state_dict().items()}
required = "encoder.final_layer_norm.weight"
if missing:
weights.pop(required)
save_file(weights, tmp_path / "model.safetensors")
if competing_index:
# LTX-2 ships both indexes in text_encoder; only the HF one owns this model
(tmp_path / "model.safetensors.index.json").write_text(
json.dumps({"weight_map": {name: "model.safetensors" for name in weights}})
)
alternate = "diffusion_pytorch_model.safetensors"
save_file({"diffusion.weight": torch.ones(2, 2)}, tmp_path / alternate)
(tmp_path / "diffusion_pytorch_model.safetensors.index.json").write_text(
json.dumps({"weight_map": {"diffusion.weight": alternate}})
)
alternate_weights = dict(checkpoint_weights_iterator(str(tmp_path)))
assert set(alternate_weights) == {"diffusion.weight"}
args = ServerArgs(
model_path="x",
component_precisions={"text_encoder": "fp32"},
component_residency={"text_encoder": "component-offload"},
)
args.pipeline_config.text_encoder_configs = (T5Config(),)
loader = TextEncoderLoader()
with mock.patch.object(
loader, "load_native", side_effect=AssertionError("fallback")
):
if missing:
with pytest.raises(ComponentCheckpointUnsupportedError, match=required):
loader.load(str(tmp_path), args, "text_encoder", "transformers")
else:
model, _ = loader.load(str(tmp_path), args, "text_encoder", "transformers")
torch.testing.assert_close(
model.state_dict()[required], weights[required], rtol=0, atol=0
)
assert all(not parameter.requires_grad for parameter in model.parameters())
class TestTextEncoderWeightDiscovery(unittest.TestCase):
def test_prepare_weights_prefers_canonical_over_fp16_variant(self):
with tempfile.TemporaryDirectory() as tmpdir:
@@ -59,22 +121,13 @@ class TestTextEncoderWeightDiscovery(unittest.TestCase):
canonical = model_dir / "model.safetensors"
variant = model_dir / "model.fp16.safetensors"
canonical.touch()
variant.touch()
weight = torch.ones(2, 2)
save_file({"weight": weight}, canonical)
save_file({"weight": weight.half()}, variant)
(
hf_folder,
weight_files,
use_safetensors,
) = TextEncoderLoader()._prepare_weights(
str(model_dir),
fall_back_to_pt=True,
allow_patterns_overrides=None,
)
self.assertEqual(hf_folder, str(model_dir))
self.assertTrue(use_safetensors)
self.assertEqual(weight_files, [str(canonical)])
state = dict(checkpoint_weights_iterator(str(model_dir)))
self.assertEqual(state["weight"].dtype, torch.float32)
torch.testing.assert_close(state["weight"], weight)
class TestTextEncoderClassResolution(unittest.TestCase):
@@ -757,7 +810,10 @@ class TestTextEncoderQuantization(unittest.TestCase):
),
):
_resolve_and_configure_encoder_quantization(
SimpleNamespace(architectures=[architecture], quant_config=None),
SimpleNamespace(
arch_config=SimpleNamespace(architectures=[architecture]),
quant_config=None,
),
component_config,
"/model/text_encoder",
"/model/text_encoder",
@@ -791,7 +847,10 @@ class TestTextEncoderQuantization(unittest.TestCase):
):
_resolve_and_configure_encoder_quantization(
SimpleNamespace(
architectures=["ThirdPartyTextEncoder"], quant_config=None
arch_config=SimpleNamespace(
architectures=["ThirdPartyTextEncoder"]
),
quant_config=None,
),
{
"quantization_config": {
@@ -812,7 +871,10 @@ class TestTextEncoderQuantization(unittest.TestCase):
):
_resolve_and_configure_encoder_quantization(
SimpleNamespace(
architectures=["ThirdPartyTextEncoder"], quant_config=None
arch_config=SimpleNamespace(
architectures=["ThirdPartyTextEncoder"]
),
quant_config=None,
),
{
"quantization_config": {
@@ -887,10 +949,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
quant_method = _RecordingQuantMethod()
model = _SRTQuantizedLinear(quant_method)
processed = _process_quantized_encoder_weights(
processed = process_model_weights_after_loading(
model,
torch.device("cpu"),
"image_encoder",
quantized_only=True,
)
self.assertEqual(processed, 1)
@@ -926,10 +988,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
quant_method = _RecordingQuantMethod()
model = _QuantizedEncoder(quant_method)
processed = _process_quantized_encoder_weights(
processed = process_model_weights_after_loading(
model,
torch.device("cpu"),
"text_encoder",
quantized_only=True,
)
self.assertEqual(processed, 1)
@@ -941,10 +1003,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
quant_method = _RecordingQuantMethod()
model = _QuantizedEncoder(quant_method)
processed = _process_quantized_encoder_weights(
processed = process_model_weights_after_loading(
model,
torch.device("cuda", torch.cuda.current_device()),
"text_encoder",
quantized_only=True,
)
self.assertEqual(processed, 1)
@@ -957,10 +1019,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase):
model = _QuantizedEncoder(_RecordingQuantMethod(error=RuntimeError("boom")))
with self.assertRaisesRegex(RuntimeError, "boom"):
_process_quantized_encoder_weights(
process_model_weights_after_loading(
model,
torch.device("cuda", torch.cuda.current_device()),
"text_encoder",
quantized_only=True,
)
self.assertEqual(model.quantized.weight.device, torch.device("cpu"))
python/sglang/multimodal_gen/test/unit/test_vae_loader.py
+20 -24
View File
@@ -27,18 +27,17 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
ComponentCheckpointUnsupportedError,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import (
_adopt_plain_weight_norm_state,
_assign_direct_gpu_vae_state,
_backfill_ltx2_audio_vae_latent_stats,
_consume_vae_checkpoint_arch_metadata,
_direct_gpu_vae_state_slots,
_match_checkpoint_dtypes,
_require_native_loader_for_quantized_vae,
_should_use_channels_last_3d,
)
from sglang.multimodal_gen.runtime.loader.utils import (
checkpoint_bytes,
keep_checkpoint_mapped,
load_model_state_dict,
)
from sglang.multimodal_gen.runtime.managers.memory_managers import (
host_memory_budget,
@@ -47,6 +46,7 @@ from sglang.multimodal_gen.runtime.models.vaes import wanvae
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.decoding_av import (
LTX2AVDecodingStage,
)
from sglang.test.test_utils import CustomTestCase
class _FakeServerArgs:
@@ -131,28 +131,26 @@ class TestKeepCheckpointMapped(unittest.TestCase):
)
class TestMatchCheckpointDtypes(unittest.TestCase):
class TestMatchCheckpointDtypes(CustomTestCase):
"""Assignment replaces a parameter, so only matching dtypes may stay mapped."""
def test_a_matching_tensor_is_left_alone(self):
loaded = {"w": torch.zeros(4, dtype=torch.float32)}
before = loaded["w"]
_match_checkpoint_dtypes(loaded, {"w": torch.zeros(4, dtype=torch.float32)})
self.assertIs(loaded["w"], before)
def test_a_mismatched_tensor_is_converted(self):
loaded = {"w": torch.zeros(4, dtype=torch.float32)}
_match_checkpoint_dtypes(loaded, {"w": torch.zeros(4, dtype=torch.bfloat16)})
self.assertEqual(loaded["w"].dtype, torch.bfloat16)
def test_a_tensor_the_module_does_not_want_is_left_alone(self):
loaded = {"extra": torch.zeros(4, dtype=torch.float32)}
before = loaded["extra"]
_match_checkpoint_dtypes(loaded, {})
self.assertIs(loaded["extra"], before)
def test_assignment_preserves_mixed_dtypes_and_matching_storage(self):
model = nn.Linear(4, 4, bias=False, dtype=torch.bfloat16)
model.register_buffer("scale", torch.zeros(4, dtype=torch.float32))
weights = {
"weight": torch.ones(4, 4, dtype=torch.float32),
"scale": torch.ones(4, dtype=torch.float32),
}
checkpoint_weight = weights["weight"]
load_model_state_dict(model, weights, assign=True)
self.assertEqual(model.weight.dtype, torch.bfloat16)
self.assertEqual(model.scale.dtype, torch.float32)
self.assertEqual(model.scale.data_ptr(), weights["scale"].data_ptr())
self.assertNotEqual(model.weight.data_ptr(), checkpoint_weight.data_ptr())
self.assertTrue(torch.equal(model.weight.float(), checkpoint_weight))
class TestPlainWeightNormCheckpoint(unittest.TestCase):
class TestPlainWeightNormCheckpoint(CustomTestCase):
def test_adopts_a_folded_weight_without_reconstructing_it(self):
module = nn.Sequential(
torch.nn.utils.parametrizations.weight_norm(
@@ -162,8 +160,7 @@ class TestPlainWeightNormCheckpoint(unittest.TestCase):
expected = torch.arange(18, dtype=torch.float32).reshape(3, 2, 3) / 19
loaded = {"0.weight": expected}
self.assertEqual(_adopt_plain_weight_norm_state(module, loaded), 1)
module.load_state_dict(loaded, strict=True)
load_model_state_dict(module, loaded)
self.assertEqual(set(module.state_dict()), {"0.weight"})
self.assertTrue(torch.equal(module[0].weight, expected))
@@ -180,8 +177,7 @@ class TestPlainWeightNormCheckpoint(unittest.TestCase):
"0.weight_v": original_state["0.parametrizations.weight.original1"].clone(),
}
self.assertEqual(_adopt_plain_weight_norm_state(module, loaded), 0)
module.load_state_dict(loaded, strict=True)
load_model_state_dict(module, loaded)
self.assertIn("0.parametrizations.weight.original0", module.state_dict())