# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Core helpers for loading model checkpoints."""

from __future__ import annotations

import math
import os
import re
import traceback
from abc import ABC, abstractmethod
from collections import defaultdict
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import contextmanager
from copy import deepcopy
from itertools import chain
from typing import TYPE_CHECKING, Any

import torch

from .distributed.sharding_utils import DtensorShardOperation, _dtensor_from_local_like
from .distributed.utils import is_dtensor
from .integrations.accelerate import get_device, offload_weight
from .utils import is_env_variable_true
from .utils.loading_report import LoadStateDictInfo
from .utils.logging import get_logger, tqdm


if TYPE_CHECKING:
    from .modeling_utils import LoadStateDictConfig, PreTrainedModel
    from .quantizers import HfQuantizer


logger = get_logger(__name__)


def build_glob_alternation(
    globs: list[WeightRenaming | WeightConverter | str],
) -> tuple[re.Pattern, dict[str, str], dict[str, str]]:
    """
    Build a single alternation regex with one named group per glob.
    """
    src_group_to_glob: dict[str, str] = {}
    tgt_group_to_glob: dict[str, str] = {}
    branches: list[str] = []
    i = 0
    for glob in globs:
        if isinstance(glob, (WeightRenaming, WeightConverter)):
            for src in glob.source_patterns:
                group_name = f"g{i}"
                src_group_to_glob[group_name] = src
                i += 1
                body = src.replace("*", r".*")
                branches.append(f"(?P<{group_name}>{body})")
                tgt_group_to_glob[group_name] = glob.target_patterns[0]  # we index with the first target
        else:
            group_name = f"g{i}"
            src_group_to_glob[group_name] = glob
            i += 1
            body = glob
            body = body.replace("*", r".*")
            branches.append(f"(?P<{group_name}>{body})")
            tgt_group_to_glob[group_name] = glob

    alternation = re.compile("|".join(branches))
    return alternation, src_group_to_glob, tgt_group_to_glob


class ConversionOps(ABC):
    """Base class for weight conversion operations."""

    def __repr__(self):
        if hasattr(self, "dim"):
            return f"{self.__class__.__name__}(dim={self.dim})"
        else:
            return f"{self.__class__.__name__}"

    @abstractmethod
    def convert(
        self, input_dict: dict[str, Any], source_patterns: list[str], target_patterns: list[str], **kwargs
    ) -> dict[str, list[torch.Tensor]]:
        raise NotImplementedError

    @property
    def reverse_op(self) -> ConversionOps:
        raise NotImplementedError


class _IdentityOp(ConversionOps):
    """Pass-through reverse op for dequantize operations.

    Dequantized weights are already in their target dtype and should be
    saved as-is without any conversion.
    """

    def convert(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
        return input_dict


class Chunk(ConversionOps):
    """Split a tensor along `dim` into equally sized chunks. Additionally, `num_shards_attribute` is a config field to read
    to know how many tensors to chunk into. Useful when concatenating an arbitrary number of tensors."""

    def __init__(self, dim: int = 0, num_shards_attribute: str | None = None):
        self.dim = dim
        self.num_shards_attribute = num_shards_attribute

    @torch.no_grad
    def convert(
        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
    ) -> dict[str, torch.Tensor]:
        if len(input_dict) > 1:
            raise ValueError("Undefined Operation encountered")
        tensors = next(iter(input_dict.values()))
        tensor = tensors[0] if isinstance(tensors, list) else tensors
        targets = self.get_target_patterns(target_patterns, **kwargs)
        num_shards = len(targets)
        chunks = tuple(chunk.contiguous() for chunk in torch.chunk(tensor, num_shards, dim=self.dim))
        return dict(zip(targets, chunks))

    def get_target_patterns(self, target_patterns: list[str], **kwargs) -> list[str]:
        if self.num_shards_attribute is None:
            return target_patterns
        # In this case we need to use the config to know how many chunks to create
        else:
            if len(target_patterns) > 1:
                raise ValueError("Undefined Operation encountered")
            subconfig = kwargs["config"].get_text_config()
            num_shards = getattr(subconfig, self.num_shards_attribute)
            target_patterns = [target_patterns[0].replace("*", str(i)) for i in range(num_shards)]
            return target_patterns

    @property
    def reverse_op(self) -> ConversionOps:
        return Concatenate(self.dim, self.num_shards_attribute)


class Concatenate(ConversionOps):
    """Concatenate tensors along `dim`. Additionally, if concatenating an aribitrary number of tensors, `num_shards_attribute` is
    a config field to read to know how many tensors to recreate when using the opposite Ops."""

    def __init__(self, dim: int = 0, num_shards_attribute: str | None = None):
        self.dim = dim
        self.num_shards_attribute = num_shards_attribute

    @torch.no_grad
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor]],
        source_patterns: list[str],
        target_patterns: list[str],
        **kwargs,
    ) -> dict[str, torch.Tensor]:
        target_pattern = self.get_target_pattern(target_patterns)
        all_tensors = []
        # Very important to keep the relative order of the source patterns here, so we iterate over them not the
        # input directly as it's unordered! Skip patterns that prior ops in the chain (e.g. `Fp8Dequantize`)
        # have already consumed and dropped from `input_dict`.
        for source_pattern in source_patterns:
            if source_pattern not in input_dict:
                continue
            # Immediately free the input_dict, so that we do not keep many copies simultaneously - otherwise we have to
            # wait for this function to return to be able to clean-up, which will not get garbage collected as fast as if
            # everything is freed right now
            tensors = input_dict.pop(source_pattern)
            if isinstance(tensors, list):
                all_tensors.extend(tensors)
            else:
                all_tensors.append(tensors)
        return {target_pattern: torch.cat(all_tensors, dim=self.dim)}

    def get_target_pattern(self, target_patterns: list[str]) -> str:
        # Here we always return the target pattern
        if len(target_patterns) > 1:
            raise ValueError("Undefined Operation encountered!")
        return target_patterns[0]

    @property
    def reverse_op(self) -> ConversionOps:
        return Chunk(self.dim, self.num_shards_attribute)


class Interleave(ConversionOps):
    """Deinterleaves a tensor along `dim` by splitting in two and transposing. Reshapes param back to its original size."""

    def __init__(self, dim: int = 0, inverse: bool = False):
        self.dim = dim
        self.inverse = inverse

    def convert(self, input_dict, source_patterns, target_patterns, **kwargs):
        tensor = next(iter(input_dict.values()))
        tensor = tensor[0] if isinstance(tensor, list) else tensor

        # Split into two in given dim and transpose to interleave along it
        shape = list(tensor.shape)
        if self.inverse:
            shape[self.dim : self.dim + 1] = [2, shape[self.dim] // 2]
        else:
            shape[self.dim : self.dim + 1] = [shape[self.dim] // 2, 2]

        tensor = tensor.reshape(shape).transpose(self.dim, self.dim + 1).reshape(tensor.shape).contiguous()
        return {target_patterns[0]: tensor}

    @property
    def reverse_op(self) -> ConversionOps:
        # can use the same dim and it will inverse it back
        return Interleave(self.dim, inverse=not self.inverse)


class MergeModulelist(ConversionOps):
    """
    Merge a list of tensors into a single tensor along the first dimension.
    We explicitly define this because for EP or TP you want to make sure you know what you are doing!

    """

    def __init__(self, dim: int = 0):
        self.dim = dim

    @torch.no_grad
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor]],
        source_patterns: list[str],
        target_patterns: list[str],
        **kwargs,
    ) -> dict[str, torch.Tensor]:
        input_size = len(input_dict)
        merged: dict[str, torch.Tensor] = {}
        for source_pattern in list(input_dict.keys()):
            # Immediately free the input dict, so that we do not keep many copies simultaneously - otherwise we have to
            # wait for this function to return to be able to clean-up, and if the size of the input_dict is larger than 1
            # (such as the MoEs' gate_proj/up_proj merging), we are wasting quite some memory
            tensors = input_dict.pop(source_pattern)
            target_pattern = self.get_target_pattern(input_size, source_pattern, target_patterns)
            # DecompressExperts pre-allocates a stacked tensor to avoid holding N individual
            # decompressed tensors simultaneously.  Pass it through to skip the redundant copy
            # that torch.stack would otherwise make.
            if isinstance(tensors, torch.Tensor):
                merged[target_pattern] = tensors
            else:
                merged[target_pattern] = torch.stack(tensors, dim=self.dim)
        return merged

    def get_target_pattern(self, input_size: int, source_pattern: str, target_patterns: list[str]) -> str:
        # Here it's a single operation, so we use the target
        if input_size == 1:
            if len(target_patterns) == 1:
                return target_patterns[0]
            else:
                raise ValueError("Undefined Operation encountered!")
        #  Here it's the first operation in a chain, so we use the source as they were replaced before in the chain
        else:
            return source_pattern

    @property
    def reverse_op(self) -> ConversionOps:
        return SplitModulelist(self.dim)


