FSDP orchestration: apply + loading/saving

.circleci/create_circleci_config.py

@@ -313,7 +313,7 @@ def job_name(self):
 torch_job = CircleCIJob(
     "torch",
     docker_image=[{"image": "huggingface/transformers-torch-light"}],
-    marker="not generate",
+    marker="not (generate or is_training_test or is_tensor_parallel_test or is_fsdp_test)",
     parallelism=6,
 )
 
@@ -404,6 +404,15 @@ def job_name(self):
     parallelism=6,
 )
 
+fsdp_ci_job = CircleCIJob(
+    "fsdp_ci",
+    additional_env={"RUN_FSDP_TESTS": True},
+    docker_image=[{"image": "huggingface/transformers-torch-light"}],
+    install_steps=["uv pip install ."],
+    marker="is_fsdp_test",
+    parallelism=6,
+)
+
 # We also include a `dummy.py` file in the files to be doc-tested to prevent edge case failure. Otherwise, the pytest
 # hangs forever during test collection while showing `collecting 0 items / 21 errors`. (To see this, we have to remove
 # the bash output redirection.)
@@ -435,7 +444,8 @@ def job_name(self):
 DOC_TESTS = [doc_test_job]
 TRAINING_CI_TESTS = [training_ci_job]
 TENSOR_PARALLEL_CI_TESTS = [tensor_parallel_ci_job]
-ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS  # fmt: skip
+FSDP_CI_TESTS = [fsdp_ci_job]
+ALL_TESTS = REGULAR_TESTS + EXAMPLES_TESTS + PIPELINE_TESTS + REPO_UTIL_TESTS + DOC_TESTS + [custom_tokenizers_job] + [exotic_models_job] + TRAINING_CI_TESTS + TENSOR_PARALLEL_CI_TESTS + FSDP_CI_TESTS  # fmt: skip
 
 
 def create_circleci_config(folder=None):

conftest.py

@@ -97,6 +97,7 @@ def pytest_configure(config):
     )
     config.addinivalue_line("markers", "training_ci: mark test for training CI validation")
     config.addinivalue_line("markers", "tensor_parallel_ci: mark test for tensor parallel CI validation")
+    config.addinivalue_line("markers", "fsdp_ci: mark test for FSDP CI validation")
 
     os.environ["DISABLE_SAFETENSORS_CONVERSION"] = "true"
     register_network_debug_plugin(config)

docs/source/en/expert_parallelism.md

