diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index eb1bafd7355f..c3d2ff5c3027 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -196,6 +196,7 @@ def maybe_distribute_model( if isinstance(distributed_config.tp_plan, dict): model.tp_plan = distributed_config.tp_plan model = apply_tensor_parallelism(model, tp_mesh) + model._tp_size = distributed_config.tp_size elif distributed_config.fsdp_size > 1: fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh diff --git a/src/transformers/integrations/peft.py b/src/transformers/integrations/peft.py index 91ebf98791ce..219fe26c7c28 100644 --- a/src/transformers/integrations/peft.py +++ b/src/transformers/integrations/peft.py @@ -17,7 +17,7 @@ from dataclasses import replace from typing import TYPE_CHECKING, Any, Literal, Optional -from safetensors import safe_open +from transformers.utils.import_utils import is_peft_greater_or_equal from .._typing import PeftConfigLike from ..conversion_mapping import get_model_conversion_mapping @@ -155,9 +155,8 @@ def load_adapter( `find_adapter_config_file` method. """ from peft import PeftType - from peft.utils.save_and_load import _maybe_shard_state_dict_for_tp - from ..modeling_utils import LoadStateDictConfig, _get_resolved_checkpoint_files, load_state_dict + from ..modeling_utils import LoadStateDictConfig, _get_resolved_checkpoint_files if local_files_only: kwargs["local_files_only"] = True @@ -274,44 +273,22 @@ def is_adapter_key(key: str) -> bool: device_map = getattr(self, "hf_device_map", {"": self.device}) - # If the model is tensor parallel, we handle the sharding of the state dict here since the logic in `self._load_pretrained_model` - # is not compatible with the way PEFT adapter should be sharded. has_tp_adapters = False for module in self.modules(): + # Legacy, pre-DTensor TP integration: PEFT stamps a `_tp_info` marker on each TP-sharded LoRA module. tp_info = getattr(module, "_tp_info", None) if tp_info is not None: has_tp_adapters = True break + # DTensor TP integration: the base model itself carries the TP plan, so any adapter injected into it will be + # TP-sharded too; no per-module PEFT marker is needed to detect this. + has_tp_adapters = has_tp_adapters or bool(getattr(self, "_tp_plan", None)) - if has_tp_adapters: - all_pointer = set() - if adapter_state_dict is not None: - merged_state_dict = adapter_state_dict - elif ( - checkpoint_files is not None - and checkpoint_files[0].endswith(".safetensors") - and adapter_state_dict is None - ): - merged_state_dict = {} - for file in checkpoint_files: - file_pointer = safe_open(file, framework="pt", device="cpu") - all_pointer.add(file_pointer) - for k in file_pointer.keys(): - merged_state_dict[k] = file_pointer.get_tensor(k) - # Checkpoints are .bin - elif checkpoint_files is not None: - merged_state_dict = {} - for ckpt_file in checkpoint_files: - merged_state_dict.update(load_state_dict(ckpt_file)) - else: - raise ValueError("Neither a state dict nor checkpoint files were found.") - - adapter_state_dict = merged_state_dict - - if any(not isinstance(v, torch.Tensor) for v in adapter_state_dict.values()): - raise ValueError("Expected all values in the adapter state dict to be tensors.") - - _maybe_shard_state_dict_for_tp(self, adapter_state_dict, adapter_name) + if has_tp_adapters and not is_peft_greater_or_equal("0.20.1", accept_dev=True): + raise ValueError( + "Loading a tensor-parallel PEFT adapter requires peft >= 0.20.1, please upgrade your peft " + "installation." + ) load_config = replace( load_config,