Add native FSDP2 module + migration

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"],

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]

src/transformers/distributed/__init__.py

@@ -19,13 +19,15 @@
 
 _import_structure = {
     "configuration_utils": ["DistributedConfig"],
+    "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"],
 }
 
 
 if TYPE_CHECKING:
     from .configuration_utils import (
         DistributedConfig,
     )
+    from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan
 
 else:
     import sys

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