[diffusion] refactor: split component_loader into component-wise files (#17820)
This commit is contained in:
Mick
@@ -0,0 +1,70 @@
|
||||
from safetensors.torch import load_file as safetensors_load_file
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
_list_safetensors_files,
|
||||
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.hf_diffusers_utils import (
|
||||
get_diffusers_component_config,
|
||||
)
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
|
||||
class AdapterLoader(ComponentLoader):
|
||||
"""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 a single safetensors state_dict.
|
||||
"""
|
||||
|
||||
component_names = ["connectors"]
|
||||
expected_library = "diffusers"
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, *args
|
||||
):
|
||||
config = get_diffusers_component_config(model_path=component_model_path)
|
||||
|
||||
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["connectors"] = component_model_path
|
||||
|
||||
model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
|
||||
|
||||
target_device = get_local_torch_device()
|
||||
default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
with set_default_torch_dtype(default_dtype), skip_init_modules():
|
||||
connector_cfg = SimpleNamespace(**config)
|
||||
model = model_cls(connector_cfg).to(
|
||||
device=target_device, dtype=default_dtype
|
||||
)
|
||||
|
||||
safetensors_list = _list_safetensors_files(component_model_path)
|
||||
if not safetensors_list:
|
||||
raise ValueError(f"No safetensors files found in {component_model_path}")
|
||||
if len(safetensors_list) != 1:
|
||||
raise ValueError(
|
||||
f"Found {len(safetensors_list)} safetensors files in {component_model_path}, expected 1"
|
||||
)
|
||||
|
||||
loaded = safetensors_load_file(safetensors_list[0])
|
||||
model.load_state_dict(loaded, strict=False)
|
||||
|
||||
return model
|
||||
105
python/sglang/multimodal_gen/runtime/loader/bridge_loader.py
Normal file
105
python/sglang/multimodal_gen/runtime/loader/bridge_loader.py
Normal file
@@ -0,0 +1,105 @@
|
||||
from copy import deepcopy
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
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.models.registry import ModelRegistry
|
||||
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
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class BridgeLoader(ComponentLoader):
|
||||
"""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 = get_diffusers_component_config(model_path=component_model_path)
|
||||
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_model_path)
|
||||
if not safetensors_list:
|
||||
raise ValueError(f"No safetensors files found in {component_model_path}")
|
||||
|
||||
default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
|
||||
|
||||
logger.info(
|
||||
"Loading %s from %s safetensors files, default_dtype: %s",
|
||||
class_name,
|
||||
len(safetensors_list),
|
||||
default_dtype,
|
||||
)
|
||||
|
||||
# Check if FSDP loading is available
|
||||
if (
|
||||
server_args.hsdp_shard_dim is not None
|
||||
and hasattr(model_cls, "_fsdp_shard_conditions")
|
||||
and model_cls._fsdp_shard_conditions
|
||||
):
|
||||
# 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=get_local_torch_device(),
|
||||
hsdp_replicate_dim=server_args.hsdp_replicate_dim,
|
||||
hsdp_shard_dim=server_args.hsdp_shard_dim,
|
||||
cpu_offload=server_args.dit_cpu_offload,
|
||||
pin_cpu_memory=server_args.pin_cpu_memory,
|
||||
fsdp_inference=server_args.use_fsdp_inference,
|
||||
default_dtype=default_dtype,
|
||||
param_dtype=torch.bfloat16,
|
||||
reduce_dtype=torch.float32,
|
||||
output_dtype=None,
|
||||
strict=False,
|
||||
)
|
||||
else:
|
||||
# Fallback to simple loading (for non-FSDP or legacy models)
|
||||
model = model_cls.from_pretrained(
|
||||
component_model_path, torch_dtype=default_dtype
|
||||
)
|
||||
model = model.to(device=get_local_torch_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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,58 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
from sglang.multimodal_gen.configs.models import ModelConfig
|
||||
from sglang.multimodal_gen.runtime.loader.text_encoder_loader import TextEncoderLoader
|
||||
from sglang.multimodal_gen.runtime.loader.utils import _clean_hf_config_inplace
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class ImageEncoderLoader(TextEncoderLoader):
|
||||
|
||||
component_names = ["image_encoder"]
|
||||
expected_library = "transformers"
|
||||
|
||||
def should_offload(self, server_args, model_config: ModelConfig | None = None):
|
||||
should_offload = server_args.image_encoder_cpu_offload
|
||||
if not should_offload:
|
||||
return False
|
||||
# _fsdp_shard_conditions is in arch_config, not directly on model_config
|
||||
arch_config = (
|
||||
getattr(model_config, "arch_config", model_config) if model_config else None
|
||||
)
|
||||
fsdp_shard_conditions = (
|
||||
getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
|
||||
)
|
||||
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
|
||||
return use_cpu_offload
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, *args
|
||||
):
|
||||
"""Load the text encoders based on the model path, and inference args."""
|
||||
# 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,
|
||||
# )
|
||||
with open(os.path.join(component_model_path, "config.json")) as f:
|
||||
model_config = json.load(f)
|
||||
_clean_hf_config_inplace(model_config)
|
||||
logger.debug("HF model config: %s", model_config)
|
||||
|
||||
encoder_config = server_args.pipeline_config.image_encoder_config
|
||||
encoder_config.update_model_arch(model_config)
|
||||
|
||||
# Always start with local device; load_model will adjust for offload if needed
|
||||
# TODO(will): add support for other dtypes
|
||||
return self.load_model(
|
||||
component_model_path,
|
||||
encoder_config,
|
||||
server_args,
|
||||
server_args.pipeline_config.image_encoder_precision,
|
||||
cpu_offload_flag=server_args.image_encoder_cpu_offload,
|
||||
)
|
||||
@@ -0,0 +1,35 @@
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
|
||||
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 SchedulerLoader(ComponentLoader):
|
||||
"""Loader for scheduler."""
|
||||
|
||||
component_names = ["scheduler"]
|
||||
expected_library = "diffusers"
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, *args
|
||||
):
|
||||
"""Load the scheduler based on the model path, and inference args."""
|
||||
config = get_diffusers_component_config(model_path=component_model_path)
|
||||
|
||||
class_name = config.pop("_class_name")
|
||||
assert (
|
||||
class_name is not None
|
||||
), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
|
||||
|
||||
scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
|
||||
scheduler = scheduler_cls(**config)
|
||||
if server_args.pipeline_config.flow_shift is not None:
|
||||
scheduler.set_shift(server_args.pipeline_config.flow_shift)
|
||||
|
||||
return scheduler
|
||||
@@ -0,0 +1,285 @@
|
||||
import dataclasses
|
||||
import glob
|
||||
import os
|
||||
from collections.abc import Generator, Iterable
|
||||
from typing import Generator, Iterable, cast
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
from torch import nn
|
||||
from torch.distributed import init_device_mesh
|
||||
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
|
||||
|
||||
from sglang.multimodal_gen.configs.models import EncoderConfig, ModelConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
|
||||
QwenImageEditPipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
from sglang.multimodal_gen.runtime.loader.fsdp_load import shard_model
|
||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
_clean_hf_config_inplace,
|
||||
set_default_torch_dtype,
|
||||
skip_init_modules,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.weight_utils import (
|
||||
filter_duplicate_safetensors_files,
|
||||
filter_files_not_needed_for_inference,
|
||||
pt_weights_iterator,
|
||||
safetensors_weights_iterator,
|
||||
)
|
||||
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 (
|
||||
get_config,
|
||||
get_diffusers_component_config,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class TextEncoderLoader(ComponentLoader):
|
||||
"""Loader for text encoders."""
|
||||
|
||||
component_names = ["text_encoder"]
|
||||
expected_library = "transformers"
|
||||
|
||||
@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 should_offload(self, server_args, model_config: ModelConfig | None = None):
|
||||
should_offload = server_args.text_encoder_cpu_offload
|
||||
if not should_offload:
|
||||
return False
|
||||
# _fsdp_shard_conditions is in arch_config, not directly on model_config
|
||||
arch_config = (
|
||||
getattr(model_config, "arch_config", model_config) if model_config else None
|
||||
)
|
||||
fsdp_shard_conditions = (
|
||||
getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
|
||||
)
|
||||
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
|
||||
return use_cpu_offload
|
||||
|
||||
def _prepare_weights(
|
||||
self,
|
||||
model_name_or_path: str,
|
||||
fall_back_to_pt: bool,
|
||||
allow_patterns_overrides: list[str] | 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)
|
||||
|
||||
is_local = os.path.isdir(model_name_or_path)
|
||||
assert is_local, "Model path must be a local directory"
|
||||
|
||||
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:
|
||||
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 use_safetensors:
|
||||
hf_weights_files = filter_duplicate_safetensors_files(
|
||||
hf_weights_files, hf_folder, index_file
|
||||
)
|
||||
else:
|
||||
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}`"
|
||||
)
|
||||
|
||||
return hf_folder, hf_weights_files, use_safetensors
|
||||
|
||||
def _get_weights_iterator(
|
||||
self, source: "Source", to_cpu: bool
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
"""get an iterator for the model weights based on the load format."""
|
||||
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
|
||||
source.model_or_path,
|
||||
source.fall_back_to_pt,
|
||||
source.allow_patterns_overrides,
|
||||
)
|
||||
if use_safetensors:
|
||||
weights_iterator = safetensors_weights_iterator(
|
||||
hf_weights_files, to_cpu=to_cpu
|
||||
)
|
||||
else:
|
||||
weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
|
||||
|
||||
# apply the prefix.
|
||||
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
|
||||
|
||||
def _get_all_weights(
|
||||
self,
|
||||
model: nn.Module,
|
||||
model_path: str,
|
||||
to_cpu: bool,
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
primary_weights = TextEncoderLoader.Source(
|
||||
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),
|
||||
)
|
||||
yield from self._get_weights_iterator(primary_weights, to_cpu)
|
||||
|
||||
secondary_weights = cast(
|
||||
Iterable[TextEncoderLoader.Source],
|
||||
getattr(model, "secondary_weights", ()),
|
||||
)
|
||||
for source in secondary_weights:
|
||||
yield from self._get_weights_iterator(source, to_cpu)
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, component_name: str
|
||||
):
|
||||
"""Load the text encoders based on the model path, and inference args."""
|
||||
# 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,
|
||||
# )
|
||||
diffusers_pretrained_config = get_config(
|
||||
component_model_path, trust_remote_code=True
|
||||
)
|
||||
model_config = get_diffusers_component_config(model_path=component_model_path)
|
||||
_clean_hf_config_inplace(model_config)
|
||||
logger.debug("HF model config: %s", model_config)
|
||||
|
||||
def is_not_first_encoder(module_name):
|
||||
return "2" in module_name
|
||||
|
||||
# TODO(mick): had to throw an exception for different text-encoder arch
|
||||
if not is_not_first_encoder(component_name):
|
||||
encoder_config = server_args.pipeline_config.text_encoder_configs[0]
|
||||
encoder_config.update_model_arch(model_config)
|
||||
for key, value in diffusers_pretrained_config.__dict__.items():
|
||||
setattr(encoder_config.arch_config, key, value)
|
||||
encoder_dtype = server_args.pipeline_config.text_encoder_precisions[0]
|
||||
else:
|
||||
assert len(server_args.pipeline_config.text_encoder_configs) == 2
|
||||
encoder_config = server_args.pipeline_config.text_encoder_configs[1]
|
||||
encoder_config.update_model_arch(model_config)
|
||||
encoder_dtype = server_args.pipeline_config.text_encoder_precisions[1]
|
||||
# TODO(will): add support for other dtypes
|
||||
return self.load_model(
|
||||
component_model_path,
|
||||
encoder_config,
|
||||
server_args,
|
||||
encoder_dtype,
|
||||
)
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
model_path: str,
|
||||
model_config: EncoderConfig,
|
||||
server_args: ServerArgs,
|
||||
dtype: str = "fp16",
|
||||
cpu_offload_flag: bool | None = None,
|
||||
):
|
||||
# Determine CPU offload behavior and target device
|
||||
|
||||
local_torch_device = get_local_torch_device()
|
||||
should_offload = self.should_offload(server_args, model_config)
|
||||
|
||||
if should_offload and not current_platform.is_mps():
|
||||
model_device = torch.device("cpu")
|
||||
else:
|
||||
model_device = local_torch_device
|
||||
|
||||
with 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 = (
|
||||
True
|
||||
if isinstance(
|
||||
server_args.pipeline_config, QwenImageEditPipelineConfig
|
||||
)
|
||||
else False
|
||||
)
|
||||
model_config.enable_image_understanding = enable_image_understanding
|
||||
model = model_cls(model_config)
|
||||
|
||||
weights_to_load = {name for name, _ in model.named_parameters()}
|
||||
loaded_weights = model.load_weights(
|
||||
self._get_all_weights(model, model_path, to_cpu=should_offload)
|
||||
)
|
||||
|
||||
# Explicitly move model to target device after loading weights
|
||||
if not should_offload:
|
||||
model = model.to(local_torch_device)
|
||||
|
||||
if should_offload:
|
||||
# Disable FSDP for MPS as it's not compatible
|
||||
if current_platform.is_mps():
|
||||
logger.info(
|
||||
"Disabling FSDP sharding for MPS platform as it's not compatible"
|
||||
)
|
||||
model = model.to(local_torch_device)
|
||||
else:
|
||||
mesh = init_device_mesh(
|
||||
current_platform.device_type,
|
||||
mesh_shape=(1, dist.get_world_size()),
|
||||
mesh_dim_names=("offload", "replicate"),
|
||||
)
|
||||
shard_model(
|
||||
model,
|
||||
cpu_offload=True,
|
||||
reshard_after_forward=True,
|
||||
mesh=mesh["offload"],
|
||||
fsdp_shard_conditions=model_config.arch_config._fsdp_shard_conditions
|
||||
or getattr(model, "_fsdp_shard_conditions", None),
|
||||
pin_cpu_memory=server_args.pin_cpu_memory,
|
||||
)
|
||||
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:
|
||||
raise ValueError(
|
||||
"Following model weights were not initialized from "
|
||||
f"checkpoint: {weights_not_loaded}"
|
||||
)
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
from copy import deepcopy
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
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,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
|
||||
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
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class TransformerLoader(ComponentLoader):
|
||||
"""Shared loader for (video/audio) DiT transformers."""
|
||||
|
||||
component_names = ["transformer", "audio_dit", "video_dit"]
|
||||
expected_library = "diffusers"
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, component_name: str
|
||||
):
|
||||
"""Load the transformer based on the model path, and inference args."""
|
||||
config = get_diffusers_component_config(model_path=component_model_path)
|
||||
hf_config = deepcopy(config)
|
||||
cls_name = config.pop("_class_name")
|
||||
if cls_name is None:
|
||||
raise ValueError(
|
||||
"Model config does not contain a _class_name attribute. "
|
||||
"Only diffusers format is supported."
|
||||
)
|
||||
|
||||
component_name = _normalize_component_type(component_name)
|
||||
server_args.model_paths[component_name] = component_model_path
|
||||
|
||||
if component_name in ("transformer", "video_dit"):
|
||||
pipeline_dit_config_attr = "dit_config"
|
||||
elif component_name in ("audio_dit",):
|
||||
pipeline_dit_config_attr = "audio_dit_config"
|
||||
else:
|
||||
raise ValueError(f"Invalid module name: {component_name}")
|
||||
# Config from Diffusers supersedes sgl_diffusion's model config
|
||||
dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr)
|
||||
dit_config.update_model_arch(config)
|
||||
|
||||
model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
|
||||
|
||||
# Find all safetensors files
|
||||
safetensors_list = _list_safetensors_files(component_model_path)
|
||||
if not safetensors_list:
|
||||
raise ValueError(f"No safetensors files found in {component_model_path}")
|
||||
|
||||
# Check if we should use custom initialization weights
|
||||
custom_weights_path = getattr(
|
||||
server_args, "init_weights_from_safetensors", None
|
||||
)
|
||||
use_custom_weights = False
|
||||
|
||||
if use_custom_weights:
|
||||
logger.info(
|
||||
"Using custom initialization weights from: %s", custom_weights_path
|
||||
)
|
||||
assert (
|
||||
custom_weights_path is not None
|
||||
), "Custom initialization weights must be provided"
|
||||
if os.path.isdir(custom_weights_path):
|
||||
safetensors_list = _list_safetensors_files(custom_weights_path)
|
||||
else:
|
||||
assert custom_weights_path.endswith(
|
||||
".safetensors"
|
||||
), "Custom initialization weights must be a safetensors file"
|
||||
safetensors_list = [custom_weights_path]
|
||||
|
||||
default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
|
||||
|
||||
logger.info(
|
||||
"Loading %s from %s safetensors files, default_dtype: %s",
|
||||
cls_name,
|
||||
len(safetensors_list),
|
||||
default_dtype,
|
||||
)
|
||||
|
||||
# Load the model using FSDP loader
|
||||
assert server_args.hsdp_shard_dim is not None
|
||||
model = maybe_load_fsdp_model(
|
||||
model_cls=model_cls,
|
||||
init_params={"config": dit_config, "hf_config": hf_config},
|
||||
weight_dir_list=safetensors_list,
|
||||
device=get_local_torch_device(),
|
||||
hsdp_replicate_dim=server_args.hsdp_replicate_dim,
|
||||
hsdp_shard_dim=server_args.hsdp_shard_dim,
|
||||
cpu_offload=server_args.dit_cpu_offload,
|
||||
pin_cpu_memory=server_args.pin_cpu_memory,
|
||||
fsdp_inference=server_args.use_fsdp_inference,
|
||||
# TODO(will): make these configurable
|
||||
default_dtype=default_dtype,
|
||||
param_dtype=torch.bfloat16,
|
||||
reduce_dtype=torch.float32,
|
||||
output_dtype=None,
|
||||
strict=False,
|
||||
)
|
||||
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
|
||||
|
||||
assert (
|
||||
next(model.parameters()).dtype == default_dtype
|
||||
), "Model dtype does not match default dtype"
|
||||
|
||||
return model
|
||||
@@ -3,12 +3,15 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Utilities for selecting and loading models."""
|
||||
import contextlib
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Any
|
||||
from typing import Any, Dict, Type
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
@@ -102,3 +105,65 @@ def hf_to_custom_state_dict(
|
||||
continue
|
||||
custom_param_sd[target_param_name] = full_tensor
|
||||
return custom_param_sd, reverse_param_names_mapping
|
||||
|
||||
|
||||
class skip_init_modules:
|
||||
def __enter__(self):
|
||||
# Save originals
|
||||
self._orig_reset = {}
|
||||
for cls in (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d):
|
||||
self._orig_reset[cls] = cls.reset_parameters
|
||||
cls.reset_parameters = lambda self: None # skip init
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
# restore originals
|
||||
for cls, orig in self._orig_reset.items():
|
||||
cls.reset_parameters = orig
|
||||
|
||||
|
||||
def _normalize_component_type(module_type: str) -> str:
|
||||
"""Normalize module types like 'text_encoder_2' -> 'text_encoder'."""
|
||||
if module_type.endswith("_2"):
|
||||
return module_type[:-2]
|
||||
return module_type
|
||||
|
||||
|
||||
def _clean_hf_config_inplace(model_config: dict) -> None:
|
||||
"""Remove common extraneous HF fields if present."""
|
||||
for key in (
|
||||
"_name_or_path",
|
||||
"transformers_version",
|
||||
"model_type",
|
||||
"tokenizer_class",
|
||||
"torch_dtype",
|
||||
):
|
||||
model_config.pop(key, None)
|
||||
|
||||
|
||||
def _list_safetensors_files(model_path: str) -> list[str]:
|
||||
"""List all .safetensors files under a directory."""
|
||||
return sorted(glob.glob(os.path.join(str(model_path), "*.safetensors")))
|
||||
|
||||
|
||||
def get_memory_usage_of_component(module) -> float | None:
|
||||
"""
|
||||
returned value is in GB, rounded to 2 decimal digits
|
||||
"""
|
||||
if not isinstance(module, nn.Module):
|
||||
return None
|
||||
BYTES_PER_GB = 1024**3
|
||||
if hasattr(module, "get_memory_footprint"):
|
||||
usage = module.get_memory_footprint() / BYTES_PER_GB
|
||||
else:
|
||||
# manually
|
||||
param_size = sum(p.numel() * p.element_size() for p in module.parameters())
|
||||
buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
|
||||
|
||||
total_size_bytes = param_size + buffer_size
|
||||
usage = total_size_bytes / (1024**3)
|
||||
|
||||
return round(usage, 2)
|
||||
|
||||
|
||||
# component name -> ComponentLoader class
|
||||
component_name_to_loader_cls: Dict[str, Type[Any]] = {}
|
||||
|
||||
112
python/sglang/multimodal_gen/runtime/loader/vae_loader.py
Normal file
112
python/sglang/multimodal_gen/runtime/loader/vae_loader.py
Normal file
@@ -0,0 +1,112 @@
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
from safetensors.torch import load_file as safetensors_load_file
|
||||
|
||||
from sglang.multimodal_gen.configs.models import ModelConfig
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
_list_safetensors_files,
|
||||
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.hf_diffusers_utils import (
|
||||
get_diffusers_component_config,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class VAELoader(ComponentLoader):
|
||||
"""Shared loader for (video/audio) VAE modules."""
|
||||
|
||||
component_names = ["vae", "audio_vae"]
|
||||
expected_library = "diffusers"
|
||||
|
||||
def should_offload(
|
||||
self, server_args: ServerArgs, model_config: ModelConfig | None = None
|
||||
):
|
||||
return server_args.vae_cpu_offload
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, component_name: str
|
||||
):
|
||||
"""Load the VAE based on the model path, and inference args."""
|
||||
config = get_diffusers_component_config(model_path=component_model_path)
|
||||
class_name = config.pop("_class_name", None)
|
||||
assert (
|
||||
class_name is not None
|
||||
), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
|
||||
|
||||
server_args.model_paths[component_name] = component_model_path
|
||||
|
||||
logger.debug("HF model config: %s", config)
|
||||
if component_name in ("vae", "video_vae"):
|
||||
pipeline_vae_config_attr = "vae_config"
|
||||
pipeline_vae_precision = "vae_precision"
|
||||
elif component_name in ("audio_vae",):
|
||||
pipeline_vae_config_attr = "audio_vae_config"
|
||||
pipeline_vae_precision = "audio_vae_precision"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported module name for VAE loader: {component_name}"
|
||||
)
|
||||
vae_config = getattr(server_args.pipeline_config, pipeline_vae_config_attr)
|
||||
vae_precision = getattr(server_args.pipeline_config, pipeline_vae_precision)
|
||||
vae_config.update_model_arch(config)
|
||||
if hasattr(vae_config, "post_init"):
|
||||
# NOTE: some post init logics are only available after updated with config
|
||||
vae_config.post_init()
|
||||
|
||||
should_offload = self.should_offload(server_args)
|
||||
target_device = self.target_device(should_offload)
|
||||
|
||||
# Check for auto_map first (custom VAE classes)
|
||||
auto_map = config.get("auto_map", {})
|
||||
auto_model_map = auto_map.get("AutoModel")
|
||||
if auto_model_map:
|
||||
module_path, cls_name = auto_model_map.rsplit(".", 1)
|
||||
custom_module_file = os.path.join(component_model_path, f"{module_path}.py")
|
||||
spec = importlib.util.spec_from_file_location("_custom", custom_module_file)
|
||||
custom_module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(custom_module)
|
||||
vae_cls = getattr(custom_module, cls_name)
|
||||
vae_dtype = PRECISION_TO_TYPE[vae_precision]
|
||||
with set_default_torch_dtype(vae_dtype):
|
||||
vae = vae_cls.from_pretrained(
|
||||
component_model_path,
|
||||
revision=server_args.revision,
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
)
|
||||
vae = vae.to(device=target_device, dtype=vae_dtype)
|
||||
return vae
|
||||
|
||||
# Load from ModelRegistry (standard VAE classes)
|
||||
with (
|
||||
set_default_torch_dtype(PRECISION_TO_TYPE[vae_precision]),
|
||||
skip_init_modules(),
|
||||
):
|
||||
vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
|
||||
vae = vae_cls(vae_config).to(target_device)
|
||||
|
||||
safetensors_list = _list_safetensors_files(component_model_path)
|
||||
assert (
|
||||
len(safetensors_list) == 1
|
||||
), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
|
||||
loaded = safetensors_load_file(safetensors_list[0])
|
||||
vae.load_state_dict(loaded, strict=False)
|
||||
|
||||
state_keys = set(vae.state_dict().keys())
|
||||
loaded_keys = set(loaded.keys())
|
||||
missing_keys = sorted(state_keys - loaded_keys)
|
||||
unexpected_keys = sorted(loaded_keys - state_keys)
|
||||
if missing_keys:
|
||||
logger.warning("VAE missing keys: %s", missing_keys)
|
||||
if unexpected_keys:
|
||||
logger.warning("VAE unexpected keys: %s", unexpected_keys)
|
||||
|
||||
return vae
|
||||
@@ -0,0 +1,39 @@
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import get_hf_config
|
||||
|
||||
|
||||
class VisionLanguageEncoderLoader(ComponentLoader):
|
||||
"""Loader for vision language encoder (typically Causal LM or Vision2Seq)."""
|
||||
|
||||
component_names = ["vision_language_encoder"]
|
||||
expected_library = "transformers"
|
||||
|
||||
def load_customized(
|
||||
self,
|
||||
component_model_path: str,
|
||||
server_args: ServerArgs,
|
||||
transformers_or_diffusers: str = "vision_language_encoder",
|
||||
) -> Any:
|
||||
if transformers_or_diffusers == "vision_language_encoder":
|
||||
from transformers import GlmImageForConditionalGeneration
|
||||
|
||||
config = get_hf_config(
|
||||
component_model_path,
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
revision=server_args.revision,
|
||||
)
|
||||
model = GlmImageForConditionalGeneration.from_pretrained(
|
||||
component_model_path,
|
||||
config=config,
|
||||
trust_remote_code=server_args.trust_remote_code,
|
||||
revision=server_args.revision,
|
||||
).to(get_local_torch_device())
|
||||
return model
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported library for VisionLanguageEncoder: {transformers_or_diffusers}"
|
||||
)
|
||||
@@ -0,0 +1,86 @@
|
||||
from safetensors.torch import load_file as safetensors_load_file
|
||||
|
||||
from sglang.multimodal_gen.configs.models import ModelConfig
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import ComponentLoader
|
||||
from sglang.multimodal_gen.runtime.loader.utils import (
|
||||
_list_safetensors_files,
|
||||
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.hf_diffusers_utils import (
|
||||
get_diffusers_component_config,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class VocoderLoader(ComponentLoader):
|
||||
component_names = ["vocoder"]
|
||||
expected_library = "diffusers"
|
||||
|
||||
def should_offload(
|
||||
self, server_args: ServerArgs, model_config: ModelConfig | None = None
|
||||
):
|
||||
return server_args.vae_cpu_offload
|
||||
|
||||
def load_customized(
|
||||
self, component_model_path: str, server_args: ServerArgs, component_name: str
|
||||
):
|
||||
config = get_diffusers_component_config(model_path=component_model_path)
|
||||
class_name = config.pop("_class_name", None)
|
||||
assert (
|
||||
class_name is not None
|
||||
), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
|
||||
|
||||
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)
|
||||
|
||||
try:
|
||||
vocoder_precision = server_args.pipeline_config.audio_vae_precision
|
||||
except AttributeError:
|
||||
vocoder_precision = "fp32"
|
||||
vocoder_dtype = PRECISION_TO_TYPE[vocoder_precision]
|
||||
|
||||
should_offload = self.should_offload(server_args)
|
||||
target_device = self.target_device(should_offload)
|
||||
|
||||
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)
|
||||
|
||||
safetensors_list = _list_safetensors_files(component_model_path)
|
||||
assert (
|
||||
len(safetensors_list) == 1
|
||||
), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
|
||||
loaded = safetensors_load_file(safetensors_list[0])
|
||||
incompatible = vocoder.load_state_dict(loaded, strict=False)
|
||||
missing_keys = []
|
||||
unexpected_keys = []
|
||||
try:
|
||||
missing_keys = incompatible.missing_keys
|
||||
unexpected_keys = incompatible.unexpected_keys
|
||||
except AttributeError:
|
||||
# Best-effort fallback in case older torch returns a tuple-like.
|
||||
try:
|
||||
missing_keys = incompatible[0]
|
||||
unexpected_keys = incompatible[1]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if missing_keys or unexpected_keys:
|
||||
logger.warning(
|
||||
"Loaded vocoder with missing_keys=%d unexpected_keys=%d",
|
||||
len(missing_keys),
|
||||
len(unexpected_keys),
|
||||
)
|
||||
return vocoder
|
||||
@@ -248,7 +248,7 @@ class ComposedPipelineBase(ABC):
|
||||
required_modules = self.required_config_modules
|
||||
logger.info("Loading required components: %s", required_modules)
|
||||
|
||||
components = {}
|
||||
loaded_components = {}
|
||||
for module_name, (
|
||||
transformers_or_diffusers,
|
||||
architecture,
|
||||
@@ -266,7 +266,7 @@ class ComposedPipelineBase(ABC):
|
||||
continue
|
||||
if loaded_modules is not None and module_name in loaded_modules:
|
||||
logger.info("Using module %s already provided", module_name)
|
||||
components[module_name] = loaded_modules[module_name]
|
||||
loaded_components[module_name] = loaded_modules[module_name]
|
||||
continue
|
||||
|
||||
# we load the module from the extra config module map if it exists
|
||||
@@ -288,8 +288,8 @@ class ComposedPipelineBase(ABC):
|
||||
)
|
||||
else:
|
||||
component_model_path = os.path.join(self.model_path, load_module_name)
|
||||
module, memory_usage = PipelineComponentLoader.load_module(
|
||||
module_name=load_module_name,
|
||||
module, memory_usage = PipelineComponentLoader.load_component(
|
||||
component_name=load_module_name,
|
||||
component_model_path=component_model_path,
|
||||
transformers_or_diffusers=transformers_or_diffusers,
|
||||
server_args=server_args,
|
||||
@@ -297,20 +297,23 @@ class ComposedPipelineBase(ABC):
|
||||
|
||||
self.memory_usages[load_module_name] = memory_usage
|
||||
|
||||
if module_name in components:
|
||||
if module_name in loaded_components:
|
||||
logger.warning("Overwriting module %s", module_name)
|
||||
components[module_name] = module
|
||||
loaded_components[module_name] = module
|
||||
|
||||
# Check if all required modules were loaded
|
||||
for module_name in required_modules:
|
||||
if module_name not in components or components[module_name] is None:
|
||||
if (
|
||||
module_name not in loaded_components
|
||||
or loaded_components[module_name] is None
|
||||
):
|
||||
raise ValueError(
|
||||
f"Required module key: {module_name} value: {components.get(module_name)} was not found in loaded modules {components.keys()}"
|
||||
f"Required module: {module_name} was not found in loaded modules: {list(loaded_components.keys())}"
|
||||
)
|
||||
|
||||
logger.debug("Memory usage of loaded modules: %s", self.memory_usages)
|
||||
|
||||
return components
|
||||
return loaded_components
|
||||
|
||||
def add_stage(self, stage_name: str, stage: PipelineStage):
|
||||
assert self.modules is not None, "No modules are registered"
|
||||
|
||||
@@ -10,7 +10,7 @@ import weakref
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import VAELoader
|
||||
from sglang.multimodal_gen.runtime.loader.vae_loader import VAELoader
|
||||
from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
|
||||
|
||||
@@ -44,7 +44,7 @@ from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
|
||||
configure_sta,
|
||||
save_mask_search_results,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader.component_loader import TransformerLoader
|
||||
from sglang.multimodal_gen.runtime.loader.transformer_loader import TransformerLoader
|
||||
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
|
||||
|
||||
Reference in New Issue
Block a user