chore: import upstream snapshot with attribution

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

@@ -0,0 +1,769 @@
+# 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)