class SplitModulelist(ConversionOps):
    """Inverse of `MergeModulelist` using explicit split sizes per group."""

    def __init__(self, dim: int = 0):
        self.dim = dim

    @torch.no_grad
    def convert(
        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
    ) -> dict[str, torch.Tensor]:
        all_tensors = {}
        for source_pattern, tensors in input_dict.items():
            tensor = tensors[0] if isinstance(tensors, list) else tensors
            # We split in the number of tensors present in the given dim
            sizes = tensor.size(self.dim)
            targets = self.get_target_patterns(input_dict, source_pattern, target_patterns, sizes)
            chunks = torch.chunk(tensor, sizes, dim=self.dim)
            # We squeeze each chunk here as well to make sure to give them their original shape
            all_tensors.update({target: chunk.squeeze() for target, chunk in zip(targets, chunks)})
        return all_tensors

    def get_target_patterns(
        self, input_dict: dict, source_pattern: str, target_patterns: list[str], sizes: int
    ) -> list[str]:
        # Here it's a single operation, so we use the target
        if len(input_dict) == 1:
            if len(target_patterns) == 1:
                return [target_patterns[0].replace("*", f"{i}") for i in range(sizes)]
            else:
                raise ValueError("Undefined Operation encountered!")
        # Here it's the last operation in a chain, so we use the source as they were replaced before in the chain
        else:
            return [source_pattern.replace("*", f"{i}") for i in range(sizes)]

    @property
    def reverse_op(self) -> ConversionOps:
        return MergeModulelist(self.dim)


class Transpose(ConversionOps):
    """
    Transposes the given tensor along dim0 and dim1.
    """

    def __init__(self, dim0: int = 0, dim1: int = 1, check_dims: bool = False):
        self.dim0 = dim0
        self.dim1 = dim1
        self.check_dims = check_dims

    @torch.no_grad
    def convert(
        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
    ) -> dict[str, torch.Tensor]:
        target_pattern = self.get_target_pattern(input_dict, source_patterns, target_patterns)
        tensors = next(iter(input_dict.values()))
        tensor = tensors[0] if isinstance(tensors, list) else tensors
        # In this case, always transpose
        if not self.check_dims:
            return {target_pattern: torch.transpose(tensor, dim0=self.dim0, dim1=self.dim1).contiguous()}
        # In this case, check the shapes before transposing
        else:
            # NOTE: this rely on the first param name, so cannot be used for many-to-one operation
            expected_shape = kwargs["model"].get_parameter(kwargs["full_layer_name"]).shape
            # The shapes are the same: do NOT transpose
            if tensor.shape == expected_shape:
                return {target_pattern: tensor}
            else:
                return {target_pattern: torch.transpose(tensor, dim0=self.dim0, dim1=self.dim1).contiguous()}

    def get_target_pattern(
        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str]
    ) -> str:
        if len(input_dict) != 1:
            raise ValueError("Undefined Operation encountered!")
        # Here it's the first operation of a chain, so return the source
        if len(target_patterns) > 1:
            if len(source_patterns) == 1:
                return source_patterns[0]
            else:
                raise ValueError("Undefined Operation encountered!")
        # Here it's the only operation, or the last operation in a chain, so we return the target
        else:
            return target_patterns[0]

    @property
    def reverse_op(self) -> ConversionOps:
        return Transpose(dim0=self.dim1, dim1=self.dim0, check_dims=self.check_dims)


class Conv3dToLinear(ConversionOps):
    """Conv3d weights → flattened Linear layout."""

    def __init__(self, in_channels: int, kernel_size: tuple[int, int, int]):
        self.in_channels = in_channels
        self.kernel_size = kernel_size

    @staticmethod
    def _get_target_pattern(
        input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str]
    ) -> str:
        if len(input_dict) != 1:
            raise ValueError("Undefined Operation encountered!")
        if len(target_patterns) > 1:
            if len(source_patterns) == 1:
                return source_patterns[0]
            else:
                raise ValueError("Undefined Operation encountered!")
        return target_patterns[0]

    @torch.no_grad
    def convert(
        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
    ) -> dict[str, torch.Tensor]:
        target_pattern = self._get_target_pattern(input_dict, source_patterns, target_patterns)
        tensors = next(iter(input_dict.values()))
        tensor = tensors[0] if isinstance(tensors, list) else tensors

        if tensor.ndim == 5:
            tensor = tensor.reshape(tensor.shape[0], -1).contiguous()
        elif tensor.ndim != 2:
            raise ValueError(f"Conv3dToLinear expects a 5D or 2D tensor, got {tensor.ndim}D")

        return {target_pattern: tensor}

    @property
    def reverse_op(self) -> ConversionOps:
        return LinearToConv3d(in_channels=self.in_channels, kernel_size=self.kernel_size)


class LinearToConv3d(ConversionOps):
    """Flattened Linear weights → Conv3d layout."""

    def __init__(self, in_channels: int, kernel_size: tuple[int, int, int]):
        self.in_channels = in_channels
        self.kernel_size = kernel_size

    @torch.no_grad
    def convert(
        self, input_dict: dict[str, torch.Tensor], source_patterns: list[str], target_patterns: list[str], **kwargs
    ) -> dict[str, torch.Tensor]:
        target_pattern = Conv3dToLinear._get_target_pattern(input_dict, source_patterns, target_patterns)
        tensors = next(iter(input_dict.values()))
        tensor = tensors[0] if isinstance(tensors, list) else tensors

        target_shape = (tensor.shape[0], self.in_channels, *self.kernel_size)
        if tensor.numel() != math.prod(target_shape):
            raise ValueError(f"Cannot reshape tensor with shape {tensor.shape} into {target_shape}")

        return {target_pattern: tensor.reshape(target_shape).contiguous()}

    @property
    def reverse_op(self) -> ConversionOps:
        return Conv3dToLinear(in_channels=self.in_channels, kernel_size=self.kernel_size)


class PermuteForRope(ConversionOps):
    """
    Applies the permutation required to convert complex RoPE weights to the split sin/cos format.
    """

    def __init__(
        self, subconfig_key: str | None = None, permute_layer_names: list[str] | None = None, inverse: bool = False
    ):
        self.subconfig_key = subconfig_key
        self.inverse = inverse
        self.permute_layer_names = permute_layer_names

    def _apply(self, tensor: torch.Tensor) -> torch.Tensor:
        dim0 = tensor.shape[0]
        config = self.config
        if self.subconfig_key is not None:
            config = getattr(self.config, self.subconfig_key, self.config)
        n_heads = getattr(config, "num_attention_heads", 1)
        half_head = dim0 // n_heads // 2

        head_shape = (2, half_head) if self.inverse else (half_head, 2)
        if tensor.ndim == 2:
            tensor = tensor.view(n_heads, *head_shape, tensor.shape[1])
            tensor = tensor.transpose(1, 2).reshape(dim0, tensor.shape[-1])
        elif tensor.ndim == 1:
            tensor = tensor.view(n_heads, *head_shape)
            tensor = tensor.transpose(1, 2).reshape(dim0)
        return tensor

    @torch.no_grad
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor] | torch.Tensor],
        source_patterns: list[str],
        target_patterns: list[str],
        config,
        **kwargs,
    ) -> dict[str, list[torch.Tensor]]:
        self.config = config
        output: dict[str, list[torch.Tensor]] = {}
        for key, tensors in input_dict.items():
            # Permute q and key weights back (skip biases) to match original RoPE implementation
            if not any(name in key for name in self.permute_layer_names):
                output[key] = tensors
                continue

            if isinstance(tensors, list):
                if len(tensors) != 1:
                    raise ValueError("PermuteForRope expects a single tensor per key.")
                tensors = tensors[0]
            output[key] = self._apply(tensors)
        return output

    @property
    def reverse_op(self) -> ConversionOps:
        return PermuteForRope(
            subconfig_key=self.subconfig_key, permute_layer_names=self.permute_layer_names, inverse=not self.inverse
        )


