WIP: initial multimodal-gen support (#12484)

Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: JiLi <leege233@gmail.com>
Co-authored-by: CHEN Xi <78632976+RubiaCx@users.noreply.github.com>
Co-authored-by: laixin <xielx@shanghaitech.edu.cn>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: jzhang38 <a1286225768@gmail.com>
Co-authored-by: BrianChen1129 <yongqichcd@gmail.com>
Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: rlsu9 <r3su@ucsd.edu>
Co-authored-by: Jinzhe Pan <48981407+eigensystem@users.noreply.github.com>
Co-authored-by: foreverpiano <pianoqwz@qq.com>
Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
Co-authored-by: PorridgeSwim <yz3883@columbia.edu>
Co-authored-by: Jiali Chen <90408393+gary-chenjl@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/loader/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -0,0 +1,670 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+import glob
+import json
+import os
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Generator, Iterable
+from copy import deepcopy
+from typing import cast
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from safetensors.torch import load_file as safetensors_load_file
+from torch.distributed import init_device_mesh
+from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
+
+from sglang.multimodal_gen.configs.models import EncoderConfig
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.fsdp_load import (
+    maybe_load_fsdp_model,
+    shard_model,
+)
+from sglang.multimodal_gen.runtime.loader.utils import set_default_torch_dtype
+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_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 ComponentLoader(ABC):
+    """Base class for loading a specific type of model component."""
+
+    def __init__(self, device=None) -> None:
+        self.device = device
+
+    @abstractmethod
+    def load(self, model_path: str, server_args: ServerArgs, module_name: str):
+        """
+        Load the component based on the model path, architecture, and inference args.
+
+        Args:
+            model_path: Path to the component model
+            server_args: ServerArgs
+
+        Returns:
+            The loaded component
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def for_module_type(
+        cls, module_type: str, transformers_or_diffusers: str
+    ) -> "ComponentLoader":
+        """
+        Factory method to create a component loader for a specific module type.
+
+        Args:
+            module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            transformers_or_diffusers: Whether the module is from transformers or diffusers
+
+        Returns:
+            A component loader for the specified module type
+        """
+        # Map of module types to their loader classes and expected library
+        module_loaders = {
+            "scheduler": (SchedulerLoader, "diffusers"),
+            "transformer": (TransformerLoader, "diffusers"),
+            "transformer_2": (TransformerLoader, "diffusers"),
+            "vae": (VAELoader, "diffusers"),
+            "text_encoder": (TextEncoderLoader, "transformers"),
+            "text_encoder_2": (TextEncoderLoader, "transformers"),
+            "tokenizer": (TokenizerLoader, "transformers"),
+            "tokenizer_2": (TokenizerLoader, "transformers"),
+            "image_processor": (ImageProcessorLoader, "transformers"),
+            "image_encoder": (ImageEncoderLoader, "transformers"),
+            "processor": (AutoProcessorLoader, "transformers"),
+        }
+
+        if module_type in module_loaders:
+            loader_cls, expected_library = module_loaders[module_type]
+            # Assert that the library matches what's expected for this module type
+            assert (
+                transformers_or_diffusers == expected_library
+            ), f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
+            return loader_cls()
+
+        # For unknown module types, use a generic loader
+        logger.warning(
+            "No specific loader found for module type: %s. Using generic loader.",
+            module_type,
+        )
+        return GenericComponentLoader(transformers_or_diffusers)
+
+
+class TextEncoderLoader(ComponentLoader):
+    """Loader for text encoders."""
+
+    @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."""
+
+    counter_before_loading_weights: float = 0.0
+    counter_after_loading_weights: float = 0.0
+
+    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)
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+        # 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(self, model_path: str, server_args: ServerArgs, module_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(model_path, trust_remote_code=True)
+        model_config = get_diffusers_config(model=model_path)
+        model_config.pop("_name_or_path", None)
+        model_config.pop("transformers_version", None)
+        model_config.pop("model_type", None)
+        model_config.pop("tokenizer_class", None)
+        model_config.pop("torch_dtype", None)
+        logger.info("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(module_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]
+        target_device = get_local_torch_device()
+        # TODO(will): add support for other dtypes
+        return self.load_model(
+            model_path,
+            encoder_config,
+            target_device,
+            server_args,
+            encoder_dtype,
+        )
+
+    def load_model(
+        self,
+        model_path: str,
+        model_config: EncoderConfig,
+        target_device: torch.device,
+        server_args: ServerArgs,
+        dtype: str = "fp16",
+    ):
+        use_cpu_offload = (
+            server_args.text_encoder_cpu_offload
+            and len(getattr(model_config, "_fsdp_shard_conditions", [])) > 0
+        )
+
+        if server_args.text_encoder_cpu_offload:
+            target_device = (
+                torch.device("mps")
+                if current_platform.is_mps()
+                else torch.device("cpu")
+            )
+
+        with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
+            with target_device:
+                architectures = getattr(model_config, "architectures", [])
+                model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
+                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=use_cpu_offload)
+            )
+            self.counter_after_loading_weights = time.perf_counter()
+            logger.info(
+                "Loading weights took %.2f seconds",
+                self.counter_after_loading_weights
+                - self.counter_before_loading_weights,
+            )
+
+            # Explicitly move model to target device after loading weights
+            model = model.to(target_device)
+
+            if use_cpu_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"
+                    )
+                else:
+                    mesh = init_device_mesh(
+                        "cuda",
+                        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._fsdp_shard_conditions,
+                        pin_cpu_memory=server_args.pin_cpu_memory,
+                    )
+            # 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 weights were not initialized from "
+                    f"checkpoint: {weights_not_loaded}"
+                )
+
+        return model.eval()
+
+
+class ImageEncoderLoader(TextEncoderLoader):
+
+    def load(self, 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(model_path, "config.json")) as f:
+            model_config = json.load(f)
+        model_config.pop("_name_or_path", None)
+        model_config.pop("transformers_version", None)
+        model_config.pop("torch_dtype", None)
+        model_config.pop("model_type", None)
+        logger.info("HF model config: %s", model_config)
+
+        encoder_config = server_args.pipeline_config.image_encoder_config
+        encoder_config.update_model_arch(model_config)
+
+        if server_args.image_encoder_cpu_offload:
+            target_device = (
+                torch.device("mps")
+                if current_platform.is_mps()
+                else torch.device("cpu")
+            )
+        else:
+            target_device = get_local_torch_device()
+        # TODO(will): add support for other dtypes
+        return self.load_model(
+            model_path,
+            encoder_config,
+            target_device,
+            server_args,
+            server_args.pipeline_config.image_encoder_precision,
+        )
+
+
+class ImageProcessorLoader(ComponentLoader):
+    """Loader for image processor."""
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load the image processor based on the model path, and inference args."""
+        logger.info("Loading image processor from %s", model_path)
+
+        image_processor = AutoImageProcessor.from_pretrained(model_path, use_fast=True)
+        logger.info("Loaded image processor: %s", image_processor.__class__.__name__)
+        return image_processor
+
+
+class AutoProcessorLoader(ComponentLoader):
+    """Loader for auto processor."""
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load the image processor based on the model path, and inference args."""
+        logger.info("Loading auto processor from %s", model_path)
+
+        processor = AutoProcessor.from_pretrained(
+            model_path,
+        )
+        logger.info("Loaded auto processor: %s", processor.__class__.__name__)
+        return processor
+
+
+class TokenizerLoader(ComponentLoader):
+    """Loader for tokenizers."""
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load the tokenizer based on the model path, and inference args."""
+        logger.info("Loading tokenizer from %s", model_path)
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path,  # "<path to model>/tokenizer"
+            # in v0, this was same string as encoder_name "ClipTextModel"
+            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
+            # other method of config?
+            padding_size="right",
+        )
+        logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__)
+        return tokenizer
+
+
+class VAELoader(ComponentLoader):
+    """Loader for VAE."""
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load the VAE based on the model path, and inference args."""
+        config = get_diffusers_config(model=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."
+
+        server_args.model_paths["vae"] = model_path
+
+        # TODO: abstract these logics
+        logger.info("HF model config: %s", config)
+        vae_config = server_args.pipeline_config.vae_config
+        vae_config.update_model_arch(config)
+
+        # NOTE: some post init logics are only available after updated with config
+        vae_config.post_init()
+
+        if server_args.vae_cpu_offload:
+            target_device = (
+                torch.device("mps")
+                if current_platform.is_mps()
+                else torch.device("cpu")
+            )
+        else:
+            target_device = get_local_torch_device()
+
+        with set_default_torch_dtype(
+            PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        ):
+            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+            vae = vae_cls(vae_config).to(target_device)
+
+        # Find all safetensors files
+        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+        # TODO(PY)
+        assert (
+            len(safetensors_list) == 1
+        ), f"Found {len(safetensors_list)} safetensors files in {model_path}"
+        loaded = safetensors_load_file(safetensors_list[0])
+        vae.load_state_dict(
+            loaded, strict=False
+        )  # We might only load encoder or decoder
+
+        return vae.eval()
+
+
+class TransformerLoader(ComponentLoader):
+    """Loader for transformer."""
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load the transformer based on the model path, and inference args."""
+        config = get_diffusers_config(model=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."
+            )
+
+        logger.info("transformer cls_name: %s", cls_name)
+        if server_args.override_transformer_cls_name is not None:
+            cls_name = server_args.override_transformer_cls_name
+            logger.info("Overriding transformer cls_name to %s", cls_name)
+
+        server_args.model_paths["transformer"] = model_path
+
+        # Config from Diffusers supersedes sgl_diffusion's model config
+        dit_config = server_args.pipeline_config.dit_config
+        dit_config.update_model_arch(config)
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+
+        # Find all safetensors files
+        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {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 = glob.glob(
+                    os.path.join(str(custom_weights_path), "*.safetensors")
+                )
+            else:
+                assert custom_weights_path.endswith(
+                    ".safetensors"
+                ), "Custom initialization weights must be a safetensors file"
+                safetensors_list = [custom_weights_path]
+
+        logger.info(
+            "Loading model from %s safetensors files: %s",
+            len(safetensors_list),
+            safetensors_list,
+        )
+
+        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        # Load the model using FSDP loader
+        logger.info("Loading %s, default_dtype: %s", cls_name, default_dtype)
+        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,
+        )
+
+        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"
+
+        model = model.eval()
+        return model
+
+
+class SchedulerLoader(ComponentLoader):
+    """Loader for scheduler."""
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load the scheduler based on the model path, and inference args."""
+        config = get_diffusers_config(model=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)
+        if server_args.pipeline_config.timesteps_scale is not None:
+            scheduler.set_timesteps_scale(server_args.pipeline_config.timesteps_scale)
+        return scheduler
+
+
+class GenericComponentLoader(ComponentLoader):
+    """Generic loader for components that don't have a specific loader."""
+
+    def __init__(self, library="transformers") -> None:
+        super().__init__()
+        self.library = library
+
+    def load(self, model_path: str, server_args: ServerArgs, *args):
+        """Load a generic component based on the model path, and inference args."""
+        logger.warning(
+            "Using generic loader for %s with library %s", model_path, self.library
+        )
+
+        if self.library == "transformers":
+            from transformers import AutoModel
+
+            model = AutoModel.from_pretrained(
+                model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+            )
+            logger.info(
+                "Loaded generic transformers model: %s", model.__class__.__name__
+            )
+            return model
+        elif self.library == "diffusers":
+            logger.warning(
+                "Generic loading for diffusers components is not fully implemented"
+            )
+
+            model_config = get_diffusers_config(model=model_path)
+            logger.info("Diffusers Model config: %s", model_config)
+            # This is a placeholder - in a real implementation, you'd need to handle this properly
+            return None
+        else:
+            raise ValueError(f"Unsupported library: {self.library}")
+
+
+class PipelineComponentLoader:
+    """
+    Utility class for loading pipeline components.
+    This replaces the chain of if-else statements in load_pipeline_module.
+    """
+
+    @staticmethod
+    def load_module(
+        module_name: str,
+        component_model_path: str,
+        transformers_or_diffusers: str,
+        server_args: ServerArgs,
+    ):
+        """
+        Load a pipeline module.
+
+        Args:
+            module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler")
+            component_model_path: Path to the component model
+            transformers_or_diffusers: Whether the module is from transformers or diffusers
+
+        Returns:
+            The loaded module
+        """
+        logger.info(
+            "Loading %s using %s from %s",
+            module_name,
+            transformers_or_diffusers,
+            component_model_path,
+        )
+
+        # Get the appropriate loader for this module type
+        loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers)
+
+        try:
+            # Load the module
+            return loader.load(component_model_path, server_args, module_name)
+        except Exception as e:
+            logger.error(
+                f"Error while loading component: {module_name}, {component_model_path=}"
+            )
+            raise e

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -0,0 +1,314 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Adapted from torchtune
+# Copyright 2024 The TorchTune Authors.
+# Copyright 2025 The sgl-diffusion Authors.
+
+import contextlib
+from collections.abc import Callable, Generator
+from itertools import chain
+from typing import Any
+
+import torch
+from torch import nn
+from torch.distributed import DeviceMesh, init_device_mesh
+from torch.distributed._tensor import distribute_tensor
+from torch.distributed.fsdp import (
+    CPUOffloadPolicy,
+    FSDPModule,
+    MixedPrecisionPolicy,
+    fully_shard,
+)
+from torch.nn.modules.module import _IncompatibleKeys
+
+from sglang.multimodal_gen.runtime.loader.utils import (
+    get_param_names_mapping,
+    hf_to_custom_state_dict,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    safetensors_weights_iterator,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import set_mixed_precision_policy
+
+logger = init_logger(__name__)
+
+
+# TODO(PY): move this to utils elsewhere
+@contextlib.contextmanager
+def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
+    """
+    Context manager to set torch's default dtype.
+
+    Args:
+        dtype (torch.dtype): The desired default dtype inside the context manager.
+
+    Returns:
+        ContextManager: context manager for setting default dtype.
+
+    Example:
+        >>> with set_default_dtype(torch.bfloat16):
+        >>>     x = torch.tensor([1, 2, 3])
+        >>>     x.dtype
+        torch.bfloat16
+
+
+    """
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)
+
+
+# TODO(PY): add compile option
+def maybe_load_fsdp_model(
+    model_cls: type[nn.Module],
+    init_params: dict[str, Any],
+    weight_dir_list: list[str],
+    device: torch.device,
+    hsdp_replicate_dim: int,
+    hsdp_shard_dim: int,
+    default_dtype: torch.dtype,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
+    cpu_offload: bool = False,
+    fsdp_inference: bool = False,
+    output_dtype: torch.dtype | None = None,
+    pin_cpu_memory: bool = True,
+) -> torch.nn.Module:
+    """
+    Load the model with FSDP if is training, else load the model without FSDP.
+    """
+    # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
+    # manually casting the inputs to the model
+    mp_policy = MixedPrecisionPolicy(
+        param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
+    )
+
+    set_mixed_precision_policy(
+        param_dtype=param_dtype,
+        reduce_dtype=reduce_dtype,
+        output_dtype=output_dtype,
+        mp_policy=mp_policy,
+    )
+
+    with set_default_dtype(default_dtype), torch.device("meta"):
+        model = model_cls(**init_params)
+
+    # Check if we should use FSDP
+    use_fsdp = fsdp_inference
+
+    # Disable FSDP for MPS as it's not compatible
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    if current_platform.is_mps():
+        use_fsdp = False
+        logger.info("Disabling FSDP for MPS platform as it's not compatible")
+
+    if use_fsdp:
+        world_size = hsdp_replicate_dim * hsdp_shard_dim
+        if not fsdp_inference:
+            hsdp_replicate_dim = world_size
+            hsdp_shard_dim = 1
+
+        device_mesh = init_device_mesh(
+            "cuda",
+            # (Replicate(), Shard(dim=0))
+            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
+            mesh_dim_names=("replicate", "shard"),
+        )
+        shard_model(
+            model,
+            cpu_offload=cpu_offload,
+            reshard_after_forward=True,
+            mp_policy=mp_policy,
+            mesh=device_mesh,
+            fsdp_shard_conditions=model._fsdp_shard_conditions,
+            pin_cpu_memory=pin_cpu_memory,
+        )
+
+    weight_iterator = safetensors_weights_iterator(weight_dir_list)
+    param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
+    load_model_from_full_model_state_dict(
+        model,
+        weight_iterator,
+        device,
+        default_dtype,
+        strict=True,
+        cpu_offload=cpu_offload,
+        param_names_mapping=param_names_mapping_fn,
+    )
+    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
+    return model
+
+
+def shard_model(
+    model,
+    *,
+    cpu_offload: bool,
+    reshard_after_forward: bool = True,
+    mp_policy: MixedPrecisionPolicy | None = MixedPrecisionPolicy(),  # noqa
+    mesh: DeviceMesh | None = None,
+    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] = [],  # noqa
+    pin_cpu_memory: bool = True,
+) -> None:
+    """
+    Utility to shard a model with FSDP using the PyTorch Distributed fully_shard API.
+
+    This method will over the model's named modules from the bottom-up and apply shard modules
+    based on whether they meet any of the criteria from shard_conditions.
+
+    Args:
+        model (TransformerDecoder): Model to shard with FSDP.
+        cpu_offload (bool): If set to True, FSDP will offload parameters, gradients, and optimizer
+            states to CPU.
+        reshard_after_forward (bool): Whether to reshard parameters and buffers after
+            the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy
+            from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy.
+        mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism.
+            Default to None.
+        fsdp_shard_conditions (List[Callable[[str, nn.Module], bool]]): A list of functions to determine
+            which modules to shard with FSDP.
+        pin_cpu_memory (bool): If set to True, FSDP will pin the CPU memory of the offloaded parameters.
+
+    Raises:
+        ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered.
+    """
+    if fsdp_shard_conditions is None or len(fsdp_shard_conditions) == 0:
+        logger.warning(
+            "The FSDP shard condition list is empty or None. No modules will be sharded in %s",
+            type(model).__name__,
+        )
+        return
+
+    fsdp_kwargs = {
+        "reshard_after_forward": reshard_after_forward,
+        "mesh": mesh,
+        "mp_policy": mp_policy,
+    }
+    if cpu_offload:
+        fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=pin_cpu_memory)
+
+    # iterating in reverse to start with
+    # lowest-level modules first
+    num_layers_sharded = 0
+    # TODO(will): don't reshard after forward for the last layer to save on the
+    # all-gather that will immediately happen Shard the model with FSDP,
+    for n, m in reversed(list(model.named_modules())):
+        if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
+            fully_shard(m, **fsdp_kwargs)
+            num_layers_sharded += 1
+
+    if num_layers_sharded == 0:
+        raise ValueError(
+            "No layer modules were sharded. Please check if shard conditions are working as expected."
+        )
+
+    # Finally shard the entire model to account for any stragglers
+    fully_shard(model, **fsdp_kwargs)
+
+
+# TODO(PY): device mesh for cfg parallel
+def load_model_from_full_model_state_dict(
+    model: FSDPModule | torch.nn.Module,
+    full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
+    device: torch.device,
+    param_dtype: torch.dtype,
+    strict: bool = False,
+    cpu_offload: bool = False,
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
+) -> _IncompatibleKeys:
+    """
+    Converting full state dict into a sharded state dict
+    and loading it into FSDP model (if training) or normal huggingface model
+    Args:
+        model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict
+        full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
+        device (torch.device): device used to move full state dict tensors
+        param_dtype (torch.dtype): dtype used to move full state dict tensors
+        strict (bool): flag to check if to load the model in strict mode
+        cpu_offload (bool): flag to check if FSDP offload is enabled
+        param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
+    Returns:
+        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
+            * **missing_keys** is a list of str containing the missing keys
+            * **unexpected_keys** is a list of str containing the unexpected keys
+
+    Raises:
+        NotImplementedError: If got FSDP with more than 1D.
+    """
+    meta_sd = model.state_dict()
+    sharded_sd = {}
+    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
+        full_sd_iterator, param_names_mapping
+    )  # type: ignore
+    for target_param_name, full_tensor in custom_param_sd.items():
+        meta_sharded_param = meta_sd.get(target_param_name)
+        if meta_sharded_param is None:
+            raise ValueError(
+                f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
+            )
+        if not hasattr(meta_sharded_param, "device_mesh"):
+            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
+            # In cases where parts of the model aren't sharded, some parameters will be plain tensors
+            sharded_tensor = full_tensor
+        else:
+            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                meta_sharded_param.device_mesh,
+                meta_sharded_param.placements,
+            )
+            if cpu_offload:
+                sharded_tensor = sharded_tensor.cpu()
+        sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)
+
+    model.reverse_param_names_mapping = reverse_param_names_mapping
+    unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
+    if unused_keys:
+        logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys)
+
+    # List of allowed parameter name patterns
+    ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"]  # Can be extended as needed
+    for new_param_name in unused_keys:
+        if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS):
+            logger.error(
+                "Unsupported new parameter: %s. Allowed patterns: %s",
+                new_param_name,
+                ALLOWED_NEW_PARAM_PATTERNS,
+            )
+            raise ValueError(
+                f"New parameter '{new_param_name}' is not supported. "
+                f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed."
+            )
+        meta_sharded_param = meta_sd.get(new_param_name)
+        if not hasattr(meta_sharded_param, "device_mesh"):
+            # Initialize with zeros
+            sharded_tensor = torch.zeros_like(
+                meta_sharded_param, device=device, dtype=param_dtype
+            )
+        else:
+            # Initialize with zeros and distribute
+            full_tensor = torch.zeros_like(
+                meta_sharded_param, device=device, dtype=param_dtype
+            )
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                meta_sharded_param.device_mesh,
+                meta_sharded_param.placements,
+            )
+            if cpu_offload:
+                sharded_tensor = sharded_tensor.cpu()
+        sharded_sd[new_param_name] = nn.Parameter(sharded_tensor)
+
+    # choose `assign=True` since we cannot call `copy_` on meta tensor
+    return model.load_state_dict(sharded_sd, strict=strict, assign=True)

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -0,0 +1,103 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for selecting and loading models."""
+import contextlib
+import re
+from collections import defaultdict
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@contextlib.contextmanager
+def set_default_torch_dtype(dtype: torch.dtype):
+    """Sets the default torch dtype to the given dtype."""
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    yield
+    torch.set_default_dtype(old_dtype)
+
+
+def get_param_names_mapping(
+    mapping_dict: dict[str, str]
+) -> Callable[[str], tuple[str, Any, Any]]:
+    """
+    Creates a mapping function that transforms parameter names using regex patterns.
+
+    Args:
+        mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns
+        param_name (str): The parameter name to be transformed
+
+    Returns:
+        Callable[[str], str]: A function that maps parameter names from source to target format
+    """
+
+    def mapping_fn(name: str) -> tuple[str, Any, Any]:
+        # Try to match and transform the name using the regex patterns in mapping_dict
+        for pattern, replacement in mapping_dict.items():
+            match = re.match(pattern, name)
+            if match:
+                merge_index = None
+                total_splitted_params = None
+                if isinstance(replacement, tuple):
+                    merge_index = replacement[1]
+                    total_splitted_params = replacement[2]
+                    replacement = replacement[0]
+                name = re.sub(pattern, replacement, name)
+                return name, merge_index, total_splitted_params
+
+        # If no pattern matches, return the original name
+        return name, None, None
+
+    return mapping_fn
+
+
+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]],
+) -> tuple[dict[str, torch.Tensor], dict[str, tuple[str, Any, Any]]]:
+    """
+    Converts a Hugging Face parameter state dictionary to a custom parameter state dictionary.
+
+    Args:
+        hf_param_sd (Dict[str, torch.Tensor]): The Hugging Face parameter state dictionary
+        param_names_mapping (Callable[[str], tuple[str, Any, Any]]): A function that maps parameter names from source to target format
+
+    Returns:
+        custom_param_sd (Dict[str, torch.Tensor]): The custom formatted parameter state dict
+        reverse_param_names_mapping (Dict[str, Tuple[str, Any, Any]]): Maps back from custom to hf
+    """
+    custom_param_sd = {}
+    to_merge_params = defaultdict(dict)  # type: ignore
+    reverse_param_names_mapping = {}
+    if isinstance(hf_param_sd, dict):
+        hf_param_sd = hf_param_sd.items()  # type: ignore
+    for source_param_name, full_tensor in hf_param_sd:  # type: ignore
+        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
+            source_param_name
+        )
+        reverse_param_names_mapping[target_param_name] = (
+            source_param_name,
+            merge_index,
+            num_params_to_merge,
+        )
+        if merge_index is not None:
+            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
+                sorted_tensors = [
+                    to_merge_params[target_param_name][i]
+                    for i in range(num_params_to_merge)
+                ]
+                full_tensor = torch.cat(sorted_tensors, dim=0)
+                del to_merge_params[target_param_name]
+            else:
+                continue
+        custom_param_sd[target_param_name] = full_tensor
+    return custom_param_sd, reverse_param_names_mapping

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -0,0 +1,238 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/model_loader/weight_utils.py
+"""Utilities for downloading and initializing model weights."""
+import hashlib
+import json
+import os
+import tempfile
+from collections.abc import Generator
+from pathlib import Path
+
+import filelock
+import huggingface_hub.constants
+import torch
+from safetensors.torch import safe_open
+from tqdm.auto import tqdm
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+# use system-level temp directory for file locks, so that multiple users
+# can share the same lock without error.
+# lock files in the temp directory will be automatically deleted when the
+# system reboots, so users will not complain about annoying lock files
+temp_dir = tempfile.gettempdir()
+
+
+def enable_hf_transfer() -> None:
+    """automatically activates hf_transfer"""
+    if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
+        try:
+            # enable hf hub transfer if available
+            import hf_transfer  # type: ignore # noqa
+
+            huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True
+        except ImportError:
+            pass
+
+
+enable_hf_transfer()
+
+
+class DisabledTqdm(tqdm):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, disable=True)
+
+
+def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
+    lock_dir = cache_dir or temp_dir
+    model_name_or_path = str(model_name_or_path)
+    os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
+    model_name = model_name_or_path.replace("/", "-")
+    hash_name = hashlib.sha256(model_name.encode()).hexdigest()
+    # add hash to avoid conflict with old users' lock files
+    lock_file_name = hash_name + model_name + ".lock"
+    # mode 0o666 is required for the filelock to be shared across users
+    lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
+    return lock
+
+
+# For models like Mistral-7B-v0.3, there are both sharded
+# safetensors files and a consolidated safetensors file.
+# Passing both of these to the weight loader functionality breaks.
+# So, we use the index_file to
+# look up which safetensors files should be used.
+def filter_duplicate_safetensors_files(
+    hf_weights_files: list[str], hf_folder: str, index_file: str
+) -> list[str]:
+    # model.safetensors.index.json is a mapping from keys in the
+    # torch state_dict to safetensors file holding that weight.
+    index_file_name = os.path.join(hf_folder, index_file)
+    if not os.path.isfile(index_file_name):
+        return hf_weights_files
+
+    # Iterate through the weight_map (weight_name: safetensors files)
+    # to identify weights that we should use.
+    with open(index_file_name) as f:
+        weight_map = json.load(f)["weight_map"]
+    weight_files_in_index = set()
+    for weight_name in weight_map:
+        weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name]))
+    # Filter out any fields that are not found in the index file.
+    hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index]
+    return hf_weights_files
+
+
+def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]:
+    """
+    Exclude files that are not needed for inference.
+
+    See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
+    """
+    blacklist = [
+        "training_args.bin",
+        "optimizer.bin",
+        "optimizer.pt",
+        "scheduler.pt",
+        "scaler.pt",
+    ]
+    hf_weights_files = [
+        f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist)
+    ]
+    return hf_weights_files
+
+
+# explicitly use pure text format, with a newline at the end
+# this makes it impossible to see the animation in the progress bar
+# but will avoid messing up with ray or multiprocessing, which wraps
+# each line of output with some prefix.
+_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501
+
+
+def safetensors_weights_iterator(
+    hf_weights_files: list[str],
+    to_cpu: bool = True,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model safetensor files."""
+    enable_tqdm = (
+        not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+    )
+    device = "cpu" if to_cpu else str(get_local_torch_device())
+    for st_file in tqdm(
+        hf_weights_files,
+        desc="Loading safetensors checkpoint shards",
+        disable=not enable_tqdm,
+        bar_format=_BAR_FORMAT,
+    ):
+        with safe_open(st_file, framework="pt", device=device) as f:
+            for name in f.keys():  # noqa: SIM118
+                param = f.get_tensor(name)
+                yield name, param
+
+
+def pt_weights_iterator(
+    hf_weights_files: list[str],
+    to_cpu: bool = True,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Iterate over the weights in the model bin/pt files."""
+    device = "cpu" if to_cpu else str(get_local_torch_device())
+    enable_tqdm = (
+        not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+    )
+    for bin_file in tqdm(
+        hf_weights_files,
+        desc="Loading pt checkpoint shards",
+        disable=not enable_tqdm,
+        bar_format=_BAR_FORMAT,
+    ):
+        state = torch.load(bin_file, map_location=device, weights_only=True)
+        yield from state.items()
+        del state
+
+
+def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+    """Default weight loader."""
+    try:
+        if param.numel() == 1 and loaded_weight.numel() == 1:
+            # Sometimes scalar values aren't considered tensors with shapes
+            # so if both param and loaded_weight are a scalar,
+            # "broadcast" instead of copy
+            param.data.fill_(loaded_weight.item())
+        else:
+            assert param.size() == loaded_weight.size(), (
+                f"Attempted to load weight ({loaded_weight.size()}) "
+                f"into parameter ({param.size()})"
+            )
+
+            param.data.copy_(loaded_weight)
+    except Exception:
+        # NOTE: This exception is added for the purpose of setting breakpoint to
+        # debug weight loading issues.
+        raise
+
+
+def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
+    """Remap the name of FP8 k/v_scale parameters.
+
+    This function handles the remapping of FP8 k/v_scale parameter names.
+    It detects if the given name ends with a suffix and attempts to remap
+    it to the expected name format in the model. If the remapped name is not
+    found in the params_dict, a warning is printed and None is returned.
+
+    Args:
+        name (str): The original loaded checkpoint parameter name.
+        params_dict (dict): Dictionary containing the model's named parameters.
+
+    Returns:
+        str: The remapped parameter name if successful, or the original name
+             if no remapping is needed.
+        None: If the remapped name is not found in params_dict.
+    """
+    if name.endswith(".kv_scale"):
+        logger.warning_once(
+            "DEPRECATED. Found kv_scale in the checkpoint. "
+            "This format is deprecated in favor of separate k_scale and "
+            "v_scale tensors and will be removed in a future release. "
+            "Functionally, we will remap kv_scale to k_scale and duplicate "
+            "k_scale to v_scale"
+        )
+        # NOTE: we remap the deprecated kv_scale to k_scale
+        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
+        if remapped_name not in params_dict:
+            logger.warning_once(
+                f"Found kv_scale in the checkpoint (e.g. {name}), "
+                "but not found the expected name in the model "
+                f"(e.g. {remapped_name}). kv_scale is "
+                "not loaded."
+            )
+            return None
+        return remapped_name
+
+    possible_scale_names = [".k_scale", ".v_scale"]
+    modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"]
+    for scale_name in possible_scale_names:
+        if name.endswith(scale_name):
+            if any(mo_scale_name in name for mo_scale_name in modelopt_scale_names):
+                remapped_name = name.replace(
+                    f".self_attn.{scale_name[1]}_proj{scale_name}",
+                    f".self_attn.attn{scale_name}",
+                )
+            else:
+                remapped_name = name.replace(scale_name, f".attn{scale_name}")
+            if remapped_name not in params_dict:
+                logger.warning_once(
+                    f"Found {scale_name} in the checkpoint (e.g. {name}), "
+                    "but not found the expected name in the model "
+                    f"(e.g. {remapped_name}). {scale_name} is "
+                    "not loaded."
+                )
+                return None
+            return remapped_name
+
+    # If there were no matches, return the untouched param name
+    return name