transformers 5.0.0__py3-none-any.whl → 5.0.0rc0__py3-none-any.whl

transformers/core_model_loading.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,24 +16,22 @@
 
 from __future__ import annotations
 
-import math
 import os
 import re
 from abc import abstractmethod
 from collections import defaultdict
-from collections.abc import Callable, MutableMapping, MutableSet
+from collections.abc import MutableMapping, MutableSet
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
-from itertools import chain
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import torch
 
-from .integrations.accelerate import get_device, offload_weight
+from .integrations.accelerate import offload_weight
 from .integrations.tensor_parallel import ALL_PARALLEL_STYLES
-from .utils import is_env_variable_true, is_torch_greater_or_equal, logging
+from .utils import is_torch_greater_or_equal, logging
 
 
 _torch_distributed_available = torch.distributed.is_available()
@@ -49,40 +48,8 @@ if TYPE_CHECKING:
 logger = logging.get_logger(__name__)
 
 
-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!
-    pattern = re.sub(r"\(\?.+\)", "", pattern)
-    # 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 build_glob_alternation(
-    globs: list[WeightRenaming | WeightConverter | str],
+    globs: list[Union[WeightRenaming, WeightConverter, str]],
 ) -> tuple[re.Pattern, dict[str, str], dict[str, str]]:
     """
     Build a single alternation regex with one named group per glob.
@@ -145,12 +112,12 @@ class Chunk(ConversionOps):
     ) -> dict[str, torch.Tensor]:
         tensors = next(iter(input_dict.values()))
         tensor = tensors[0] if isinstance(tensors, list) else tensors
-        targets = self.get_target_patterns(input_dict, target_patterns)
+        targets = self.get_target_pattern(input_dict, target_patterns)
         sizes = len(targets)
         chunks = torch.chunk(tensor, sizes, dim=self.dim)
         return dict(zip(targets, chunks))
 
-    def get_target_patterns(self, input_dict: dict, target_patterns: list[str]) -> list[str]:
+    def get_target_pattern(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!")
@@ -177,7 +144,7 @@ class Concatenate(ConversionOps):
     ) -> 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
+        # Very important to keep the relative order of the source patterms here, so we iterate over them not the
         # input directly as it's unordered!
         for source_pattern in source_patterns:
             tensors = input_dict[source_pattern]
@@ -277,44 +244,6 @@ class SplitModulelist(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):
-        self.dim0 = dim0
-        self.dim1 = dim1
-
-    @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
-        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)
-
-
 class PermuteForRope(ConversionOps):
     """
     Applies the permutation required to convert complex RoPE weights to the split sin/cos format.
@@ -349,217 +278,18 @@ class PermuteForRope(ConversionOps):
         return output
 
 
-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)
-
-
-class Force16BytesAlignment(ConversionOps):
-    """
-    Ensures that the given tensor is 16-bytes aligned in memory and clones it if not.
-    This garantees 16-bytes alignmenet for kernels / implementations that use TMA or SIMD instructions like torch._grouped_mm.
-    """
-
-    @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
-        tensor = tensor.clone() if tensor.data_ptr() % 16 != 0 else tensor
-        return {target_pattern: tensor}
-
-    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 Force16BytesAlignment()
-
-
 @dataclass(slots=True)
 class WeightTransform:
-    source_patterns: str | list[str] = field(init=True)
-    target_patterns: str | list[str] = field(init=True)
+    source_patterns: Union[str, list[str]] = field(init=True)
+    target_patterns: Union[str, list[str]] = field(init=True)
     compiled_sources: re.Pattern = field(init=False)
 
-    distributed_operation: TensorParallelLayer | None = None
-    quantization_operation: ConversionOps | None = None
+    distributed_operation: Optional[TensorParallelLayer] = None
+    quantization_operation: Optional[ConversionOps] = None
 
     collected_tensors: dict[str, list[Future]] = field(default_factory=lambda: defaultdict(list), init=False)
     layer_targets: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set), init=False)
 