class VisionFuseAndPermuteForRope(ConversionOps):
    """
    Applies the permutation required to convert complex RoPE weights to the split sin/cos format on fused QKV.
    Same as calling `PermuteForRope() + Concatenate()` but lets us call `Permute` only on a subset of chunked tensors.

    NOTE: this conversion applies only to a vision backbone in multimodal models, because it checks `config.vision_config`
    """

    def __init__(self, dim: int = 0, permute_layer_names: list[str] | None = None, inverse: bool = True):
        logger.warning(
            "`VisionUnfuseAndPermuteForRope` is deprecated and will be removed in v5.20. Use `PermuteForRope()` and `Concatenate()` "
            "consecutively instead to permute back and fuse projections."
        )
        self.dim = dim
        self.permute_layer_names = permute_layer_names
        self.inverse = inverse
        self.concat_op = Concatenate(dim=dim)
        self.permute_op = PermuteForRope(
            subconfig_key="vision_config", permute_layer_names=permute_layer_names, inverse=inverse
        )

    @torch.no_grad
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor]],
        source_patterns: list[str],
        target_patterns: list[str],
        config,
        **kwargs,
    ) -> dict[str, torch.Tensor]:
        input_dict = self.permute_op.convert(
            input_dict=input_dict,
            source_patterns=source_patterns,
            target_patterns=target_patterns,
            config=config,
        )
        input_dict = self.concat_op.convert(
            input_dict=input_dict,
            source_patterns=source_patterns,
            target_patterns=target_patterns,
        )
        return input_dict

    def get_target_pattern(self, target_patterns: list[str]) -> str:
        # Here we always return the target pattern
        if len(target_patterns) > 1:
            raise ValueError("Undefined Operation encountered!")
        return target_patterns[0]

    @property
    def reverse_op(self) -> ConversionOps:
        return VisionUnfuseAndPermuteForRope(self.dim, self.permute_layer_names, inverse=not self.inverse)


class VisionUnfuseAndPermuteForRope(ConversionOps):
    """
    Applies the permutation required to convert complex RoPE weights to the split sin/cos format on fused QKV.
    Same as calling `Chunk() + PermuteForRope()` but lets us call `Permute` only on a subset of chunked tensors.

    NOTE: this conversion applies only to a vision backbone in multimodal models, because it checks `config.vision_config`
    """

    def __init__(self, dim: int = 0, permute_layer_names: list[str] | None = None, inverse: bool = False):
        logger.warning(
            "`VisionUnfuseAndPermuteForRope` is deprecated and will be removed in v5.20. Use `Chunk()` and `PermuteForRope()` "
            "consecutively instead to unfuse projections and permute for block-split RoPE"
        )
        self.dim = dim
        self.permute_layer_names = permute_layer_names
        self.inverse = inverse
        self.chunk_op = Chunk(dim=dim)
        self.permute_op = PermuteForRope(
            subconfig_key="vision_config", permute_layer_names=permute_layer_names, inverse=inverse
        )

    @torch.no_grad
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor]],
        source_patterns: list[str],
        target_patterns: list[str],
        config,
        **kwargs,
    ) -> dict[str, torch.Tensor]:
        input_dict = self.chunk_op.convert(
            input_dict=input_dict,
            source_patterns=source_patterns,
            target_patterns=target_patterns,
        )
        input_dict = self.permute_op.convert(
            input_dict=input_dict,
            source_patterns=source_patterns,
            target_patterns=target_patterns,
            config=config,
        )
        return input_dict

    def get_target_patterns(self, input_dict: dict, target_patterns: list[str]) -> list[str]:
        # Here we always return the target patterns
        if len(input_dict) > 1 or len(target_patterns) == 1:
            raise ValueError("Undefined Operation encountered!")
        return target_patterns

    @property
    def reverse_op(self) -> ConversionOps:
        return VisionFuseAndPermuteForRope(self.dim, self.permute_layer_names, inverse=not self.inverse)


class ErnieFuseAndSplitTextVisionExperts(ConversionOps):
    r"""
    Special operation that splits a module list over all keys and fuses over the number of original modules.

    Example with 2 original modules "Gate" and "Up" with 2 target keys "Text" and "Vision":

                 ModuleList 1            ModuleList 2
                [   Gate    ]            [   Up    ]
                |           |            |         |
          [Gate_Text] [Gate_Vision] [Up_Text]   [Up_Vision]
              \                  \  /                /
               \                 \ /                /
                \               /  \               /
                 \            /     \             /
                 [GateUp_Text]      [GateUp_Vision]

    The splits are equal and are defined by the amount of target keys.
    The final fusions are defined by the amount of original module lists.
    """

    def __init__(self, stack_dim: int = 0, concat_dim: int = 1):
        self.stack_dim = stack_dim
        self.concat_dim = concat_dim

    def split_list_into_chunks(self, tensor_list: list[torch.Tensor], chunks: int = 2):
        split_size = math.ceil(len(tensor_list) / chunks)  # best effort split size
        return [tensor_list[i * split_size : (i + 1) * split_size] for i in range(chunks)]

    @torch.no_grad()
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor]],
        source_patterns: list[str],
        target_patterns: list[str],
        config,
        **kwargs,
    ) -> dict[str, list[torch.Tensor]]:
        valid_keys = input_dict.keys()
        split_and_fused = defaultdict(list)
        for key in source_patterns:
            if key not in valid_keys:
                raise ValueError(
                    f"Expected pattern {key} in collected tensors but only found tensors for: {valid_keys}"
                )

            tensors = input_dict.get(key, [])
            split_tensor_lists = self.split_list_into_chunks(tensors, chunks=len(target_patterns))
            stacked_tensors = (torch.stack(tensor_group, dim=self.stack_dim) for tensor_group in split_tensor_lists)
            for idx, tensor_group in enumerate(stacked_tensors):
                split_and_fused[target_patterns[idx]].append(tensor_group)

        for k, v in split_and_fused.items():
            split_and_fused[k] = torch.cat(v, dim=self.concat_dim)

        return split_and_fused

    @property
    def reverse_op(self) -> ConversionOps:
        return ErnieSplitAndDecoupleTextVisionExperts(stack_dim=self.stack_dim, concat_dim=self.concat_dim)


class ErnieSplitAndDecoupleTextVisionExperts(ConversionOps):
    r"""
    Special operation that splits a fused module list over all original modules and
    then decouples them into a mixed module list each over all keys.

    Example with 2 original modules "Gate" and "Up" with 2 target keys "Text" and "Vision":

                    [GateUp_Text]     [GateUp_Vision]
                  /              \   /             \
                 /                \ /               \
                /                / \                 \
               /                /   \                 \
          [Gate_Text] [Gate_Vision] [Up_Text]   [Up_Vision]
                |           |            |         |
                [   Gate    ]            [   Up    ]
                 ModuleList 1            ModuleList 2

    The splits are equal and are defined by the amount of original module lists.
    The final decoupled module lists are defined by the amount of keys.
    """

    def __init__(self, stack_dim: int = 0, concat_dim: int = 1):
        self.stack_dim = stack_dim
        self.concat_dim = concat_dim

    @torch.no_grad()
    def convert(
        self,
        input_dict: dict[str, list[torch.Tensor]],
        source_patterns: list[str],
        target_patterns: list[str],
        config,
        **kwargs,
    ) -> dict[str, list[torch.Tensor]]:
        fused_modules = len(target_patterns)
        valid_keys = input_dict.keys()
        split_tensors = []
        for key in source_patterns:
            if key not in valid_keys:
                raise ValueError(
                    f"Expected pattern {key} in collected tensors but only found tensors for: {valid_keys}"
                )

            # Assuming that we get single sized lists here to index with 0
            split_tensors.append(input_dict[key][0].chunk(fused_modules, dim=self.concat_dim))

        decoupled = {}
        for idx, key in enumerate(target_patterns):
            tensor_groups = [
                list(torch.unbind(tensor_group[idx], dim=self.stack_dim)) for tensor_group in split_tensors
            ]
            tensor_list = list(chain.from_iterable(tensor_groups))
            targets = [key.replace("*", f"{i}") for i in range(len(tensor_list))]
            decoupled |= dict(zip(targets, tensor_list))

        return decoupled

    @property
    def reverse_op(self) -> ConversionOps:
        return ErnieFuseAndSplitTextVisionExperts(stack_dim=self.stack_dim, concat_dim=self.concat_dim)


def process_target_pattern(pattern: str) -> tuple[str, str | None]:
    """
    Process a target pattern for reverse mapping (when targets become sources).

    This handles several edge cases in checkpoint conversion mappings:
    - Removes `^` prefix and `$` suffix (start/end of string anchors)
    - Removes negative lookahead/lookbehind assertions
    - Detects capturing groups and replaces them with `\\1` backreference

    Args:
        pattern: The target pattern to process for reverse mapping.

    Returns:
        A tuple of (processed_pattern, captured_group) where captured_group is
        the original capturing group found (e.g., "(encoder|decoder)") or None.
    """
    # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping
    pattern = pattern.removeprefix("^")
    # Some mapping contains `$` to notify end of string when matching -> remove it during reverse mapping
    pattern = pattern.removesuffix("$")
    # Remove negative lookahead/behind if any. This is ugly but needed for reverse mapping of
    # Qwen2.5, Sam3, Ernie4.5 VL MoE! It needs to be non greedy in case there are several
    pattern = re.sub(r"\(\?.+?\)?\)", "", pattern)
    # Remove the backslash for literal dots
    pattern = pattern.replace(r"\.", ".")
    # Allow capturing groups in patterns, i.e. to add/remove a prefix to all keys (e.g. timm_wrapper, sam3)
    capturing_group_match = re.search(r"\(.+?\)", pattern)
    captured_group = None
    if capturing_group_match:
        captured_group = capturing_group_match.group(0)
        pattern = pattern.replace(captured_group, r"\1", 1)
    return pattern, captured_group


