sgl-project--sglang/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

# 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 sglang-diffusion Authors.

from collections import Counter, defaultdict
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.configs.models.fsdp import is_module_list_entry_in
from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
    attach_bitsandbytes_4bit_quant_states,
    build_bitsandbytes_4bit_quant_states,
    split_bitsandbytes_4bit_state,
)
from sglang.multimodal_gen.runtime.loader.utils import (
    get_param_names_mapping,
    hf_to_custom_state_dict,
    set_default_torch_dtype,
)
from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
from sglang.multimodal_gen.runtime.loader.weight_utils import (
    safetensors_weights_iterator,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import set_mixed_precision_policy
from sglang.srt.utils import is_npu

_is_npu = is_npu()

logger = init_logger(__name__)

_QUANTIZED_DTYPES = (
    torch.uint8,
    torch.float8_e4m3fn,
    torch.float8_e5m2,
    torch.int8,
)
_DTYPE_MISMATCH_EXAMPLE_LIMIT = 3


def _is_bitsandbytes_quant_config(quant_config: Any | None) -> bool:
    if quant_config is None:
        return False
    quant_name_getter = getattr(type(quant_config), "get_name", None)
    return bool(callable(quant_name_getter) and quant_name_getter() == "bitsandbytes")


def _format_dtype_mismatch_summary(
    mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]],
    mismatch_examples: dict[tuple[torch.dtype, torch.dtype], list[str]],
) -> str:
    parts: list[str] = []
    for (checkpoint_dtype, target_dtype), count in mismatch_counts.items():
        examples = mismatch_examples[(checkpoint_dtype, target_dtype)]
        part = f"{checkpoint_dtype}->{target_dtype} x{count}"
        if examples:
            part += f" (e.g. {', '.join(examples)})"
        parts.append(part)
    return "; ".join(parts)


def _make_param_like(
    actual_param: torch.nn.Parameter, tensor: torch.Tensor
) -> torch.nn.Parameter:
    cls = actual_param.__class__
    # nn.Parameter defaults to requires_grad=True, which is illegal for non-floating/complex dtypes (e.g., int8/FP8
    # quantized weights).
    try:
        new_param = cls.__new__(cls, tensor, requires_grad=False)
    except TypeError:
        try:
            new_param = cls.__new__(cls, tensor)
        except TypeError:
            new_param = nn.Parameter(tensor, requires_grad=False)
    new_param.__dict__.update(actual_param.__dict__)
    new_param.requires_grad = False
    return new_param


def _get_param_for_weight_loading(
    model: torch.nn.Module,
    param_dict: dict[str, torch.nn.Parameter],
    param_name: str,
) -> torch.nn.Parameter | None:
    actual_param = param_dict.get(param_name)
    if actual_param is not None and getattr(actual_param, "weight_loader", None):
        return actual_param

    pre_fsdp_weight_loader_params = getattr(model, "_pre_fsdp_weight_loader_params", {})
    pre_fsdp_param = pre_fsdp_weight_loader_params.get(param_name)
    if pre_fsdp_param is not None:
        return pre_fsdp_param

    return actual_param


def _make_class_name_shard_condition(class_names: set[str]):
    def shard_condition(n: str, m: nn.Module) -> bool:
        return type(m).__name__ in class_names

    return shard_condition


def _is_common_numbered_block(n: str, m: nn.Module) -> bool:
    return is_module_list_entry_in(
        n,
        (
            "blocks",
            "layers",
            "double_blocks",
            "single_blocks",
            "refiner_blocks",
            "noise_refiner",
            "context_refiner",
            "transformer_blocks",
            "single_transformer_blocks",
        ),
    )


def _resolve_fsdp_shard_conditions(
    model: torch.nn.Module,
    fsdp_shard_conditions: list[Callable[[str, nn.Module], bool]] | None,
) -> tuple[list[Callable[[str, nn.Module], bool]], str]:
    if fsdp_shard_conditions:
        return fsdp_shard_conditions, "explicit"

    block_class_names = set(getattr(model, "_repeated_blocks", []) or [])
    block_class_names.update(getattr(model, "_no_split_modules", []) or [])
    if block_class_names:
        return [_make_class_name_shard_condition(block_class_names)], "block-class"

    return [_is_common_numbered_block], "common-numbered-block"