-    def __setattr__(self, name: str, value: Any):
-        if name == "source_patterns":
-            if isinstance(value, str):
-                value = [value]
-            normalized = []
-            for pattern in value:
-                if r"\1" in pattern:
-                    pattern = pattern.replace(r"\1", r"(.+)")
-                normalized.append(pattern)
-            object.__setattr__(self, name, normalized)
-            self._rebuild_compiled_sources()
-            return
-        if name == "target_patterns":
-            if isinstance(value, str):
-                value = [value]
-            normalized = []
-            for pattern in value:
-                # 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!
-                pattern = re.sub(r"\(\?.+\)", "", pattern)
-                # Allow capturing groups in patterns, i.e. to add/remove a prefix to all keys (e.g. timm_wrapper, sam3)
-                if r"(.+)" in pattern:
-                    pattern = pattern.replace(r"(.+)", r"\1")
-                normalized.append(pattern)
-            object.__setattr__(self, name, normalized)
-            return
-        object.__setattr__(self, name, value)
-
-    def _rebuild_compiled_sources(self):
-        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})")
-        object.__setattr__(self, "compiled_sources", re.compile("|".join(branches)))
-
     def __post_init__(self):
         if isinstance(self.source_patterns, str):
             self.source_patterns = [self.source_patterns]
@@ -569,45 +299,39 @@ class WeightTransform:
         # Due to how our `_checkpoint_conversion_mapping` mappings are written, we need 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
+            # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping
+            pattern = pattern.removeprefix("^")
+            # Remove negative lookahead if any. This is ugly but needed for reverse mapping of Qwen2.5 and Sam3!
+            pattern = re.sub(r"\(\?!.+\)", "", pattern)
+            # Allow capturing groups in patterns, i.e. to add/remove a prefix to all keys (e.g. timm_wrapper, sam3)
+            if r"(.+)" in pattern:
+                pattern = pattern.replace(r"(.+)", r"\1")
+            self.target_patterns[i] = pattern
 
         # 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):
             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)
+                pattern = pattern.replace(r"\1", r"(.+)")
             self.source_patterns[i] = pattern
 
         # Construct the regex we will use to rename keys from the sources to the targets
-        self._rebuild_compiled_sources()
+        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 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 rename_source_key(self, source_key: str) -> tuple[str, str | re.Pattern]:
+    def reset(self) -> None:
+        """Clean-up the collected tensors to make sure we don't keep references to past tensors in memory."""
+        self.collected_tensors = defaultdict(list)
+
+    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.
@@ -618,19 +342,19 @@ class WeightTransform:
         match_object = self.compiled_sources.search(source_key)
         if match_object is None:
             return source_key, None
-
         # 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)
+        # # 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 = source_key.replace(match_object.group(0), replacement)
+
         return renamed_key, source_pattern_that_matched
 
     def reverse_transform(self) -> WeightTransform:
@@ -651,32 +375,6 @@ class WeightTransform:
 
         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 set(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]
-            # Sync loading
-            elif callable(tensors[0]):
-                tensors = [func() for func in tensors]
-            # Add them to the new dictionary
-            collected_tensors[key] = tensors
-
-        return collected_tensors
-
 
 @dataclass(slots=True)
 class WeightRenaming(WeightTransform):
@@ -688,22 +386,22 @@ class WeightRenaming(WeightTransform):
         model=None,
         config=None,
         hf_quantizer=None,
-        missing_keys: MutableSet[str] | None = None,
-        conversion_errors: MutableMapping[str, str] | None = None,
+        missing_keys: Optional[MutableSet[str]] = None,
+        misc: Optional[MutableMapping[str, str]] = 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()
+        # Collect the tensor if using threading
+        for pattern, futures in self.collected_tensors.items():
+            self.collected_tensors[pattern] = (
+                futures if isinstance(futures[0], torch.Tensor) else [future.result() for future in futures]
+            )
 
         # 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]]}
+        collected_tensors = {target_key: self.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, conversion_errors, (len(collected_tensors), layer_name), self.quantization_operation
-            ):
+            with log_to_misc(layer_name, misc, (self.collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
@@ -714,14 +412,7 @@ class WeightRenaming(WeightTransform):
                     missing_keys=missing_keys,
                 )
 
-        return collected_tensors, conversion_errors
-
-
-# List of classes that are known to be able to use m:n
-_INTERNAL_MANY_TO_MANY_CONVERSIONS = (
-    ErnieFuseAndSplitTextVisionExperts,
-    ErnieSplitAndDecoupleTextVisionExperts,
-)
+        return collected_tensors, misc
 
 
 @dataclass(slots=True)
@@ -729,12 +420,11 @@ class WeightConverter(WeightTransform):
     operations: list[ConversionOps] = field(default_factory=list, repr=False)
 
     def __post_init__(self):
+        WeightTransform.__post_init__(self)
         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."
-                )
+            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.")
 
@@ -744,15 +434,18 @@ class WeightConverter(WeightTransform):
         model=None,
         config=None,
         hf_quantizer=None,
-        missing_keys: MutableSet[str] | None = None,
-        conversion_errors: MutableMapping[str, str] | None = None,
+        missing_keys: Optional[MutableSet[str]] = None,
+        misc: Optional[MutableMapping[str, str]] = 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()
+        # Collect all tensors if using threading
+        for pattern, futures in self.collected_tensors.items():
+            self.collected_tensors[pattern] = (
+                futures if isinstance(futures[0], torch.Tensor) else [future.result() for future in futures]
+            )
 
+        collected_tensors = self.collected_tensors
         for op in self.operations:
-            with log_conversion_errors(layer_name, conversion_errors, (len(collected_tensors), layer_name), op):
+            with log_to_misc(layer_name, misc, (collected_tensors, layer_name), op):
                 collected_tensors = op.convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
@@ -769,19 +462,11 @@ class WeightConverter(WeightTransform):
         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
-
+        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()}
         if hf_quantizer is not None and self.quantization_operation is not None:
-            with log_conversion_errors(
-                layer_name, conversion_errors, (len(collected_tensors), layer_name), self.quantization_operation
-            ):
+            with log_to_misc(layer_name, misc, (collected_tensors, layer_name), self.quantization_operation):
                 collected_tensors = self.quantization_operation.convert(
                     collected_tensors,
                     source_patterns=self.source_patterns,
@@ -791,7 +476,7 @@ class WeightConverter(WeightTransform):
                     model=model,
                     missing_keys=missing_keys,
                 )
-        return collected_tensors, conversion_errors
+        return collected_tensors, misc
 
 
 # For I/O bound operations (i.e. here reading files), it is better to have fewer threads, e.g. 4 is a good default.
@@ -800,46 +485,25 @@ class WeightConverter(WeightTransform):
 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
+def _materialize_copy(tensor, device=None, dtype=None):
     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
-) -> Future | Callable:
-    """Materialize a tensor from file asynchronously if `thread_pool` is provided, or return a Callable that will
-    load the tensor synchronously when called."""
-
+def spawn_materialize(thread_pool, tensor, device=None, dtype=None) -> Future:
     def _job():
         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
-
+    return thread_pool.submit(_job)
 
-def spawn_tp_materialize(
-    thread_pool: ThreadPoolExecutor | None, tensor: torch.Tensor, sharding_method, tensor_idx, device=None, dtype=None
-) -> Future | Callable:
-    """Materialize and shard a tensor (according to the TP-plan) from file asynchronously if `thread_pool` is provided, or
-    return a Callable that will load the tensor synchronously when called."""
 
+def spawn_tp_materialize(thread_pool, tensor, sharding_method, tensor_idx, dtype=None) -> Future:
     def _job():
-        return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype)
+        return sharding_method.shard_tensor(tensor, param_casting_dtype=dtype, tensor_idx=tensor_idx)[0]
 
-    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
+    return thread_pool.submit(_job)
 
 
 def dot_natural_key(s: str):
@@ -852,22 +516,18 @@ def dot_natural_key(s: str):
 
 
 @contextmanager
-def log_conversion_errors(
+def log_to_misc(
     first_target_key: str,
-    conversion_errors: MutableMapping[str, str] | None,
+    misc: MutableMapping[str, str],
     extras: Any = None,
-    op: list[ConversionOps] | ConversionOps | None = None,
+    op: Union[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 catched later to skip the parameters that raised the original Exception."""
+    # A simple helper to handle errors with contextual messages.
     try:
         yield
     except Exception as e:
-        # During reverse mapping, we do not log and skip errors
-        if conversion_errors is None:
-            raise e
 
-        def _format_op_name(curr_op: list[ConversionOps] | ConversionOps | None) -> str | None:
+        def _format_op_name(curr_op: Union[list[ConversionOps], ConversionOps, None]) -> Optional[str]:
             if curr_op is None:
                 return None
             if isinstance(curr_op, (list, tuple, set)):
@@ -879,21 +539,19 @@ def log_conversion_errors(
 
         op_name = _format_op_name(op)
         if isinstance(extras, tuple) and len(extras) == 2:
-            length, target_keys = extras
+            values, target_keys = extras
             descriptor = f"{op_name} " if op_name else ""
-            conversion_errors[first_target_key] = (
-                f"{e}\nError: {descriptor}on tensors destined for {target_keys}. Ckpt contains: {length}"
+            misc[first_target_key] = (
+                f"{e}\nError: {descriptor}on tensors destined for {target_keys}. Ckpt contains: {len(values)}"
             )
         elif isinstance(extras, str):
             suffix = f" via {op_name}" if op_name else ""
-            conversion_errors[first_target_key] = f"{e}\nError{suffix} when processing parameter {extras}"
+            misc[first_target_key] = f"{e}\nError{suffix} when processing parameter {extras}"
         elif extras is None and op_name:
-            conversion_errors[first_target_key] = f"{op_name}: {e}"
+            misc[first_target_key] = f"{op_name}: {e}"
         else:
-            conversion_errors[first_target_key] = f"{extras} |Error: {e}"
-
-        # Raise a specific Exception that we can catch easily
-        raise SkipParameters()
+            misc[first_target_key] = f"{extras} |Error: {e}"
+        raise SkipLayer()
 
 
 def set_param_for_module(
@@ -902,20 +560,22 @@ def set_param_for_module(
     param_value: torch.Tensor,
     mismatch_keys: MutableSet[tuple[str, torch.Size, torch.Size]],
     missing_keys: MutableSet[str],
+    misc: MutableMapping[str, Any],
     unexpected_keys: MutableSet[str],
-    distributed_operation: TensorParallelLayer | None,
+    distributed_operation: Optional[TensorParallelLayer],
     hf_quantizer: HfQuantizer,
 ):
-    module_path, _, param_name = target_name.rpartition(".")
-    module_obj = model.get_submodule(module_path) if module_path else model
-
-    ref = getattr(module_obj, param_name)
-    if ref is None:
-        unexpected_keys.add(target_name)
-    else:
-        if not isinstance(param_value, torch.nn.Parameter):
-            if distributed_operation is not None:
-                if getattr(distributed_operation, "use_dtensor", False):
+    with log_to_misc(target_name, misc, target_name):
+        module_path, _, param_name = target_name.rpartition(".")
+        module_obj = model.get_submodule(module_path) if module_path else model
+
+        ref = getattr(module_obj, param_name)
+        if ref is None:
+            unexpected_keys.add(target_name)
+        else:
+            use_dtensor = hasattr(distributed_operation, "use_dtensor") and distributed_operation.use_dtensor
+            if not isinstance(param_value, torch.nn.Parameter):
+                if distributed_operation is not None:
                     param_value = DTensor.from_local(
                         param_value,
                         distributed_operation.device_mesh,
@@ -924,17 +584,20 @@ def set_param_for_module(
                         shape=ref.size(),
                         stride=ref.stride(),
                     )
-            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)
-        missing_keys.discard(target_name)
-        if ref is not None and ref.shape != param_value.shape and hf_quantizer is None:
-            mismatch_keys.add((target_name, param_value.shape, ref.shape))
-        else:
-            # super important otherwise _init_weight will re-init the param
-            param_value._is_hf_initialized = True
-            setattr(module_obj, param_name, param_value)
+                    if not use_dtensor:
+                        # we convert to local
+                        param_value = param_value.to_local()
+                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)
+            missing_keys.discard(target_name)
+            if ref is not None and ref.shape != param_value.shape and hf_quantizer is None:
+                mismatch_keys.add((target_name, param_value.shape, ref.shape))
+            else:
+                # 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(
@@ -956,9 +619,8 @@ def offload_and_maybe_resave_param(
     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)."""
+class SkipLayer(Exception):
+    """Control-flow sentinel: abort processing of the current layer only."""
 
     pass
 
@@ -1004,10 +666,15 @@ def rename_source_key(
 def convert_and_load_state_dict_in_model(
     model: PreTrainedModel,
     state_dict: dict[str, Any],
-    load_config: Any,
+    weight_mapping: list[WeightConverter | WeightRenaming] | None,
     tp_plan: dict[str, str] | None,
+    hf_quantizer: HfQuantizer | None,
+    dtype: torch.dtype | None = None,
+    device_map: dict | None = None,
     dtype_plan: dict | None = None,
+    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
     disk_offload_index: dict | None = None,
+    disk_offload_folder: str | None = None,
 ):
     r"""
     We build a mapping from the keys obtained by renaming each of the checkpoint keys according to the weight_mapping rules.
@@ -1021,7 +688,7 @@ def convert_and_load_state_dict_in_model(
                 target_patterns=["q", "k","v"],
                 operations=[Chunk(dim=0, chunks=3)]),
                 collected_tensors={
-                    "qkv": [Future]},
+                    "qkv": [Future, Future, 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"},
@@ -1097,32 +764,26 @@ def convert_and_load_state_dict_in_model(
     """
     prefix = model.base_model_prefix
     tp_plan = tp_plan or {}
-    device_map = load_config.device_map or {"": "cpu"}
-    hf_quantizer = load_config.hf_quantizer
-    dtype = load_config.dtype
-    device_mesh = load_config.device_mesh
-    disk_offload_folder = load_config.disk_offload_folder
-    offload_buffers = load_config.offload_buffers
+    device_map = device_map or {"": "cpu"}
+    # Here, we first sort by number of submodules, then length of the full string, to make sure to match correctly
+    device_map_regex = re.compile(
+        "|".join(rf"({k})" for k in sorted(device_map.keys(), key=lambda x: (x.count("."), len(x)), reverse=True))
+    )
     dtype_plan = dtype_plan or {}
-    weight_mapping = load_config.weight_mapping or []
+    weight_mapping = weight_mapping or []
     meta_model_state_dict = model.state_dict()
-    model_buffers = {k for k, _ in model.named_buffers()}
-
     missing_keys = set(meta_model_state_dict.keys())
-    conversion_errors = {}
+
+    misc = {}
     mismatch_keys = set()
     unexpected_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
-    if is_env_variable_true("HF_DEACTIVATE_ASYNC_LOAD") or "disk" in device_map.values():
-        thread_pool = None
-    else:
-        thread_pool = ThreadPoolExecutor(max_workers=GLOBAL_WORKERS)
+    # Global thread_pool
+    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] = {}
+
+    param_name_to_load: dict[str, Union[WeightRenaming | WeightConverter]] = {}
 
     # build '(?P<g0>.*.*\\.block_sparse_moe\\..*)' and group to source {'g0': '*.block_sparse_moe.'}
     # and target to source {'g0': '*.mlp.'}. This allows us to quickly find which pattern matched.
@@ -1165,40 +826,41 @@ def convert_and_load_state_dict_in_model(
             if hf_quantizer and hf_quantizer.pre_quantized and original_key != renamed_key:
                 # 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]]
+                    _dtype = dtype_plan[matched_dtype_pattern.group()]
             elif empty_param is not None and empty_param.dtype != _dtype:
                 _dtype = empty_param.dtype  # usually correct when initializing
 
-            # 4. Handle TP sharding or device_map placement
-            future_or_tensor = None
+            # 4. Handle TP sharding or device_map placement -> scheduled materialization
+            future = None
             if device_mesh:
                 if matched_tp_pattern := tp_plan_alt.search(renamed_key):
                     matched_tp_pattern = tp_plan_by_group_name[matched_tp_pattern.lastgroup]
                     if getattr(mapping, "distributed_operation", None) is None:
                         tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]].__class__
                         mapping.distributed_operation = tp_layer(
-                            device_mesh=device_mesh, rank=device_mesh.get_local_rank(), empty_param=empty_param.clone()
+                            device_mesh=device_mesh, rank=device_map[""].index, empty_param=empty_param.clone()
                         )
                     shard_index = len(mapping.collected_tensors.get(original_key, []))
-                    future_or_tensor = spawn_tp_materialize(
+                    future = spawn_tp_materialize(
                         thread_pool,
                         tensor,
                         mapping.distributed_operation,
                         shard_index,
-                        device_map[""],
                         _dtype,
                     )
 
-            if future_or_tensor is None:
-                param_device = get_device(device_map, renamed_key, valid_torch_device=True)
-                future_or_tensor = spawn_materialize(thread_pool, tensor, param_device, _dtype)
+            if future is None:
+                device_match = device_map_regex.match(renamed_key)
+                param_device = device_map[device_match.group()] if device_match else device_map.get("", "cpu")
+                # If disk, we need to materialize on cpu first
+                param_device = "cpu" if param_device == "disk" else param_device
+                future = spawn_materialize(thread_pool, tensor, param_device, _dtype)
 
-            mapping.add_tensor(renamed_key, original_key, source_pattern, future_or_tensor)
+            mapping.add_tensor(renamed_key, original_key, source_pattern, future)
         elif source_pattern is not None:  # add all target keys as unexpected
             mapping = pattern_to_converter[source_pattern]
             for k in mapping.target_patterns:
@@ -1206,57 +868,52 @@ def convert_and_load_state_dict_in_model(
         else:
             unexpected_keys.add(renamed_key)
 
-    try:
-        total_entries = len(param_name_to_load)
-        with logging.tqdm(total=total_entries, desc="Loading weights") as pbar:
-            for first_param_name, mapping in param_name_to_load.items():
-                pbar.update(1)
-                pbar.set_postfix({"Materializing param": first_param_name})
-                pbar.refresh()
-                try:
-                    realized_value, conversion_errors = mapping.convert(
-                        first_param_name,
-                        model=model,
-                        config=model.config,
-                        hf_quantizer=hf_quantizer,
-                        missing_keys=missing_keys,
-                        conversion_errors=conversion_errors,
-                    )
-                    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, missing_keys, disk_offload_folder, disk_offload_index, mapping
-                            )
-                        else:
-                            set_param_for_module(
-                                model,
-                                target_name,
-                                param,
-                                mismatch_keys,
-                                missing_keys,
-                                unexpected_keys,
-                                mapping.distributed_operation,
-                                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 interupted, to avoid wasting time on exit
-            thread_pool.shutdown(wait=False, cancel_futures=True)
+    total_entries = len(param_name_to_load)
+    with logging.tqdm(total=total_entries, desc="Loading weights") as pbar:
+        for first_param_name, mapping in param_name_to_load.items():
+            pbar.update(1)
+            pbar.set_postfix({"Materializing param": first_param_name})
+            pbar.refresh()
+            try:
+                realized_value, misc = mapping.convert(
+                    first_param_name,
+                    model=model,
+                    config=model.config,
+                    hf_quantizer=hf_quantizer,
+                    missing_keys=missing_keys,
+                    misc=misc,
+                )
+                for target_name, param in realized_value.items():
+                    param = param[0] if isinstance(param, list) else param
+                    device_match = device_map_regex.match(target_name)
+                    param_device = device_map[device_match.group()] if device_match else device_map.get("", "cpu")
+                    # Offloading support
+                    if param_device == "disk":
+                        disk_offload_index = offload_and_maybe_resave_param(
+                            target_name, param, missing_keys, disk_offload_folder, disk_offload_index, mapping
+                        )
+                    else:
+                        set_param_for_module(
+                            model,
+                            target_name,
+                            param,
+                            mismatch_keys,
+                            missing_keys,
+                            misc,
+                            unexpected_keys,
+                            mapping.distributed_operation,
+                            hf_quantizer,
+                        )
+
+                # Cleanup the tensors
+                mapping.reset()
+            except SkipLayer:
+                continue
 
     # Keep the current weight conversion mapping for later saving (in case it was coming directly from the user)
     model._weight_conversions = weight_mapping
-    return missing_keys, unexpected_keys, mismatch_keys, disk_offload_index, conversion_errors
+    thread_pool.shutdown(wait=False)
+    return missing_keys, unexpected_keys, mismatch_keys, disk_offload_index, misc
 
 
 def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]):
@@ -1303,7 +960,7 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch
     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)
+        realized_value, misc = 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