def process_source_pattern(source_pattern: str, target_pattern: str) -> str:
    """
    Process a source pattern for reverse mapping (when sources become targets).
    This is useful because usually if the original source (so now the target in reverse mode) had a `^` or `$`
    to restrict to start/end of string, we should do the same in reverse mode. This is why this method is conditioned
    on the target pattern, we want to do it only for pairs (source, target) when the original source (so the current target
    in reverse mode) had it.
    """
    if target_pattern.startswith("^"):
        source_pattern = f"^{source_pattern}" if not source_pattern.startswith("^") else source_pattern
    if target_pattern.endswith("$"):
        source_pattern = f"{source_pattern}$" if not source_pattern.endswith("$") else source_pattern

    return source_pattern


class WeightTransform:
    # Restrict the attributes that can be attached
    __slots__ = (
        "source_patterns",
        "target_patterns",
        "compiled_sources",
        "distributed_operation",
        "quantization_operation",
        "collected_tensors",
        "layer_targets",
        "_original_source_patterns",
        "_original_target_patterns",
        "_was_used",
        "scope_prefix",
        "base_model_prefix",
    )

    def __init__(self, source_patterns: str | list[str], target_patterns: str | list[str]):
        self.source_patterns: list[str] = source_patterns
        self.target_patterns: list[str] = target_patterns
        # Those are needed to be able to reverse correctly the transform, as the patterns may be processed
        self._original_source_patterns = self.source_patterns.copy()
        self._original_target_patterns = self.target_patterns.copy()

        # Init fields that will be used during conversion
        self.distributed_operation: Any = None
        self.quantization_operation: ConversionOps | None = None
        self.collected_tensors: dict[str, list[Future]] = defaultdict(list)
        self.layer_targets: dict[str, set[str]] = defaultdict(set)

        # Flag to notice if the Transform was used
        self._was_used = False

        # Optional scope_prefix/base_model_prefix. When used, the transform will only match and apply to keys containing
        # either `base_model_prefix.scope_prefix.` or `scope_prefix.` prefixes
        self.scope_prefix: str | None = None
        self.base_model_prefix: str | None = None

        # We need to process a few exceptions here when instantiating the reverse mapping (i.e. the targets become
        # sources, and sources become targets). The issues lie in the sources usually, so here we need to check the
        # targets for the reversed mapping

        # Process target_patterns: detect capturing groups and replace with \1
        # Store the original capturing group patterns for reverse mapping
        target_capturing_groups: list[str] = []
        for i, pattern in enumerate(self.target_patterns):
            self.target_patterns[i], captured_group = process_target_pattern(pattern)
            if captured_group is not None:
                target_capturing_groups.append(captured_group)

        # Validate that we only have one unique capturing group pattern across all targets
        # This ensures deterministic reverse mapping when sources have \1 backreferences
        unique_capturing_groups = set(target_capturing_groups)
        if len(unique_capturing_groups) > 1:
            raise ValueError(
                f"Multiple different capturing groups found in target_patterns: {unique_capturing_groups}. "
                f"All target patterns must use the same capturing group pattern."
            )
        unique_capturing_group = unique_capturing_groups.pop() if unique_capturing_groups else None

        # We also need to check capturing groups in the sources during reverse mapping (e.g. timm_wrapper, sam3)
        for i, pattern in enumerate(self.source_patterns):
            # Replace capturing groups
            if r"\1" in pattern:
                if unique_capturing_group is None:
                    raise ValueError(
                        f"Source pattern '{pattern}' contains \\1 backreference, but no capturing groups "
                        f"found in target_patterns."
                    )
                # Use the unique capturing group from target_patterns for all sources
                pattern = pattern.replace(r"\1", unique_capturing_group, 1)
            # Potentially process a bit more for consistency - only if they are consistent pairs, i.e. the length is the same
            if len(self.source_patterns) == len(self.target_patterns):
                pattern = process_source_pattern(pattern, self._original_target_patterns[i])
            self.source_patterns[i] = pattern

        # Construct the regex we will use to rename keys from the sources to the targets
        branches = []
        for i, source_pattern in enumerate(self.source_patterns):
            group_name = f"g{i}"
            pattern = source_pattern.replace("*.", r".*\.")
            branches.append(f"(?P<{group_name}>{pattern})")
        self.compiled_sources = re.compile("|".join(branches))

    def __repr__(self):
        return f"{self.__class__.__name__}(source_patterns={self.source_patterns}, target_patterns={self.target_patterns})"

    def __setattr__(self, name, value):
        if name in ("source_patterns", "target_patterns"):
            # We do not allow to re-set the patterns, as they are linked between each other and changing one
            # without the other can mess-up with the capturing groups/compiled sources
            if hasattr(self, name):
                raise ValueError(f"Cannot assign to field {name}, you should create a new instance")
            # Switch str to list
            elif isinstance(value, str):
                value = [value]
        object.__setattr__(self, name, value)

    def add_tensor(self, target_key: str, source_key: str, source_pattern: str, future: Future):
        self.collected_tensors[source_pattern].append(future)
        self.layer_targets[target_key].add(source_key)

    def _scoped_match(self, source_key: str) -> tuple[str | None, str, re.Match[str]] | None:
        """
        Strip `scope_prefix` (if any) from `source_key`, then match `compiled_sources` against the
        remaining suffix.

        Returns `(prefix_dot, key_to_match, match_object)` on match, else `None`. `prefix_dot` is
        the prefix consumed from `source_key`: either `f"{scope_prefix}."` or that same string with
        one `base_model_prefix` level stripped or prepended when the former didn't match.
        `None` when `scope_prefix` is unset.
        """
        key_to_match = source_key
        prefix = None
        if self.scope_prefix is not None:
            scope_prefix = f"{self.scope_prefix}." if self.scope_prefix != "" else ""
            base_model_prefix = f"{self.base_model_prefix}." if self.base_model_prefix != "" else ""
            # First, try to match the longest sequence, i.e. base_model_prefix + scope_prefix
            if source_key.startswith(base_model_prefix + scope_prefix):
                prefix = base_model_prefix + scope_prefix
            # Then, try to strip the base_model_prefix, in case we load a ForXXX model from BaseModel weights
            elif source_key.startswith(scope_prefix):
                prefix = scope_prefix
            # In this case, no match is ever possible
            else:
                return None
            key_to_match = source_key.removeprefix(prefix)

        match_object = self.compiled_sources.search(key_to_match)
        if match_object is None:
            return None
        return (prefix, key_to_match, match_object)

    def rename_source_key(self, source_key: str) -> tuple[str, str | None]:
        """
        Return a tuple (renamed_key, source_pattern_producing_the_match).
        Try renaming `source_key` according to the source and target patterns of the current WeightTransform.
        In case of a one-to-many transform, i.e. we have several target patterns, the matching source pattern
        will be replaced by the first of all the target patterns (they are then correctly expanded in the Operations).
        """
        matched = self._scoped_match(source_key)
        if matched is None:
            return source_key, None

        prefix_dot, key_to_match, match_object = matched

        # We have a match, so the Transform was used
        self._was_used = True

        # Find the source that produced the match (it's the first group that matched, as the search stops after first branch match)
        matching_group_name = next(name for name, val in match_object.groupdict().items() if val is not None)
        source_pattern_that_matched = self.source_patterns[int(matching_group_name[1:])]
        # If we matched, we always replace with the first target pattern, in case we have several (one to many transform)
        replacement = self.target_patterns[0]
        # Allow capturing groups in patterns, i.e. to add a prefix to all keys (e.g. timm_wrapper, sam3)
        if r"\1" in replacement:
            # The index of the internal group we need to replace is the index of the matched named group as it comes
            # inside that matched named group
            replaced_group_idx = self.compiled_sources.groupindex[matching_group_name] + 1
            replacement = replacement.replace(r"\1", match_object.group(replaced_group_idx))
        renamed_key = key_to_match.replace(match_object.group(0), replacement, 1)
        if prefix_dot is not None:
            renamed_key = prefix_dot + renamed_key
        return renamed_key, source_pattern_that_matched

    def reverse_transform(self) -> WeightTransform:
        """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations."""
        # TODO: check this and relax when quantizer have `reverse_op`
        if self.quantization_operation is not None:
            raise ValueError("Cannot reverse the transform with TP or quantization")

        kwargs = {}
        # Add the reverse ops if applicable (it needs to be provided at __init__)
        if hasattr(self, "operations"):
            # All reverse ops, in reverse order
            kwargs["operations"] = [op.reverse_op for op in self.operations[::-1]]

        reverse_transform = self.__class__(
            source_patterns=self._original_target_patterns, target_patterns=self._original_source_patterns, **kwargs
        )
        reverse_transform.scope_prefix = self.scope_prefix
        reverse_transform.base_model_prefix = self.base_model_prefix
        return reverse_transform

    def materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
        """
        Materialize all the tensors that were saved in `self.collected_tensors`. This function removes them from the
        internal attribute to avoid keeping them in memory during the different `self.convert` operations, and return
        a new dictionary (otherwise we use more memory than needed during loading).

        We basically have 3 cases here:
        - async loading (default): the tensors are Future instances that we need to wait for
        - sync loading: the tensors are Callable, we need to call the Callable to actually load them from disk
        - saving: the tensors are already torch.Tensor instances (the existing model weights)
        """
        collected_tensors = {}
        for key in list(self.collected_tensors.keys()):
            # Remove from internal attribute
            tensors = self.collected_tensors.pop(key)
            # Async loading
            if isinstance(tensors[0], Future):
                tensors = [future.result() for future in tensors if future.result() is not None]
            # Sync loading
            elif callable(tensors[0]):
                tensors = [func() for func in tensors]
                # Some may be None for some distributed setups
                tensors = [tensor for tensor in tensors if tensor is not None]
            # Add them to the new dictionary
            collected_tensors[key] = tensors

        return collected_tensors

    def was_used(self) -> bool:
        """
        Return whether the current Transform matched any weights during loading/saving. This is needed as some
        weight renaming transforms are not bijective, i.e. if we drop/add full parts of a name with PrefixChange, we
        lose some information that we cannot get back if we don't know if the Transform was used before already (say we
        have a prefix to drop, we need to know whether the checkpoints we loaded before contained the said prefix or not
        before adding it back, or not, during saving).
        """
        return self._was_used


class WeightRenaming(WeightTransform):
    # Special case of WeightTransform that only renames keys without any conversion.

    # Needs to be empty, otherwise the class will not be slotted
    __slots__ = ()

    def convert(
        self,
        layer_name: str,
        model=None,
        config=None,
        hf_quantizer=None,
        loading_info: LoadStateDictInfo | None = None,
    ):
        # Collect the tensors here - we use a new dictionary to avoid keeping them in memory in the internal
        # attribute during the whole process
        collected_tensors = self.materialize_tensors()

        # Perform renaming op (for a simple WeightRenaming, `self.source_patterns` and `self.target_patterns` can
        # only be of length 1, and are actually the full key names - we also have only 1 single related tensor)
        target_key = self.target_patterns[0]
        collected_tensors = {target_key: collected_tensors[self.source_patterns[0]]}

        if hf_quantizer is not None and self.quantization_operation is not None:
            with log_conversion_errors(
                layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation
            ):
                collected_tensors = self.quantization_operation.convert(
                    collected_tensors,
                    source_patterns=self.source_patterns,
                    target_patterns=self.target_patterns,
                    full_layer_name=target_key,
                    model=model,
                    config=config,
                    missing_keys=loading_info.missing_keys if loading_info else None,
                )

        return collected_tensors


class GroupWeightRename(WeightRenaming):
    """
    Applies a list of paired WeightRenaming transforms, activated lazily by the first ("guard")
    source pattern.  Use this when two renames share an intermediate name (e.g. `norm0→norm1`
    and `norm1→norm2`) so that loading an already-converted checkpoint (which has `norm1`
    and `norm2` but no `norm0`) does not incorrectly re-apply the renames.

    NOTE: order `source_patterns` so that the one that is absent in an already-converted checkpoint
    comes first.  The state dict is iterated in sorted key order, so the guard pattern must be
    lexicographically smaller than the dependent patterns. Otherwise the dependents will be
    skipped on the first pass and never retried.
    """

    __slots__ = ("_active",)

    def __init__(self, source_patterns: list[str], target_patterns: list[str]):
        if len(source_patterns) != len(target_patterns):
            raise ValueError(
                "GroupWeightRename requires N:N length matching, but found "
                f"len(source_patterns)={len(source_patterns)} != len(target_patterns)={len(target_patterns)}"
            )
        super().__init__(source_patterns=source_patterns, target_patterns=target_patterns)
        self._active = None  # None = undecided; True = guard was seen

    def rename_source_key(self, source_key: str) -> tuple[str, str | None]:
        matched = self._scoped_match(source_key)
        if matched is None:
            return source_key, None

        prefix_dot, key_to_match, match_object = matched
        matching_group_name = next(name for name, val in match_object.groupdict().items() if val is not None)
        group_index = int(matching_group_name[1:])

        if group_index == 0:
            # Guard pattern matched — activate the group for subsequent keys
            self._active = True
        elif not self._active:
            # Dependent pattern matched before the guard was ever seen → skip
            return source_key, None

        self._was_used = True
        replacement = self.target_patterns[group_index]
        if re.search(r"\\\d", replacement):
            group_start = self.compiled_sources.groupindex[matching_group_name]
            replacement = re.sub(
                r"\\(\d+)",
                lambda m: match_object.group(group_start + int(m.group(1))),
                replacement,
            )
        renamed_key = key_to_match.replace(match_object.group(0), replacement, 1)
        if prefix_dot is not None:
            renamed_key = prefix_dot + renamed_key
        return renamed_key, self.source_patterns[group_index]

    def reverse_transform(self) -> GroupWeightRename:
        # Swap source↔target then sort pairs by (new) source pattern so the guard
        # (lex-smallest) stays first, preserving the ordering contract.
        pairs = sorted(zip(self._original_target_patterns, self._original_source_patterns))
        reversed_srcs = [p[0] for p in pairs]
        reversed_tgts = [p[1] for p in pairs]
        result = GroupWeightRename(source_patterns=reversed_srcs, target_patterns=reversed_tgts)
        result.scope_prefix = self.scope_prefix
        result.base_model_prefix = self.base_model_prefix
        return result


class PrefixChange(WeightRenaming):
    """
    Special case of WeightRenaming, used to simplify adding/removing full parts of a weight name. The regexes
    that are needed for such operations are complex, so this is a much easier API for such cases.
    """

    __slots__ = (
        "prefix_to_add",
        "prefix_to_remove",
        "model_prefix",
    )

    def __init__(
        self, prefix_to_add: str | None = None, prefix_to_remove: str | None = None, model_prefix: str | None = None
    ):
        if (prefix_to_add is None) ^ (prefix_to_remove is not None):
            raise ValueError("You must provide only one of `prefix_to_add` and `prefix_to_remove`")

        self.prefix_to_add = prefix_to_add
        self.prefix_to_remove = prefix_to_remove
        self.model_prefix = "" if model_prefix is None else model_prefix
        prefix = rf"{self.model_prefix}\." if self.model_prefix != "" else ""

        if prefix_to_add is not None:
            super().__init__(
                # We use a lookbehind to avoid adding the prefix if we detect that it's already present
                source_patterns=rf"^{prefix}(?:(?!{prefix_to_add}\.))(.+)$",
                target_patterns=rf"{prefix}{prefix_to_add}\.\1",
            )
        else:
            super().__init__(source_patterns=rf"^{prefix}{prefix_to_remove}\.(.+)$", target_patterns=rf"{prefix}\1")

    def reverse_transform(self) -> WeightTransform:
        """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations."""
        # TODO: check this and relax when quantizer have `reverse_op`
        if self.quantization_operation is not None:
            raise ValueError("Cannot reverse the transform with TP or quantization")

        # Only one of the 2 can ever be used, so 1 is always None
        result = PrefixChange(
            prefix_to_add=self.prefix_to_remove, prefix_to_remove=self.prefix_to_add, model_prefix=self.model_prefix
        )
        result.scope_prefix = self.scope_prefix
        result.base_model_prefix = self.base_model_prefix
        return result


# List of classes that are known to be able to use m:n
_INTERNAL_MANY_TO_MANY_CONVERSIONS = (
    ErnieFuseAndSplitTextVisionExperts,
    ErnieSplitAndDecoupleTextVisionExperts,
)