def _maybe_dequantize_fp8(
    full_tensor: torch.Tensor,
    target_dtype: torch.dtype,
    target_param_name: str,
    param_sd: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Auto-dequantize an FP8 checkpoint weight when the model parameter expects a higher-precision type.

    Some modules (e.g. AdaLayerNormZero) don't accept quant_config, so their
    parameters remain in higher precision even when the checkpoint stores FP8
    weights.  In that case we multiply by the per-tensor weight_scale to
    recover the original unquantized value.
    """
    if not (
        full_tensor.dtype == torch.float8_e4m3fn and target_dtype != torch.float8_e4m3fn
    ):
        return full_tensor

    scale_key = target_param_name.rsplit(".", 1)[0] + ".weight_scale"
    scale_tensor = param_sd.get(scale_key)
    if scale_tensor is not None:
        full_tensor = full_tensor.to(torch.float32) * scale_tensor.float()
        logger.debug(
            "Auto-dequantized FP8 weight %s using %s",
            target_param_name,
            scale_key,
        )
    return full_tensor


# 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,
    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,
    strict: bool = True,
    weight_load_plan: WeightLoadPlan | None = None,
) -> torch.nn.Module:
    """Load a model with optional FSDP (Fully Sharded Data Parallel) support.

    Args:
        param_dtype: Data type for model parameters, also used for:
            - Model initialization context (set_default_torch_dtype)
            - FSDP mixed precision policy
            - Weight loading and casting
        reduce_dtype: Data type for gradient reduction in FSDP mixed precision.
        strict: If True, enforce strict state dict loading (all keys must match).
        weight_load_plan: Optional checkpoint/postprocess device plan for this load.
    """
    # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
    # manually casting the inputs to the model

    # 1. prepare for loading
    default_torch_dtype = param_dtype if param_dtype else torch.bfloat16
    mp_policy = MixedPrecisionPolicy(
        default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
    )

    set_mixed_precision_policy(
        param_dtype=default_torch_dtype,
        reduce_dtype=reduce_dtype,
        output_dtype=output_dtype,
        mp_policy=mp_policy,
    )

    with set_default_torch_dtype(default_torch_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
    if current_platform.is_mps():
        use_fsdp = False
        logger.info("Disabling FSDP for MPS platform as it's not compatible")

    weight_load_plan = weight_load_plan or WeightLoadPlan(checkpoint_load_device=device)
    defer_cpu_offload = bool(
        cpu_offload and weight_load_plan.defer_component_cpu_offload
    )
    if defer_cpu_offload and use_fsdp:
        logger.warning(
            "Ignoring deferred CPU offload for FSDP loading; keeping the existing "
            "FSDP offload policy."
        )
        defer_cpu_offload = False
    load_cpu_offload = bool(cpu_offload and not defer_cpu_offload)
    weight_postprocess_device = weight_load_plan.weight_postprocess_device
    if use_fsdp and weight_postprocess_device is not None:
        logger.warning("Ignoring weight postprocess device override for FSDP loading.")
        weight_postprocess_device = None

    if use_fsdp:
        model._pre_fsdp_weight_loader_params = {
            n: p
            for n, p in model.named_parameters()
            if getattr(p, "weight_loader", None)
        }
        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(
            current_platform.device_type,
            # (Replicate(), Shard(dim=0))
            mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
            mesh_dim_names=("replicate", "shard"),
        )
        shard_model(
            model,
            cpu_offload=load_cpu_offload,
            reshard_after_forward=True,
            mp_policy=mp_policy,
            mesh=device_mesh,
            fsdp_shard_conditions=getattr(model, "_fsdp_shard_conditions", None),
            pin_cpu_memory=pin_cpu_memory,
        )

    param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)

    # 2. load model from disk
    weight_iterator = safetensors_weights_iterator(weight_dir_list)
    preprocess_loaded_state_dict = getattr(model, "preprocess_loaded_state_dict", None)
    if preprocess_loaded_state_dict is not None:
        weight_iterator = preprocess_loaded_state_dict(weight_iterator)
    bnb_quant_states = None
    if _is_bitsandbytes_quant_config(init_params.get("quant_config")):
        normal_weights, raw_quant_state = split_bitsandbytes_4bit_state(weight_iterator)
        bnb_quant_states = build_bitsandbytes_4bit_quant_states(
            [name for name, _ in normal_weights],
            raw_quant_state,
            device,
            param_names_mapping_fn,
        )
        weight_iterator = iter(normal_weights)
    load_model_from_full_model_state_dict(
        model,
        weight_iterator,
        weight_load_plan.checkpoint_load_device,
        param_dtype,
        strict=strict,
        cpu_offload=load_cpu_offload,
        param_names_mapping=param_names_mapping_fn,
    )
    if bnb_quant_states:
        attach_bitsandbytes_4bit_quant_states(
            dict(model.named_parameters()), bnb_quant_states
        )

    # 3. postprocessing
    if weight_postprocess_device is not None:
        # move to device to perform postprocessing
        model.to(weight_postprocess_device)

    for _, module in model.named_modules():
        quant_method = getattr(module, "quant_method", None)
        if quant_method is not None and hasattr(
            quant_method, "process_weights_after_loading"
        ):
            if _is_npu and not isinstance(quant_method, UnquantizedLinearMethod):
                # Activate the NZ format for storing weights,
                # which is a specific optimization for Ascend NPU
                torch.npu.config.allow_internal_format = True
            quant_method.process_weights_after_loading(module)
            if _is_npu:
                torch.npu.empty_cache()
    model.post_load_weights()

    for n, p in chain(model.named_parameters(), model.named_buffers()):
        if p.is_meta:
            raise RuntimeError(f"Unexpected param or buffer {n} on meta device.")
        # Avoid unintended computation graph accumulation during inference
        if isinstance(p, torch.nn.Parameter):
            p.requires_grad = False

    # 4. deferred cpu offload
    if defer_cpu_offload:
        model.to("cpu")

    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]] | None = None,
    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.

    """
    fsdp_shard_conditions, condition_source = _resolve_fsdp_shard_conditions(
        model, fsdp_shard_conditions
    )
    if condition_source != "explicit":
        logger.warning(
            "Using %s FSDP shard condition fallback for %s",
            condition_source,
            type(model).__name__,
        )

    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]):  # type: ignore
            fully_shard(m, **fsdp_kwargs)
            num_layers_sharded += 1

    if num_layers_sharded == 0:
        raise ValueError(
            f"No layer modules were sharded in {type(model).__name__}. "
            f"FSDP shard condition source: {condition_source}."
        )

    # Finally shard the entire model to account for any stragglers
    fully_shard(model, **fsdp_kwargs)
    logger.info(
        "Applied FSDP to %d submodules in %s using %s shard conditions",
        num_layers_sharded,
        type(model).__name__,
        condition_source,
    )


