From 799ac94d45b37cfbea8e180ae47e1846e1af526b Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 17 Jun 2026 03:51:49 +0000 Subject: [PATCH 01/24] add distributed config --- src/transformers/__init__.py | 1 + src/transformers/configuration_utils.py | 7 + .../distributed/configuration_utils.py | 142 ++++++++---------- src/transformers/modeling_utils.py | 17 ++- tests/test_distributed_config.py | 97 ++++++++++++ 5 files changed, 182 insertions(+), 82 deletions(-) create mode 100644 tests/test_distributed_config.py diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index fa6dbb438702..9c442318800b 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -104,6 +104,7 @@ "debug_utils": [], "dependency_versions_check": [], "dependency_versions_table": [], + "distributed": [], "dynamic_module_utils": [], "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"], "feature_extraction_utils": ["BatchFeature", "FeatureExtractionMixin"], diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 2b6fdce06e28..d9fcdad17256 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -149,6 +149,9 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): naming of attributes. - **base_model_tp_plan** (`dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a tensor parallel plan applied to the sub-module when `model.tensor_parallel` is called. + - **base_model_fsdp_plan** (`dict[Any, str]`) -- A dict that maps sub-modules of a base model to an FSDP2 + sharding strategy (e.g. `"free_full_weight"` / `"keep_full_weight"`). Keys can be wildcard module paths + (e.g. `"layers.*"`) or tuples of paths (grouped into a single `fully_shard` call). - **base_model_pp_plan** (`dict[str, tuple[list[str]]]`) -- A dict that maps child-modules of a base model to a pipeline parallel plan that enables users to place the child-module on the appropriate device. @@ -220,6 +223,7 @@ class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin): keys_to_ignore_at_inference: ClassVar[list[str]] = [] attribute_map: ClassVar[dict[str, str]] = {} base_model_tp_plan: ClassVar[dict[str, Any] | None] = None + base_model_fsdp_plan: ClassVar[dict[Any, str] | None] = None base_model_pp_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None base_model_ep_plan: ClassVar[dict[str, Sequence[list[str]]] | None] = None _auto_class: ClassVar[str | None] = None @@ -1022,6 +1026,9 @@ def to_dict(self) -> dict[str, Any]: # Pop "kwargs" since they are unpacked and set in the post init output.pop("kwargs", None) + if "distributed_config" in output and hasattr(output["distributed_config"], "to_dict"): + output["distributed_config"] = output["distributed_config"].to_dict() + def to_list(value): if isinstance(value, tuple): value = [to_list(item) for item in value] diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index 7726d9f3290d..9307398f6d37 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -12,99 +12,79 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import json import os -from dataclasses import dataclass -from typing import Any +from dataclasses import asdict, dataclass + +from ..utils import is_torch_available + + +if is_torch_available(): + import torch @dataclass class DistributedConfig: """ - Base class for distributed configs + Configuration for native distributed training (FSDP2 + TP). + + Args: + tp_size (`int`, *optional*): + Number of devices for tensor parallelism. If `None` and `fsdp_size` is set, defaults to 1. + tp_plan (`dict`, *optional*): + Tensor parallel sharding plan. Leave as `None` to select the model's SP/TP and EP plan + (see ``select_parallel_plan``). Set explicitly to override. + enable_sequence_parallel (`bool`, *optional*, defaults to `False`): + Select ``base_model_sp_plan`` (dense-only) or ``base_model_sp_ep_plan`` (with EP). + enable_expert_parallel (`bool`, *optional*, defaults to `False`): + Select ``base_model_tp_ep_plan`` or ``base_model_sp_ep_plan`` when combined with SP flag. + fsdp_size (`int`, *optional*): + Number of devices for FSDP (data parallelism). If `None` and `tp_size` is set, defaults to 1. + fsdp_cpu_offload (`bool`, *optional*, defaults to `False`): + Whether to enable CPU offloading for FSDP2. + fsdp_mixed_precision (`bool`, *optional*, defaults to `False`): + Whether to enable mixed precision for FSDP2. """ + tp_size: int | None = None + tp_plan: dict[str, str] | None = None + enable_sequence_parallel: bool = False enable_expert_parallel: bool = False - # TODO: add tp_plan, pp_plan, device_mesh etc.. + fsdp_size: int | None = None + fsdp_cpu_offload: bool = False + fsdp_mixed_precision: bool = False + + def __post_init__(self): + if self.tp_size is None and self.fsdp_size is None: + return + + if self.tp_size is None: + self.tp_size = 1 + if self.fsdp_size is None: + self.fsdp_size = 1 + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + if self.tp_size * self.fsdp_size != world_size: + raise RuntimeError( + f"tp_size ({self.tp_size}) * fsdp_size ({self.fsdp_size}) is not equal to world_size ({world_size})" + ) @classmethod - def from_dict(cls, config_dict, **kwargs): - """ - Constructs a DistributedConfig instance from a dictionary of parameters. - Args: - config_dict (Dict[str, Any]): Dictionary containing configuration parameters. - **kwargs: Additional keyword arguments to override dictionary values. - Returns: - DistributedConfig: Instance of DistributedConfig constructed from the dictionary. - """ - config = cls(**config_dict) - to_remove = [] - for key, value in kwargs.items(): - if hasattr(config, key): - setattr(config, key, value) - to_remove.append(key) - for key in to_remove: - kwargs.pop(key, None) - return config - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file + def from_dict(cls, config_dict: dict, **kwargs) -> "DistributedConfig": + merged = {**config_dict, **kwargs} + valid_keys = {f.name for f in cls.__dataclass_fields__.values()} + return cls(**{k: v for k, v in merged.items() if k in valid_keys}) + + def to_dict(self) -> dict: + return asdict(self) + + def to_json_string(self) -> str: + return json.dumps(self.to_dict(), indent=2) + "\n" + def to_json_file(self, json_file_path: str | os.PathLike): - """ - Save this instance to a JSON file. - Args: - json_file_path (`str` or `os.PathLike`): - Path to the JSON file in which this configuration instance's parameters will be saved. - use_diff (`bool`, *optional*, defaults to `True`): - If set to `True`, only the difference between the config instance and the default - `QuantizationConfig()` is serialized to JSON file. - """ - with open(json_file_path, "w", encoding="utf-8") as writer: - config_dict = self.to_dict() - json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" - - writer.write(json_string) - - def to_dict(self) -> dict[str, Any]: - """ - Serializes this instance to a Python dictionary. Returns: - `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. - """ - return copy.deepcopy(self.__dict__) - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__ - def __iter__(self): - """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" - yield from copy.deepcopy(self.__dict__).items() - - # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ + with open(json_file_path, "w", encoding="utf-8") as f: + f.write(self.to_json_string()) + def __repr__(self): return f"{self.__class__.__name__} {self.to_json_string()}" - - def to_json_string(self): - """ - Serializes this instance to a JSON formatted string. - Returns: - str: JSON formatted string representing the configuration instance. - """ - return json.dumps(self.__dict__, indent=2) + "\n" - - def update(self, **kwargs): - """ - Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, - returning all the unused kwargs. - Args: - kwargs (`Dict[str, Any]`): - Dictionary of attributes to tentatively update this class. - Returns: - `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. - """ - to_remove = [] - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - to_remove.append(key) - - # Remove all the attributes that were updated, without modifying the input dict - unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} - return unused_kwargs diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 1b8c93328f43..5574e6fb9231 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1251,6 +1251,10 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH # models, this attribute is currently defined in respective model code. For base models, it comes from # `config.base_model_pp_plan` during `post_init`. _pp_plan: dict[str, tuple[str, str]] = None + # FSDP2 sharding plan of the form `{"layers.*": "free_full_weight"}`. For top-level models, this attribute is + # defined on the head class (e.g. `*ForCausalLM`). For base models, it comes from `config.base_model_fsdp_plan` + # during `post_init`. + _fsdp_plan: dict[str, str] = None # Advanced functionalities support supports_gradient_checkpointing: bool = False @@ -1391,13 +1395,18 @@ def post_init(self): correctly in the case of composite models (that is, the top level model should know about those properties from its children). """ # Attach the different parallel plans and tied weight keys to the top-most model, so that everything is - # easily available + # easily available. Seed with the class-level FSDP plan before the instance attribute shadows it. + cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} self._tp_plan, self._ep_plan, self._pp_plan = {}, {}, {} + self._fsdp_plan = dict(cls_fsdp_plan) # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} self._ep_plan = self.config.base_model_ep_plan.copy() if self.config.base_model_ep_plan is not None else {} + self._fsdp_plan = ( + self.config.base_model_fsdp_plan.copy() if self.config.base_model_fsdp_plan is not None else {} + ) # Current submodel should register its tied weights self.all_tied_weights_keys = self.get_expanded_tied_weights_keys(all_submodels=False) # Current submodel should register its `_keep_in_fp32_modules` @@ -1421,6 +1430,8 @@ def post_init(self): self._tp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) if plan := getattr(module, "_pp_plan", None): self._pp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + if plan := getattr(module, "_fsdp_plan", None): + self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) # Always attach the keys of the children (if the children's config says to NOT tie, then it's empty) if tied_keys := getattr(module, "all_tied_weights_keys", None): self.all_tied_weights_keys.update({f"{name}.{k}": f"{name}.{v}" for k, v in tied_keys.copy().items()}) @@ -1457,6 +1468,10 @@ def tp_plan(self) -> dict[str, str]: return self._ep_plan return self._tp_plan + @property + def fsdp_plan(self) -> dict[str, str]: + return self._fsdp_plan + @property def pp_plan(self) -> dict[str, tuple[str, str]]: return self._pp_plan diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py new file mode 100644 index 000000000000..d6b2b0ab2bd5 --- /dev/null +++ b/tests/test_distributed_config.py @@ -0,0 +1,97 @@ +import json +import tempfile + +from transformers.distributed import DistributedConfig + + +class TestDistributedConfig: + def test_2d_parallelism(self): + dc = DistributedConfig(tp_size=2, fsdp_size=2) + assert dc.tp_size == 2 + assert dc.fsdp_size == 2 + assert dc.tp_plan is None + assert dc.fsdp_cpu_offload is False + assert dc.fsdp_mixed_precision is False + + def test_tp_only_defaults_fsdp_to_1(self): + dc = DistributedConfig(tp_size=4) + assert dc.tp_size == 4 + assert dc.fsdp_size == 1 + assert dc.tp_plan is None + + def test_fsdp_only_defaults_tp_to_1(self): + dc = DistributedConfig(fsdp_size=4) + assert dc.tp_size == 1 + assert dc.fsdp_size == 4 + assert dc.fsdp_cpu_offload is False + assert dc.fsdp_mixed_precision is False + assert dc.tp_plan is None + + def test_empty_config(self): + dc = DistributedConfig() + assert dc.tp_size is None + assert dc.fsdp_size is None + assert dc.tp_plan is None + assert dc.fsdp_cpu_offload is False + assert dc.fsdp_mixed_precision is False + + def test_from_dict(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) + assert dc.tp_size == 2 + assert dc.fsdp_size == 4 + assert dc.tp_plan is None + + def test_from_dict_ignores_unknown_keys(self): + dc = DistributedConfig.from_dict({"tp_size": 2, "unknown_key": 42}) + assert dc.tp_size == 2 + assert not hasattr(dc, "unknown_key") + + def test_from_dict_kwargs_override(self): + dc = DistributedConfig.from_dict({"tp_size": 2}, tp_size=8) + assert dc.tp_size == 8 + + def test_to_dict(self): + dc = DistributedConfig(tp_size=2, fsdp_size=4) + d = dc.to_dict() + assert d == { + "tp_size": 2, + "tp_plan": None, + "enable_sequence_parallel": False, + "enable_expert_parallel": False, + "fsdp_size": 4, + "fsdp_cpu_offload": False, + "fsdp_mixed_precision": False, + } + + def test_to_dict_is_a_copy(self): + dc = DistributedConfig(tp_plan={"layer": "colwise"}) + d = dc.to_dict() + d["tp_plan"]["layer"] = "rowwise" + assert dc.tp_plan["layer"] == "colwise" + + def test_to_json_string(self): + dc = DistributedConfig(tp_size=2, fsdp_size=2) + s = dc.to_json_string() + parsed = json.loads(s) + assert parsed["tp_size"] == 2 + assert parsed["fsdp_size"] == 2 + + def test_to_json_file(self): + dc = DistributedConfig(tp_size=4) + with tempfile.NamedTemporaryFile(mode="r", suffix=".json", delete=False) as f: + dc.to_json_file(f.name) + f.seek(0) + parsed = json.load(f) + assert parsed["tp_size"] == 4 + assert parsed["tp_plan"] is None + + def test_roundtrip_dict(self): + original = DistributedConfig(tp_size=2, fsdp_size=4) + restored = DistributedConfig.from_dict(original.to_dict()) + assert original == restored + + def test_repr(self): + dc = DistributedConfig(tp_size=2) + r = repr(dc) + assert "DistributedConfig" in r + assert '"tp_size": 2' in r From 22d4b52bd6e7f8a4635a080965587ff08b7454bf Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 17 Jun 2026 03:56:34 +0000 Subject: [PATCH 02/24] Add native FSDP2 module and migrate FSDP imports (Phase A PR-2). Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep integrations/fsdp.py as a backward-compatible re-export, and update core call sites to import from transformers.distributed.fsdp. --- src/transformers/distributed/__init__.py | 2 + src/transformers/distributed/fsdp.py | 277 ++++++++++++++++++++ src/transformers/generation/utils.py | 19 +- src/transformers/integrations/accelerate.py | 2 +- src/transformers/integrations/fsdp.py | 85 +----- src/transformers/trainer.py | 2 +- src/transformers/trainer_seq2seq.py | 2 +- tests/test_distributed_fsdp.py | 63 +++++ 8 files changed, 369 insertions(+), 83 deletions(-) create mode 100644 src/transformers/distributed/fsdp.py create mode 100644 tests/test_distributed_fsdp.py diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py index ba6db8358d2b..e179a2306156 100644 --- a/src/transformers/distributed/__init__.py +++ b/src/transformers/distributed/__init__.py @@ -19,6 +19,7 @@ _import_structure = { "configuration_utils": ["DistributedConfig"], + "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"], } @@ -26,6 +27,7 @@ from .configuration_utils import ( DistributedConfig, ) + from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan else: import sys diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py new file mode 100644 index 000000000000..95f9b3403614 --- /dev/null +++ b/src/transformers/distributed/fsdp.py @@ -0,0 +1,277 @@ +# Copyright 2024 The HuggingFace 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. +from __future__ import annotations + +import inspect +import os +from typing import TYPE_CHECKING, Any + +from ..utils import is_torch_available, is_torch_greater_or_equal, logging, strtobool +from ..utils.quantization_config import QuantizationMethod +from ..integrations.tensor_parallel import replace_layer_number_by_wildcard + + +if TYPE_CHECKING: + import torch.nn as nn + + from .configuration_utils import DistributedConfig + +if is_torch_available(): + import torch + +if is_torch_available() and is_torch_greater_or_equal("2.6"): + from torch.distributed._composable.fsdp import fully_shard + from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy + +logger = logging.get_logger(__name__) + + +def is_fsdp_enabled() -> bool: + """Check if FSDP is active via Accelerate (env var based) — covers FSDP1 only.""" + if not is_torch_available(): + return False + + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 + and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 + ) + + +def is_fsdp_managed_module(module: nn.Module) -> bool: + """Check if a module is managed by FSDP (1 or 2).""" + if not is_torch_available(): + return False + if not torch.distributed.is_available(): + return False + + # FSDP2: attribute set by apply_fsdp2() + if getattr(module, "_is_fsdp_managed_module", False): + return True + # FSDP1: wrapped by FullyShardedDataParallel + try: + from torch.distributed.fsdp import FullyShardedDataParallel + except ImportError: + return False + return isinstance(module, FullyShardedDataParallel) + + +def _get_fsdp_policy_kwargs(distributed_config: DistributedConfig | None) -> dict[str, Any]: + """Build ``fully_shard`` policy kwargs from ``DistributedConfig`` runtime flags.""" + if distributed_config is None: + return {} + + fsdp_policy_kwargs = {} + if distributed_config.fsdp_cpu_offload: + fsdp_policy_kwargs["offload_policy"] = CPUOffloadPolicy() + if distributed_config.fsdp_mixed_precision: + fsdp_policy_kwargs["mp_policy"] = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) + return fsdp_policy_kwargs + + +def is_norm_and_head_pair(modules: list[tuple[str, Any]]) -> bool: + """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" + if len(modules) != 2: + return False + module_names = [name for name, _ in modules] + has_final_norm = any(name == "norm" or name.endswith(".norm") for name in module_names) + has_output_head = any( + name in {"lm_head", "embed_tokens"} or name.endswith((".lm_head", ".embed_tokens")) for name in module_names + ) + return has_final_norm and has_output_head + + +def _resolve_tied_embed_lm_head_plan( + fsdp_plan: dict[str, str], + *, + tie_word_embeddings: bool, +) -> dict[str, str]: + """ + Rewrite the plan so tied embed/lm_head weights are wrapped once. + Example: + {"model.embed_tokens": "free_full_weight", + "model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "lm_head": "keep_full_weight"} + -> + {"model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "model.embed_tokens": "keep_full_weight"} + """ + if not tie_word_embeddings: + return fsdp_plan + + embed_module_key = next( + (key for key in fsdp_plan if key == "embed_tokens" or key.endswith(".embed_tokens")), + None, + ) + if embed_module_key is None: + return fsdp_plan + + adapted_plan = {} + for key, sharding_strategy in fsdp_plan.items(): + if key == embed_module_key: + continue + if key == "lm_head" and sharding_strategy == "keep_full_weight": + adapted_plan[embed_module_key] = sharding_strategy + else: + adapted_plan[key] = sharding_strategy + return adapted_plan + + +def expand_fsdp_plan(model, fsdp_plan: dict[str, str]) -> list[tuple[str, nn.Module, str]]: + """Expand plan keys into ``(module_name, module, sharding_strategy)`` shard targets.""" + module_lookup = dict(model.named_modules()) + shard_targets: list[tuple[str, nn.Module, str]] = [] + + for plan_key, sharding_strategy in fsdp_plan.items(): + if plan_key in module_lookup: + shard_targets.append((plan_key, module_lookup[plan_key], sharding_strategy)) + continue + + for module_name, module in module_lookup.items(): + if replace_layer_number_by_wildcard(module_name) == plan_key: + shard_targets.append((module_name, module, sharding_strategy)) + + return shard_targets + + +def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) -> None: + """ + Verify the FSDP plan of the model, log a warning if plan keys were not applied or strategies are invalid. + """ + if not fsdp_plan: + return + + name_lookup = dict.fromkeys(module_names) + unused_rules: dict[str, str] = {} + invalid_strategies: dict[str, str] = {} + + for key, strategy in fsdp_plan.items(): + if strategy not in {"free_full_weight", "keep_full_weight"}: + invalid_strategies[key] = strategy + continue + if key not in name_lookup and not any(replace_layer_number_by_wildcard(name) == key for name in name_lookup): + unused_rules[key] = strategy + + if invalid_strategies: + logger.warning(f"The following FSDP entries have unknown strategies: {invalid_strategies}") + if unused_rules: + logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}") + + +def apply_fully_shard_data_parallel(model, fsdp_mesh): + """ + Apply FSDP2 (fully_shard) to a model. + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for FSDP support") + + if not is_torch_greater_or_equal("2.6"): + raise OSError("FSDP2 requires torch>=2.6") + + fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {}) + if not fsdp_plan: + raise ValueError( + f"{type(model).__name__} does not have a FSDP2 plan declared. Set " + "`base_model_fsdp_plan` on the config and `_fsdp_plan` on the head class." + ) + + distributed_config = getattr(model.config, "distributed_config", None) + fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) + tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) + + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, tie_word_embeddings=tie_word_embeddings) + shard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) + + reshard_targets = [] + no_reshard_targets = [] + for module_name, module, sharding_strategy in shard_targets: + if sharding_strategy == "keep_full_weight": + no_reshard_targets.append((module_name, module)) + else: + reshard_targets.append((module_name, module)) + + for module_name, module in reshard_targets: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) + logger.debug(f"Applied fully_shard to {module_name} (reshard=True)") + + # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) + # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. + if is_norm_and_head_pair(no_reshard_targets): + module_names = [name for name, _ in no_reshard_targets] + modules = [module for _, module in no_reshard_targets] + fully_shard(modules, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) + logger.debug(f"Grouped tail {module_names} (reshard=False)") + else: + for module_name, module in no_reshard_targets: + fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) + logger.debug(f"Applied fully_shard to {module_name} (reshard=False)") + + # Apply FSDP2 to the root module + fully_shard(model, mesh=fsdp_mesh, **fsdp_policy_kwargs) + + logger.info(f"FSDP2 applied to model via _fsdp_plan: {len(fsdp_plan)} entries") + + # Used by generation code to detect FSDP and enable synced_gpus. + model._is_fsdp_managed_module = True + + if tie_word_embeddings and hasattr(model, "tie_weights"): + model.tie_weights() + + return model + + +# ========================= PEFT compatibility ========================= +# TODO(3outeille): make sure new FSDP works with PEFT +def get_fsdp_ckpt_kwargs(): + """ + Returns checkpoint kwargs for FSDP model saving. + + Checks if the `adapter_only` parameter is supported by `save_fsdp_model` from accelerate + and returns the appropriate kwargs. + """ + from accelerate.utils import save_fsdp_model + + if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): + return {"adapter_only": True} + else: + return {} + + +def update_fsdp_plugin_peft(model, accelerator): + """ + Updates the FSDP plugin for PEFT LoRA/QLoRA compatibility. + + When using FSDP with PEFT LoRA, the auto wrap policy needs to be updated to additionally wrap + LoRA trainable layers separately. When using FSDP with QLoRA, the mixed precision policy needs + to be updated to use the quantization storage data type. + """ + from peft import PeftConfig + from peft.utils.other import fsdp_auto_wrap_policy + + if isinstance(model.active_peft_config, PeftConfig): + accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) + if ( + getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES + and model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point + ): + accelerator.state.fsdp_plugin.set_mixed_precision( + model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True + ) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 02808e34b82f..ed14b60fbfa8 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -33,6 +33,7 @@ QuantizedCache, StaticCache, ) +from ..distributed.fsdp import is_fsdp_managed_module from ..dynamic_module_utils import ( check_python_requirements, get_cached_module_file, @@ -40,7 +41,6 @@ resolve_trust_remote_code, ) from ..integrations.deepspeed import is_deepspeed_zero3_enabled -from ..integrations.fsdp import is_fsdp_managed_module from ..masking_utils import create_masks_for_generate from ..tokenization_python import ExtensionsTrie from ..utils import ( @@ -118,6 +118,18 @@ logger = logging.get_logger(__name__) + +def _get_dist_world_size() -> int: + if ( + dist.is_available() + and hasattr(dist, "is_initialized") + and dist.is_initialized() + and hasattr(dist, "get_world_size") + ): + return dist.get_world_size() + return 1 + + if is_accelerate_available(): from accelerate.hooks import AlignDevicesHook, add_hook_to_module @@ -2165,7 +2177,7 @@ def _extract_generation_mode_kwargs( "assistant_model": assistant_model, "streamer": streamer, } - world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 # type: ignore + world_size = _get_dist_world_size() generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2612,7 +2624,8 @@ def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, # The following logic allows an early break if all peers finished generating their sequence this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) # send 0.0 if we finished, 1.0 otherwise - dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # type: ignore + if hasattr(dist, "all_reduce") and hasattr(dist, "ReduceOp"): + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # did all peers finish? the reduced sum will be 0.0 then if this_peer_finished_flag.item() == 0.0: return False diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index 142489c4c21a..63d600e3e88a 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -34,7 +34,7 @@ ) from ..utils.quantization_config import QuantizationMethod from .deepspeed import is_deepspeed_zero3_enabled -from .fsdp import is_fsdp_enabled +from ..distributed.fsdp import is_fsdp_enabled if is_torch_available(): diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py index 7cda7ad55acc..086d5f06cb0a 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/integrations/fsdp.py @@ -11,82 +11,13 @@ # 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. -from __future__ import annotations +"""Backward-compatible re-exports. Prefer ``transformers.distributed.fsdp``.""" -import inspect -import os -from typing import TYPE_CHECKING +from ..distributed.fsdp import ( + get_fsdp_ckpt_kwargs, + is_fsdp_enabled, + is_fsdp_managed_module, + update_fsdp_plugin_peft, +) -from ..utils import is_torch_available, strtobool -from ..utils.quantization_config import QuantizationMethod - - -if TYPE_CHECKING: - from torch import nn - - -def is_fsdp_managed_module(module: nn.Module) -> bool: - if not is_torch_available(): - return False - - import torch - - if not torch.distributed.is_available(): - return False - - import torch.distributed.fsdp - - return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr( - module, "_is_fsdp_managed_module", False - ) - - -def is_fsdp_enabled(): - if is_torch_available(): - import torch - - return ( - torch.distributed.is_available() - and torch.distributed.is_initialized() - and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 - and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 - ) - - return False - - -def get_fsdp_ckpt_kwargs(): - """ - Returns checkpoint kwargs for FSDP model saving. - - Checks if the `adapter_only` parameter is supported by `save_fsdp_model` from accelerate - and returns the appropriate kwargs. - """ - from accelerate.utils import save_fsdp_model - - if "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): - return {"adapter_only": True} - else: - return {} - - -def update_fsdp_plugin_peft(model, accelerator): - """ - Updates the FSDP plugin for PEFT LoRA/QLoRA compatibility. - - When using FSDP with PEFT LoRA, the auto wrap policy needs to be updated to additionally wrap - LoRA trainable layers separately. When using FSDP with QLoRA, the mixed precision policy needs - to be updated to use the quantization storage data type. - """ - from peft import PeftConfig - from peft.utils.other import fsdp_auto_wrap_policy - - if isinstance(model.active_peft_config, PeftConfig): - accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) - if ( - getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES - and model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point - ): - accelerator.state.fsdp_plugin.set_mixed_precision( - model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True - ) +__all__ = ["get_fsdp_ckpt_kwargs", "is_fsdp_enabled", "is_fsdp_managed_module", "update_fsdp_plugin_peft"] diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 672cb7151dfa..c159d69de295 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -66,7 +66,7 @@ is_deepspeed_available, propagate_args_to_deepspeed, ) -from .integrations.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft +from .distributed.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .integrations.liger import apply_liger_kernel from .integrations.neftune import activate_neftune, deactivate_neftune from .integrations.peft import MIN_PEFT_VERSION diff --git a/src/transformers/trainer_seq2seq.py b/src/transformers/trainer_seq2seq.py index ada588adbd21..4f4045ae574f 100644 --- a/src/transformers/trainer_seq2seq.py +++ b/src/transformers/trainer_seq2seq.py @@ -24,7 +24,7 @@ from .generation.configuration_utils import GenerationConfig from .integrations.deepspeed import is_deepspeed_zero3_enabled -from .integrations.fsdp import is_fsdp_managed_module +from .distributed.fsdp import is_fsdp_managed_module from .trainer import Trainer from .utils import is_datasets_available, logging diff --git a/tests/test_distributed_fsdp.py b/tests/test_distributed_fsdp.py new file mode 100644 index 000000000000..fce25521ae90 --- /dev/null +++ b/tests/test_distributed_fsdp.py @@ -0,0 +1,63 @@ +import torch +from torch import nn + +from transformers.distributed.fsdp import ( + _resolve_tied_embed_lm_head_plan, + expand_fsdp_plan, + is_fsdp_managed_module, + is_norm_and_head_pair, + verify_fsdp_plan, +) + + +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = nn.Embedding(8, 4) + self.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + self.norm = nn.LayerNorm(4) + self.lm_head = nn.Linear(4, 8) + + +class TestDistributedFSDP: + def test_verify_fsdp_plan_warns_on_unknown_strategy(self, caplog): + verify_fsdp_plan(list(TinyModel().named_modules()), {"layers.*": "bad_strategy"}) + assert "unknown strategies" in caplog.text + + def test_verify_fsdp_plan_warns_on_unused_key(self, caplog): + verify_fsdp_plan(list(TinyModel().named_modules()), {"missing.*": "free_full_weight"}) + assert "not applied to any module" in caplog.text + + def test_expand_fsdp_plan_wildcard(self): + model = TinyModel() + targets = expand_fsdp_plan(model, {"layers.*": "free_full_weight"}) + assert len(targets) == 2 + assert all(strategy == "free_full_weight" for _, _, strategy in targets) + + def test_expand_fsdp_plan_exact_key(self): + model = TinyModel() + targets = expand_fsdp_plan(model, {"norm": "keep_full_weight"}) + assert targets == [("norm", model.norm, "keep_full_weight")] + + def test_resolve_tied_embed_lm_head_plan(self): + plan = { + "model.embed_tokens": "free_full_weight", + "model.layers.*": "free_full_weight", + "model.norm": "keep_full_weight", + "lm_head": "keep_full_weight", + } + adapted = _resolve_tied_embed_lm_head_plan(plan, tie_word_embeddings=True) + assert "lm_head" not in adapted + assert adapted["model.embed_tokens"] == "keep_full_weight" + assert adapted["model.layers.*"] == "free_full_weight" + + def test_is_norm_and_head_pair(self): + model = TinyModel() + assert is_norm_and_head_pair([("norm", model.norm), ("lm_head", model.lm_head)]) + assert not is_norm_and_head_pair([("norm", model.norm)]) + + def test_is_fsdp_managed_module_flag(self): + model = TinyModel() + assert not is_fsdp_managed_module(model) + model._is_fsdp_managed_module = True + assert is_fsdp_managed_module(model) From 4bfd1a6432dbd18f542707a1169014b8dc474f7c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 17 Jun 2026 04:19:56 +0000 Subject: [PATCH 03/24] linting --- src/transformers/distributed/fsdp.py | 2 +- src/transformers/integrations/accelerate.py | 2 +- src/transformers/integrations/fsdp.py | 1 + src/transformers/trainer.py | 2 +- src/transformers/trainer_seq2seq.py | 2 +- tests/test_distributed_fsdp.py | 1 - 6 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 95f9b3403614..abf9a7a29df3 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -17,9 +17,9 @@ import os from typing import TYPE_CHECKING, Any +from ..integrations.tensor_parallel import replace_layer_number_by_wildcard from ..utils import is_torch_available, is_torch_greater_or_equal, logging, strtobool from ..utils.quantization_config import QuantizationMethod -from ..integrations.tensor_parallel import replace_layer_number_by_wildcard if TYPE_CHECKING: diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py index 63d600e3e88a..69df6a618bb2 100644 --- a/src/transformers/integrations/accelerate.py +++ b/src/transformers/integrations/accelerate.py @@ -26,6 +26,7 @@ from safetensors import safe_open from safetensors.torch import save_file +from ..distributed.fsdp import is_fsdp_enabled from ..utils import ( is_accelerate_available, is_torch_available, @@ -34,7 +35,6 @@ ) from ..utils.quantization_config import QuantizationMethod from .deepspeed import is_deepspeed_zero3_enabled -from ..distributed.fsdp import is_fsdp_enabled if is_torch_available(): diff --git a/src/transformers/integrations/fsdp.py b/src/transformers/integrations/fsdp.py index 086d5f06cb0a..6cf34db2cadd 100644 --- a/src/transformers/integrations/fsdp.py +++ b/src/transformers/integrations/fsdp.py @@ -20,4 +20,5 @@ update_fsdp_plugin_peft, ) + __all__ = ["get_fsdp_ckpt_kwargs", "is_fsdp_enabled", "is_fsdp_managed_module", "update_fsdp_plugin_peft"] diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index c159d69de295..84ea3447db48 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -55,6 +55,7 @@ from .configuration_utils import PreTrainedConfig from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator from .debug_utils import DebugOption, DebugUnderflowOverflow +from .distributed.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .feature_extraction_sequence_utils import SequenceFeatureExtractor from .feature_extraction_utils import FeatureExtractionMixin from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend @@ -66,7 +67,6 @@ is_deepspeed_available, propagate_args_to_deepspeed, ) -from .distributed.fsdp import get_fsdp_ckpt_kwargs, update_fsdp_plugin_peft from .integrations.liger import apply_liger_kernel from .integrations.neftune import activate_neftune, deactivate_neftune from .integrations.peft import MIN_PEFT_VERSION diff --git a/src/transformers/trainer_seq2seq.py b/src/transformers/trainer_seq2seq.py index 4f4045ae574f..c3907ed2556c 100644 --- a/src/transformers/trainer_seq2seq.py +++ b/src/transformers/trainer_seq2seq.py @@ -22,9 +22,9 @@ from torch import nn from torch.utils.data import Dataset +from .distributed.fsdp import is_fsdp_managed_module from .generation.configuration_utils import GenerationConfig from .integrations.deepspeed import is_deepspeed_zero3_enabled -from .distributed.fsdp import is_fsdp_managed_module from .trainer import Trainer from .utils import is_datasets_available, logging diff --git a/tests/test_distributed_fsdp.py b/tests/test_distributed_fsdp.py index fce25521ae90..93a270ce8a2c 100644 --- a/tests/test_distributed_fsdp.py +++ b/tests/test_distributed_fsdp.py @@ -1,4 +1,3 @@ -import torch from torch import nn from transformers.distributed.fsdp import ( From 9487bddaac3a7885ff3de1728b1f2da8f35dea72 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 17 Jun 2026 04:21:25 +0000 Subject: [PATCH 04/24] unecessary --- tests/test_distributed_fsdp.py | 62 ---------------------------------- 1 file changed, 62 deletions(-) delete mode 100644 tests/test_distributed_fsdp.py diff --git a/tests/test_distributed_fsdp.py b/tests/test_distributed_fsdp.py deleted file mode 100644 index 93a270ce8a2c..000000000000 --- a/tests/test_distributed_fsdp.py +++ /dev/null @@ -1,62 +0,0 @@ -from torch import nn - -from transformers.distributed.fsdp import ( - _resolve_tied_embed_lm_head_plan, - expand_fsdp_plan, - is_fsdp_managed_module, - is_norm_and_head_pair, - verify_fsdp_plan, -) - - -class TinyModel(nn.Module): - def __init__(self): - super().__init__() - self.embed_tokens = nn.Embedding(8, 4) - self.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) - self.norm = nn.LayerNorm(4) - self.lm_head = nn.Linear(4, 8) - - -class TestDistributedFSDP: - def test_verify_fsdp_plan_warns_on_unknown_strategy(self, caplog): - verify_fsdp_plan(list(TinyModel().named_modules()), {"layers.*": "bad_strategy"}) - assert "unknown strategies" in caplog.text - - def test_verify_fsdp_plan_warns_on_unused_key(self, caplog): - verify_fsdp_plan(list(TinyModel().named_modules()), {"missing.*": "free_full_weight"}) - assert "not applied to any module" in caplog.text - - def test_expand_fsdp_plan_wildcard(self): - model = TinyModel() - targets = expand_fsdp_plan(model, {"layers.*": "free_full_weight"}) - assert len(targets) == 2 - assert all(strategy == "free_full_weight" for _, _, strategy in targets) - - def test_expand_fsdp_plan_exact_key(self): - model = TinyModel() - targets = expand_fsdp_plan(model, {"norm": "keep_full_weight"}) - assert targets == [("norm", model.norm, "keep_full_weight")] - - def test_resolve_tied_embed_lm_head_plan(self): - plan = { - "model.embed_tokens": "free_full_weight", - "model.layers.*": "free_full_weight", - "model.norm": "keep_full_weight", - "lm_head": "keep_full_weight", - } - adapted = _resolve_tied_embed_lm_head_plan(plan, tie_word_embeddings=True) - assert "lm_head" not in adapted - assert adapted["model.embed_tokens"] == "keep_full_weight" - assert adapted["model.layers.*"] == "free_full_weight" - - def test_is_norm_and_head_pair(self): - model = TinyModel() - assert is_norm_and_head_pair([("norm", model.norm), ("lm_head", model.lm_head)]) - assert not is_norm_and_head_pair([("norm", model.norm)]) - - def test_is_fsdp_managed_module_flag(self): - model = TinyModel() - assert not is_fsdp_managed_module(model) - model._is_fsdp_managed_module = True - assert is_fsdp_managed_module(model) From 588884e6fce683d1abb0f8b89eba24e2c63ec2fb Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 17 Jun 2026 04:22:22 +0000 Subject: [PATCH 05/24] copyright edit --- src/transformers/distributed/fsdp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index abf9a7a29df3..078659717ba7 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -1,4 +1,4 @@ -# Copyright 2024 The HuggingFace Team. All rights reserved. +# Copyright 2026 The HuggingFace 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. From 8cc48a0678b44e931485ab1a1322ba4cf4ad802c Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 17 Jun 2026 04:31:35 +0000 Subject: [PATCH 06/24] revert --- src/transformers/generation/utils.py | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index ed14b60fbfa8..749074c3d892 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -118,18 +118,6 @@ logger = logging.get_logger(__name__) - -def _get_dist_world_size() -> int: - if ( - dist.is_available() - and hasattr(dist, "is_initialized") - and dist.is_initialized() - and hasattr(dist, "get_world_size") - ): - return dist.get_world_size() - return 1 - - if is_accelerate_available(): from accelerate.hooks import AlignDevicesHook, add_hook_to_module @@ -2177,7 +2165,7 @@ def _extract_generation_mode_kwargs( "assistant_model": assistant_model, "streamer": streamer, } - world_size = _get_dist_world_size() + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 # type: ignore generation_mode_kwargs["synced_gpus"] = ( (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and world_size > 1 if synced_gpus is None @@ -2624,8 +2612,7 @@ def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, # The following logic allows an early break if all peers finished generating their sequence this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) # send 0.0 if we finished, 1.0 otherwise - if hasattr(dist, "all_reduce") and hasattr(dist, "ReduceOp"): - dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) # type: ignore # did all peers finish? the reduced sum will be 0.0 then if this_peer_finished_flag.item() == 0.0: return False From 79457b3f739aaa0c119b82ab5e4b3230cf9d0a91 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 23 Jun 2026 05:21:12 +0000 Subject: [PATCH 07/24] fix --- src/transformers/modeling_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 7cd363c4982e..6ebdb91acdb8 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1400,9 +1400,7 @@ def post_init(self): """ # Attach the different parallel plans and tied weight keys to the top-most model, so that everything is # easily available. Seed with the class-level FSDP plan before the instance attribute shadows it. - cls_fsdp_plan = getattr(self, "_fsdp_plan", None) or {} - self._tp_plan, self._ep_plan, self._pp_plan = {}, {}, {} - self._fsdp_plan = dict(cls_fsdp_plan) + self._tp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = {}, {}, {}, {} # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config if self.base_model is self: self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else {} From f3e80210ecfe1f71c057b7da5065b37e4e5828dd Mon Sep 17 00:00:00 2001 From: 3outeille Date: Tue, 23 Jun 2026 05:21:54 +0000 Subject: [PATCH 08/24] fix --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 6ebdb91acdb8..9fae0dd8504e 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1399,7 +1399,7 @@ def post_init(self): correctly in the case of composite models (that is, the top level model should know about those properties from its children). """ # Attach the different parallel plans and tied weight keys to the top-most model, so that everything is - # easily available. Seed with the class-level FSDP plan before the instance attribute shadows it. + # easily available. self._tp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = {}, {}, {}, {} # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config if self.base_model is self: From ea8243f8d326e9dd1a1af6af04e3d18355e1e3ae Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 04:25:11 +0000 Subject: [PATCH 09/24] remove redundant test file --- tests/test_distributed_config.py | 97 -------------------------------- 1 file changed, 97 deletions(-) delete mode 100644 tests/test_distributed_config.py diff --git a/tests/test_distributed_config.py b/tests/test_distributed_config.py deleted file mode 100644 index d6b2b0ab2bd5..000000000000 --- a/tests/test_distributed_config.py +++ /dev/null @@ -1,97 +0,0 @@ -import json -import tempfile - -from transformers.distributed import DistributedConfig - - -class TestDistributedConfig: - def test_2d_parallelism(self): - dc = DistributedConfig(tp_size=2, fsdp_size=2) - assert dc.tp_size == 2 - assert dc.fsdp_size == 2 - assert dc.tp_plan is None - assert dc.fsdp_cpu_offload is False - assert dc.fsdp_mixed_precision is False - - def test_tp_only_defaults_fsdp_to_1(self): - dc = DistributedConfig(tp_size=4) - assert dc.tp_size == 4 - assert dc.fsdp_size == 1 - assert dc.tp_plan is None - - def test_fsdp_only_defaults_tp_to_1(self): - dc = DistributedConfig(fsdp_size=4) - assert dc.tp_size == 1 - assert dc.fsdp_size == 4 - assert dc.fsdp_cpu_offload is False - assert dc.fsdp_mixed_precision is False - assert dc.tp_plan is None - - def test_empty_config(self): - dc = DistributedConfig() - assert dc.tp_size is None - assert dc.fsdp_size is None - assert dc.tp_plan is None - assert dc.fsdp_cpu_offload is False - assert dc.fsdp_mixed_precision is False - - def test_from_dict(self): - dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) - assert dc.tp_size == 2 - assert dc.fsdp_size == 4 - assert dc.tp_plan is None - - def test_from_dict_ignores_unknown_keys(self): - dc = DistributedConfig.from_dict({"tp_size": 2, "unknown_key": 42}) - assert dc.tp_size == 2 - assert not hasattr(dc, "unknown_key") - - def test_from_dict_kwargs_override(self): - dc = DistributedConfig.from_dict({"tp_size": 2}, tp_size=8) - assert dc.tp_size == 8 - - def test_to_dict(self): - dc = DistributedConfig(tp_size=2, fsdp_size=4) - d = dc.to_dict() - assert d == { - "tp_size": 2, - "tp_plan": None, - "enable_sequence_parallel": False, - "enable_expert_parallel": False, - "fsdp_size": 4, - "fsdp_cpu_offload": False, - "fsdp_mixed_precision": False, - } - - def test_to_dict_is_a_copy(self): - dc = DistributedConfig(tp_plan={"layer": "colwise"}) - d = dc.to_dict() - d["tp_plan"]["layer"] = "rowwise" - assert dc.tp_plan["layer"] == "colwise" - - def test_to_json_string(self): - dc = DistributedConfig(tp_size=2, fsdp_size=2) - s = dc.to_json_string() - parsed = json.loads(s) - assert parsed["tp_size"] == 2 - assert parsed["fsdp_size"] == 2 - - def test_to_json_file(self): - dc = DistributedConfig(tp_size=4) - with tempfile.NamedTemporaryFile(mode="r", suffix=".json", delete=False) as f: - dc.to_json_file(f.name) - f.seek(0) - parsed = json.load(f) - assert parsed["tp_size"] == 4 - assert parsed["tp_plan"] is None - - def test_roundtrip_dict(self): - original = DistributedConfig(tp_size=2, fsdp_size=4) - restored = DistributedConfig.from_dict(original.to_dict()) - assert original == restored - - def test_repr(self): - dc = DistributedConfig(tp_size=2) - r = repr(dc) - assert "DistributedConfig" in r - assert '"tp_size": 2' in r From c384fcdc262b43d2610b415158d7569c592c6f6f Mon Sep 17 00:00:00 2001 From: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Date: Wed, 24 Jun 2026 13:44:58 +0900 Subject: [PATCH 10/24] Update src/transformers/distributed/fsdp.py naming Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/distributed/fsdp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 078659717ba7..b9933a7e66a0 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -176,7 +176,7 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}") -def apply_fully_shard_data_parallel(model, fsdp_mesh): +def apply_fully_sharded_data_parallel(model, fsdp_mesh): """ Apply FSDP2 (fully_shard) to a model. """ From 962581633523ff096fba2cf4abe1e66504b97844 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 05:37:07 +0000 Subject: [PATCH 11/24] avoid looping, just look at dict --- src/transformers/distributed/fsdp.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index b9933a7e66a0..db6b594459ab 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -124,14 +124,11 @@ def _resolve_tied_embed_lm_head_plan( if embed_module_key is None: return fsdp_plan - adapted_plan = {} - for key, sharding_strategy in fsdp_plan.items(): - if key == embed_module_key: - continue - if key == "lm_head" and sharding_strategy == "keep_full_weight": - adapted_plan[embed_module_key] = sharding_strategy - else: - adapted_plan[key] = sharding_strategy + adapted_plan = fsdp_plan.copy() + adapted_plan.pop(embed_module_key, None) + if fsdp_plan.get("lm_head") == "keep_full_weight": + adapted_plan.pop("lm_head", None) + adapted_plan[embed_module_key] = "keep_full_weight" return adapted_plan From 59bcec56ffe354770c95ddd75c321275a5ae80f3 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 05:37:28 +0000 Subject: [PATCH 12/24] expand_fsdp returns reshard_targets, no_reshard_targets right away --- src/transformers/distributed/fsdp.py | 32 +++++++++++++++------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index db6b594459ab..9de56ab88281 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -132,21 +132,31 @@ def _resolve_tied_embed_lm_head_plan( return adapted_plan -def expand_fsdp_plan(model, fsdp_plan: dict[str, str]) -> list[tuple[str, nn.Module, str]]: - """Expand plan keys into ``(module_name, module, sharding_strategy)`` shard targets.""" +def expand_fsdp_plan( + model, fsdp_plan: dict[str, str] +) -> tuple[list[tuple[str, nn.Module]], list[tuple[str, nn.Module]]]: + """Expand plan keys into reshard and no-reshard ``(module_name, module)`` shard targets.""" module_lookup = dict(model.named_modules()) - shard_targets: list[tuple[str, nn.Module, str]] = [] + reshard_targets: list[tuple[str, nn.Module]] = [] + no_reshard_targets: list[tuple[str, nn.Module]] = [] for plan_key, sharding_strategy in fsdp_plan.items(): if plan_key in module_lookup: - shard_targets.append((plan_key, module_lookup[plan_key], sharding_strategy)) + target = (plan_key, module_lookup[plan_key]) + if sharding_strategy == "keep_full_weight": + no_reshard_targets.append(target) + else: + reshard_targets.append(target) continue for module_name, module in module_lookup.items(): if replace_layer_number_by_wildcard(module_name) == plan_key: - shard_targets.append((module_name, module, sharding_strategy)) + if sharding_strategy == "keep_full_weight": + no_reshard_targets.append((module_name, module)) + else: + reshard_targets.append((module_name, module)) - return shard_targets + return reshard_targets, no_reshard_targets def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) -> None: @@ -195,15 +205,7 @@ def apply_fully_sharded_data_parallel(model, fsdp_mesh): tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, tie_word_embeddings=tie_word_embeddings) - shard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) - - reshard_targets = [] - no_reshard_targets = [] - for module_name, module, sharding_strategy in shard_targets: - if sharding_strategy == "keep_full_weight": - no_reshard_targets.append((module_name, module)) - else: - reshard_targets.append((module_name, module)) + reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) for module_name, module in reshard_targets: fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) From ebf3585ceeee4096f7124a52abf3dc6565d52c47 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 05:58:29 +0000 Subject: [PATCH 13/24] better _resolve_tied_embed_lm_head_plan --- src/transformers/distributed/fsdp.py | 31 ++++++++++++++-------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 9de56ab88281..4324c05436c4 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -99,8 +99,7 @@ def is_norm_and_head_pair(modules: list[tuple[str, Any]]) -> bool: def _resolve_tied_embed_lm_head_plan( fsdp_plan: dict[str, str], - *, - tie_word_embeddings: bool, + model: nn.Module, ) -> dict[str, str]: """ Rewrite the plan so tied embed/lm_head weights are wrapped once. @@ -114,24 +113,26 @@ def _resolve_tied_embed_lm_head_plan( "model.norm": "keep_full_weight", "model.embed_tokens": "keep_full_weight"} """ - if not tie_word_embeddings: + tied_keys = getattr(model, "all_tied_weights_keys", None) or {} + if not tied_keys: return fsdp_plan - embed_module_key = next( - (key for key in fsdp_plan if key == "embed_tokens" or key.endswith(".embed_tokens")), - None, - ) - if embed_module_key is None: + target_param, source_param = next(iter(tied_keys.items())) + head_module = target_param.rsplit(".", 1)[0] + embed_module = source_param.rsplit(".", 1)[0] + + if embed_module not in fsdp_plan: return fsdp_plan - + adapted_plan = fsdp_plan.copy() - adapted_plan.pop(embed_module_key, None) - if fsdp_plan.get("lm_head") == "keep_full_weight": - adapted_plan.pop("lm_head", None) - adapted_plan[embed_module_key] = "keep_full_weight" + adapted_plan.pop(embed_module, None) + + if fsdp_plan.get(head_module) == "keep_full_weight": + adapted_plan.pop(head_module, None) + adapted_plan[embed_module] = "keep_full_weight" + return adapted_plan - def expand_fsdp_plan( model, fsdp_plan: dict[str, str] ) -> tuple[list[tuple[str, nn.Module]], list[tuple[str, nn.Module]]]: @@ -204,7 +205,7 @@ def apply_fully_sharded_data_parallel(model, fsdp_mesh): fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) - adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, tie_word_embeddings=tie_word_embeddings) + adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) for module_name, module in reshard_targets: From e9693251751a7476971c7425b71643996cfaa198 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 06:08:32 +0000 Subject: [PATCH 14/24] cleaning --- src/transformers/distributed/fsdp.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 4324c05436c4..e2492d6319c5 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -117,13 +117,9 @@ def _resolve_tied_embed_lm_head_plan( if not tied_keys: return fsdp_plan - target_param, source_param = next(iter(tied_keys.items())) - head_module = target_param.rsplit(".", 1)[0] - embed_module = source_param.rsplit(".", 1)[0] - - if embed_module not in fsdp_plan: - return fsdp_plan - + head_param, embed_param = next(iter(tied_keys.items())) + head_module = head_param.rsplit(".", 1)[0] + embed_module = embed_param.rsplit(".", 1)[0] adapted_plan = fsdp_plan.copy() adapted_plan.pop(embed_module, None) From 2376965ee7c763ea8288a1cf97107e118cc4a7c5 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 06:10:24 +0000 Subject: [PATCH 15/24] ruff --- src/transformers/distributed/fsdp.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index e2492d6319c5..a055c7931633 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -122,13 +122,14 @@ def _resolve_tied_embed_lm_head_plan( embed_module = embed_param.rsplit(".", 1)[0] adapted_plan = fsdp_plan.copy() adapted_plan.pop(embed_module, None) - + if fsdp_plan.get(head_module) == "keep_full_weight": adapted_plan.pop(head_module, None) adapted_plan[embed_module] = "keep_full_weight" - + return adapted_plan + def expand_fsdp_plan( model, fsdp_plan: dict[str, str] ) -> tuple[list[tuple[str, nn.Module]], list[tuple[str, nn.Module]]]: From 0f62c45e9c73febb52a30f8ab633e87903764b79 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 08:58:57 +0000 Subject: [PATCH 16/24] more robust detection of embed and lm_head --- src/transformers/distributed/fsdp.py | 42 ++++++++++++++++++++-------- 1 file changed, 31 insertions(+), 11 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index a055c7931633..c07dab560197 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -85,15 +85,27 @@ def _get_fsdp_policy_kwargs(distributed_config: DistributedConfig | None) -> dic return fsdp_policy_kwargs -def is_norm_and_head_pair(modules: list[tuple[str, Any]]) -> bool: - """Match the canonical tail pair: one final norm + the output head (or tied embedding).""" +def _get_input_output_embeddings(model: nn.Module) -> tuple[nn.Module | None, nn.Module | None]: + input_embed = None + output_head = None + if hasattr(model, "get_input_embeddings"): + input_embed = model.get_input_embeddings() + if hasattr(model, "get_output_embeddings"): + output_head = model.get_output_embeddings() + return input_embed, output_head + + +def is_norm_and_head_pair(modules: list[tuple[str, nn.Module]], model: nn.Module) -> bool: if len(modules) != 2: return False + input_embed, output_head = _get_input_output_embeddings(model) + head_modules = {module for module in (input_embed, output_head) if module is not None} + module_names = [name for name, _ in modules] + module_objects = [module for _, module in modules] + has_final_norm = any(name == "norm" or name.endswith(".norm") for name in module_names) - has_output_head = any( - name in {"lm_head", "embed_tokens"} or name.endswith((".lm_head", ".embed_tokens")) for name in module_names - ) + has_output_head = any(module in head_modules for module in module_objects) return has_final_norm and has_output_head @@ -117,9 +129,14 @@ def _resolve_tied_embed_lm_head_plan( if not tied_keys: return fsdp_plan - head_param, embed_param = next(iter(tied_keys.items())) - head_module = head_param.rsplit(".", 1)[0] - embed_module = embed_param.rsplit(".", 1)[0] + input_embed, output_head = _get_input_output_embeddings(model) + name_by_module = {module: name for name, module in model.named_modules()} + embed_module = name_by_module.get(input_embed) + head_module = name_by_module.get(output_head) + + if embed_module is None or head_module is None: + return fsdp_plan + adapted_plan = fsdp_plan.copy() adapted_plan.pop(embed_module, None) @@ -131,7 +148,8 @@ def _resolve_tied_embed_lm_head_plan( def expand_fsdp_plan( - model, fsdp_plan: dict[str, str] + model: nn.Module, + fsdp_plan: dict[str, str], ) -> tuple[list[tuple[str, nn.Module]], list[tuple[str, nn.Module]]]: """Expand plan keys into reshard and no-reshard ``(module_name, module)`` shard targets.""" module_lookup = dict(model.named_modules()) @@ -181,7 +199,9 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) logger.warning(f"The following FSDP rules were not applied to any module: {unused_rules}") -def apply_fully_sharded_data_parallel(model, fsdp_mesh): +def apply_fully_sharded_data_parallel( + model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh +) -> nn.Module: """ Apply FSDP2 (fully_shard) to a model. """ @@ -211,7 +231,7 @@ def apply_fully_sharded_data_parallel(model, fsdp_mesh): # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. - if is_norm_and_head_pair(no_reshard_targets): + if is_norm_and_head_pair(no_reshard_targets, model): module_names = [name for name, _ in no_reshard_targets] modules = [module for _, module in no_reshard_targets] fully_shard(modules, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) From da302ade32f20bebadb84ff9ad93ef78f6597bc8 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 13:48:28 +0000 Subject: [PATCH 17/24] cleaning --- src/transformers/distributed/fsdp.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index c07dab560197..c452aa3bbcfe 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -95,17 +95,19 @@ def _get_input_output_embeddings(model: nn.Module) -> tuple[nn.Module | None, nn return input_embed, output_head -def is_norm_and_head_pair(modules: list[tuple[str, nn.Module]], model: nn.Module) -> bool: +def is_norm_and_head_pair(no_reshard_targets: list[tuple[str, nn.Module]], model: nn.Module) -> bool: if len(modules) != 2: return False input_embed, output_head = _get_input_output_embeddings(model) head_modules = {module for module in (input_embed, output_head) if module is not None} - module_names = [name for name, _ in modules] - module_objects = [module for _, module in modules] + names, modules = [], [] + for name, module in no_reshard_targets: + names.append(name) + modules.append(module) - has_final_norm = any(name == "norm" or name.endswith(".norm") for name in module_names) - has_output_head = any(module in head_modules for module in module_objects) + has_final_norm = any(name == "norm" or name.endswith(".norm") for name in names) + has_output_head = any(module in head_modules for module in modules) return has_final_norm and has_output_head @@ -232,12 +234,14 @@ def apply_fully_sharded_data_parallel( # Optimization: when the keep buffer is exactly the (final_norm, lm_head/embed) # tail pair, bundle them into one fully_shard so that we dont need to do all-gather during backward pass. if is_norm_and_head_pair(no_reshard_targets, model): - module_names = [name for name, _ in no_reshard_targets] - modules = [module for _, module in no_reshard_targets] + names, modules = [], [] + for name, module in no_reshard_targets: + names.append(name) + modules.append(module) fully_shard(modules, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) - logger.debug(f"Grouped tail {module_names} (reshard=False)") + logger.debug(f"Grouped tail {names} (reshard=False)") else: - for module_name, module in no_reshard_targets: + for name, module in no_reshard_targets: fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) logger.debug(f"Applied fully_shard to {module_name} (reshard=False)") From dfc665cd67a1a553f491d9fb4388dcfc7e4b86d4 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 13:50:40 +0000 Subject: [PATCH 18/24] ruff --- src/transformers/distributed/fsdp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index c452aa3bbcfe..1584b1064749 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -96,7 +96,7 @@ def _get_input_output_embeddings(model: nn.Module) -> tuple[nn.Module | None, nn def is_norm_and_head_pair(no_reshard_targets: list[tuple[str, nn.Module]], model: nn.Module) -> bool: - if len(modules) != 2: + if len(no_reshard_targets) != 2: return False input_embed, output_head = _get_input_output_embeddings(model) head_modules = {module for module in (input_embed, output_head) if module is not None} From 446fd6e987c6d401ccc1a8ddfa09c88120f15961 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 13:56:20 +0000 Subject: [PATCH 19/24] typo --- src/transformers/distributed/fsdp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 1584b1064749..26c78fe282f2 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -243,7 +243,7 @@ def apply_fully_sharded_data_parallel( else: for name, module in no_reshard_targets: fully_shard(module, mesh=fsdp_mesh, reshard_after_forward=False, **fsdp_policy_kwargs) - logger.debug(f"Applied fully_shard to {module_name} (reshard=False)") + logger.debug(f"Applied fully_shard to {name} (reshard=False)") # Apply FSDP2 to the root module fully_shard(model, mesh=fsdp_mesh, **fsdp_policy_kwargs) From 5aeaff749cce72951a1f19949e3b33d92c01ced7 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 14:06:56 +0000 Subject: [PATCH 20/24] cleaner --- src/transformers/distributed/fsdp.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 26c78fe282f2..2fb0468cd11c 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -160,19 +160,20 @@ def expand_fsdp_plan( for plan_key, sharding_strategy in fsdp_plan.items(): if plan_key in module_lookup: - target = (plan_key, module_lookup[plan_key]) - if sharding_strategy == "keep_full_weight": - no_reshard_targets.append(target) - else: - reshard_targets.append(target) - continue - - for module_name, module in module_lookup.items(): - if replace_layer_number_by_wildcard(module_name) == plan_key: - if sharding_strategy == "keep_full_weight": - no_reshard_targets.append((module_name, module)) - else: - reshard_targets.append((module_name, module)) + # model.norm, lm_head etc. + targets = [(plan_key, module_lookup[plan_key])] + else: + # model.layers.* + targets = [ + (module_name, module) + for module_name, module in module_lookup.items() + if replace_layer_number_by_wildcard(module_name) == plan_key + ] + + if sharding_strategy == "keep_full_weight": + no_reshard_targets.extend(targets) + else: + reshard_targets.extend(targets) return reshard_targets, no_reshard_targets From 819ff1498e2b7b7bd915d3ae49458a1384682e71 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Wed, 24 Jun 2026 14:10:02 +0000 Subject: [PATCH 21/24] cleaner --- src/transformers/distributed/fsdp.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 2fb0468cd11c..38baed939b60 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -192,8 +192,7 @@ def verify_fsdp_plan(module_names: list[str], fsdp_plan: dict[str, str] | None) for key, strategy in fsdp_plan.items(): if strategy not in {"free_full_weight", "keep_full_weight"}: invalid_strategies[key] = strategy - continue - if key not in name_lookup and not any(replace_layer_number_by_wildcard(name) == key for name in name_lookup): + elif key not in name_lookup and not any(replace_layer_number_by_wildcard(name) == key for name in name_lookup): unused_rules[key] = strategy if invalid_strategies: From 5c95559e848c59a711f76bcc5ec223d90309664e Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 3 Jul 2026 04:04:45 +0000 Subject: [PATCH 22/24] expand_fsdp_plan iterate over modules --- src/transformers/distributed/fsdp.py | 24 +++++++----------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 38baed939b60..051aef86e364 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -154,26 +154,16 @@ def expand_fsdp_plan( fsdp_plan: dict[str, str], ) -> tuple[list[tuple[str, nn.Module]], list[tuple[str, nn.Module]]]: """Expand plan keys into reshard and no-reshard ``(module_name, module)`` shard targets.""" - module_lookup = dict(model.named_modules()) reshard_targets: list[tuple[str, nn.Module]] = [] no_reshard_targets: list[tuple[str, nn.Module]] = [] - for plan_key, sharding_strategy in fsdp_plan.items(): - if plan_key in module_lookup: - # model.norm, lm_head etc. - targets = [(plan_key, module_lookup[plan_key])] - else: - # model.layers.* - targets = [ - (module_name, module) - for module_name, module in module_lookup.items() - if replace_layer_number_by_wildcard(module_name) == plan_key - ] - - if sharding_strategy == "keep_full_weight": - no_reshard_targets.extend(targets) - else: - reshard_targets.extend(targets) + for module_name, module in model.named_modules(): + plan_key = module_name if module_name in fsdp_plan else replace_layer_number_by_wildcard(module_name) + if plan_key in fsdp_plan: + if fsdp_plan[plan_key] == "keep_full_weight": + no_reshard_targets.append((module_name, module)) + else: + reshard_targets.append((module_name, module)) return reshard_targets, no_reshard_targets From e4613e670535e581983a3e6c13d13e48271d5f39 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 3 Jul 2026 05:01:45 +0000 Subject: [PATCH 23/24] comment about tie embedding --- src/transformers/distributed/fsdp.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 051aef86e364..d3870485b54d 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -243,8 +243,7 @@ def apply_fully_sharded_data_parallel( # Used by generation code to detect FSDP and enable synced_gpus. model._is_fsdp_managed_module = True - if tie_word_embeddings and hasattr(model, "tie_weights"): - model.tie_weights() + #NOTE(3outeille): No need to tie the word embeddings here, it will be done _finalize_model_loading in modeling_utils.py return model From d558f9991c45da6f214cd53aec0a8f25c181b6c1 Mon Sep 17 00:00:00 2001 From: 3outeille Date: Fri, 3 Jul 2026 05:06:53 +0000 Subject: [PATCH 24/24] add comment tied embedding --- src/transformers/distributed/fsdp.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index d3870485b54d..1eacfa209317 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -212,7 +212,6 @@ def apply_fully_sharded_data_parallel( distributed_config = getattr(model.config, "distributed_config", None) fsdp_policy_kwargs = _get_fsdp_policy_kwargs(distributed_config) - tie_word_embeddings = getattr(model.config, "tie_word_embeddings", False) adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, model) reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan) @@ -243,7 +242,7 @@ def apply_fully_sharded_data_parallel( # Used by generation code to detect FSDP and enable synced_gpus. model._is_fsdp_managed_module = True - #NOTE(3outeille): No need to tie the word embeddings here, it will be done _finalize_model_loading in modeling_utils.py + # NOTE(3outeille): No need to tie the word embeddings here, it will be done _finalize_model_loading in modeling_utils.py return model