class WeightConverter(WeightTransform):
    __slots__ = ("operations", "force_cpu")

    def __init__(
        self,
        source_patterns: str | list[str],
        target_patterns: str | list[str],
        operations: list[ConversionOps],
        force_cpu: bool = False,
    ):
        super().__init__(source_patterns, target_patterns)
        self.operations: list[ConversionOps] = operations
        self.force_cpu = force_cpu

        if bool(len(self.source_patterns) - 1) + bool(len(self.target_patterns) - 1) >= 2:
            # We allow many-to-many only if we use an internal operation that can handle it
            if not any(isinstance(op, _INTERNAL_MANY_TO_MANY_CONVERSIONS) for op in self.operations):
                raise ValueError(
                    f"source keys={self.source_patterns}, target_patterns={self.target_patterns} but you can only have one to many, one to one or many to one."
                )
        if not self.operations:
            raise ValueError("WeightConverter requires at least one operation.")

    def convert(
        self,
        layer_name: str,
        model=None,
        config=None,
        hf_quantizer=None,
        loading_info: LoadStateDictInfo | None = None,
    ):
        # Collect the tensors here - we use a new dictionary to avoid keeping them in memory in the internal
        # attribute during the whole process
        collected_tensors = self.materialize_tensors()

        for op in self.operations:
            with log_conversion_errors(layer_name, loading_info, (len(collected_tensors), layer_name), op):
                collected_tensors = op.convert(
                    collected_tensors,
                    source_patterns=self.source_patterns,
                    target_patterns=self.target_patterns,
                    # Additional kwargs, usually not used
                    full_layer_name=layer_name,
                    model=model,
                    config=config,
                    missing_keys=loading_info.missing_keys if loading_info else None,
                )

        # Tensors are returned from ops with the target patterns, we need to expand them to full name.
        # This means we need to grab the prefix and suffix to add to every target key
        full_name = layer_name
        if "*." in layer_name:
            full_name = layer_name.replace("*.", "0.")

        try:
            prefix, _, suffix = next(full_name.partition(k) for k in collected_tensors.keys() if k in full_name)
            # Rename the tensors
            collected_tensors = {prefix + k + suffix: v for k, v in collected_tensors.items()}
        # some quantizers need to already rename in `convert` as they cannot only rely on prefix and suffix
        except StopIteration:
            pass

        if hf_quantizer is not None and self.quantization_operation is not None:
            with log_conversion_errors(
                layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation
            ):
                collected_tensors = self.quantization_operation.convert(
                    collected_tensors,
                    source_patterns=self.source_patterns,
                    target_patterns=self.target_patterns,
                    full_layer_name=layer_name,
                    config=config,
                    model=model,
                    missing_keys=loading_info.missing_keys if loading_info else None,
                )
        return collected_tensors


# For I/O bound operations (i.e. here reading files), it is better to have fewer threads, e.g. 4 is a good default.
# Having too many is actually harming performances quite a lot, i.e. using 16 can sometimes lead to taking TWICE
# as much time to load the same model
GLOBAL_WORKERS = min(4, os.cpu_count() or 4)


def _materialize_copy(tensor: torch.Tensor, device=None, dtype=None) -> torch.Tensor:
    # This slicing is what actually loads the tensor from the safetensors slice object
    tensor = tensor[...]
    if dtype is not None or device is not None:
        tensor = tensor.to(device=device, dtype=dtype)
    return tensor


def spawn_materialize(
    thread_pool: ThreadPoolExecutor | None,
    tensor: torch.Tensor,
    device=None,
    dtype=None,
    sharding_op: DtensorShardOperation | None = None,
    tensor_idx: int | None = None,
) -> Future | Callable:
    """Materialize (and optionally shard) a tensor, asynchronously if a thread pool is provided.

    When ``sharding_op`` is given the tensor is sharded according to its DTensor placements;
    otherwise it is simply copied to *device*/*dtype*. Without a thread pool a deferred callable
    is returned instead of a Future.
    """

    def _job():
        if sharding_op is not None:
            return sharding_op.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype)
        return _materialize_copy(tensor, device, dtype)

    if thread_pool is not None:
        return thread_pool.submit(_job)
    else:
        # Return the Callable here, not the Tensor itself, so we actually delay loading to avoid saturating cpu
        # memory during Conversion
        return _job


def dot_natural_key(s: str):
    """
    Sort key for state-dict names: split on `"."` and sort digits numerically and strings alphabetically. It emits a
    tuple at each point to sort ints first and strings second to avoid int-string comparison failures.
    """
    parts = []
    for part in s.split("."):
        if part.isdigit():
            parts.append((0, int(part)))
        else:
            # This will remove all trailing digit characters, as `rstrip` actually considers it as a set of chars
            text_part = part.rstrip("0123456789")
            trailing_digits = part[len(text_part) :]
            # Sort numeric suffixes numerically, so `shard_2` precedes `shard_11` for example
            if trailing_digits != "":
                parts.append((1, text_part, int(trailing_digits)))
            else:
                parts.append((1, text_part))
    return parts


@contextmanager
def log_conversion_errors(
    first_target_key: str,
    loading_info: LoadStateDictInfo | None,
    extras: Any = None,
    op: list[ConversionOps] | ConversionOps | None = None,
):
    """Catch all exceptions during `convert` calls, and log the errors for later. Re-raise a `SkipParameters` exception
    that will be caught later to skip the parameters that raised the original Exception."""
    try:
        yield
    except Exception as e:
        # During reverse mapping, we do not log and skip errors
        if loading_info is None:
            raise e

        def _format_op_name(curr_op: list[ConversionOps] | ConversionOps | None) -> str | None:
            if curr_op is None:
                return None
            if isinstance(curr_op, (list, tuple, set)):
                names = [o.__class__.__name__ for o in curr_op if o is not None]
                if not names:
                    return None
                return ", ".join(names)
            return curr_op.__class__.__name__

        op_name = _format_op_name(op)

        tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
        if isinstance(extras, tuple) and len(extras) == 2:
            length, target_keys = extras
            descriptor = f"{op_name} " if op_name else ""
            loading_info.conversion_errors[first_target_key] = (
                f"{tb_str}{e}\nError: {descriptor}on tensors destined for {target_keys}. Ckpt contains: {length}"
            )
        elif isinstance(extras, str):
            suffix = f" via {op_name}" if op_name else ""
            loading_info.conversion_errors[first_target_key] = (
                f"{tb_str}{e}\nError{suffix} when processing parameter {extras}"
            )
        elif extras is None and op_name:
            loading_info.conversion_errors[first_target_key] = f"{op_name}: {e}"
        else:
            loading_info.conversion_errors[first_target_key] = f"{extras} |Error: {e}"

        # Raise a specific Exception that we can catch easily
        raise SkipParameters()


@torch.no_grad()
def set_param_for_module(
    model: PreTrainedModel,
    target_name: str,
    param_value: torch.Tensor,
    loading_info: LoadStateDictInfo,
    hf_quantizer: HfQuantizer,
):
    module_path, _, param_name = target_name.rpartition(".")
    module_obj = model.get_submodule(module_path) if module_path else model

    if param_name == torch.nn.modules.module._EXTRA_STATE_KEY_SUFFIX:
        module_obj.set_extra_state(param_value)
        loading_info.missing_keys.discard(target_name)
        return

    ref = getattr(module_obj, param_name)
    if ref is None:
        loading_info.unexpected_keys.add(target_name)
    else:
        if not isinstance(param_value, torch.nn.Parameter) and not is_dtensor(ref):
            if param_name not in module_obj._buffers:
                param_value = torch.nn.Parameter(param_value, requires_grad=param_value.is_floating_point())

        # Remove from missing keys (it's either mismatched, or all good)
        loading_info.missing_keys.discard(target_name)

        # For DTensor parameters, compare against the local shard loaded on this rank.
        expected_shape = ref._local_tensor.shape if is_dtensor(ref) else ref.shape

        if ref is not None and param_value.shape != expected_shape and hf_quantizer is None:
            loading_info.mismatched_keys.add((target_name, param_value.shape, expected_shape))
        else:
            if is_dtensor(ref):
                local_param = param_value.detach() if isinstance(param_value, torch.nn.Parameter) else param_value
                dtensor_param = _dtensor_from_local_like(local_param, ref)
                param_value = torch.nn.Parameter(
                    dtensor_param, requires_grad=ref.requires_grad and dtensor_param.is_floating_point()
                )
            # super important otherwise _init_weight will re-init the param
            param_value._is_hf_initialized = True
            setattr(module_obj, param_name, param_value)


def offload_and_maybe_resave_param(
    target_name: str,
    param: torch.Tensor,
    loading_info: LoadStateDictInfo,
    disk_offload_folder: str,
    disk_offload_index: dict,
    applied_ops: WeightConverter | WeightRenaming,
) -> dict:
    """Takes care of correctly offloading `param`. If it's not already present in the `disk_offload_index`, or if any
    WeightConverter operations have been applied, it will resave the new parameter. Otherwise, it will use the original
    `disk_offload_index` for this given param."""
    # We need to remove from missing keys
    loading_info.missing_keys.discard(target_name)
    # If not already offloaded, or if we applied any special Operation except Renaming, we need to re-save
    if target_name not in disk_offload_index or isinstance(applied_ops, WeightConverter):
        disk_offload_index = offload_weight(param, target_name, disk_offload_folder, disk_offload_index)
    return disk_offload_index


class SkipParameters(Exception):
    """Control-flow sentinel: abort processing of the current parameters only (that were supposed to be created
    by a WeightConverter)."""

    pass