@@ -22,11 +22,16 @@ rendered properly in your Markdown viewer.
 Enable expert parallelism with the [`DistributedConfig`] class and the `enable_expert_parallel` argument.
 
 ```py
+import os
+
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.distributed.configuration_utils import DistributedConfig
 
-distributed_config = DistributedConfig(enable_expert_parallel=True)
+distributed_config = DistributedConfig(
+    tp_size=int(os.environ["WORLD_SIZE"]),
+    enable_expert_parallel=True,
+)
 
 model = AutoModelForCausalLM.from_pretrained(
     "openai/gpt-oss-120b",

docs/source/en/model_doc/minimax_m3_vl.md

@@ -165,8 +165,10 @@ model = AutoModelForImageTextToText.from_pretrained(
     dtype=torch.bfloat16,
     # Dequantize the native MXFP8 weights to bf16 at load (the speed win); needs even TP/EP sharding.
     quantization_config=FineGrainedFP8Config(dequantize=True),
-    tp_plan="auto",
-    distributed_config=DistributedConfig(enable_expert_parallel=True),
+    distributed_config=DistributedConfig(
+        tp_size=int(os.environ["WORLD_SIZE"]),
+        enable_expert_parallel=True,
+    ),
     attn_implementation="kernels-staging/msa@v0",  # MSA block-sparse attention kernel
 )
 model.eval()

docs/source/en/pr_checks.md

@@ -123,6 +123,7 @@ Tests are split across parallel CI jobs, and each job picks up files by path pat
 - `pipelines_torch`: pipeline tests
 - `tests_training_ci`: training loop tests
 - `tests_tensor_parallel_ci`: tensor parallel tests
+- `tests_fsdp_ci`: FSDP tests
 
 ### Slow tests
 

docs/source/ro/pr_checks.md

@@ -123,6 +123,7 @@ Testele sunt împărțite pe job-uri CI paralele, iar fiecare job preia fișiere
 - `pipelines_torch`: teste de pipeline
 - `tests_training_ci`: teste pentru loop-ul de antrenare
 - `tests_tensor_parallel_ci`: teste de tensor parallelism
+- `tests_fsdp_ci`: teste FSDP
 
 ### Testele slow
 

pyproject.toml

@@ -83,6 +83,7 @@ markers = [
     "generate: marks tests that use the GenerationTesterMixin",
     "is_training_test: marks tests that use the TrainingTesterMixin (deselect with '-m \"not is_training_test\"')",
     "is_tensor_parallel_test: marks tests that use the TensorParallelTesterMixin (deselect with '-m \"not is_tensor_parallel_test\"')",
+    "is_fsdp_test: marks tests that use the FSDPTesterMixin (deselect with '-m \"not is_fsdp_test\"')",
 ]
 log_cli = 1
 log_cli_level = "WARNING"

src/transformers/configuration_utils.py

@@ -1040,9 +1040,6 @@ 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]
@@ -1189,6 +1186,7 @@ def _remove_keys_not_serialized(self, d: dict[str, Any]) -> None:
             "ignore_keys_at_rope_validation",
             "base_model_tp_plan",
             "base_model_pp_plan",
+            "distributed_config",
         ]:
             d.pop(key_to_remove, None)
 

src/transformers/distributed/__init__.py

@@ -20,6 +20,15 @@
 _import_structure = {
     "configuration_utils": ["DistributedConfig"],
     "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"],
+    "mixin": ["DistributedMixin"],
+    "utils": [
+        "distribute_model",
+        "gather_full_state_dict",
+        "initialize_fully_sharded_data_parallelism",
+        "load_optimizer_distributed",
+        "save_model_checkpoint_distributed",
+        "save_optimizer_distributed",
+    ],
 }
 
 
@@ -28,6 +37,15 @@
         DistributedConfig,
     )
     from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan
+    from .mixin import DistributedMixin
+    from .utils import (
+        distribute_model,
+        gather_full_state_dict,
+        initialize_fully_sharded_data_parallelism,
+        load_optimizer_distributed,
+        save_model_checkpoint_distributed,
+        save_optimizer_distributed,
+    )
 
 else:
     import sys

src/transformers/distributed/configuration_utils.py

@@ -63,12 +63,35 @@ def __post_init__(self):
         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})"
-                )
+        if self.tp_size > 1 and self.fsdp_size > 1:
+            raise ValueError(
+                "FSDP+TP is not supported yet. "
+                "Use DistributedConfig(fsdp_size=N) or DistributedConfig(tp_size=N), not both. "
+                "2D support will come soon."
+            )
+
+    def validate(self) -> None:
+        """Validate against the live process group. Call before distributed load/train."""
+        if self.tp_size is None and self.fsdp_size is None:
+            return
+
+        if self.tp_size <= 1 and self.fsdp_size <= 1:
+            return
+
+        if not is_torch_available():
+            raise RuntimeError("PyTorch is required to use DistributedConfig.")
+
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            raise RuntimeError(
+                "torch.distributed must be initialized before using DistributedConfig with tp_size > 1 or "
+                "fsdp_size > 1. Call dist.init_process_group(...) first, or launch with torchrun."
+            )
+
+        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: dict, **kwargs) -> "DistributedConfig":

src/transformers/distributed/fsdp.py

@@ -30,7 +30,7 @@
 if is_torch_available():
     import torch
 
-if is_torch_available() and is_torch_greater_or_equal("2.6"):
+if is_torch_available() and is_torch_greater_or_equal("2.7"):
     from torch.distributed._composable.fsdp import fully_shard
     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy
 
@@ -191,7 +191,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_sharded_data_parallel(
+def apply_fully_sharded_data_parallelism(
     model: nn.Module, fsdp_mesh: torch.distributed.device_mesh.DeviceMesh
 ) -> nn.Module:
     """
@@ -200,8 +200,8 @@ def apply_fully_sharded_data_parallel(
     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")
+    if not is_torch_greater_or_equal("2.7"):
+        raise OSError("FSDP2 requires torch>=2.7")
 
     fsdp_plan = dict(getattr(model, "_fsdp_plan", None) or {})
     if not fsdp_plan: