Merge branch 'main' into alm-base-model-class

Author: eustlb

.github/dependabot.yml

@@ -0,0 +1,11 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    cooldown:
+      default-days: 7
+    groups:
+      actions:
+        patterns: ["*"]

setup.py

@@ -325,7 +325,7 @@ def run(self):
 
     setup(
         name="transformers",
-        version="5.8.0.dev0",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
+        version="5.10.0.dev0",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
         author="The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/transformers/graphs/contributors)",
         author_email="transformers@huggingface.co",
         description="Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.",

src/transformers/__init__.py

@@ -18,7 +18,7 @@
 # to defer the actual importing for when the objects are requested. This way `import transformers` provides the names
 # in the namespace without actually importing anything (and especially none of the backends).
 
-__version__ = "5.8.0.dev0"
+__version__ = "5.10.0.dev0"
 
 import importlib
 import sys

src/transformers/conversion_mapping.py

@@ -1174,9 +1174,14 @@ def get_model_conversion_mapping(
 
         is_root_model = module_name == ""
         if not is_root_model:
-            # Scope each transform so it only matches keys under this sub-module's prefix.
+            # Scope each transform so it only matches keys under this sub-module's prefix - but we still allow to
+            # arbitrary add/remove base_model_prefix to load ForXXX model from BaseModel and the opposite
+            # Note that we need 2 removeprefix calls here, as only one level of nesting would not have the ending dot to module_name
+            scope_prefix = module_name.removeprefix(model.base_model_prefix)
+            scope_prefix = module_name.removeprefix(".")
             for transform in conversions:
-                transform.scope_prefix = module_name
+                transform.scope_prefix = scope_prefix
+                transform.base_model_prefix = model.base_model_prefix
         weight_conversions.extend(conversions)
 
         seen_identifiers[class_name].append(module_name)

src/transformers/core_model_loading.py

@@ -38,12 +38,12 @@
 
 
 _torch_distributed_available = torch.distributed.is_available()
+if _torch_distributed_available:
+    from torch.distributed.tensor import DTensor
 
 if TYPE_CHECKING:
     from .modeling_utils import LoadStateDictConfig, PreTrainedModel
     from .quantizers import HfQuantizer
-elif _torch_distributed_available:
-    from torch.distributed.tensor import DTensor
 
 logger = get_logger(__name__)
 
@@ -602,6 +602,7 @@ class WeightTransform:
         "_original_target_patterns",
         "_was_used",
         "scope_prefix",
+        "base_model_prefix",
     )
 
     def __init__(self, source_patterns: str | list[str], target_patterns: str | list[str]):
@@ -619,9 +620,11 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list
 
         # Flag to notice if the Transform was used
         self._was_used = False
-        # Optional prefix scope: when set, this transform only applies to keys starting with
-        # `scope_prefix + "."`, stripping / re-attaching the prefix around the pattern match.
+
+        # 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
@@ -689,23 +692,34 @@ def add_tensor(self, target_key: str, source_key: str, source_pattern: str, futu
 
     def _scoped_match(self, source_key: str) -> tuple[str | None, str, re.Match[str]] | None:
         """
-        Apply `scope_prefix` stripping (if any), then match `compiled_sources` against the suffix.
+        Strip `scope_prefix` (if any) from `source_key`, then match `compiled_sources` against the
+        remaining suffix.
 
-        Returns `(prefix_dot, key_to_match, match_object)` when a branch matches, where `prefix_dot` is `None`
-        if `scope_prefix` is unset, else `f"{scope_prefix}."`. Returns `None` when out of scope or unmatched.
+        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.
         """
-        prefix_dot = None
         key_to_match = source_key
+        prefix = None
         if self.scope_prefix is not None:
-            prefix_dot = self.scope_prefix + "."
-            if not source_key.startswith(prefix_dot):
+            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[len(prefix_dot) :]
+            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_dot, key_to_match, match_object)
+        return (prefix, key_to_match, match_object)
 
     def rename_source_key(self, source_key: str) -> tuple[str, str | None]:
         """
@@ -762,6 +776,7 @@ def reverse_transform(self) -> WeightTransform:
             **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]]:
@@ -784,7 +799,9 @@ def materialize_tensors(self) -> dict[str, list[torch.Tensor]]:
                 tensors = [future.result() for future in tensors if future.result() is not None]
             # Sync loading
             elif callable(tensors[0]):
-                tensors = [tensor for func in tensors if (tensor := func()) is not None]
+                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
 
@@ -884,6 +901,7 @@ def reverse_transform(self) -> WeightTransform:
             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
 
 
@@ -1003,9 +1021,10 @@ def _job():
 
     if thread_pool is not None:
         return thread_pool.submit(_job)
-    # Return the Callable here, not the Tensor itself, so we actually delay loading
-    # to avoid saturating cpu memory during Conversion
-    return _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):
@@ -1142,7 +1161,7 @@ def rename_source_key(
     source_key: str,
     weight_renamings: list[WeightRenaming],
     weight_converters: list[WeightConverter],
-    prefix: str | None = None,
+    base_model_prefix: str | None = None,
     meta_state_dict: dict | None = None,
 ) -> tuple[str, str | None]:
     """
@@ -1160,10 +1179,10 @@ def rename_source_key(
             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.
-        prefix (`str`, *optional*):
-            Base-model prefix to add or strip when both `prefix` and `meta_state_dict` are given.
+        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 `prefix` should be added or stripped.
+            Meta state dict used to decide whether `base_model_prefix` should be added or stripped.
 
     Returns:
         `tuple[str, str | None]`: The renamed key and the matched converter's source pattern
@@ -1183,15 +1202,15 @@ def rename_source_key(
         if source_pattern is not None:
             break
 
-    # 3. check if we need to add or remove prefix if necessary (only during loading, not saving)
-    if prefix is not None and meta_state_dict is not None:
+    # 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(prefix)
-            and meta_state_dict.get(re.sub(f"^{prefix}.", "", renamed_key, count=1)) is not None
+            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"^{prefix}.", "", renamed_key, count=1)
-        elif meta_state_dict.get(f"{prefix}.{renamed_key}") is not None:
-            renamed_key = f"{prefix}.{renamed_key}"
+            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, source_pattern
 
@@ -1289,23 +1308,11 @@ def convert_and_load_state_dict_in_model(
     ```
 
     """
-    prefix = model.base_model_prefix
+    base_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
-
-    if load_config.device_map is not None:
-        device_map = load_config.device_map
-    elif device_mesh is not None:
-        if device_mesh.device_type == "cpu":
-            device_map = {"": torch.device("cpu")}
-        else:
-            device_map = {
-                "": torch.device(device_mesh.device_type, getattr(torch, device_mesh.device_type).current_device())
-            }
-    else:
-        device_map = {"": "cpu"}
     disk_offload_folder = load_config.disk_offload_folder
     offload_buffers = load_config.offload_buffers
     dtype_plan = load_config.dtype_plan or {}
@@ -1349,12 +1356,12 @@ def convert_and_load_state_dict_in_model(
     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, prefix=prefix, meta_state_dict=meta_model_state_dict
+            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, [], [], prefix=prefix, meta_state_dict=meta_model_state_dict
+                original_key, [], [], base_model_prefix=base_model_prefix, meta_state_dict=meta_model_state_dict
             )
 
         # 2. finally, collect the tensor into the proper converter

src/transformers/integrations/accelerate.py

@@ -469,7 +469,9 @@ def accelerate_disk_offload(
 
         # Update the weight names according to the `weight_mapping`
         weight_renaming_map = {
-            rename_source_key(k, renamings, [], prefix=model.base_model_prefix, meta_state_dict=meta_state_dict)[0]: k
+            rename_source_key(
+                k, renamings, [], base_model_prefix=model.base_model_prefix, meta_state_dict=meta_state_dict
+            )[0]: k
             for k in weight_map
         }
 

src/transformers/integrations/deepspeed.py

@@ -352,7 +352,7 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     # Preserve metadata from the original state dict
     metadata = getattr(state_dict, "_metadata", None)
 
-    prefix = model.base_model_prefix
+    base_model_prefix = model.base_model_prefix
 
     # Build a meta state dict for matching - only keys/shapes, no actual tensor data
     # This minimizes memory since we don't duplicate the model's parameters
@@ -368,7 +368,7 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
         new_state_dict = {}
         for original_key, tensor in state_dict.items():
             renamed_key, _ = rename_source_key(
-                original_key, renamings, [], prefix=prefix, meta_state_dict=model_state_dict
+                original_key, renamings, [], base_model_prefix=base_model_prefix, meta_state_dict=model_state_dict
             )
             if renamed_key in model_state_dict:
                 new_state_dict[renamed_key] = tensor
@@ -389,7 +389,7 @@ def _apply_weight_conversions_to_state_dict(model, state_dict, weight_mapping):
     for original_key in sorted_keys:
         tensor = state_dict.pop(original_key)
         renamed_key, source_pattern = rename_source_key(
-            original_key, renamings, converters, prefix=prefix, meta_state_dict=model_state_dict
+            original_key, renamings, converters, base_model_prefix=base_model_prefix, meta_state_dict=model_state_dict
         )
 
         # Only process if the renamed key is in the model's state dict

src/transformers/integrations/flash_attention.py

@@ -12,8 +12,9 @@
 def get_target_dtype(query: torch.Tensor, module: torch.nn.Module) -> torch.dtype:
     """If the query is in float32, return a target dtype compatible with flash attention. Return None otherwise."""
     if query.dtype == torch.float32:
-        if torch.is_autocast_enabled("cuda"):
-            return torch.get_autocast_dtype("cuda")
+        device_type = query.device.type
+        if torch.is_autocast_enabled(device_type):
+            return torch.get_autocast_dtype(device_type)
         # Handle the case where the model is quantized
         elif hasattr(module.config, "_is_quantized"):
             return module.config.dtype

src/transformers/modeling_rope_utils.py

@@ -719,7 +719,9 @@ def convert_rope_params_to_dict(self, **kwargs):
         partial_rotary_factor = kwargs.get("partial_rotary_factor", getattr(self, "partial_rotary_factor", None))
         if partial_rotary_factor is not None:
             self.rope_parameters.setdefault("partial_rotary_factor", partial_rotary_factor)
-            self.ignore_keys_at_rope_validation = self.ignore_keys_at_rope_validation | {"partial_rotary_factor"}
+            self.ignore_keys_at_rope_validation = set(self.ignore_keys_at_rope_validation or []) | {
+                "partial_rotary_factor"
+            }
 
         self.standardize_rope_params()
         return kwargs

src/transformers/modeling_utils.py

@@ -4188,6 +4188,12 @@ def from_pretrained(
 
         if distributed_config is not None:
             device_mesh = init_device_mesh(distributed_config)
+            # When using any distributed setup, we simply set the device_map here to the current rank so that we correctly
+            # load the params on the right device
+            if device_mesh.device_type == "cpu":
+                device_map = {"": torch.device("cpu")}
+            else:
+                device_map = {"": torch.device(device_mesh.device_type, int(os.environ["LOCAL_RANK"]))}
         else:
             # Accelerate path
             if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")):
@@ -4342,10 +4348,9 @@ def from_pretrained(
 
         if distributed_config is not None:
             model = distribute_model(model, distributed_config, device_mesh)
-        else:
-            # Accelerate path: auto device mapping
-            if device_map is not None:
-                device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
+        elif device_map is not None:
+            # Expand device_map if it was passed as a `str`, i.e. `device_map="auto"`
+            device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
 
         # Finalize model weight initialization
         active_tp_plan = (
@@ -4907,7 +4912,7 @@ def mark_tied_weights_as_initialized(self, loading_info):
         later as they will be tied (overwritten) anyway.
         This is very important as most embeddings are tied, and they are huge params (vocabularies are often 256k), so
         running inits on them is very costly."""
-        for tied_param in self.all_tied_weights_keys.keys():
+        for tied_param in getattr(self, "all_tied_weights_keys", {}).keys():
             param = self.get_parameter(tied_param)
             setattr(param, "_is_hf_initialized", True)
 
@@ -5041,7 +5046,7 @@ def get_total_byte_count(
 
     total_byte_count = defaultdict(lambda: 0)
     tied_param_names = model.all_tied_weights_keys.keys()
-    tp_plan = model._tp_plan if _is_torch_distributed_initialized() else []
+    tp_plan = model.tp_plan if _is_torch_distributed_initialized() else []
 
     for param_name, device in accelerator_device_map.items():
         # Skip if the parameter has already been accounted for (tied weights)