def rename_source_key(
    source_key: str,
    weight_renamings: list[WeightRenaming],
    weight_converters: list[WeightConverter],
    base_model_prefix: str | None = None,
    meta_state_dict: dict | None = None,
    reverse: bool = False,
) -> tuple[str, str | None]:
    """
    Rename a checkpoint key by first applying all `WeightRenaming`s, then at most one `WeightConverter`.
    This means that the `WeightConverter` must specify the final name for the key, but some `WeightRenaming`s can act on other
    parts of that key beforehand. If `reverse` is True, i.e. when reverting all the `WeightTransform`s, the opposite is performed
    to be coherent and correctly respect the `scope_prefix` of all `WeightTransform`s: first try to match 1 `WeightConverter`, and only
    then try to apply all `WeightRenaming`s. Indeed, in reverse mode, all reverse transforms should be applied in the opposite order to
    be consistent.
    Note that we proceed in this way because there is no need for a Converter-then-Rename order because Converters act only on specific
    leaf patterns, so no subsequent Renamings should ever target their output.

    Args:
        source_key (`str`):
            The original checkpoint key to rename.
        weight_renamings (`list[WeightRenaming]`):
            Applied in order; every matching renaming fires (they may chain).
        weight_converters (`list[WeightConverter]`):
            Applied after all renamings; at most one may match. Subsequent converters are skipped.
        base_model_prefix (`str`, *optional*):
            Base-model prefix to add or strip when both `base_model_prefix` and `meta_state_dict` are given.
        meta_state_dict (`dict`, *optional*):
            Meta state dict used to decide whether `base_model_prefix` should be added or stripped.
        reverse (`bool`, *optional*):
            This specifies if we are reverting all the `WeightTransform`s (saving back original format).

    Returns:
        `tuple[str, str | None]`: The renamed key and the matched converter's source pattern
        (or `None` if no converter matched).
    """
    renamed_key = source_key
    converter_source_pattern = None
    # 1. If `reverse` is False: apply all renamings in turns (if multiple match, it's the responsibility of the mappings to make sure
    # they are coherent).
    # Else, first apply the `WeightConverter` if any
    first_iterable = weight_renamings if not reverse else weight_converters
    for transform in first_iterable:
        renamed_key, source_pattern = transform.rename_source_key(renamed_key)
        # Only break after match is we are using the `WeightConverter`s here, i.e. `reverse` is True
        if reverse and source_pattern is not None:
            converter_source_pattern = source_pattern
            break

    # 2. If `reverse` is False: apply renaming through weight conversions on the key if we have any WeightConverter (here we stop after
    # the first match, as we assume only 1 converter can match any source key).
    # Else, apply all the `WeightRenaming`s
    second_iterable = weight_converters if not reverse else weight_renamings
    for transform in second_iterable:
        renamed_key, source_pattern = transform.rename_source_key(renamed_key)
        # Only break after match is we are using the `WeightConverter`s here, i.e. `reverse` is False
        if not reverse and source_pattern is not None:
            converter_source_pattern = source_pattern
            break

    # 3. check if we need to add or remove base_model_prefix if necessary (only during loading, not saving)
    if base_model_prefix is not None and meta_state_dict is not None:
        if (
            renamed_key.startswith(base_model_prefix)
            and meta_state_dict.get(re.sub(f"^{base_model_prefix}.", "", renamed_key, count=1)) is not None
        ):
            renamed_key = re.sub(f"^{base_model_prefix}.", "", renamed_key, count=1)
        elif meta_state_dict.get(f"{base_model_prefix}.{renamed_key}") is not None:
            renamed_key = f"{base_model_prefix}.{renamed_key}"

    return renamed_key, converter_source_pattern


def _add_unmatched_checkpoint_key(
    key: str,
    model: PreTrainedModel,
    loading_info: LoadStateDictInfo,
) -> None:
    """Classify a checkpoint key that does not match the current model state dict.

    During pipeline-parallel loading, each stage receives keys for the entire model even though its local state dict
    contains only that stage's parameters. A key owned by another stage is therefore recorded as intentionally skipped
    instead of unexpected. Without the key, it would be considered unexpected.
    """
    stage = getattr(model, "_pp_stage", None)
    if stage is None:
        loading_info.unexpected_keys.add(key)
        return

    base_model = getattr(model, model.base_model_prefix)
    owner_rank = stage.find_rank_for_key(key, len(base_model.layers), model.base_model_prefix)
    owned_by_another_stage = owner_rank is not None and owner_rank != stage.pp_rank

    if owned_by_another_stage:
        loading_info.skipped_pp_keys.add(key)
    else:
        loading_info.unexpected_keys.add(key)