# TODO(mick): need refactor, to move out checkpoint-specific adjustments
def load_model_from_full_model_state_dict(
    model: FSDPModule | torch.nn.Module,
    full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
    checkpoint_load_device: torch.device,
    param_dtype: torch.dtype | None,
    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
        checkpoint_load_device (torch.device): device used to move full state dict tensors
        param_dtype (torch.dtype): dtype used to move full state dict tensors. If none, respect original dtype from checkpoint
        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

    """
    meta_sd = model.state_dict()
    param_dict = dict(model.named_parameters())

    # map names from checkpoint to customized names
    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
        full_sd_iterator,
        param_names_mapping,
        valid_target_names=set(meta_sd.keys()),
    )  # type: ignore

    is_fsdp_model = isinstance(model, FSDPModule) or any(
        hasattr(p, "device_mesh") for p in meta_sd.values()
    )

    # sort parameter names to ensure all ranks process parameters in the same order
    sorted_param_names = sorted(custom_param_sd.keys())

    sharded_sd = {}
    skipped_checkpoint_keys: list[str] = []
    non_quantized_dtype_mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]] = (
        Counter()
    )
    non_quantized_dtype_mismatch_examples: dict[
        tuple[torch.dtype, torch.dtype], list[str]
    ] = defaultdict(list)
    quantized_dtype_mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]] = (
        Counter()
    )
    quantized_dtype_mismatch_examples: dict[
        tuple[torch.dtype, torch.dtype], list[str]
    ] = defaultdict(list)

    # shard from loaded state_dict, custom_param_sd -> sharded_sd
    for target_param_name in sorted_param_names:
        full_tensor = custom_param_sd[target_param_name]
        meta_sharded_param = meta_sd.get(target_param_name)

        if meta_sharded_param is None:
            # For FSDP models, ensure all ranks process parameters consistently
            if strict or is_fsdp_model:
                raise ValueError(
                    f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
                )
            else:
                skipped_checkpoint_keys.append(target_param_name)
                continue

        # use meta param dtype so quantized params (e.g. FP8) keep their dtype;
        # for non-quantized models meta dtype equals param_dtype anyway
        if meta_sharded_param is None:
            # for nunchaku, some scales are patched later
            target_dtype = full_tensor.dtype
        else:
            target_dtype = meta_sharded_param.dtype

        full_tensor = _maybe_dequantize_fp8(
            full_tensor, target_dtype, target_param_name, custom_param_sd
        )

        if full_tensor.dtype != target_dtype:
            mismatch_key = (full_tensor.dtype, target_dtype)
            if (
                full_tensor.dtype in _QUANTIZED_DTYPES
                or target_dtype in _QUANTIZED_DTYPES
            ):
                quantized_dtype_mismatch_counts[mismatch_key] += 1
                if (
                    len(quantized_dtype_mismatch_examples[mismatch_key])
                    < _DTYPE_MISMATCH_EXAMPLE_LIMIT
                ):
                    quantized_dtype_mismatch_examples[mismatch_key].append(
                        target_param_name
                    )
            else:
                non_quantized_dtype_mismatch_counts[mismatch_key] += 1
                if (
                    len(non_quantized_dtype_mismatch_examples[mismatch_key])
                    < _DTYPE_MISMATCH_EXAMPLE_LIMIT
                ):
                    non_quantized_dtype_mismatch_examples[mismatch_key].append(
                        target_param_name
                    )

        if not hasattr(meta_sharded_param, "device_mesh"):
            full_tensor = full_tensor.to(
                device=checkpoint_load_device, dtype=target_dtype
            )
            actual_param = _get_param_for_weight_loading(
                model, param_dict, target_param_name
            )
            weight_loader = (
                getattr(actual_param, "weight_loader", None)
                if actual_param is not None
                else None
            )
            if weight_loader is not None:
                assert actual_param is not None
                sharded_tensor = torch.empty_like(
                    meta_sharded_param,
                    device=checkpoint_load_device,
                    dtype=target_dtype,
                )
                # Preserve requires_grad flag to avoid errors with non-floating dtypes
                requires_grad = getattr(meta_sharded_param, "requires_grad", False)
                temp_param = _make_param_like(actual_param, sharded_tensor)
                if not (
                    sharded_tensor.is_floating_point() or sharded_tensor.is_complex()
                ):
                    requires_grad = False
                temp_param.requires_grad = requires_grad
                try:
                    weight_loader(temp_param, full_tensor)
                except AssertionError as exc:
                    raise AssertionError(
                        "Failed to shard/load parameter "
                        f"{target_param_name}: full_tensor.shape={tuple(full_tensor.shape)}, "
                        f"meta_sharded_param.shape={tuple(meta_sharded_param.shape)}, "
                        f"temp_param.shape={tuple(temp_param.shape)}, "
                        f"param_cls={type(actual_param).__name__}"
                    ) from exc
                sharded_tensor = temp_param.data
            else:
                # In cases where parts of the model aren't sharded, some parameters will be plain tensors
                sharded_tensor = full_tensor

            # Important: `cpu_offload` is intended for FSDP-managed parameter movement.
            # If a parameter is not sharded into a DTensor (i.e., no `device_mesh`), FSDP
            # will NOT manage it. Offloading it here would leave CPU parameters that
            # later participate in GPU kernels (e.g., conv/embedding), causing device/dtype
            # mismatches like "Input type (CUDABFloat16Type) and weight type (CPUBFloat16Type)".
            #
            # Therefore:
            # - For non-FSDP models, keep the historical behavior (allow CPU offload).
            # - For FSDP models, do NOT offload non-sharded parameters here.
            if cpu_offload and not is_fsdp_model:
                sharded_tensor = sharded_tensor.cpu()
        else:
            full_tensor = full_tensor.to(
                device=checkpoint_load_device, dtype=target_dtype
            )
            actual_param = _get_param_for_weight_loading(
                model, param_dict, target_param_name
            )
            weight_loader = (
                getattr(actual_param, "weight_loader", None)
                if actual_param is not None
                else None
            )
            if weight_loader is not None:
                assert actual_param is not None
                tp_sharded_tensor = torch.empty(
                    tuple(actual_param.shape),
                    device=checkpoint_load_device,
                    dtype=target_dtype,
                )
                temp_param = _make_param_like(actual_param, tp_sharded_tensor)
                if not (
                    tp_sharded_tensor.is_floating_point()
                    or tp_sharded_tensor.is_complex()
                ):
                    temp_param.requires_grad = False
                try:
                    weight_loader(temp_param, full_tensor)
                except AssertionError as exc:
                    raise AssertionError(
                        "Failed to TP-shard/load FSDP parameter "
                        f"{target_param_name}: full_tensor.shape={tuple(full_tensor.shape)}, "
                        f"meta_sharded_param.shape={tuple(meta_sharded_param.shape)}, "
                        f"temp_param.shape={tuple(temp_param.shape)}, "
                        f"param_cls={type(actual_param).__name__}"
                    ) from exc
                full_tensor = temp_param.data
            sharded_tensor = distribute_tensor(
                full_tensor,
                meta_sharded_param.device_mesh,
                meta_sharded_param.placements,
            )
            if cpu_offload:
                sharded_tensor = sharded_tensor.to("cpu")

        actual_param = param_dict.get(target_param_name)
        if actual_param is not None:
            sharded_sd[target_param_name] = _make_param_like(
                actual_param, sharded_tensor
            )
        else:
            sharded_sd[target_param_name] = nn.Parameter(
                sharded_tensor, requires_grad=False
            )

    model.reverse_param_names_mapping = reverse_param_names_mapping

    if non_quantized_dtype_mismatch_counts:
        logger.debug(
            "Casting checkpoint tensors to target dtype during load: %s",
            _format_dtype_mismatch_summary(
                non_quantized_dtype_mismatch_counts,
                non_quantized_dtype_mismatch_examples,
            ),
            main_process_only=True,
            local_main_process_only=True,
        )

    if quantized_dtype_mismatch_counts:
        logger.warning(
            "Dtype mismatches detected for quantized parameters during load: %s",
            _format_dtype_mismatch_summary(
                quantized_dtype_mismatch_counts,
                quantized_dtype_mismatch_examples,
            ),
            main_process_only=True,
            local_main_process_only=True,
        )

    if skipped_checkpoint_keys:
        logger.warning(
            "Checkpoint keys not loaded (no matching model parameter) %s",
            (
                skipped_checkpoint_keys[:20]
                if len(skipped_checkpoint_keys) > 20
                else skipped_checkpoint_keys
            ),
        )
        if len(skipped_checkpoint_keys) > 20:
            logger.warning(
                "... and %d more skipped keys.",
                len(skipped_checkpoint_keys) - 20,
            )

    # parameters in nn.Module that doesn't exist in safetensor files
    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)

    # Legacy allowlist for parameter families synthesized after loading.
    # New formats should declare missing_param_init on the parameter instead.
    LEGACY_ALLOWED_NEW_PARAM_PATTERNS = [
        "gate_compress",
        "wcscales",
        "wtscale",
        "input_scale",
        "weight_scale",
        "bias",
        "norm_q",
        "norm_k",
        "weight_scale",
    ]
    for new_param_name in unused_keys:
        meta_sharded_param = meta_sd.get(new_param_name)
        meta_sharded_param_dtype = meta_sharded_param.dtype
        actual_param = param_dict.get(new_param_name)
        missing_param_init = (
            getattr(actual_param, "missing_param_init", None)
            if actual_param is not None
            else None
        )

        if missing_param_init == "error":
            raise ValueError(
                f"Required checkpoint parameter '{new_param_name}' was not loaded. "
                "This usually indicates a checkpoint/model-arch mismatch or a "
                "broken weight-name mapping."
            )

        if missing_param_init is None and not any(
            pattern in new_param_name for pattern in LEGACY_ALLOWED_NEW_PARAM_PATTERNS
        ):
            logger.error(
                "Unsupported new parameter: %s. Allowed legacy patterns: %s",
                new_param_name,
                LEGACY_ALLOWED_NEW_PARAM_PATTERNS,
            )
            raise ValueError(
                f"New parameter '{new_param_name}' is not supported. "
                "Checkpoint-specific synthesized parameters should either match "
                f"{LEGACY_ALLOWED_NEW_PARAM_PATTERNS} or declare missing_param_init."
            )

        if missing_param_init == "ones" or any(
            p in new_param_name
            for p in (
                "wcscales",
                "wtscale",
                "input_scale",
                "weight_scale",
                "norm_q",
                "norm_k",
            )
        ):
            init_like = torch.ones_like
        elif missing_param_init == "zeros" or missing_param_init is None:
            init_like = torch.zeros_like
        else:
            raise ValueError(
                f"Unsupported missing_param_init={missing_param_init!r} for {new_param_name}"
            )

        if not hasattr(meta_sharded_param, "device_mesh"):
            sharded_tensor = init_like(
                meta_sharded_param,
                device=checkpoint_load_device,
                dtype=meta_sharded_param_dtype,
            )
            if cpu_offload and not is_fsdp_model:
                sharded_tensor = sharded_tensor.cpu()
        else:
            full_tensor = init_like(
                meta_sharded_param,
                device=checkpoint_load_device,
                dtype=meta_sharded_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)