def convert_and_load_state_dict_in_model(
    model: PreTrainedModel,
    state_dict: dict[str, Any],
    load_config: LoadStateDictConfig,
    disk_offload_index: dict | None = None,
):
    r"""
    We build a mapping from the keys obtained by renaming each of the checkpoint keys according to the weight_mapping rules.
    Then we load the tensors into the model, applying any conversion operations as needed.

    The `param_name_to_load` will look like this:
    {
        "model.layers.0.attention.q.weight": # Notice here there is only the first key of the target keys
            WeightConverter(
                source_patterns=["qkv"],
                target_patterns=["q", "k","v"],
                operations=[Chunk(dim=0, chunks=3)]),
                collected_tensors={
                    "qkv": [Future]},
                layer_targets={
                    "model.layers.0.attention.q.weight": {"model.layers.0.attention.qkv.weight"},
                    "model.layers.0.attention.k.weight": {"model.layers.0.attention.qkv.weight"},
                    "model.layers.0.attention.v.weight": {"model.layers.0.attention.qkv.weight"},
                }
            ),
        ...
    }

    We make sure that the keys are the full keys. The only "nit" here is that 1 key can map to multiple target keys (e.g. qkv -> q, k, v).
    In that case the weight converter will take care of doing the appropriate renaming.

    For example for:
    ```python
    WeightConverter(
        source_patterns=["mlp.experts.*.gate_proj.weight","mlp.experts.*.up_proj.weight"],
        target_patterns="mlp.experts.gate_up_proj",
        operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
    )
    ```
    we would have the following collected tensors:
    ```python
    collected_tensors = {
        "mlp.experts.*.gate_proj.weight": [Future, Future, Future, Future, Future, Future, Future, Future],
        "mlp.experts.*.up_proj.weight": [Future, Future, Future, Future, Future, Future, Future, Future],
    }
    ```
    The first op, `MergeModulelist`, would stack the 8 tensors of each source but will not "rename" them into the fused target name.
    The second op, `Concatenate`, would then rename the fused tensor into the final target name.

    If we want to split `qkv` we would have:
    ```python
    collected_tensors = {
        "attention.qkv.weight": [Future], # here its the full SOURCE keys.
    }
    ```
    The `Chunk` operation would then split the single tensor into 3 and rename them accordingly and update the collected tensors to:
    ```python
    realized_values = {
        "attention.q.weight": [Tensor],
        "attention.k.weight": [Tensor],
        "attention.v.weight": [Tensor],
    }
    ```

    Now that this is done, we can quantize / dequantize accordingly the collected_tensors.

    For some quantization methods, we need to gather different tensors:

    ```python
    # for "medmekk/llama-3.2-1b-float8-torchao"
    WeightConverter(
        source_patterns=[":qdata", ":scale"],
        target_patterns="",
        operations=[TorchaoDeserialize()],
    )
    ```
    This will collect all tensors that have the same prefix, but end with `:qdata` or `:scale`. This will give us:
    ```python
    all_weight_mapping = {
        "model.layers.13.self_attn.o_proj.weight": WeightConverter(
            source_patterns=[":qdata", ":scale"],
            target_patterns="",
            operations=[TorchaoDeserialize()],
            collected_tensors={
                ":qdata": [Future],
                ":scale": [Future],
            },
        ...
    }
    ```

    """
    base_model_prefix = model.base_model_prefix
    device_map = load_config.device_map or {"": "cpu"}
    hf_quantizer = load_config.hf_quantizer
    dtype = load_config.dtype
    disk_offload_folder = load_config.disk_offload_folder
    offload_buffers = load_config.offload_buffers
    dtype_plan = load_config.dtype_plan or {}
    weight_mapping = load_config.weight_mapping or []
    meta_model_state_dict = model.state_dict()
    model_buffers = {k for k, _ in model.named_buffers()}

    # We start from all missing keys, and we will remove/add them from the proper containers as loading advances
    loading_info = LoadStateDictInfo(
        missing_keys=set(meta_model_state_dict.keys()),
        unexpected_keys=set(),
        mismatched_keys=set(),
        conversion_errors={},
        error_msgs=[],
        skipped_pp_keys=set(),
    )

    # We use threading by default, if not explicitly deactivated via env variable. If we have to offload,
    # we cannot use it either to control the memory as we are under memory constraints, so we need to be sequential.
    # When doing on-the-fly quantization, we also use sync loading to avoid worker threads loading full-precision
    # tensors to GPU faster than the main thread can quantize them, which would cause a large memory spike.
    has_on_the_fly_quantization = hf_quantizer is not None and not hf_quantizer.pre_quantized
    if (
        is_env_variable_true("HF_DEACTIVATE_ASYNC_LOAD")
        or "disk" in device_map.values()
        or has_on_the_fly_quantization
    ):
        thread_pool = None
    else:
        thread_pool = ThreadPoolExecutor(max_workers=GLOBAL_WORKERS)

    renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)]
    converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)]
    param_name_to_load: dict[str, WeightRenaming | WeightConverter] = {}

    if dtype_plan != {}:
        dtype_policy_alt, dtype_policy_by_group_name, _ = build_glob_alternation(list(dtype_plan.keys()))

    pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns}

    state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
    for original_key, tensor in state_dict:
        # 1. Rename the key according to all renaming and weight conversion patterns.
        renamed_key, source_pattern = rename_source_key(
            original_key, renamings, converters, base_model_prefix, meta_model_state_dict
        )
        if renamed_key not in meta_model_state_dict and original_key in meta_model_state_dict:
            # Key should probably not have been renamed but we might need the `prefix` to be added.
            renamed_key, source_pattern = rename_source_key(
                original_key, [], [], base_model_prefix=base_model_prefix, meta_state_dict=meta_model_state_dict
            )

        # 2. finally, collect the tensor into the proper converter
        if renamed_key in meta_model_state_dict:
            empty_param = meta_model_state_dict.get(renamed_key)
            # If we enter here, we have a WeightConverter operation to perform
            if source_pattern is not None:
                new_converter = deepcopy(pattern_to_converter[source_pattern])
                # each target key gets its own converter instance
                mapping = param_name_to_load.setdefault(renamed_key, new_converter)
            # Otherwise, only potential renaming
            else:
                mapping = param_name_to_load.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
                source_pattern = original_key

            # 3. Handle dtype casting
            needs_quantization = (
                hf_quantizer
                and not hf_quantizer.pre_quantized
                and hf_quantizer.param_needs_quantization(model, renamed_key)
            )
            if needs_quantization:
                mapping.quantization_operation = hf_quantizer.get_quantize_ops()

            _dtype = dtype
            if (
                hf_quantizer
                and hf_quantizer.pre_quantized
                and (
                    original_key != renamed_key
                    or not (
                        tensor.get_dtype().startswith(("F", "BF"))
                        if hasattr(tensor, "get_dtype")
                        else tensor.is_floating_point()
                    )
                )
            ):
                # if the key was renamed as it is not available in the state dict otherwise, it means that we are deserializing it,
                # so we need to make sure to load the tensor with the same dtype from the checkpoint
                # TODO: make the condition more srict for native fp8 model such as qwen2moe fp8
                _dtype = None
            elif dtype_plan != {} and dtype_policy_alt.search(renamed_key):
                matched_dtype_pattern = dtype_policy_alt.search(renamed_key)
                if matched_dtype_pattern is not None:
                    _dtype = dtype_plan[dtype_policy_by_group_name[matched_dtype_pattern.lastgroup]]
            elif empty_param is not None and empty_param.dtype != _dtype:
                _dtype = empty_param.dtype  # usually correct when initializing

            # Per-expert sharding (EP) needs `tensor_idx` = the expert index so the
            # distributed op selects whole experts. The signal is a `MergeModulelist`
            # in the chain; it isn't always `operations[0]` (e.g. an FP8 quantizer
            # prepends a scale-decode op), so scan the whole chain rather than just the head.
            tensor_idx = (
                len(mapping.collected_tensors.get(source_pattern, []))
                if isinstance(mapping, WeightConverter)
                and any(isinstance(op, MergeModulelist) for op in mapping.operations)
                else None
            )

            # 4. Handle DTensor sharding or device_map placement
            param_device = get_device(device_map, renamed_key, valid_torch_device=True)
            sharding_op = None
            if is_dtensor(empty_param):
                sharding_op = DtensorShardOperation(empty_param)

            # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations
            # directly on the device, as it will completely blow up the memory during the ops memory spike. So defer to "cpu", then
            # accelerate will take care of putting back on correct device after loading
            # Note that we only do it with `device_map` but not with `tp_plan`, as tp will perform local sharding before, so memory
            # spike during conversion ops should be fine
            if sharding_op is None and isinstance(mapping, WeightConverter) and mapping.force_cpu:
                param_device = "cpu"

            future_or_tensor = spawn_materialize(
                thread_pool,
                tensor,
                param_device,
                _dtype,
                sharding_op=sharding_op,
                tensor_idx=tensor_idx,
            )

            mapping.add_tensor(renamed_key, original_key, source_pattern, future_or_tensor)
        elif source_pattern is not None:  # add all target keys as unexpected
            mapping = pattern_to_converter[source_pattern]
            for k in mapping.target_patterns:
                _add_unmatched_checkpoint_key(
                    renamed_key.replace(mapping.target_patterns[0], k),
                    model,
                    loading_info,
                )
        else:
            _add_unmatched_checkpoint_key(renamed_key, model, loading_info)

    try:
        for first_param_name, mapping in tqdm(param_name_to_load.items(), desc="Loading weights"):
            try:
                realized_value = mapping.convert(
                    first_param_name,
                    model=model,
                    config=model.config,
                    hf_quantizer=hf_quantizer,
                    loading_info=loading_info,
                )
                for target_name, param in realized_value.items():
                    param = param[0] if isinstance(param, list) else param
                    param_device = get_device(device_map, target_name)
                    # Offloading support
                    if param_device == "disk" and (target_name not in model_buffers or offload_buffers):
                        disk_offload_index = offload_and_maybe_resave_param(
                            target_name, param, loading_info, disk_offload_folder, disk_offload_index, mapping
                        )
                    else:
                        set_param_for_module(
                            model,
                            target_name,
                            param,
                            loading_info,
                            hf_quantizer,
                        )

                # Cleanup all the tensors that were gathered before next iteration
                del realized_value

            except SkipParameters:
                continue

    # Close the pool, independently of whether the code was interrupted or finished successfully
    finally:
        if thread_pool is not None:
            # `cancel_futures=True` in case the program was interrupted, to avoid wasting time on exit
            thread_pool.shutdown(wait=False, cancel_futures=True)

    # Keep the current weight conversion mapping for later saving (in case it was coming directly from the user), but
    # only if it was used, i.e. it matched any weight from the checkpoints
    model_specific_conversions = [conversion for conversion in weight_mapping if conversion.was_used()]
    model._weight_conversions = model_specific_conversions

    return loading_info, disk_offload_index


def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]):
    """
    Revert the conversion mapping that was used to load the model with `from_pretrained`, or the default one
    if the model was created in another way and is part of the default mappings.
    """
    weight_conversions = getattr(model, "_weight_conversions", None)
    # In this case, the model was not created with `from_pretrained` -> let's check if it's in the hardcoded
    # mappings, and recreate the mapping from there if it is
    if weight_conversions is None:
        from .conversion_mapping import get_model_conversion_mapping

        # Do not resave with the legacy renaming, if present
        weight_conversions = get_model_conversion_mapping(model, add_legacy=False)
        # If the model had no `_weight_conversions` attached, drop any PrefixChange transform - this is because the
        # model was almost surely instantiated from scratch (at least not from `from_pretrained`), and PrefixChange with
        # `prefix_to_remove` would otherwise add a unwanted prefix (as we dont have any information about whether the prefix
        # was there or not during load)
        weight_conversions = [x for x in weight_conversions if not isinstance(x, PrefixChange)]
        weight_conversions = weight_conversions if len(weight_conversions) > 0 else None

    # We did not find any operations to perform -> quick escape
    if weight_conversions is None:
        return state_dict

    # Important: we need to revert the order here, so that potential conversions from submodels are performed first
    weight_conversions = weight_conversions[::-1]
    # Reverse all Transforms
    reverse_weight_conversions = [conversion.reverse_transform() for conversion in weight_conversions]
    renamings = [entry for entry in reverse_weight_conversions if isinstance(entry, WeightRenaming)]
    converters = [entry for entry in reverse_weight_conversions if isinstance(entry, WeightConverter)]
    pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns}

    conversion_mapping: dict[str, WeightTransform] = {}
    state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0]))
    for original_key, tensor in state_dict:
        # Rename the key according to all renaming pattern and optional weight converter patterns
        renamed_key, source_pattern = rename_source_key(original_key, renamings, converters, reverse=True)
        if source_pattern is not None:
            new_converter = deepcopy(pattern_to_converter[source_pattern])
            # each target key gets its own converter instance
            mapping = conversion_mapping.setdefault(renamed_key, new_converter)
        else:
            mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
            source_pattern = original_key

        mapping.add_tensor(renamed_key, original_key, source_pattern, tensor)

    new_state_dict = {}
    for first_param_name, reversed_converter in conversion_mapping.items():
        # Apply the reverse converter
        realized_value = reversed_converter.convert(first_param_name, model=model, config=model.config)
        for target_name, param in realized_value.items():
            param = param[0] if isinstance(param, list) else param
            new_state_dict[target_name] = param

    return new_state_dict