From 7c11394079b9155a76e5e559edf870bcdab1353a Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 05:08:49 +0000
Subject: [PATCH 01/33] Add FSDP orchestration: mesh init,
 distribute-before-load, and DCP save.

Wire distributed_config from_pretrained/save_pretrained alongside the legacy tp_plan path, add distributed/utils.py for mesh orchestration and checkpoint I/O, and extend sharding_utils with DTensor gather/optimizer fusion helpers needed by save/load.
---
 src/transformers/distributed/__init__.py      |  16 +
 .../distributed/sharding_utils.py             | 130 +++++++-
 src/transformers/distributed/utils.py         | 282 ++++++++++++++++++
 src/transformers/modeling_utils.py            | 260 ++++++++++------
 tests/utils/test_modeling_utils.py            |  40 +++
 5 files changed, 635 insertions(+), 93 deletions(-)
 create mode 100644 src/transformers/distributed/utils.py

diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py
index e179a2306156..318fd2961360 100644
--- a/src/transformers/distributed/__init__.py
+++ b/src/transformers/distributed/__init__.py
@@ -20,6 +20,14 @@
 _import_structure = {
     "configuration_utils": ["DistributedConfig"],
     "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"],
+    "utils": [
+        "init_device_mesh",
+        "distribute_model",
+        "gather_full_state_dict",
+        "save_model_checkpoint_distributed",
+        "save_optimizer_distributed",
+        "load_optimizer_distributed",
+    ],
 }
 
 
@@ -28,6 +36,14 @@
         DistributedConfig,
     )
     from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan
+    from .utils import (
+        distribute_model,
+        gather_full_state_dict,
+        init_device_mesh,
+        load_optimizer_distributed,
+        save_model_checkpoint_distributed,
+        save_optimizer_distributed,
+    )
 
 else:
     import sys
diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py
index 365713113cc9..fb0d30368c2a 100644
--- a/src/transformers/distributed/sharding_utils.py
+++ b/src/transformers/distributed/sharding_utils.py
@@ -25,9 +25,10 @@
 
 if is_torch_available():
     import torch
-    from torch.distributed.tensor import DTensor
+    from torch.distributed._functional_collectives import wait_tensor
+    from torch.distributed.tensor import DTensor, Replicate
     from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
-    from torch.distributed.tensor.placement_types import Shard
+    from torch.distributed.tensor.placement_types import Shard, _StridedShard
 
     # torch < 2.10 names as an underscore before `local_shard_size_and_offset`: alias it the non-underscored version
     if not hasattr(Shard, "local_shard_size_and_offset") and hasattr(Shard, "_local_shard_size_and_offset"):
@@ -330,3 +331,128 @@ def _dtensor_from_local_like(local_tensor: torch.Tensor, ref: DTensor) -> DTenso
         shape=ref.shape,
         stride=tuple(ref.stride()),
     )
+
+
+def _replicate_dtensor(tensor: DTensor) -> DTensor:
+    """All-gather a DTensor to fully Replicate, handling _StridedShard."""
+    mesh = tensor.device_mesh
+    placements = tensor.placements
+    replicate_all = tuple(Replicate() for _ in range(mesh.ndim))
+
+    if not any(isinstance(p, _StridedShard) for p in placements):
+        return tensor.redistribute(placements=replicate_all)
+
+    with torch.no_grad():
+        local = tensor._local_tensor
+        for i in reversed(range(mesh.ndim)):
+            p = placements[i]
+            if p.is_replicate():
+                continue
+            logical_shape = list(tensor.shape)
+            for j, pj in enumerate(placements[:i]):
+                if not pj.is_replicate():
+                    size, _ = Shard.local_shard_size_and_offset(
+                        logical_shape[pj.dim], mesh.size(j), mesh.get_local_rank(j)
+                    )
+                    logical_shape[pj.dim] = size
+            local = p._to_replicate_tensor(local, mesh, i, logical_shape)
+            local = wait_tensor(local)
+        return DTensor.from_local(local, mesh, replicate_all, run_check=False)
+
+
+def _find_strided_shard_placement_from_fused_params(placements):
+    for i, p in enumerate(placements):
+        if not isinstance(p, _StridedShard):
+            continue
+        has_partner_on_same_dim = any(
+            j != i and getattr(other, "dim", None) == p.dim for j, other in enumerate(placements)
+        )
+        if not has_partner_on_same_dim:
+            return p
+    return None
+
+
+def _split_fused_dtensor(dt, dim, n_pieces):
+    local_pieces = list(dt._local_tensor.chunk(n_pieces, dim=dim))
+    if len(local_pieces) != n_pieces:
+        raise RuntimeError(
+            f"Cannot split DTensor of shape {tuple(dt.shape)} into {n_pieces} pieces along dim {dim}: "
+            f"got {len(local_pieces)} chunks instead."
+        )
+    new_placements = tuple(
+        Shard(p.dim) if (isinstance(p, _StridedShard) and p.dim == dim) else p for p in dt.placements
+    )
+    return [
+        DTensor.from_local(lp.contiguous(), dt.device_mesh, new_placements, run_check=False) for lp in local_pieces
+    ]
+
+
+def _merge_unfused_dtensors(pieces, dim, target_placements):
+    local_cat = torch.cat([p._local_tensor for p in pieces], dim=dim).contiguous()
+    return DTensor.from_local(local_cat, pieces[0].device_mesh, target_placements, run_check=False)
+
+
+def get_fusion_metadata(optimizer_state_dict):
+    """Inspect the optimizer state dict and return metadata describing every
+    fused param that needs to be split for DCP."""
+    optimizer_state = optimizer_state_dict.get("state", {})
+    fusion_metadata: dict[str, dict] = {}
+
+    for param_name, optimizer_fields in optimizer_state.items():
+        dtensor = next((v for v in optimizer_fields.values() if isinstance(v, DTensor)), None)
+        if dtensor is None:
+            continue
+        strided_shard_placement = _find_strided_shard_placement_from_fused_params(dtensor.placements)
+        if strided_shard_placement is None:
+            continue
+        fusion_metadata[param_name] = {
+            "chunk_dim": strided_shard_placement.dim,
+            "placements": tuple(dtensor.placements),
+            "unfused_keys": [f"{param_name}.{i}" for i in range(strided_shard_placement.split_factor)],
+        }
+
+    return fusion_metadata
+
+
+def unfuse_optimizer_state(optimizer_state_dict, fusion_metadata):
+    """Apply `fusion_metadata` to split each fused param into plain-Shard pieces in place."""
+    optimizer_state = optimizer_state_dict.get("state", {})
+
+    for param_name, metadata in fusion_metadata.items():
+        optimizer_fields = optimizer_state[param_name]
+
+        for unfused_key in metadata["unfused_keys"]:
+            optimizer_state[unfused_key] = {}
+
+        for optim_param_name, value in optimizer_fields.items():
+            if isinstance(value, DTensor):
+                chunks = _split_fused_dtensor(value, metadata["chunk_dim"], len(metadata["unfused_keys"]))
+            else:
+                chunks = [value] * len(metadata["unfused_keys"])
+            for unfused_key, chunk in zip(metadata["unfused_keys"], chunks):
+                optimizer_state[unfused_key][optim_param_name] = chunk
+
+        del optimizer_state[param_name]
+
+
+def fuse_optimizer_state(optimizer_state_dict, fusion_metadata):
+    """Fuse the optimizer state dict back in place using the fusion info."""
+    optimizer_state = optimizer_state_dict.get("state", {})
+
+    for param_name, metadata in fusion_metadata.items():
+        first_piece = optimizer_state[metadata["unfused_keys"][0]]
+
+        merged_fields = {}
+        for optim_param_name, value in first_piece.items():
+            chunks = [optimizer_state[unfused_key][optim_param_name] for unfused_key in metadata["unfused_keys"]]
+            if isinstance(value, DTensor):
+                merged_fields[optim_param_name] = _merge_unfused_dtensors(
+                    chunks, metadata["chunk_dim"], metadata["placements"]
+                )
+            else:
+                merged_fields[optim_param_name] = value
+
+        optimizer_state[param_name] = merged_fields
+
+        for unfused_key in metadata["unfused_keys"]:
+            del optimizer_state[unfused_key]
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
new file mode 100644
index 000000000000..a9b571aa85fc
--- /dev/null
+++ b/src/transformers/distributed/utils.py
@@ -0,0 +1,282 @@
+# Copyright 2025 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 os
+from typing import TYPE_CHECKING
+
+from ..utils import is_torch_available, is_torch_greater_or_equal, logging
+from .fsdp import apply_fully_sharded_data_parallel
+from .sharding_utils import (
+    _find_strided_shard_placement_from_fused_params,
+    _replicate_dtensor,
+    fuse_optimizer_state,
+    get_fusion_metadata,
+    unfuse_optimizer_state,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+if TYPE_CHECKING:
+    import torch.nn as nn
+
+    from .configuration_utils import DistributedConfig
+
+if is_torch_available():
+    import torch
+    import torch.distributed.checkpoint as dcp
+    from torch.distributed.checkpoint.state_dict import (
+        get_model_state_dict,
+        get_optimizer_state_dict,
+        set_optimizer_state_dict,
+    )
+    from torch.distributed.tensor import DTensor
+
+    if is_torch_greater_or_equal("2.7"):
+        from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter
+
+
+def _ensure_torch_distributed(device_type: str):
+    """Initialize torch.distributed if not already initialized."""
+    if not torch.distributed.is_initialized():
+        try:
+            rank = int(os.environ["RANK"])
+            local_rank = int(os.environ["LOCAL_RANK"])
+            world_size = int(os.environ["WORLD_SIZE"])
+
+            backend_map = {
+                "cuda": "nccl",
+                "cpu": "gloo",
+                "xpu": "xccl",
+                "hpu": "hccl",
+                "neuron": "neuron",
+                "tpu": "tpu_dist",
+            }
+            backend = backend_map.get(device_type)
+
+            # Bind the accelerator before init so the process group is created with a
+            # device_id, otherwise collectives like barrier() warn (and may spin up an
+            # extra NCCL comm) about the missing device binding.
+            device_id = None
+            if device_type != "cpu":
+                getattr(torch, device_type).set_device(local_rank)
+                device_id = torch.device(device_type, local_rank)
+            torch.distributed.init_process_group(
+                backend=backend, rank=rank, world_size=world_size, device_id=device_id
+            )
+        except Exception as e:
+            raise OSError(
+                "We tried to initialize torch.distributed for you, but it failed. Make "
+                "sure you init torch distributed in your script to use distributed training."
+            ) from e
+
+
+def _distributed_barrier():
+    """Barrier bound to the current accelerator device.
+
+    Passing `device_ids` is required when the process group was initialized without a
+    `device_id`; with it, the call is a no-op compared to plain `barrier()`. Safe to call
+    when torch.distributed has not been initialized — returns immediately.
+    """
+    if not torch.distributed.is_initialized():
+        return
+    device_type = torch._C._get_accelerator().type
+    if device_type != "cpu":
+        torch.distributed.barrier(device_ids=[getattr(torch, device_type).current_device()])
+    else:
+        torch.distributed.barrier()
+
+
+def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed.device_mesh.DeviceMesh:
+    if not is_torch_greater_or_equal("2.5"):
+        raise OSError("Distributed training with DistributedConfig requires `torch>=2.5`.")
+
+    if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.6"):
+        raise OSError("FSDP2 requires `torch>=2.6`.")
+
+    device_type = torch._C._get_accelerator().type
+    _ensure_torch_distributed(device_type)
+
+    world_size = torch.distributed.get_world_size()
+    if device_type != "cpu":
+        getattr(torch, device_type).set_device(int(os.environ.get("LOCAL_RANK", 0)))
+
+    tp_size = distributed_config.tp_size
+    fsdp_size = distributed_config.fsdp_size
+
+    assert world_size == tp_size * fsdp_size, (
+        f"world_size ({world_size}) must be equal to tp_size ({tp_size}) * fsdp_size ({fsdp_size})"
+    )
+
+    dims, names = [], []
+    if fsdp_size > 1:
+        dims.append(fsdp_size)
+        names.append("fsdp")
+    if tp_size > 1:
+        dims.append(tp_size)
+        names.append("tp")
+
+    # Build the N-dimensional device mesh
+    mesh = torch.distributed.init_device_mesh(device_type, tuple(dims), mesh_dim_names=tuple(names))
+    # If N > 1, create a flattened sub-mesh so all-reduces across the world mesh ae done in one collective
+    if len(dims) > 1:
+        mesh._flatten("_".join(names))
+
+    return mesh
+
+
+def distribute_model(model, distributed_config: DistributedConfig, device_mesh) -> nn.Module:
+    """Apply TP and/or FSDP2 to `model` based on the mesh dims in `device_mesh`."""
+    model.config.distributed_config = distributed_config
+    model.device_mesh = device_mesh
+    mesh_dim_names = device_mesh.mesh_dim_names or ()
+    if "tp" in mesh_dim_names:
+        from .tensor_parallel import apply_tensor_parallel
+
+        tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh
+        model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan)
+    if "fsdp" in mesh_dim_names:
+        fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh
+        model = apply_fully_sharded_data_parallel(model, fsdp_mesh)
+    return model
+
+
+def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
+    """Gather all sharded params to full plain tensors for saving.
+
+    Handles FSDP unshard and TP DTensor gather.
+    Streams one parameter at a time to avoid holding all full tensors on GPU.
+    Only rank 0 accumulates the result; other ranks return ``{}``.
+    """
+    is_rank0 = torch.distributed.get_rank() == 0
+    state_dict = get_model_state_dict(model)
+
+    result = {}
+    for key, tensor in state_dict.items():
+        if not isinstance(tensor, DTensor):
+            if is_rank0:
+                result[key] = tensor.detach().to(device="cpu", copy=True).contiguous()
+            continue
+
+        with torch.no_grad():
+            full = _replicate_dtensor(tensor).to_local()
+            if is_rank0:
+                result[key] = full.detach().to(device="cpu", copy=True).contiguous()
+            del full
+    return result
+
+
+def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
+    """Save model parameters as standard HF-format sharded safetensors using
+    DCP + HuggingFaceStorageWriter with consolidation enabled.
+
+    Every rank first writes its own shard in parallel under
+    `<checkpoint_dir>/sharded/`, then a consolidation pass reads those shards
+    and emits HF-compatible `model-*-of-N.safetensors` (+ index) at
+    `<checkpoint_dir>/`. The result is a directory `from_pretrained` reads
+    through its normal path — no special flag needed at load time.
+
+    DTensors carrying an uncomposed `_StridedShard` placement (e.g. fused
+    gate||up MoE weights) are replicated to a full tensor on every rank
+    before the save, otherwise DCP cannot encode that placement.
+    """
+    if not is_torch_greater_or_equal("2.7"):
+        raise OSError("Distributed checkpoint saving requires `torch>=2.7`.")
+
+    state_dict = get_model_state_dict(model)
+    for key, value in list(state_dict.items()):
+        if (
+            isinstance(value, DTensor)
+            and _find_strided_shard_placement_from_fused_params(value.placements) is not None
+        ):
+            state_dict[key] = _replicate_dtensor(value)
+
+    dcp.save(
+        state_dict,
+        storage_writer=HuggingFaceStorageWriter(
+            path=checkpoint_dir,
+            save_distributed=True,
+            enable_consolidation=True,
+        ),
+    )
+    # Wait for rank 0 to finish writing the HF safetensors so other
+    # ranks don't return (and hit `from_pretrained`) before the files exist.
+    _distributed_barrier()
+
+
+def has_mixed_tensor_and_dtensor_params(params) -> bool:
+    has_dtensor = False
+    has_tensor = False
+    for param in params:
+        if isinstance(param, DTensor):
+            has_dtensor = True
+        elif isinstance(param, torch.Tensor):
+            has_tensor = True
+
+        if has_dtensor and has_tensor:
+            return True
+    return False
+
+
+def maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) -> None:
+    """
+    When get_optimizer_state_dict() or set_optimizer_state_dict() runs on an optimizer with no state yet,
+    PyTorch first materializes that state by doing a no-op step() with zero gradients. If an optimizer
+    group mixes regular tensors and DTensors, the batched foreach/fused optimizer kernels cannot process
+    that mixed group, so we turn those kernels off for such groups before distributed optimizer save/
+    load.
+    """
+    for i, param_group in enumerate(optimizer.param_groups):
+        if has_mixed_tensor_and_dtensor_params(param_group.get("params", ())):
+            logger.warning_once(
+                f"Param group {i} mixes regular tensors and DTensors; disabling foreach/fused "
+                "optimizer kernels for that group so distributed optimizer save/load can materialize state."
+            )
+            param_group["foreach"] = False
+            if "fused" in param_group:
+                param_group["fused"] = False
+
+
+def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
+    """Save optimizer state via DCP.
+
+    Params whose DTensors carry a lonely `_StridedShard` placement (e.g. Mixtral
+    `gate_up_proj`) are locally split into plain-`Shard` pieces at the boundary
+    so DCP only ever sees DTensors it can encode as one contiguous chunk per
+    rank.
+    """
+    maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer)
+    optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
+    fusion_metadata = get_fusion_metadata(optimizer_state_dict)
+    unfuse_optimizer_state(optimizer_state_dict, fusion_metadata)
+    dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
+
+
+def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
+    """Load optimizer state via DCP.
+
+    Symmetric to `save_optimizer_distributed`: build the unfused template, let DCP fill
+    it, then merge fused params back to their original `_StridedShard` form
+    before handing the state_dict back to the optimizer.
+    """
+    maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer)
+    optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
+    fusion_metadata = get_fusion_metadata(optimizer_state_dict)
+    unfuse_optimizer_state(optimizer_state_dict, fusion_metadata)
+    dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
+    fuse_optimizer_state(optimizer_state_dict, fusion_metadata)
+    set_optimizer_state_dict(model, optimizer, optimizer_state_dict)
+    maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer)
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index c3890ac42536..1eb08322cc17 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -52,6 +52,15 @@
     revert_weight_conversion,
 )
 from .distributed import DistributedConfig
+from .distributed.utils import (
+    _distributed_barrier,
+    gather_full_state_dict,
+    init_device_mesh,
+    save_model_checkpoint_distributed,
+)
+from .distributed.utils import (
+    distribute_model as distribute_model_distributed,
+)
 from .dynamic_module_utils import custom_object_save
 from .generation import CompileConfig, GenerationConfig
 from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled, is_fsdp_enabled
@@ -209,6 +218,12 @@ def _get_torch_distributed_world_size() -> int:
     return torch.distributed.get_world_size()
 
 
+def _get_torch_distributed_rank() -> int:
+    if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_rank"):
+        return 0
+    return torch.distributed.get_rank()
+
+
 def is_local_dist_rank_0():
     return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
 
@@ -3357,6 +3372,7 @@ def save_pretrained(
         token: str | bool | None = None,
         save_peft_format: bool = True,
         save_original_format: bool = True,
+        distributed_checkpoint: bool = False,
         **kwargs,
     ):
         """
@@ -3432,9 +3448,13 @@ def save_pretrained(
             return
 
         os.makedirs(save_directory, exist_ok=True)
+        save_on_this_rank = is_main_process
+        if _is_torch_distributed_initialized() and getattr(self, "device_mesh", None) is not None:
+            # DTensor gather materializes a full plain state dict only on global rank 0.
+            save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
         save_directory_path = os.fspath(save_directory)
 
-        if push_to_hub:
+        if push_to_hub and save_on_this_rank:
             commit_message = kwargs.pop("commit_message", None)
             repo_id = kwargs.pop("repo_id", save_directory_path.split(os.path.sep)[-1])
             create_pr = kwargs.pop("create_pr", False)
@@ -3459,46 +3479,72 @@ def save_pretrained(
 
         # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
         # loaded from the Hub.
-        if self._auto_class is not None:
+        if self._auto_class is not None and save_on_this_rank:
             custom_object_save(self, save_directory, config=self.config)
 
-        # Save the config
-        if is_main_process:
-            if not _hf_peft_config_loaded:
-                model_to_save.config.save_pretrained(save_directory)
-            if self.can_generate():
-                model_to_save.generation_config.save_pretrained(save_directory)
+        # Don't persist distributed_config in saved config — it's runtime-only
+        # (otherwise AutoConfig absorbs it on reload, preventing from_pretrained from seeing it as a kwarg).
+        # Keep a runtime copy around because TP/FSDP save helpers still rely on it after config serialization.
+        distributed_config = getattr(model_to_save.config, "distributed_config", None)
+        if distributed_config is not None:
+            del model_to_save.config.distributed_config
 
-            if _hf_peft_config_loaded:
-                logger.info(
-                    "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved."
-                )
-                state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict)
+        # Save the config
+        try:
+            if save_on_this_rank:
+                if not _hf_peft_config_loaded:
+                    model_to_save.config.save_pretrained(save_directory)
+                if self.can_generate():
+                    model_to_save.generation_config.save_pretrained(save_directory)
 
-                if save_peft_format:
+                if _hf_peft_config_loaded:
                     logger.info(
-                        "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`."
+                        "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved."
                     )
-                    peft_state_dict = {}
-                    for key, value in state_dict.items():
-                        peft_state_dict[f"base_model.model.{key}"] = value
-                    state_dict = peft_state_dict
+                    state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict)
 
-                active_adapter = self.active_adapters()
+                    if save_peft_format:
+                        logger.info(
+                            "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`."
+                        )
+                        peft_state_dict = {}
+                        for key, value in state_dict.items():
+                            peft_state_dict[f"base_model.model.{key}"] = value
+                        state_dict = peft_state_dict
 
-                if len(active_adapter) > 1:
-                    raise ValueError(
-                        "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one "
-                        "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`"
-                    )
-                active_adapter = active_adapter[0]
+                    active_adapter = self.active_adapters()
+
+                    if len(active_adapter) > 1:
+                        raise ValueError(
+                            "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one "
+                            "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`"
+                        )
+                    active_adapter = active_adapter[0]
 
-                current_peft_config = self.peft_config[active_adapter]
-                current_peft_config.save_pretrained(save_directory)
+                    current_peft_config = self.peft_config[active_adapter]
+                    current_peft_config.save_pretrained(save_directory)
+        finally:
+            if distributed_config is not None:
+                model_to_save.config.distributed_config = distributed_config
 
-        # Get the model state_dict
+        if distributed_checkpoint:
+            if _is_torch_distributed_initialized() and getattr(self, "device_mesh", None) is None:
+                raise ValueError(
+                    "save_pretrained(distributed_checkpoint=True) requires the model to have been "
+                    "initialized with a distributed_config (device_mesh is None)."
+                )
+            save_model_checkpoint_distributed(self, save_directory)
+            return
+
+        # Get the model state_dict (handles FSDP unshard + TP gather in one call)
+        used_distributed_gather = False
         if state_dict is None:
-            state_dict = model_to_save.state_dict()
+            if getattr(self, "device_mesh", None) is not None:
+                # Pass self (not model_to_save) so device_mesh/tp_size/tp_plan are available
+                state_dict = gather_full_state_dict(self)
+                used_distributed_gather = True
+            else:
+                state_dict = model_to_save.state_dict()
 
         # if any model parameters are offloaded, we need to know it for later
         is_offloaded = False
@@ -3524,8 +3570,8 @@ def save_pretrained(
                 if ignore_key in state_dict:
                     del state_dict[ignore_key]
 
-        # If model was sharded with TP, gather full tensors for saving
-        if self._tp_size is not None:
+        # If model was sharded with legacy TP, gather full tensors for saving
+        if self._tp_size is not None and not used_distributed_gather:
             state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size)
 
         # Remove tied weights as safetensors do not handle them
@@ -3569,54 +3615,55 @@ def save_pretrained(
                 filename.startswith(weights_no_suffix)
                 and os.path.isfile(full_filename)
                 and filename not in state_dict_split.filename_to_tensors
-                and is_main_process
+                and save_on_this_rank
                 and reg.fullmatch(filename_no_suffix) is not None
             ):
                 os.remove(full_filename)
 
         # Save the model
-        for shard_file, tensor_names in logging.tqdm(
-            state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
-        ):
-            filename = os.path.join(save_directory, shard_file)
-            shard_state_dict = {}
-            for tensor_name in tensor_names:
-                # Get the tensor, and remove it from state_dict to avoid keeping the ref
-                tensor = state_dict.pop(tensor_name)
-
-                # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
-                # but it would otherwise not be contained in the saved shard if we were to simply move the file
-                # or something
-                if is_offloaded and tensor.device.type == "meta":
-                    tensor = load_offloaded_parameter(model_to_save, tensor_name)
-
-                # only do contiguous after it's permuted correctly in case of TP
-                shard_state_dict[tensor_name] = tensor.contiguous()
-
-            # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
-            # so it's not possible for now....
-            # Write the shard to disk
-            safe_save_file(shard_state_dict, filename, metadata=metadata)
-            # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
-            del shard_state_dict
-
-        if index is None:
-            path_to_weights = os.path.join(save_directory, weights_name)
-            logger.info(f"Model weights saved in {path_to_weights}")
-        else:
-            save_index_file = SAFE_WEIGHTS_INDEX_NAME
-            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
-            # Save the index as well
-            with open(save_index_file, "w", encoding="utf-8") as f:
-                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-                f.write(content)
-            logger.info(
-                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
-                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
-                f"index located at {save_index_file}."
-            )
+        if save_on_this_rank:
+            for shard_file, tensor_names in logging.tqdm(
+                state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
+            ):
+                filename = os.path.join(save_directory, shard_file)
+                shard_state_dict = {}
+                for tensor_name in tensor_names:
+                    # Get the tensor, and remove it from state_dict to avoid keeping the ref
+                    tensor = state_dict.pop(tensor_name)
+
+                    # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
+                    # but it would otherwise not be contained in the saved shard if we were to simply move the file
+                    # or something
+                    if is_offloaded and tensor.device.type == "meta":
+                        tensor = load_offloaded_parameter(model_to_save, tensor_name)
+
+                    # only do contiguous after it's permuted correctly in case of TP
+                    shard_state_dict[tensor_name] = tensor.contiguous()
+
+                # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
+                # so it's not possible for now....
+                # Write the shard to disk
+                safe_save_file(shard_state_dict, filename, metadata=metadata)
+                # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
+                del shard_state_dict
+
+            if index is None:
+                path_to_weights = os.path.join(save_directory, weights_name)
+                logger.info(f"Model weights saved in {path_to_weights}")
+            else:
+                save_index_file = SAFE_WEIGHTS_INDEX_NAME
+                save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
+                # Save the index as well
+                with open(save_index_file, "w", encoding="utf-8") as f:
+                    content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+                    f.write(content)
+                logger.info(
+                    f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
+                    f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
+                    f"index located at {save_index_file}."
+                )
 
-        if push_to_hub:
+        if push_to_hub and save_on_this_rank:
             # Eventually create an empty model card
             model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
 
@@ -3632,6 +3679,13 @@ def save_pretrained(
                 create_pr=create_pr,
             )
 
+        # `gather_full_state_dict` concentrates the full state on rank 0 only;
+        # other ranks then loop over an empty shard list and would race ahead
+        # of rank 0's safetensors writes. Barrier so any subsequent
+        # `from_pretrained` on this path sees the consolidated files.
+        if used_distributed_gather:
+            _distributed_barrier()
+
     @wraps(PushToHubMixin.push_to_hub)
     def push_to_hub(self, *args, **kwargs):
         tags = self.model_tags if self.model_tags is not None else []
@@ -4138,8 +4192,17 @@ def from_pretrained(
         kernel_config = kwargs.pop("kernel_config", None)
         key_mapping = kwargs.pop("key_mapping", None)
 
-        if distributed_config is not None and tp_plan is None:
-            tp_plan = "auto"
+        if distributed_config is not None:
+            if device_map is not None:
+                raise ValueError(
+                    "`distributed_config` and `device_map` are mutually exclusive. "
+                    "`distributed_config` handles device placement natively via torch.distributed."
+                )
+            # NOTE(3outeille): support quantization (fp4/fp8) with distributed training later
+            if quantization_config is not None:
+                raise ValueError(
+                    "Quantization is not currently supported with distributed training. Please disable quantization or distributed_config."
+                )
 
         # Not used anymore -- remove them from the kwargs
         for name in ["mirror", "_fast_init", "low_cpu_mem_usage", "from_tf", "from_flax", "offload_state_dict"]:
@@ -4170,17 +4233,26 @@ def from_pretrained(
                 "`state_dict` cannot be passed together with a model name or a `gguf_file`. Use one of the two loading strategies."
             )
 
-        if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")):
-            logger.info(
-                "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. "
-                "If your plan is to load the model on each device, you should set device_map={"
-                ": PartialState().process_index} where PartialState comes from accelerate library"
-            )
+        if distributed_config is not None:
+            device_mesh = init_device_mesh(distributed_config)
+            # When using any distributed setup, we simply set the device_map here to the current rank so that we correctly
+            # load the params on the right device
+            if device_mesh.device_type == "cpu":
+                device_map = {"": torch.device("cpu")}
+            else:
+                device_map = {"": torch.device(device_mesh.device_type, int(os.environ["LOCAL_RANK"]))}
+        else:
+            if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")):
+                logger.info(
+                    "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. "
+                    "If your plan is to load the model on each device, you should set device_map={"
+                    ": PartialState().process_index} where PartialState comes from accelerate library"
+                )
 
-        if tp_plan is not None or tp_size is not None:  # TP warnings, and setup
-            device_map, device_mesh, tp_size = initialize_tensor_parallelism(
-                tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map
-            )
+            if tp_plan is not None or tp_size is not None:  # TP warnings, and setup
+                device_map, device_mesh, tp_size = initialize_tensor_parallelism(
+                    tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map
+                )
 
         if gguf_file is not None and not is_accelerate_available():
             raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")
@@ -4193,7 +4265,8 @@ def from_pretrained(
             download_kwargs_with_commit,
             **adapter_kwargs,
         )
-        device_map = check_and_set_device_map(device_map)  # warn, error and fix the device map
+        if distributed_config is None:
+            device_map = check_and_set_device_map(device_map)  # warn, error and fix the device map
 
         user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class}
         if from_pipeline is not None:
@@ -4333,11 +4406,16 @@ def from_pretrained(
         # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively
         weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer)
 
-        if _torch_distributed_available and device_mesh is not None:  # add hooks to nn.Modules: no weights
+        if distributed_config is not None:
+            # Tie weights before sharding so `apply_fully_shard_data_parallel` /
+            # `apply_tensor_parallel` see the shared-parameter graph and can route tied
+            # entries (e.g. `lm_head` -> `embed_tokens`) correctly. `_finalize_model_loading`
+            # re-runs `tie_weights` after the checkpoint is loaded to handle missing-key edge cases.
+            model.tie_weights()
+            model = distribute_model_distributed(model, distributed_config, device_mesh)
+        elif _torch_distributed_available and device_mesh is not None:  # legacy TP hooks: no weights yet
             model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size)
-
-        # Prepare the full device map
-        if device_map is not None:
+        elif device_map is not None:
             device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
 
         # Finalize model weight initialization
diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py
index 4331d0987bc7..6260fe8e4868 100644
--- a/tests/utils/test_modeling_utils.py
+++ b/tests/utils/test_modeling_utils.py
@@ -127,6 +127,7 @@
         T5ForConditionalGeneration,
     )
     from transformers.conversion_mapping import MergeModulelist, WeightConverter, get_model_conversion_mapping
+    from transformers.distributed import DistributedConfig
     from transformers.modeling_utils import (
         FLASH_ATTN_KERNEL_FALLBACK,
         _find_disjoint,
@@ -436,6 +437,45 @@ def test_get_total_byte_count_does_not_require_process_group(self):
         self.assertIn(torch.device("cpu"), total_byte_count)
         self.assertGreater(total_byte_count[torch.device("cpu")], 0)
 
+    def test_model_from_pretrained_fsdp_distributes_before_loading(self):
+        model = GPT2LMHeadModel(GPT2Config(n_layer=1, n_head=2, n_embd=8, n_positions=8, n_ctx=8, vocab_size=32))
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir)
+            call_order = []
+
+            fake_mesh = mock.Mock()
+            fake_mesh.mesh_dim_names = ("fsdp",)
+            fake_mesh.ndim = 1
+            fake_mesh.device_type = "cpu"
+
+            def fake_distribute_model(model, distributed_config, device_mesh):
+                call_order.append("distribute")
+                self.assertIs(device_mesh, fake_mesh)
+                return model
+
+            def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None):
+                call_order.append("load")
+                self.assertIs(load_config.device_mesh, fake_mesh)
+                return mock.Mock(), None
+
+            with (
+                patch("transformers.modeling_utils.init_device_mesh", return_value=fake_mesh),
+                patch(
+                    "transformers.modeling_utils.distribute_model_distributed",
+                    side_effect=fake_distribute_model,
+                ),
+                patch.object(GPT2LMHeadModel, "_load_pretrained_model", side_effect=fake_load_pretrained_model),
+                patch.object(
+                    GPT2LMHeadModel,
+                    "_finalize_model_loading",
+                    side_effect=lambda model, load_config, loading_info: loading_info,
+                ),
+            ):
+                GPT2LMHeadModel.from_pretrained(tmp_dir, distributed_config=DistributedConfig(fsdp_size=2))
+
+        self.assertEqual(call_order, ["distribute", "load"])
+
     def test_hub_retry(self):
         @hub_retry(max_attempts=2)
         def test_func():

From 00eb1166685205e8b2c75ab33eee83dc77b18fca Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 05:17:42 +0000
Subject: [PATCH 02/33] add fsdp plan to 2 models for now

---
 src/transformers/models/llama/configuration_llama.py         | 5 +++++
 src/transformers/models/llama/modeling_llama.py              | 1 +
 src/transformers/models/qwen3_moe/configuration_qwen3_moe.py | 5 +++++
 src/transformers/models/qwen3_moe/modeling_qwen3_moe.py      | 1 +
 src/transformers/models/qwen3_moe/modular_qwen3_moe.py       | 2 ++
 5 files changed, 14 insertions(+)

diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py
index 6960a6970592..9d84c0c16d3c 100644
--- a/src/transformers/models/llama/configuration_llama.py
+++ b/src/transformers/models/llama/configuration_llama.py
@@ -60,6 +60,11 @@ class LlamaConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
+    base_model_fsdp_plan = {
+        "embed_tokens": "free_full_weight",
+        "layers.*": "free_full_weight",
+        "norm": "keep_full_weight",
+    }
 
     vocab_size: int = 32000
     hidden_size: int = 4096
diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py
index 9d659c7c6f08..d56a47c3651c 100644
--- a/src/transformers/models/llama/modeling_llama.py
+++ b/src/transformers/models/llama/modeling_llama.py
@@ -430,6 +430,7 @@ class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
index 9a7d4b4c8b5b..ca040eedbec4 100644
--- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
@@ -81,6 +81,11 @@ class Qwen3MoeConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
+    base_model_fsdp_plan = {
+        "embed_tokens": "free_full_weight",
+        "layers.*": "free_full_weight",
+        "norm": "keep_full_weight",
+    }
 
     vocab_size: int = 151936
     hidden_size: int = 2048
diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
index ddf84fc575b7..7f4c7404356d 100644
--- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
@@ -611,6 +611,7 @@ class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/qwen3_moe/modular_qwen3_moe.py b/src/transformers/models/qwen3_moe/modular_qwen3_moe.py
index 0fd5b451959c..6453f1754a6a 100644
--- a/src/transformers/models/qwen3_moe/modular_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/modular_qwen3_moe.py
@@ -95,6 +95,8 @@ class Qwen3MoeModel(MixtralModel):
 
 
 class Qwen3MoeForCausalLM(MixtralForCausalLM):
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
+
     def __init__(self, config):
         super().__init__(config)
         self.model = Qwen3MoeModel(config)

From be296ddf9c24275e34b82be9c83975df18279362 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 05:31:05 +0000
Subject: [PATCH 03/33] add tests fsdp mixin

---
 pyproject.toml                     |   1 +
 src/transformers/modeling_utils.py |   2 +-
 src/transformers/testing_utils.py  |  17 +
 tests/causal_lm_tester.py          |   8 +-
 tests/test_fsdp_mixin.py           | 674 +++++++++++++++++++++++++++++
 5 files changed, 700 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_fsdp_mixin.py

diff --git a/pyproject.toml b/pyproject.toml
index f9b1a5f17a33..1ff95cdd7904 100644
--- a/pyproject.toml
+++ b/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"
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 9f579a5ab64a..1f738b8e1fa2 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3486,7 +3486,7 @@ def save_pretrained(
 
         # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
         # loaded from the Hub.
-        if self.is_remote_code() and save_on_this_rank::
+        if self.is_remote_code() and save_on_this_rank:
             custom_object_save(self, save_directory, config=self.config)
 
         # Don't persist distributed_config in saved config — it's runtime-only
diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py
index 48b342e00559..22093439d9dd 100644
--- a/src/transformers/testing_utils.py
+++ b/src/transformers/testing_utils.py
@@ -316,6 +316,7 @@ def parse_int_from_env(key, default=None):
 _run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False)
 _run_training_tests = parse_flag_from_env("RUN_TRAINING_TESTS", default=True)
 _run_tensor_parallel_tests = parse_flag_from_env("RUN_TENSOR_PARALLEL_TESTS", default=True)
+_run_fsdp_tests = parse_flag_from_env("RUN_FSDP_TESTS", default=True)
 
 
 def is_staging_test(test_case):
@@ -398,6 +399,22 @@ def is_tensor_parallel_test(test_case):
             return pytest.mark.is_tensor_parallel_test()(test_case)
 
 
+def is_fsdp_test(test_case):
+    """
+    Decorator marking a test as an FSDP test. If RUN_FSDP_TESTS is set to a falsy value, those tests will be
+    skipped.
+    """
+    if not _run_fsdp_tests:
+        return unittest.skip(reason="test is fsdp test")(test_case)
+    else:
+        try:
+            import pytest  # We don't need a hard dependency on pytest in the main library
+        except ImportError:
+            return test_case
+        else:
+            return pytest.mark.is_fsdp_test()(test_case)
+
+
 def slow(test_case):
     """
     Decorator marking a test as slow.
diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py
index 3bb99dddd7ed..1d0a066d67d5 100644
--- a/tests/causal_lm_tester.py
+++ b/tests/causal_lm_tester.py
@@ -39,6 +39,7 @@
     torch_device,
 )
 from .test_pipeline_mixin import PipelineTesterMixin
+from .test_fsdp_mixin import FSDPTesterMixin
 from .test_tensor_parallel_mixin import TensorParallelTesterMixin
 from .test_training_mixin import TrainingTesterMixin
 
@@ -271,7 +272,12 @@ def prepare_config_and_inputs_for_common(self):
 
 @require_torch
 class CausalLMModelTest(
-    ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TrainingTesterMixin, TensorParallelTesterMixin
+    ModelTesterMixin,
+    GenerationTesterMixin,
+    PipelineTesterMixin,
+    TrainingTesterMixin,
+    TensorParallelTesterMixin,
+    FSDPTesterMixin,
 ):
     model_tester_class = None
     all_model_classes = None
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
new file mode 100644
index 000000000000..b046980fe294
--- /dev/null
+++ b/tests/test_fsdp_mixin.py
@@ -0,0 +1,674 @@
+# 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.
+# 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.
+
+"""FSDP tester mixin for model tests."""
+
+import json
+import logging
+import os
+import socket
+import sys
+import tempfile
+import time
+import traceback
+from abc import ABC, abstractmethod
+
+from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, is_torch_available
+from transformers.testing_utils import (
+    backend_device_count,
+    backend_empty_cache,
+    backend_torch_accelerator_module,
+    init_test_logger,
+    is_fsdp_test,
+    require_torch_greater_or_equal,
+)
+from transformers.trainer_utils import set_seed
+
+
+logger = logging.getLogger("transformers.training_test")
+
+
+if is_torch_available():
+    import torch
+    import torch.distributed as dist
+    import torch.multiprocessing as mp
+    from torch.nn.parallel import DistributedDataParallel as DDP
+
+    from transformers.distributed import DistributedConfig
+    from transformers.distributed.fsdp import _resolve_tied_embed_lm_head_plan, expand_fsdp_plan
+    from transformers.distributed.utils import (
+        gather_full_state_dict,
+        load_optimizer_distributed,
+        save_optimizer_distributed,
+    )
+
+
+# =============================================================================
+# Constants
+# =============================================================================
+
+BATCH_SIZE = 2
+SEQ_LEN = 64
+NUM_STEPS = 20
+LR = 3e-4
+SEED = 42
+FSDP_TOP_MODEL_NAMES = {
+    # Models with `base_model_fsdp_plan` + head-class `_fsdp_plan` on this branch.
+    "llama",
+    "qwen3_moe",
+}
+
+
+# =============================================================================
+# Distributed helpers (top-level for pickling by mp.spawn)
+# =============================================================================
+
+
+def _get_distributed_device_type():
+    device_type = torch._C._get_accelerator().type
+    return "cpu" if device_type == "mps" else device_type
+
+
+def _get_distributed_backend():
+    backend_map = {"cpu": "gloo", "cuda": "nccl", "xpu": "xccl", "hpu": "hccl"}
+    return backend_map.get(_get_distributed_device_type(), "gloo")
+
+
+def _get_rank_device(rank):
+    device_type = _get_distributed_device_type()
+    if device_type == "cpu":
+        return torch.device("cpu")
+    return torch.device(device_type, rank)
+
+
+def _set_rank_device(rank):
+    accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
+    if accelerator_module is not None and hasattr(accelerator_module, "set_device"):
+        accelerator_module.set_device(rank)
+
+
+def _get_accelerator_rng_state():
+    accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
+    if accelerator_module is None or not hasattr(accelerator_module, "get_rng_state"):
+        return None
+    return accelerator_module.get_rng_state()
+
+
+def _set_accelerator_rng_state(rng_state):
+    accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
+    if rng_state is not None and accelerator_module is not None and hasattr(accelerator_module, "set_rng_state"):
+        accelerator_module.set_rng_state(rng_state)
+
+
+def _get_available_fsdp_workers():
+    if _get_distributed_device_type() == "cpu":
+        return os.cpu_count() or 1
+    return backend_device_count(_get_distributed_device_type())
+
+
+def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_size, port, results_file):
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = str(port)
+
+    _set_determinism(SEED)
+
+    dist.init_process_group(backend=_get_distributed_backend(), rank=rank, world_size=world_size)
+    _set_rank_device(rank)
+
+    if rank == 0:
+        start_time = time.perf_counter()
+        print(f"[FSDP] Starting test: {test_name}", flush=True)
+
+    error = None
+    try:
+        func(rank, *func_args, **func_kwargs)
+    except Exception as e:
+        error = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
+
+    error_flag = torch.tensor([1 if error else 0], device=_get_rank_device(rank))
+    dist.all_reduce(error_flag, op=dist.ReduceOp.MAX)
+    any_failed = error_flag.item() > 0
+
+    if rank == 0:
+        elapsed = time.perf_counter() - start_time
+        status = "FAIL" if any_failed else "PASS"
+        output_stream = sys.stderr if any_failed else sys.stdout
+        print(f"[FSDP] {status} test: {test_name} ({elapsed:.1f}s)", file=output_stream, flush=True)
+        with open(results_file, "w") as f:
+            json.dump({"error": error or ("Failed on another rank" if any_failed else None)}, f)
+
+    backend_empty_cache(_get_distributed_device_type())
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+def _set_determinism(seed):
+    torch.use_deterministic_algorithms(True)
+    if _get_distributed_device_type() == "cuda" and torch.cuda.is_available():
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+    set_seed(seed)
+
+
+# =============================================================================
+# Training & comparison helpers (top-level for pickling)
+# =============================================================================
+
+
+def _build_repeated_training_batches(config, device, num_steps):
+    """Create one deterministic batch and reuse it across steps."""
+    generator = torch.Generator(device=device)
+    generator.manual_seed(SEED)
+    input_ids = torch.randint(0, config.vocab_size, (BATCH_SIZE, SEQ_LEN), device=device, generator=generator)
+    labels = input_ids.clone()
+    return [(input_ids, labels)] * num_steps
+
+
+def _create_shared_tmpdir(rank):
+    if rank == 0:
+        tmpdir_obj = tempfile.TemporaryDirectory()
+        tmpdir = tmpdir_obj.name
+        tmpdir_list = [tmpdir]
+    else:
+        tmpdir_obj = None
+        tmpdir_list = [None]
+    dist.broadcast_object_list(tmpdir_list, src=0)
+    return tmpdir_list[0], tmpdir_obj
+
+
+def _gather_ddp_state_dict(model):
+    # Only rank 0 returns data to match gather_full_state_dict semantics, so the
+    # downstream DDP-vs-FSDP comparison runs once on rank 0 instead of N times.
+    if dist.get_rank() != 0:
+        return {}
+    return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()}
+
+
+def _fsdp_target_module_names(model) -> set[str]:
+    """Return module names that should receive a ``fully_shard`` call (excluding root)."""
+    adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, model)
+    reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan)
+    return {name for name, _ in reshard_targets + no_reshard_targets}
+
+
+def _save_init_pretrained(rank, config, dtype):
+    """Save a deterministic initial model to a shared tmpdir for from_pretrained loading."""
+    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
+    if rank == 0:
+        set_seed(SEED)
+        model = AutoModelForCausalLM.from_config(config).to(dtype)
+        model.save_pretrained(tmpdir)
+        del model
+    dist.barrier()
+    return tmpdir, tmpdir_obj
+
+
+def _save_training_state(model, optimizer, training_state_dir):
+    """Save optimizer (canonical DCP path) plus per-rank RNG for resume."""
+    save_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim"))
+    rng = {"cpu": torch.get_rng_state()}
+    accel = _get_accelerator_rng_state()
+    if accel is not None:
+        rng["accel"] = accel
+    torch.save(rng, os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt"))
+
+
+def _load_training_state(model, optimizer, training_state_dir):
+    """Inverse of `_save_training_state`."""
+    load_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim"))
+    rng = torch.load(os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt"), weights_only=False)
+    torch.set_rng_state(rng["cpu"])
+    if "accel" in rng:
+        _set_accelerator_rng_state(rng["accel"])
+
+
+def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
+    _set_determinism(SEED)
+    model = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype).to(device)
+    # MoE/conditional-routing variants may not use all params on every step, and DDP would otherwise fail.
+    ddp_kwargs = {"find_unused_parameters": True}
+    if device.type != "cpu":
+        ddp_kwargs["device_ids"] = [rank]
+    ddp_model = DDP(model, **ddp_kwargs)
+    ddp_model.train()
+    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr)
+
+    losses, grad_norms = [], []
+    for input_ids, labels in batches:
+        optimizer.zero_grad()
+        output = ddp_model(input_ids=input_ids, labels=labels, use_cache=False)
+        loss = output.loss
+        loss.backward()
+        grad_norm = torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=float("inf"))
+        optimizer.step()
+
+        losses.append(loss.detach().item())
+        grad_norms.append(grad_norm)
+
+    state_dict = _gather_ddp_state_dict(ddp_model)
+
+    del optimizer, ddp_model, model
+    backend_empty_cache(_get_distributed_device_type())
+    dist.barrier()
+
+    return losses, grad_norms, state_dict
+
+
+def train_fsdp2(
+    rank,
+    batches,
+    lr,
+    dtype,
+    init_model_dir,
+    checkpoint_step,
+    fsdp_cpu_offload=False,
+    fsdp_mixed_precision=False,
+):
+    # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save
+    _set_determinism(SEED)
+    distributed_config = DistributedConfig(
+        fsdp_size=dist.get_world_size(),
+        fsdp_cpu_offload=fsdp_cpu_offload,
+        fsdp_mixed_precision=fsdp_mixed_precision,
+    )
+    pre_ckpt_model = AutoModelForCausalLM.from_pretrained(
+        init_model_dir, torch_dtype=dtype, distributed_config=distributed_config
+    )
+    pre_ckpt_model.train()
+    pre_ckpt_optimizer = torch.optim.Adam(pre_ckpt_model.parameters(), lr=lr)
+
+    pre_ckpt_losses, pre_ckpt_grad_norms = [], []
+    for step in range(0, checkpoint_step):
+        input_ids, labels = batches[step]
+        pre_ckpt_optimizer.zero_grad()
+        output = pre_ckpt_model(input_ids=input_ids, labels=labels, use_cache=False)
+        loss = output.loss
+        loss.backward()
+        grad_norm = torch.nn.utils.clip_grad_norm_(pre_ckpt_model.parameters(), max_norm=float("inf"))
+        pre_ckpt_optimizer.step()
+
+        pre_ckpt_losses.append(loss.detach().item())
+        pre_ckpt_grad_norms.append(grad_norm)
+
+    # -- Phase 2: Save checkpoint, then load into a fresh model
+    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
+    try:
+        model_dir = os.path.join(tmpdir, "model")
+        training_state_dir = os.path.join(tmpdir, "training_state")
+
+        pre_ckpt_model.save_pretrained(model_dir, is_main_process=(rank == 0))
+        _save_training_state(pre_ckpt_model, pre_ckpt_optimizer, training_state_dir)
+        dist.barrier()
+
+        # Intentionally scramble RNG to prove checkpoint restore works
+        _set_determinism(SEED + 1234)
+        resumed_model = AutoModelForCausalLM.from_pretrained(
+            model_dir, torch_dtype=dtype, distributed_config=distributed_config
+        )
+        resumed_model.train()
+        resumed_optimizer = torch.optim.Adam(resumed_model.parameters(), lr=lr)
+
+        _load_training_state(resumed_model, resumed_optimizer, training_state_dir)
+        dist.barrier()
+    finally:
+        if rank == 0:
+            tmpdir_obj.cleanup()
+
+    # -- Phase 3: Post-checkpoint run -- continue training the remaining steps from the resumed model
+    post_ckpt_losses, post_ckpt_grad_norms = [], []
+    for step in range(checkpoint_step, len(batches)):
+        input_ids, labels = batches[step]
+        resumed_optimizer.zero_grad()
+        output = resumed_model(input_ids=input_ids, labels=labels, use_cache=False)
+        loss = output.loss
+        loss.backward()
+        grad_norm = torch.nn.utils.clip_grad_norm_(resumed_model.parameters(), max_norm=float("inf"))
+        resumed_optimizer.step()
+
+        post_ckpt_losses.append(loss.detach().item())
+        post_ckpt_grad_norms.append(grad_norm)
+
+    combined_losses = pre_ckpt_losses + post_ckpt_losses
+    combined_grad_norms = pre_ckpt_grad_norms + post_ckpt_grad_norms
+    combined_state_dict = gather_full_state_dict(resumed_model)
+
+    return combined_losses, combined_grad_norms, combined_state_dict
+
+
+# =============================================================================
+# Distributed test implementations (top-level for pickling by mp.spawn)
+# =============================================================================
+
+
+def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
+    """Train FSDP2 model, save via save_pretrained, load via from_pretrained, compare state dicts."""
+    init_test_logger()
+
+    device = _get_rank_device(rank)
+    config = config_class.from_dict(config_dict)
+
+    batches = _build_repeated_training_batches(config, device, 3)
+
+    distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
+
+    init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32)
+    try:
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(init_tmpdir, distributed_config=distributed_config)
+        dist.barrier()
+    finally:
+        if rank == 0 and init_tmpdir_obj is not None:
+            init_tmpdir_obj.cleanup()
+    model.train()
+    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+
+    for input_ids, labels in batches:
+        optimizer.zero_grad()
+        output = model(input_ids=input_ids, labels=labels, use_cache=False)
+        output.loss.backward()
+        optimizer.step()
+
+    state_dict_before = gather_full_state_dict(model)
+
+    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
+    try:
+        model.save_pretrained(tmpdir, is_main_process=(rank == 0))
+        dist.barrier()
+
+        new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
+        dist.barrier()
+    finally:
+        if rank == 0:
+            tmpdir_obj.cleanup()
+
+    state_dict_after = gather_full_state_dict(new_model)
+
+    for key in state_dict_before:
+        assert key in state_dict_after, f"Key {key} missing after load"
+        torch.testing.assert_close(
+            state_dict_before[key],
+            state_dict_after[key],
+            rtol=0,
+            atol=0,
+            msg=f"Weight mismatch for {key} after save/load",
+        )
+
+    if rank == 0:
+        logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
+
+
+def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings):
+    """Verify that apply_fully_sharded_data_parallel wraps exactly the right modules."""
+    init_test_logger()
+
+    config = config_class.from_dict(config_dict)
+    config.tie_word_embeddings = tie_word_embeddings
+
+    init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32)
+    try:
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(
+            init_tmpdir, distributed_config=DistributedConfig(fsdp_size=dist.get_world_size())
+        )
+        dist.barrier()
+    finally:
+        if rank == 0 and init_tmpdir_obj is not None:
+            init_tmpdir_obj.cleanup()
+
+    expected_targets = {""} | _fsdp_target_module_names(model)
+
+    actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")}
+
+    if rank == 0:
+        logger.debug(f"  Weights tied: {config.tie_word_embeddings}")
+        logger.debug(f"  Expected FSDP targets: {sorted(expected_targets)}")
+        logger.debug(f"  Actual FSDP targets:   {sorted(actual_targets)}")
+
+    missing = expected_targets - actual_targets
+    extra = actual_targets - expected_targets
+    assert not missing and not extra, (
+        f"FSDP target mismatch.\n"
+        f"  Missing (expected but not wrapped): {sorted(missing)}\n"
+        f"  Extra (wrapped but not expected):   {sorted(extra)}"
+    )
+
+    if rank == 0:
+        logger.debug(f"  FSDP sharding structure OK ({len(actual_targets)} targets)")
+
+
+def _test_fsdp2_plan_vs_ddp_impl(
+    rank, config_class, config_dict, tie_word_embeddings, policy_options=None, dtype=None
+):
+    """Validate DDP-vs-FSDP2 trace matching using the model's declared FSDP plan."""
+    init_test_logger()
+
+    if dtype is None:
+        dtype = torch.float32
+
+    policy_options = policy_options or []
+    assert "mixed_precision" not in policy_options, (
+        "Use the mixed-precision specific tests when enabling mixed_precision policy."
+    )
+
+    device = _get_rank_device(rank)
+    config = config_class.from_dict(config_dict)
+    config.tie_word_embeddings = tie_word_embeddings
+
+    fsdp_cpu_offload = "cpu_offload" in policy_options
+    fsdp_mixed_precision = "mixed_precision" in policy_options
+    test_label = f"FSDP2{'+policies' if policy_options else ''}"
+
+    checkpoint_step = NUM_STEPS // 2
+    init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype)
+    batches = _build_repeated_training_batches(config, device, NUM_STEPS)
+    try:
+        ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir)
+
+        fsdp_losses, fsdp_grad_norms, fsdp_state_dict = train_fsdp2(
+            rank,
+            batches,
+            LR,
+            dtype,
+            init_model_dir=init_model_dir,
+            checkpoint_step=checkpoint_step,
+            fsdp_cpu_offload=fsdp_cpu_offload,
+            fsdp_mixed_precision=fsdp_mixed_precision,
+        )
+    finally:
+        if rank == 0 and init_tmpdir_obj is not None:
+            init_tmpdir_obj.cleanup()
+
+    for step in range(len(ddp_losses)):
+        torch.testing.assert_close(
+            torch.tensor(ddp_losses[step]),
+            torch.tensor(fsdp_losses[step]),
+            rtol=1e-5,
+            atol=1e-5,
+            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}",
+        )
+        torch.testing.assert_close(
+            torch.tensor(ddp_grad_norms[step]),
+            torch.tensor(fsdp_grad_norms[step]),
+            rtol=1e-5,
+            atol=1e-5,
+            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}",
+        )
+
+    for key in ddp_state_dict:
+        assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict"
+        torch.testing.assert_close(
+            ddp_state_dict[key],
+            fsdp_state_dict[key],
+            rtol=1e-5,
+            atol=1e-5,
+            msg=f"Weight mismatch for {key}: DDP vs {test_label}",
+        )
+
+    if rank == 0:
+        logger.debug(f"DDP and {test_label} comparison checks passed.")
+
+
+# =============================================================================
+# Mixin class
+# =============================================================================
+
+
+class FSDPTesterMixin(ABC):
+    fsdp_nproc_per_node: int = 2
+    skip_fsdp_tests: bool = False
+
+    @property
+    @abstractmethod
+    def model_tester(self):
+        """The model tester instance (e.g., CausalLMModelTester)."""
+        ...
+
+    def _skip_if_fsdp_disabled(self):
+        if self.skip_fsdp_tests:
+            self.skipTest("FSDP tests disabled for this model (skip_fsdp_tests=True)")
+
+    def _skip_if_insufficient_devices(self):
+        self._skip_if_fsdp_disabled()
+        available_workers = _get_available_fsdp_workers()
+        if available_workers < self.fsdp_nproc_per_node:
+            self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}")
+
+    def _get_fsdp_model_name(self):
+        module_parts = self.__class__.__module__.split(".")
+        if len(module_parts) >= 3 and module_parts[0] == "tests" and module_parts[1] == "models":
+            return module_parts[2]
+        return None
+
+    def _skip_if_fsdp_model_not_selected(self):
+        model_name = self._get_fsdp_model_name()
+        if model_name not in FSDP_TOP_MODEL_NAMES:
+            model_label = model_name or self.__class__.__module__
+            self.skipTest(
+                "FSDP mixin coverage is currently limited to models with declared FSDP plans "
+                f"(skipping {model_label})."
+            )
+
+    def _create_model_on_meta(self, config):
+        """Instantiate a model on the meta device (no memory allocated)."""
+        auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM]
+        for auto_cls in auto_classes:
+            try:
+                with torch.device("meta"):
+                    return auto_cls.from_config(config)
+            except Exception:
+                continue
+        self.skipTest(f"Cannot instantiate model with any Auto class for config {type(config).__name__}")
+
+    def _get_tiny_config(self):
+        """Get config class and serialized dict for passing to spawned processes."""
+        config = self.model_tester.get_config()
+        config.vocab_size = 256
+        config.hidden_size = 64
+        config.intermediate_size = 128
+        if hasattr(config, "ffn_config"):
+            if hasattr(config.ffn_config, "ffn_hidden_size"):
+                config.ffn_config.ffn_hidden_size = config.hidden_size
+            if hasattr(config.ffn_config, "hidden_size"):
+                config.ffn_config.hidden_size = config.intermediate_size
+        if hasattr(config, "num_attention_heads"):
+            config.num_attention_heads = 4
+        if hasattr(config, "num_key_value_heads"):
+            config.num_key_value_heads = 4
+        if hasattr(config, "moe_intermediate_size"):
+            config.moe_intermediate_size = 32
+        if hasattr(config, "vocab_size_per_layer_input"):
+            config.vocab_size_per_layer_input = config.vocab_size
+        config_dict = config.to_diff_dict()
+        return type(config), config_dict
+
+    def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs):
+        self._skip_if_fsdp_model_not_selected()
+        self._skip_if_insufficient_devices()
+
+        config_class, config_dict = self._get_tiny_config()
+        func_args = (config_class, config_dict, *test_args)
+
+        results_file = tempfile.mktemp(suffix=".json")
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            port = s.getsockname()[1]
+
+        try:
+            mp.spawn(
+                _fsdp_global_wrapper,
+                args=(test_name, test_impl, func_args, test_kwargs, self.fsdp_nproc_per_node, port, results_file),
+                nprocs=self.fsdp_nproc_per_node,
+            )
+
+            with open(results_file) as f:
+                result = json.load(f)
+        finally:
+            if os.path.exists(results_file):
+                os.unlink(results_file)
+
+        if result["error"] is not None:
+            self.fail(f"FSDP test '{test_name}' failed:\n{result['error']}")
+
+    @is_fsdp_test
+    def test_fsdp_plan_declared(self):
+        """The model exposes a non-empty `_fsdp_plan` derived from config + class-level overrides."""
+        self._skip_if_fsdp_disabled()
+        self._skip_if_fsdp_model_not_selected()
+        start_time = time.perf_counter()
+        logger.info("[FSDP] Starting test: test_fsdp_plan_declared")
+        status = "FAIL"
+        try:
+            config = self.model_tester.get_config()
+            model = self._create_model_on_meta(config)
+
+            plan = getattr(model, "_fsdp_plan", None)
+            self.assertTrue(plan, f"No _fsdp_plan declared for {type(model).__name__}")
+            status = "PASS"
+        finally:
+            logger.info("[FSDP] %s test: test_fsdp_plan_declared (%.1fs)", status, time.perf_counter() - start_time)
+
+    @is_fsdp_test
+    @require_torch_greater_or_equal("2.6")
+    def test_fsdp2_sharding_structure_untied(self):
+        self._run_fsdp2_distributed_test(
+            "test_fsdp2_sharding_structure_untied", _test_fsdp2_sharding_structure_impl, False
+        )
+
+    @is_fsdp_test
+    @require_torch_greater_or_equal("2.6")
+    def test_fsdp2_sharding_structure_tied(self):
+        self._run_fsdp2_distributed_test(
+            "test_fsdp2_sharding_structure_tied", _test_fsdp2_sharding_structure_impl, True
+        )
+
+    @is_fsdp_test
+    @require_torch_greater_or_equal("2.6")
+    def test_fsdp2_save_load(self):
+        self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl)
+
+    @is_fsdp_test
+    @require_torch_greater_or_equal("2.6")
+    def test_fsdp2_plan_vs_ddp_untied(self):
+        self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False)
+
+    @is_fsdp_test
+    @require_torch_greater_or_equal("2.6")
+    def test_fsdp2_plan_vs_ddp_tied(self):
+        self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True)

From 05900d6051d6b876fde2866e07d04f2e4d615b89 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 05:39:06 +0000
Subject: [PATCH 04/33] linting

---
 src/transformers/distributed/__init__.py | 4 ++--
 tests/causal_lm_tester.py                | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py
index 318fd2961360..dff6acb66217 100644
--- a/src/transformers/distributed/__init__.py
+++ b/src/transformers/distributed/__init__.py
@@ -21,12 +21,12 @@
     "configuration_utils": ["DistributedConfig"],
     "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"],
     "utils": [
-        "init_device_mesh",
         "distribute_model",
         "gather_full_state_dict",
+        "init_device_mesh",
+        "load_optimizer_distributed",
         "save_model_checkpoint_distributed",
         "save_optimizer_distributed",
-        "load_optimizer_distributed",
     ],
 }
 
diff --git a/tests/causal_lm_tester.py b/tests/causal_lm_tester.py
index 1d0a066d67d5..2efe6d1bcb70 100644
--- a/tests/causal_lm_tester.py
+++ b/tests/causal_lm_tester.py
@@ -30,6 +30,7 @@
 )
 
 from .test_configuration_common import ConfigTester
+from .test_fsdp_mixin import FSDPTesterMixin
 from .test_modeling_common import (
     GenerationTesterMixin,
     ModelTesterMixin,
@@ -39,7 +40,6 @@
     torch_device,
 )
 from .test_pipeline_mixin import PipelineTesterMixin
-from .test_fsdp_mixin import FSDPTesterMixin
 from .test_tensor_parallel_mixin import TensorParallelTesterMixin
 from .test_training_mixin import TrainingTesterMixin
 

From fc2423bd2262b2446782da6bf41ab5b085bb983e Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 06:15:26 +0000
Subject: [PATCH 05/33] refactor test fsdp mixin

---
 tests/test_fsdp_mixin.py | 414 ++++++++++++++++++++-------------------
 1 file changed, 216 insertions(+), 198 deletions(-)

diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index b046980fe294..c3e13b2a70a8 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -24,6 +24,8 @@
 import traceback
 from abc import ABC, abstractmethod
 
+from parameterized import parameterized
+
 from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, is_torch_available
 from transformers.testing_utils import (
     backend_device_count,
@@ -63,6 +65,9 @@
 NUM_STEPS = 20
 LR = 3e-4
 SEED = 42
+DDP_FSDP_RTOL = 1e-5
+DDP_FSDP_ATOL = 1e-5
+
 FSDP_TOP_MODEL_NAMES = {
     # Models with `base_model_fsdp_plan` + head-class `_fsdp_plan` on this branch.
     "llama",
@@ -75,6 +80,12 @@
 # =============================================================================
 
 
+def _find_free_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
 def _get_distributed_device_type():
     device_type = torch._C._get_accelerator().type
     return "cpu" if device_type == "mps" else device_type
@@ -117,6 +128,28 @@ def _get_available_fsdp_workers():
     return backend_device_count(_get_distributed_device_type())
 
 
+def _set_determinism(seed):
+    torch.use_deterministic_algorithms(True)
+    if _get_distributed_device_type() == "cuda" and torch.cuda.is_available():
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+    set_seed(seed)
+
+
+def _create_shared_tmpdir(rank):
+    if rank == 0:
+        tmpdir_obj = tempfile.TemporaryDirectory()
+        tmpdir_list = [tmpdir_obj.name]
+    else:
+        tmpdir_obj = None
+        tmpdir_list = [None]
+    dist.broadcast_object_list(tmpdir_list, src=0)
+    return tmpdir_list[0], tmpdir_obj
+
+
 def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_size, port, results_file):
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["RANK"] = str(rank)
@@ -125,7 +158,6 @@ def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_si
     os.environ["MASTER_PORT"] = str(port)
 
     _set_determinism(SEED)
-
     dist.init_process_group(backend=_get_distributed_backend(), rank=rank, world_size=world_size)
     _set_rank_device(rank)
 
@@ -156,19 +188,8 @@ def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_si
     dist.destroy_process_group()
 
 
-def _set_determinism(seed):
-    torch.use_deterministic_algorithms(True)
-    if _get_distributed_device_type() == "cuda" and torch.cuda.is_available():
-        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-        torch.backends.cuda.matmul.allow_tf32 = False
-        torch.backends.cudnn.allow_tf32 = False
-    set_seed(seed)
-
-
 # =============================================================================
-# Training & comparison helpers (top-level for pickling)
+# Training helpers (top-level for pickling)
 # =============================================================================
 
 
@@ -181,16 +202,19 @@ def _build_repeated_training_batches(config, device, num_steps):
     return [(input_ids, labels)] * num_steps
 
 
-def _create_shared_tmpdir(rank):
-    if rank == 0:
-        tmpdir_obj = tempfile.TemporaryDirectory()
-        tmpdir = tmpdir_obj.name
-        tmpdir_list = [tmpdir]
-    else:
-        tmpdir_obj = None
-        tmpdir_list = [None]
-    dist.broadcast_object_list(tmpdir_list, src=0)
-    return tmpdir_list[0], tmpdir_obj
+def _run_training_steps(model, optimizer, batches, *, track_grad_norms=True):
+    """Forward/backward/step over batches. Returns (losses, grad_norms)."""
+    losses, grad_norms = [], []
+    for input_ids, labels in batches:
+        optimizer.zero_grad()
+        loss = model(input_ids=input_ids, labels=labels, use_cache=False).loss
+        loss.backward()
+        if track_grad_norms:
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=float("inf"))
+            grad_norms.append(grad_norm)
+        optimizer.step()
+        losses.append(loss.detach().item())
+    return losses, grad_norms
 
 
 def _gather_ddp_state_dict(model):
@@ -220,6 +244,25 @@ def _save_init_pretrained(rank, config, dtype):
     return tmpdir, tmpdir_obj
 
 
+def _load_fsdp_model_from_init_checkpoint(rank, config, dtype, *, distributed_config=None, tie_word_embeddings=None):
+    """Save deterministic init weights, load FSDP model, cleanup init tmpdir."""
+    if tie_word_embeddings is not None:
+        config.tie_word_embeddings = tie_word_embeddings
+
+    init_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype)
+    try:
+        _set_determinism(SEED)
+        kwargs = {}
+        if distributed_config is not None:
+            kwargs["distributed_config"] = distributed_config
+        model = AutoModelForCausalLM.from_pretrained(init_dir, **kwargs)
+        dist.barrier()
+        return model
+    finally:
+        if rank == 0 and init_tmpdir_obj is not None:
+            init_tmpdir_obj.cleanup()
+
+
 def _save_training_state(model, optimizer, training_state_dir):
     """Save optimizer (canonical DCP path) plus per-rank RNG for resume."""
     save_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim"))
@@ -239,6 +282,75 @@ def _load_training_state(model, optimizer, training_state_dir):
         _set_accelerator_rng_state(rng["accel"])
 
 
+def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config, lr):
+    """Save model+optimizer, scramble RNG, reload and restore training state."""
+    rank = dist.get_rank()
+    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
+    try:
+        model_dir = os.path.join(tmpdir, "model")
+        training_state_dir = os.path.join(tmpdir, "training_state")
+
+        pre_model.save_pretrained(model_dir, is_main_process=(rank == 0))
+        _save_training_state(pre_model, pre_optimizer, training_state_dir)
+        dist.barrier()
+
+        # Intentionally scramble RNG to prove checkpoint restore works
+        _set_determinism(SEED + 1234)
+        resumed_model = AutoModelForCausalLM.from_pretrained(
+            model_dir, torch_dtype=dtype, distributed_config=distributed_config
+        )
+        resumed_model.train()
+        resumed_optimizer = torch.optim.Adam(resumed_model.parameters(), lr=lr)
+        _load_training_state(resumed_model, resumed_optimizer, training_state_dir)
+        dist.barrier()
+        return resumed_model, resumed_optimizer
+    finally:
+        if rank == 0:
+            tmpdir_obj.cleanup()
+
+
+def _assert_state_dicts_equal(before, after, *, rtol, atol, msg_prefix=""):
+    for key in before:
+        assert key in after, f"{msg_prefix}Key {key} missing after load"
+        torch.testing.assert_close(
+            before[key],
+            after[key],
+            rtol=rtol,
+            atol=atol,
+            msg=f"{msg_prefix}Weight mismatch for {key}",
+        )
+
+
+def _assert_ddp_fsdp_match(
+    ddp_losses, fsdp_losses, ddp_grad_norms, fsdp_grad_norms, ddp_state, fsdp_state, test_label
+):
+    for step in range(len(ddp_losses)):
+        torch.testing.assert_close(
+            torch.tensor(ddp_losses[step]),
+            torch.tensor(fsdp_losses[step]),
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}",
+        )
+        torch.testing.assert_close(
+            torch.tensor(ddp_grad_norms[step]),
+            torch.tensor(fsdp_grad_norms[step]),
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}",
+        )
+
+    for key in ddp_state:
+        assert key in fsdp_state, f"Key {key} missing from {test_label} state dict"
+        torch.testing.assert_close(
+            ddp_state[key],
+            fsdp_state[key],
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Weight mismatch for {key}: DDP vs {test_label}",
+        )
+
+
 def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
     _set_determinism(SEED)
     model = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype).to(device)
@@ -250,18 +362,7 @@ def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
     ddp_model.train()
     optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr)
 
-    losses, grad_norms = [], []
-    for input_ids, labels in batches:
-        optimizer.zero_grad()
-        output = ddp_model(input_ids=input_ids, labels=labels, use_cache=False)
-        loss = output.loss
-        loss.backward()
-        grad_norm = torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=float("inf"))
-        optimizer.step()
-
-        losses.append(loss.detach().item())
-        grad_norms.append(grad_norm)
-
+    losses, grad_norms = _run_training_steps(ddp_model, optimizer, batches)
     state_dict = _gather_ddp_state_dict(ddp_model)
 
     del optimizer, ddp_model, model
@@ -281,75 +382,38 @@ def train_fsdp2(
     fsdp_cpu_offload=False,
     fsdp_mixed_precision=False,
 ):
-    # -- Phase 1: Pre-checkpoint run -- train only the first `checkpoint_step` steps, then save
-    _set_determinism(SEED)
     distributed_config = DistributedConfig(
         fsdp_size=dist.get_world_size(),
         fsdp_cpu_offload=fsdp_cpu_offload,
         fsdp_mixed_precision=fsdp_mixed_precision,
     )
+
+    # Phase 1: Pre-checkpoint run
+    _set_determinism(SEED)
     pre_ckpt_model = AutoModelForCausalLM.from_pretrained(
         init_model_dir, torch_dtype=dtype, distributed_config=distributed_config
     )
     pre_ckpt_model.train()
     pre_ckpt_optimizer = torch.optim.Adam(pre_ckpt_model.parameters(), lr=lr)
+    pre_ckpt_losses, pre_ckpt_grad_norms = _run_training_steps(
+        pre_ckpt_model, pre_ckpt_optimizer, batches[:checkpoint_step]
+    )
 
-    pre_ckpt_losses, pre_ckpt_grad_norms = [], []
-    for step in range(0, checkpoint_step):
-        input_ids, labels = batches[step]
-        pre_ckpt_optimizer.zero_grad()
-        output = pre_ckpt_model(input_ids=input_ids, labels=labels, use_cache=False)
-        loss = output.loss
-        loss.backward()
-        grad_norm = torch.nn.utils.clip_grad_norm_(pre_ckpt_model.parameters(), max_norm=float("inf"))
-        pre_ckpt_optimizer.step()
-
-        pre_ckpt_losses.append(loss.detach().item())
-        pre_ckpt_grad_norms.append(grad_norm)
-
-    # -- Phase 2: Save checkpoint, then load into a fresh model
-    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
-    try:
-        model_dir = os.path.join(tmpdir, "model")
-        training_state_dir = os.path.join(tmpdir, "training_state")
-
-        pre_ckpt_model.save_pretrained(model_dir, is_main_process=(rank == 0))
-        _save_training_state(pre_ckpt_model, pre_ckpt_optimizer, training_state_dir)
-        dist.barrier()
-
-        # Intentionally scramble RNG to prove checkpoint restore works
-        _set_determinism(SEED + 1234)
-        resumed_model = AutoModelForCausalLM.from_pretrained(
-            model_dir, torch_dtype=dtype, distributed_config=distributed_config
-        )
-        resumed_model.train()
-        resumed_optimizer = torch.optim.Adam(resumed_model.parameters(), lr=lr)
-
-        _load_training_state(resumed_model, resumed_optimizer, training_state_dir)
-        dist.barrier()
-    finally:
-        if rank == 0:
-            tmpdir_obj.cleanup()
-
-    # -- Phase 3: Post-checkpoint run -- continue training the remaining steps from the resumed model
-    post_ckpt_losses, post_ckpt_grad_norms = [], []
-    for step in range(checkpoint_step, len(batches)):
-        input_ids, labels = batches[step]
-        resumed_optimizer.zero_grad()
-        output = resumed_model(input_ids=input_ids, labels=labels, use_cache=False)
-        loss = output.loss
-        loss.backward()
-        grad_norm = torch.nn.utils.clip_grad_norm_(resumed_model.parameters(), max_norm=float("inf"))
-        resumed_optimizer.step()
-
-        post_ckpt_losses.append(loss.detach().item())
-        post_ckpt_grad_norms.append(grad_norm)
+    # Phase 2: Save checkpoint, then load into a fresh model
+    resumed_model, resumed_optimizer = _checkpoint_and_resume(
+        pre_ckpt_model, pre_ckpt_optimizer, dtype, distributed_config, lr
+    )
 
-    combined_losses = pre_ckpt_losses + post_ckpt_losses
-    combined_grad_norms = pre_ckpt_grad_norms + post_ckpt_grad_norms
-    combined_state_dict = gather_full_state_dict(resumed_model)
+    # Phase 3: Post-checkpoint run
+    post_ckpt_losses, post_ckpt_grad_norms = _run_training_steps(
+        resumed_model, resumed_optimizer, batches[checkpoint_step:]
+    )
 
-    return combined_losses, combined_grad_norms, combined_state_dict
+    return (
+        pre_ckpt_losses + post_ckpt_losses,
+        pre_ckpt_grad_norms + post_ckpt_grad_norms,
+        gather_full_state_dict(resumed_model),
+    )
 
 
 # =============================================================================
@@ -363,27 +427,13 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
 
     device = _get_rank_device(rank)
     config = config_class.from_dict(config_dict)
-
     batches = _build_repeated_training_batches(config, device, 3)
-
     distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
 
-    init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32)
-    try:
-        _set_determinism(SEED)
-        model = AutoModelForCausalLM.from_pretrained(init_tmpdir, distributed_config=distributed_config)
-        dist.barrier()
-    finally:
-        if rank == 0 and init_tmpdir_obj is not None:
-            init_tmpdir_obj.cleanup()
+    model = _load_fsdp_model_from_init_checkpoint(rank, config, torch.float32, distributed_config=distributed_config)
     model.train()
     optimizer = torch.optim.Adam(model.parameters(), lr=LR)
-
-    for input_ids, labels in batches:
-        optimizer.zero_grad()
-        output = model(input_ids=input_ids, labels=labels, use_cache=False)
-        output.loss.backward()
-        optimizer.step()
+    _run_training_steps(model, optimizer, batches, track_grad_norms=False)
 
     state_dict_before = gather_full_state_dict(model)
 
@@ -391,24 +441,19 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
     try:
         model.save_pretrained(tmpdir, is_main_process=(rank == 0))
         dist.barrier()
-
         new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
         dist.barrier()
     finally:
         if rank == 0:
             tmpdir_obj.cleanup()
 
-    state_dict_after = gather_full_state_dict(new_model)
-
-    for key in state_dict_before:
-        assert key in state_dict_after, f"Key {key} missing after load"
-        torch.testing.assert_close(
-            state_dict_before[key],
-            state_dict_after[key],
-            rtol=0,
-            atol=0,
-            msg=f"Weight mismatch for {key} after save/load",
-        )
+    _assert_state_dicts_equal(
+        state_dict_before,
+        gather_full_state_dict(new_model),
+        rtol=0,
+        atol=0,
+        msg_prefix="After save/load: ",
+    )
 
     if rank == 0:
         logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
@@ -419,25 +464,19 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor
     init_test_logger()
 
     config = config_class.from_dict(config_dict)
-    config.tie_word_embeddings = tie_word_embeddings
-
-    init_tmpdir, init_tmpdir_obj = _save_init_pretrained(rank, config, torch.float32)
-    try:
-        _set_determinism(SEED)
-        model = AutoModelForCausalLM.from_pretrained(
-            init_tmpdir, distributed_config=DistributedConfig(fsdp_size=dist.get_world_size())
-        )
-        dist.barrier()
-    finally:
-        if rank == 0 and init_tmpdir_obj is not None:
-            init_tmpdir_obj.cleanup()
+    model = _load_fsdp_model_from_init_checkpoint(
+        rank,
+        config,
+        torch.float32,
+        tie_word_embeddings=tie_word_embeddings,
+        distributed_config=DistributedConfig(fsdp_size=dist.get_world_size()),
+    )
 
     expected_targets = {""} | _fsdp_target_module_names(model)
-
     actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")}
 
     if rank == 0:
-        logger.debug(f"  Weights tied: {config.tie_word_embeddings}")
+        logger.debug(f"  Weights tied: {tie_word_embeddings}")
         logger.debug(f"  Expected FSDP targets: {sorted(expected_targets)}")
         logger.debug(f"  Actual FSDP targets:   {sorted(actual_targets)}")
 
@@ -480,7 +519,6 @@ def _test_fsdp2_plan_vs_ddp_impl(
     batches = _build_repeated_training_batches(config, device, NUM_STEPS)
     try:
         ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir)
-
         fsdp_losses, fsdp_grad_norms, fsdp_state_dict = train_fsdp2(
             rank,
             batches,
@@ -495,31 +533,15 @@ def _test_fsdp2_plan_vs_ddp_impl(
         if rank == 0 and init_tmpdir_obj is not None:
             init_tmpdir_obj.cleanup()
 
-    for step in range(len(ddp_losses)):
-        torch.testing.assert_close(
-            torch.tensor(ddp_losses[step]),
-            torch.tensor(fsdp_losses[step]),
-            rtol=1e-5,
-            atol=1e-5,
-            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}",
-        )
-        torch.testing.assert_close(
-            torch.tensor(ddp_grad_norms[step]),
-            torch.tensor(fsdp_grad_norms[step]),
-            rtol=1e-5,
-            atol=1e-5,
-            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}",
-        )
-
-    for key in ddp_state_dict:
-        assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict"
-        torch.testing.assert_close(
-            ddp_state_dict[key],
-            fsdp_state_dict[key],
-            rtol=1e-5,
-            atol=1e-5,
-            msg=f"Weight mismatch for {key}: DDP vs {test_label}",
-        )
+    _assert_ddp_fsdp_match(
+        ddp_losses,
+        fsdp_losses,
+        ddp_grad_norms,
+        fsdp_grad_norms,
+        ddp_state_dict,
+        fsdp_state_dict,
+        test_label,
+    )
 
     if rank == 0:
         logger.debug(f"DDP and {test_label} comparison checks passed.")
@@ -530,6 +552,25 @@ def _test_fsdp2_plan_vs_ddp_impl(
 # =============================================================================
 
 
+def _shrink_config_for_fsdp(config):
+    config.vocab_size = 256
+    config.hidden_size = 64
+    config.intermediate_size = 128
+    if hasattr(config, "ffn_config"):
+        if hasattr(config.ffn_config, "ffn_hidden_size"):
+            config.ffn_config.ffn_hidden_size = config.hidden_size
+        if hasattr(config.ffn_config, "hidden_size"):
+            config.ffn_config.hidden_size = config.intermediate_size
+    if hasattr(config, "num_attention_heads"):
+        config.num_attention_heads = 4
+    if hasattr(config, "num_key_value_heads"):
+        config.num_key_value_heads = 4
+    if hasattr(config, "moe_intermediate_size"):
+        config.moe_intermediate_size = 32
+    if hasattr(config, "vocab_size_per_layer_input"):
+        config.vocab_size_per_layer_input = config.vocab_size
+
+
 class FSDPTesterMixin(ABC):
     fsdp_nproc_per_node: int = 2
     skip_fsdp_tests: bool = False
@@ -579,24 +620,8 @@ def _create_model_on_meta(self, config):
     def _get_tiny_config(self):
         """Get config class and serialized dict for passing to spawned processes."""
         config = self.model_tester.get_config()
-        config.vocab_size = 256
-        config.hidden_size = 64
-        config.intermediate_size = 128
-        if hasattr(config, "ffn_config"):
-            if hasattr(config.ffn_config, "ffn_hidden_size"):
-                config.ffn_config.ffn_hidden_size = config.hidden_size
-            if hasattr(config.ffn_config, "hidden_size"):
-                config.ffn_config.hidden_size = config.intermediate_size
-        if hasattr(config, "num_attention_heads"):
-            config.num_attention_heads = 4
-        if hasattr(config, "num_key_value_heads"):
-            config.num_key_value_heads = 4
-        if hasattr(config, "moe_intermediate_size"):
-            config.moe_intermediate_size = 32
-        if hasattr(config, "vocab_size_per_layer_input"):
-            config.vocab_size_per_layer_input = config.vocab_size
-        config_dict = config.to_diff_dict()
-        return type(config), config_dict
+        _shrink_config_for_fsdp(config)
+        return type(config), config.to_diff_dict()
 
     def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs):
         self._skip_if_fsdp_model_not_selected()
@@ -606,10 +631,7 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k
         func_args = (config_class, config_dict, *test_args)
 
         results_file = tempfile.mktemp(suffix=".json")
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(("", 0))
-            port = s.getsockname()[1]
-
+        port = _find_free_port()
         try:
             mp.spawn(
                 _fsdp_global_wrapper,
@@ -644,31 +666,27 @@ def test_fsdp_plan_declared(self):
         finally:
             logger.info("[FSDP] %s test: test_fsdp_plan_declared (%.1fs)", status, time.perf_counter() - start_time)
 
-    @is_fsdp_test
+    @parameterized.expand(["untied", "tied"])
     @require_torch_greater_or_equal("2.6")
-    def test_fsdp2_sharding_structure_untied(self):
-        self._run_fsdp2_distributed_test(
-            "test_fsdp2_sharding_structure_untied", _test_fsdp2_sharding_structure_impl, False
-        )
-
     @is_fsdp_test
-    @require_torch_greater_or_equal("2.6")
-    def test_fsdp2_sharding_structure_tied(self):
+    def test_fsdp2_sharding_structure(self, label):
         self._run_fsdp2_distributed_test(
-            "test_fsdp2_sharding_structure_tied", _test_fsdp2_sharding_structure_impl, True
+            f"test_fsdp2_sharding_structure_{label}",
+            _test_fsdp2_sharding_structure_impl,
+            label == "tied",
         )
 
-    @is_fsdp_test
     @require_torch_greater_or_equal("2.6")
+    @is_fsdp_test
     def test_fsdp2_save_load(self):
         self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl)
 
-    @is_fsdp_test
+    @parameterized.expand(["untied", "tied"])
     @require_torch_greater_or_equal("2.6")
-    def test_fsdp2_plan_vs_ddp_untied(self):
-        self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_untied", _test_fsdp2_plan_vs_ddp_impl, False)
-
     @is_fsdp_test
-    @require_torch_greater_or_equal("2.6")
-    def test_fsdp2_plan_vs_ddp_tied(self):
-        self._run_fsdp2_distributed_test("test_fsdp2_plan_vs_ddp_tied", _test_fsdp2_plan_vs_ddp_impl, True)
+    def test_fsdp2_plan_vs_ddp(self, label):
+        self._run_fsdp2_distributed_test(
+            f"test_fsdp2_plan_vs_ddp_{label}",
+            _test_fsdp2_plan_vs_ddp_impl,
+            label == "tied",
+        )

From 5bbd8201fe2d2f0f24cea3a24c785dd4378ddcc2 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 06:36:32 +0000
Subject: [PATCH 06/33] test fsdp mixin cleaning

---
 tests/test_fsdp_mixin.py | 162 ++++++++++++++++-----------------------
 1 file changed, 67 insertions(+), 95 deletions(-)

diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index c3e13b2a70a8..2136ca2bacbc 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -68,8 +68,8 @@
 DDP_FSDP_RTOL = 1e-5
 DDP_FSDP_ATOL = 1e-5
 
-FSDP_TOP_MODEL_NAMES = {
-    # Models with `base_model_fsdp_plan` + head-class `_fsdp_plan` on this branch.
+# Set to None to run distributed FSDP tests for every model with a plan.
+FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {
     "llama",
     "qwen3_moe",
 }
@@ -309,48 +309,6 @@ def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config,
             tmpdir_obj.cleanup()
 
 
-def _assert_state_dicts_equal(before, after, *, rtol, atol, msg_prefix=""):
-    for key in before:
-        assert key in after, f"{msg_prefix}Key {key} missing after load"
-        torch.testing.assert_close(
-            before[key],
-            after[key],
-            rtol=rtol,
-            atol=atol,
-            msg=f"{msg_prefix}Weight mismatch for {key}",
-        )
-
-
-def _assert_ddp_fsdp_match(
-    ddp_losses, fsdp_losses, ddp_grad_norms, fsdp_grad_norms, ddp_state, fsdp_state, test_label
-):
-    for step in range(len(ddp_losses)):
-        torch.testing.assert_close(
-            torch.tensor(ddp_losses[step]),
-            torch.tensor(fsdp_losses[step]),
-            rtol=DDP_FSDP_RTOL,
-            atol=DDP_FSDP_ATOL,
-            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}",
-        )
-        torch.testing.assert_close(
-            torch.tensor(ddp_grad_norms[step]),
-            torch.tensor(fsdp_grad_norms[step]),
-            rtol=DDP_FSDP_RTOL,
-            atol=DDP_FSDP_ATOL,
-            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}",
-        )
-
-    for key in ddp_state:
-        assert key in fsdp_state, f"Key {key} missing from {test_label} state dict"
-        torch.testing.assert_close(
-            ddp_state[key],
-            fsdp_state[key],
-            rtol=DDP_FSDP_RTOL,
-            atol=DDP_FSDP_ATOL,
-            msg=f"Weight mismatch for {key}: DDP vs {test_label}",
-        )
-
-
 def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
     _set_determinism(SEED)
     model = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype).to(device)
@@ -447,13 +405,16 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
         if rank == 0:
             tmpdir_obj.cleanup()
 
-    _assert_state_dicts_equal(
-        state_dict_before,
-        gather_full_state_dict(new_model),
-        rtol=0,
-        atol=0,
-        msg_prefix="After save/load: ",
-    )
+    state_dict_after = gather_full_state_dict(new_model)
+    for key in state_dict_before:
+        assert key in state_dict_after, f"After save/load: Key {key} missing after load"
+        torch.testing.assert_close(
+            state_dict_before[key],
+            state_dict_after[key],
+            rtol=0,
+            atol=0,
+            msg=f"After save/load: Weight mismatch for {key}",
+        )
 
     if rank == 0:
         logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
@@ -533,15 +494,31 @@ def _test_fsdp2_plan_vs_ddp_impl(
         if rank == 0 and init_tmpdir_obj is not None:
             init_tmpdir_obj.cleanup()
 
-    _assert_ddp_fsdp_match(
-        ddp_losses,
-        fsdp_losses,
-        ddp_grad_norms,
-        fsdp_grad_norms,
-        ddp_state_dict,
-        fsdp_state_dict,
-        test_label,
-    )
+    for step in range(len(ddp_losses)):
+        torch.testing.assert_close(
+            torch.tensor(ddp_losses[step]),
+            torch.tensor(fsdp_losses[step]),
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}",
+        )
+        torch.testing.assert_close(
+            torch.tensor(ddp_grad_norms[step]),
+            torch.tensor(fsdp_grad_norms[step]),
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}",
+        )
+
+    for key in ddp_state_dict:
+        assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict"
+        torch.testing.assert_close(
+            ddp_state_dict[key],
+            fsdp_state_dict[key],
+            rtol=DDP_FSDP_RTOL,
+            atol=DDP_FSDP_ATOL,
+            msg=f"Weight mismatch for {key}: DDP vs {test_label}",
+        )
 
     if rank == 0:
         logger.debug(f"DDP and {test_label} comparison checks passed.")
@@ -552,25 +529,6 @@ def _test_fsdp2_plan_vs_ddp_impl(
 # =============================================================================
 
 
-def _shrink_config_for_fsdp(config):
-    config.vocab_size = 256
-    config.hidden_size = 64
-    config.intermediate_size = 128
-    if hasattr(config, "ffn_config"):
-        if hasattr(config.ffn_config, "ffn_hidden_size"):
-            config.ffn_config.ffn_hidden_size = config.hidden_size
-        if hasattr(config.ffn_config, "hidden_size"):
-            config.ffn_config.hidden_size = config.intermediate_size
-    if hasattr(config, "num_attention_heads"):
-        config.num_attention_heads = 4
-    if hasattr(config, "num_key_value_heads"):
-        config.num_key_value_heads = 4
-    if hasattr(config, "moe_intermediate_size"):
-        config.moe_intermediate_size = 32
-    if hasattr(config, "vocab_size_per_layer_input"):
-        config.vocab_size_per_layer_input = config.vocab_size
-
-
 class FSDPTesterMixin(ABC):
     fsdp_nproc_per_node: int = 2
     skip_fsdp_tests: bool = False
@@ -591,24 +549,22 @@ def _skip_if_insufficient_devices(self):
         if available_workers < self.fsdp_nproc_per_node:
             self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}")
 
-    def _get_fsdp_model_name(self):
-        module_parts = self.__class__.__module__.split(".")
-        if len(module_parts) >= 3 and module_parts[0] == "tests" and module_parts[1] == "models":
-            return module_parts[2]
-        return None
-
-    def _skip_if_fsdp_model_not_selected(self):
-        model_name = self._get_fsdp_model_name()
-        if model_name not in FSDP_TOP_MODEL_NAMES:
-            model_label = model_name or self.__class__.__module__
+    def _skip_if_fsdp_distributed_not_enabled(self):
+        # Trigger tests only if the model has an FSDP plan (base_model_fsdp_plan)
+        config = self.model_tester.get_config()
+        if not (hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None):
+            self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)")
+        
+        # Only top-N models are tested, set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests.
+        if FSDP_DISTRIBUTED_TEST_MODEL_TYPES is not None and config.model_type not in FSDP_DISTRIBUTED_TEST_MODEL_TYPES:
             self.skipTest(
-                "FSDP mixin coverage is currently limited to models with declared FSDP plans "
-                f"(skipping {model_label})."
+                f"FSDP distributed tests are not enabled for model_type={config.model_type!r} "
+                f"(enabled: {sorted(FSDP_DISTRIBUTED_TEST_MODEL_TYPES)}). Set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests."
             )
 
     def _create_model_on_meta(self, config):
         """Instantiate a model on the meta device (no memory allocated)."""
-        auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM]
+        auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM] #TODO(3outeille): why AutoModelForSeq2SeqLM ?
         for auto_cls in auto_classes:
             try:
                 with torch.device("meta"):
@@ -620,12 +576,27 @@ def _create_model_on_meta(self, config):
     def _get_tiny_config(self):
         """Get config class and serialized dict for passing to spawned processes."""
         config = self.model_tester.get_config()
-        _shrink_config_for_fsdp(config)
+        config.vocab_size = 256
+        config.hidden_size = 64
+        config.intermediate_size = 128
+        if hasattr(config, "ffn_config"):
+            if hasattr(config.ffn_config, "ffn_hidden_size"):
+                config.ffn_config.ffn_hidden_size = config.hidden_size
+            if hasattr(config.ffn_config, "hidden_size"):
+                config.ffn_config.hidden_size = config.intermediate_size
+        if hasattr(config, "num_attention_heads"):
+            config.num_attention_heads = 4
+        if hasattr(config, "num_key_value_heads"):
+            config.num_key_value_heads = 4
+        if hasattr(config, "moe_intermediate_size"):
+            config.moe_intermediate_size = 32
+        if hasattr(config, "vocab_size_per_layer_input"):
+            config.vocab_size_per_layer_input = config.vocab_size
         return type(config), config.to_diff_dict()
 
     def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_kwargs):
-        self._skip_if_fsdp_model_not_selected()
         self._skip_if_insufficient_devices()
+        self._skip_if_fsdp_distributed_not_enabled()
 
         config_class, config_dict = self._get_tiny_config()
         func_args = (config_class, config_dict, *test_args)
@@ -652,12 +623,13 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k
     def test_fsdp_plan_declared(self):
         """The model exposes a non-empty `_fsdp_plan` derived from config + class-level overrides."""
         self._skip_if_fsdp_disabled()
-        self._skip_if_fsdp_model_not_selected()
         start_time = time.perf_counter()
         logger.info("[FSDP] Starting test: test_fsdp_plan_declared")
         status = "FAIL"
         try:
             config = self.model_tester.get_config()
+            if not (hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None):
+                self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)")
             model = self._create_model_on_meta(config)
 
             plan = getattr(model, "_fsdp_plan", None)

From b6d0b67abd02396b417b543412e77b1dcf8dbcea Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 06:53:31 +0000
Subject: [PATCH 07/33] remove fsdp policy in tests + trim down further

---
 tests/test_fsdp_mixin.py | 71 +++++++++++-----------------------------
 1 file changed, 20 insertions(+), 51 deletions(-)

diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 2136ca2bacbc..9fff23238621 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -337,14 +337,8 @@ def train_fsdp2(
     dtype,
     init_model_dir,
     checkpoint_step,
-    fsdp_cpu_offload=False,
-    fsdp_mixed_precision=False,
 ):
-    distributed_config = DistributedConfig(
-        fsdp_size=dist.get_world_size(),
-        fsdp_cpu_offload=fsdp_cpu_offload,
-        fsdp_mixed_precision=fsdp_mixed_precision,
-    )
+    distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
 
     # Phase 1: Pre-checkpoint run
     _set_determinism(SEED)
@@ -453,28 +447,17 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor
         logger.debug(f"  FSDP sharding structure OK ({len(actual_targets)} targets)")
 
 
-def _test_fsdp2_plan_vs_ddp_impl(
-    rank, config_class, config_dict, tie_word_embeddings, policy_options=None, dtype=None
-):
+def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embeddings, dtype=None):
     """Validate DDP-vs-FSDP2 trace matching using the model's declared FSDP plan."""
     init_test_logger()
 
     if dtype is None:
         dtype = torch.float32
 
-    policy_options = policy_options or []
-    assert "mixed_precision" not in policy_options, (
-        "Use the mixed-precision specific tests when enabling mixed_precision policy."
-    )
-
     device = _get_rank_device(rank)
     config = config_class.from_dict(config_dict)
     config.tie_word_embeddings = tie_word_embeddings
 
-    fsdp_cpu_offload = "cpu_offload" in policy_options
-    fsdp_mixed_precision = "mixed_precision" in policy_options
-    test_label = f"FSDP2{'+policies' if policy_options else ''}"
-
     checkpoint_step = NUM_STEPS // 2
     init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype)
     batches = _build_repeated_training_batches(config, device, NUM_STEPS)
@@ -487,8 +470,6 @@ def _test_fsdp2_plan_vs_ddp_impl(
             dtype,
             init_model_dir=init_model_dir,
             checkpoint_step=checkpoint_step,
-            fsdp_cpu_offload=fsdp_cpu_offload,
-            fsdp_mixed_precision=fsdp_mixed_precision,
         )
     finally:
         if rank == 0 and init_tmpdir_obj is not None:
@@ -500,28 +481,28 @@ def _test_fsdp2_plan_vs_ddp_impl(
             torch.tensor(fsdp_losses[step]),
             rtol=DDP_FSDP_RTOL,
             atol=DDP_FSDP_ATOL,
-            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, {test_label}={fsdp_losses[step]}",
+            msg=f"Loss mismatch at step {step}: DDP={ddp_losses[step]}, FSDP2={fsdp_losses[step]}",
         )
         torch.testing.assert_close(
             torch.tensor(ddp_grad_norms[step]),
             torch.tensor(fsdp_grad_norms[step]),
             rtol=DDP_FSDP_RTOL,
             atol=DDP_FSDP_ATOL,
-            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, {test_label}={fsdp_grad_norms[step]}",
+            msg=f"Grad norm mismatch at step {step}: DDP={ddp_grad_norms[step]}, FSDP2={fsdp_grad_norms[step]}",
         )
 
     for key in ddp_state_dict:
-        assert key in fsdp_state_dict, f"Key {key} missing from {test_label} state dict"
+        assert key in fsdp_state_dict, f"Key {key} missing from FSDP2 state dict"
         torch.testing.assert_close(
             ddp_state_dict[key],
             fsdp_state_dict[key],
             rtol=DDP_FSDP_RTOL,
             atol=DDP_FSDP_ATOL,
-            msg=f"Weight mismatch for {key}: DDP vs {test_label}",
+            msg=f"Weight mismatch for {key}: DDP vs FSDP2",
         )
 
     if rank == 0:
-        logger.debug(f"DDP and {test_label} comparison checks passed.")
+        logger.debug("DDP and FSDP2 comparison checks passed.")
 
 
 # =============================================================================
@@ -531,7 +512,7 @@ def _test_fsdp2_plan_vs_ddp_impl(
 
 class FSDPTesterMixin(ABC):
     fsdp_nproc_per_node: int = 2
-    skip_fsdp_tests: bool = False
+    #TODO(3outeille): do we put the CONSTANTS in the mixin class ?
 
     @property
     @abstractmethod
@@ -539,22 +520,20 @@ def model_tester(self):
         """The model tester instance (e.g., CausalLMModelTester)."""
         ...
 
-    def _skip_if_fsdp_disabled(self):
-        if self.skip_fsdp_tests:
-            self.skipTest("FSDP tests disabled for this model (skip_fsdp_tests=True)")
-
     def _skip_if_insufficient_devices(self):
-        self._skip_if_fsdp_disabled()
         available_workers = _get_available_fsdp_workers()
         if available_workers < self.fsdp_nproc_per_node:
             self.skipTest(f"Need at least {self.fsdp_nproc_per_node} FSDP workers, have {available_workers}")
 
-    def _skip_if_fsdp_distributed_not_enabled(self):
-        # Trigger tests only if the model has an FSDP plan (base_model_fsdp_plan)
+    def _has_fsdp_plan(self) -> bool:
         config = self.model_tester.get_config()
-        if not (hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None):
+        return hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None
+
+    def _skip_if_fsdp_distributed_not_enabled(self):
+        if not self._has_fsdp_plan():
             self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)")
-        
+
+        config = self.model_tester.get_config()
         # Only top-N models are tested, set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests.
         if FSDP_DISTRIBUTED_TEST_MODEL_TYPES is not None and config.model_type not in FSDP_DISTRIBUTED_TEST_MODEL_TYPES:
             self.skipTest(
@@ -622,21 +601,11 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k
     @is_fsdp_test
     def test_fsdp_plan_declared(self):
         """The model exposes a non-empty `_fsdp_plan` derived from config + class-level overrides."""
-        self._skip_if_fsdp_disabled()
-        start_time = time.perf_counter()
-        logger.info("[FSDP] Starting test: test_fsdp_plan_declared")
-        status = "FAIL"
-        try:
-            config = self.model_tester.get_config()
-            if not (hasattr(config, "base_model_fsdp_plan") and config.base_model_fsdp_plan is not None):
-                self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)")
-            model = self._create_model_on_meta(config)
-
-            plan = getattr(model, "_fsdp_plan", None)
-            self.assertTrue(plan, f"No _fsdp_plan declared for {type(model).__name__}")
-            status = "PASS"
-        finally:
-            logger.info("[FSDP] %s test: test_fsdp_plan_declared (%.1fs)", status, time.perf_counter() - start_time)
+        if not self._has_fsdp_plan():
+            self.skipTest("Model does not have an FSDP plan (base_model_fsdp_plan)")
+
+        model = self._create_model_on_meta(self.model_tester.get_config())
+        self.assertTrue(model._fsdp_plan, f"No _fsdp_plan declared for {type(model).__name__}")
 
     @parameterized.expand(["untied", "tied"])
     @require_torch_greater_or_equal("2.6")

From ea3612314c2189a9b9fbe78d4bceb37671237bec Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 07:18:01 +0000
Subject: [PATCH 08/33] test fsdp clean

---
 tests/test_fsdp_mixin.py | 226 ++++++++++++++++-----------------------
 1 file changed, 90 insertions(+), 136 deletions(-)

diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 9fff23238621..84c63ea032da 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -23,6 +23,7 @@
 import time
 import traceback
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 
 from parameterized import parameterized
 
@@ -80,12 +81,6 @@
 # =============================================================================
 
 
-def _find_free_port():
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("", 0))
-        return s.getsockname()[1]
-
-
 def _get_distributed_device_type():
     device_type = torch._C._get_accelerator().type
     return "cpu" if device_type == "mps" else device_type
@@ -109,19 +104,6 @@ def _set_rank_device(rank):
         accelerator_module.set_device(rank)
 
 
-def _get_accelerator_rng_state():
-    accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
-    if accelerator_module is None or not hasattr(accelerator_module, "get_rng_state"):
-        return None
-    return accelerator_module.get_rng_state()
-
-
-def _set_accelerator_rng_state(rng_state):
-    accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
-    if rng_state is not None and accelerator_module is not None and hasattr(accelerator_module, "set_rng_state"):
-        accelerator_module.set_rng_state(rng_state)
-
-
 def _get_available_fsdp_workers():
     if _get_distributed_device_type() == "cpu":
         return os.cpu_count() or 1
@@ -139,7 +121,8 @@ def _set_determinism(seed):
     set_seed(seed)
 
 
-def _create_shared_tmpdir(rank):
+@contextmanager
+def _distributed_tmpdir(rank):
     if rank == 0:
         tmpdir_obj = tempfile.TemporaryDirectory()
         tmpdir_list = [tmpdir_obj.name]
@@ -147,7 +130,23 @@ def _create_shared_tmpdir(rank):
         tmpdir_obj = None
         tmpdir_list = [None]
     dist.broadcast_object_list(tmpdir_list, src=0)
-    return tmpdir_list[0], tmpdir_obj
+    try:
+        yield tmpdir_list[0]
+    finally:
+        if rank == 0 and tmpdir_obj is not None:
+            tmpdir_obj.cleanup()
+
+
+@contextmanager
+def _deterministic_init_model_dir(rank, config, dtype):
+    with _distributed_tmpdir(rank) as model_dir:
+        if rank == 0:
+            set_seed(SEED)
+            model = AutoModelForCausalLM.from_config(config).to(dtype)
+            model.save_pretrained(model_dir)
+            del model
+        dist.barrier()
+        yield model_dir
 
 
 def _fsdp_global_wrapper(rank, test_name, func, func_args, func_kwargs, world_size, port, results_file):
@@ -217,59 +216,15 @@ def _run_training_steps(model, optimizer, batches, *, track_grad_norms=True):
     return losses, grad_norms
 
 
-def _gather_ddp_state_dict(model):
-    # Only rank 0 returns data to match gather_full_state_dict semantics, so the
-    # downstream DDP-vs-FSDP comparison runs once on rank 0 instead of N times.
-    if dist.get_rank() != 0:
-        return {}
-    return {k: v.clone().detach().cpu() for k, v in model.module.state_dict().items()}
-
-
-def _fsdp_target_module_names(model) -> set[str]:
-    """Return module names that should receive a ``fully_shard`` call (excluding root)."""
-    adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, model)
-    reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan)
-    return {name for name, _ in reshard_targets + no_reshard_targets}
-
-
-def _save_init_pretrained(rank, config, dtype):
-    """Save a deterministic initial model to a shared tmpdir for from_pretrained loading."""
-    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
-    if rank == 0:
-        set_seed(SEED)
-        model = AutoModelForCausalLM.from_config(config).to(dtype)
-        model.save_pretrained(tmpdir)
-        del model
-    dist.barrier()
-    return tmpdir, tmpdir_obj
-
-
-def _load_fsdp_model_from_init_checkpoint(rank, config, dtype, *, distributed_config=None, tie_word_embeddings=None):
-    """Save deterministic init weights, load FSDP model, cleanup init tmpdir."""
-    if tie_word_embeddings is not None:
-        config.tie_word_embeddings = tie_word_embeddings
-
-    init_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype)
-    try:
-        _set_determinism(SEED)
-        kwargs = {}
-        if distributed_config is not None:
-            kwargs["distributed_config"] = distributed_config
-        model = AutoModelForCausalLM.from_pretrained(init_dir, **kwargs)
-        dist.barrier()
-        return model
-    finally:
-        if rank == 0 and init_tmpdir_obj is not None:
-            init_tmpdir_obj.cleanup()
-
-
 def _save_training_state(model, optimizer, training_state_dir):
     """Save optimizer (canonical DCP path) plus per-rank RNG for resume."""
     save_optimizer_distributed(model, optimizer, os.path.join(training_state_dir, "optim"))
     rng = {"cpu": torch.get_rng_state()}
-    accel = _get_accelerator_rng_state()
-    if accel is not None:
-        rng["accel"] = accel
+    accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
+    if accelerator_module is not None and hasattr(accelerator_module, "get_rng_state"):
+        accel_rng = accelerator_module.get_rng_state()
+        if accel_rng is not None:
+            rng["accel"] = accel_rng
     torch.save(rng, os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt"))
 
 
@@ -279,14 +234,15 @@ def _load_training_state(model, optimizer, training_state_dir):
     rng = torch.load(os.path.join(training_state_dir, f"rng_rank{dist.get_rank()}.pt"), weights_only=False)
     torch.set_rng_state(rng["cpu"])
     if "accel" in rng:
-        _set_accelerator_rng_state(rng["accel"])
+        accelerator_module = backend_torch_accelerator_module(_get_distributed_device_type())
+        if accelerator_module is not None and hasattr(accelerator_module, "set_rng_state"):
+            accelerator_module.set_rng_state(rng["accel"])
 
 
 def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config, lr):
     """Save model+optimizer, scramble RNG, reload and restore training state."""
     rank = dist.get_rank()
-    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
-    try:
+    with _distributed_tmpdir(rank) as tmpdir:
         model_dir = os.path.join(tmpdir, "model")
         training_state_dir = os.path.join(tmpdir, "training_state")
 
@@ -304,9 +260,6 @@ def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config,
         _load_training_state(resumed_model, resumed_optimizer, training_state_dir)
         dist.barrier()
         return resumed_model, resumed_optimizer
-    finally:
-        if rank == 0:
-            tmpdir_obj.cleanup()
 
 
 def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
@@ -321,7 +274,11 @@ def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
     optimizer = torch.optim.Adam(ddp_model.parameters(), lr=lr)
 
     losses, grad_norms = _run_training_steps(ddp_model, optimizer, batches)
-    state_dict = _gather_ddp_state_dict(ddp_model)
+    if dist.get_rank() != 0:
+        state_dict = {}
+    else:
+        # Only rank 0 returns data to match gather_full_state_dict semantics.
+        state_dict = {k: v.clone().detach().cpu() for k, v in ddp_model.module.state_dict().items()}
 
     del optimizer, ddp_model, model
     backend_empty_cache(_get_distributed_device_type())
@@ -374,44 +331,38 @@ def train_fsdp2(
 
 
 def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
-    """Train FSDP2 model, save via save_pretrained, load via from_pretrained, compare state dicts."""
+    """Save FSDP2 model via save_pretrained, load via from_pretrained, compare state dicts."""
     init_test_logger()
 
-    device = _get_rank_device(rank)
     config = config_class.from_dict(config_dict)
-    batches = _build_repeated_training_batches(config, device, 3)
     distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
 
-    model = _load_fsdp_model_from_init_checkpoint(rank, config, torch.float32, distributed_config=distributed_config)
-    model.train()
-    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
-    _run_training_steps(model, optimizer, batches, track_grad_norms=False)
-
-    state_dict_before = gather_full_state_dict(model)
-
-    tmpdir, tmpdir_obj = _create_shared_tmpdir(rank)
-    try:
-        model.save_pretrained(tmpdir, is_main_process=(rank == 0))
-        dist.barrier()
-        new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
         dist.barrier()
-    finally:
-        if rank == 0:
-            tmpdir_obj.cleanup()
 
-    state_dict_after = gather_full_state_dict(new_model)
-    for key in state_dict_before:
-        assert key in state_dict_after, f"After save/load: Key {key} missing after load"
-        torch.testing.assert_close(
-            state_dict_before[key],
-            state_dict_after[key],
-            rtol=0,
-            atol=0,
-            msg=f"After save/load: Weight mismatch for {key}",
-        )
+        state_dict_before = gather_full_state_dict(model)
+
+        with _distributed_tmpdir(rank) as tmpdir:
+            model.save_pretrained(tmpdir, is_main_process=(rank == 0))
+            dist.barrier()
+            new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
+            dist.barrier()
+
+        state_dict_after = gather_full_state_dict(new_model)
+        for key in state_dict_before:
+            assert key in state_dict_after, f"After save/load: Key {key} missing after load"
+            torch.testing.assert_close(
+                state_dict_before[key],
+                state_dict_after[key],
+                rtol=0,
+                atol=0,
+                msg=f"After save/load: Weight mismatch for {key}",
+            )
 
-    if rank == 0:
-        logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
+        if rank == 0:
+            logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
 
 
 def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings):
@@ -419,32 +370,34 @@ def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_wor
     init_test_logger()
 
     config = config_class.from_dict(config_dict)
-    model = _load_fsdp_model_from_init_checkpoint(
-        rank,
-        config,
-        torch.float32,
-        tie_word_embeddings=tie_word_embeddings,
-        distributed_config=DistributedConfig(fsdp_size=dist.get_world_size()),
-    )
+    config.tie_word_embeddings = tie_word_embeddings
+    distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
 
-    expected_targets = {""} | _fsdp_target_module_names(model)
-    actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")}
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
+        dist.barrier()
 
-    if rank == 0:
-        logger.debug(f"  Weights tied: {tie_word_embeddings}")
-        logger.debug(f"  Expected FSDP targets: {sorted(expected_targets)}")
-        logger.debug(f"  Actual FSDP targets:   {sorted(actual_targets)}")
-
-    missing = expected_targets - actual_targets
-    extra = actual_targets - expected_targets
-    assert not missing and not extra, (
-        f"FSDP target mismatch.\n"
-        f"  Missing (expected but not wrapped): {sorted(missing)}\n"
-        f"  Extra (wrapped but not expected):   {sorted(extra)}"
-    )
+        adapted_fsdp_plan = _resolve_tied_embed_lm_head_plan(model._fsdp_plan, model)
+        reshard_targets, no_reshard_targets = expand_fsdp_plan(model, adapted_fsdp_plan)
+        expected_targets = {""} | {name for name, _ in reshard_targets + no_reshard_targets}
+        actual_targets = {name for name, module in model.named_modules() if type(module).__name__.startswith("FSDP")}
 
-    if rank == 0:
-        logger.debug(f"  FSDP sharding structure OK ({len(actual_targets)} targets)")
+        if rank == 0:
+            logger.debug(f"  Weights tied: {tie_word_embeddings}")
+            logger.debug(f"  Expected FSDP targets: {sorted(expected_targets)}")
+            logger.debug(f"  Actual FSDP targets:   {sorted(actual_targets)}")
+
+        missing = expected_targets - actual_targets
+        extra = actual_targets - expected_targets
+        assert not missing and not extra, (
+            f"FSDP target mismatch.\n"
+            f"  Missing (expected but not wrapped): {sorted(missing)}\n"
+            f"  Extra (wrapped but not expected):   {sorted(extra)}"
+        )
+
+        if rank == 0:
+            logger.debug(f"  FSDP sharding structure OK ({len(actual_targets)} targets)")
 
 
 def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embeddings, dtype=None):
@@ -459,9 +412,9 @@ def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embed
     config.tie_word_embeddings = tie_word_embeddings
 
     checkpoint_step = NUM_STEPS // 2
-    init_model_dir, init_tmpdir_obj = _save_init_pretrained(rank, config, dtype)
     batches = _build_repeated_training_batches(config, device, NUM_STEPS)
-    try:
+
+    with _deterministic_init_model_dir(rank, config, dtype) as init_model_dir:
         ddp_losses, ddp_grad_norms, ddp_state_dict = train_ddp(rank, batches, LR, device, dtype, init_model_dir)
         fsdp_losses, fsdp_grad_norms, fsdp_state_dict = train_fsdp2(
             rank,
@@ -471,9 +424,6 @@ def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embed
             init_model_dir=init_model_dir,
             checkpoint_step=checkpoint_step,
         )
-    finally:
-        if rank == 0 and init_tmpdir_obj is not None:
-            init_tmpdir_obj.cleanup()
 
     for step in range(len(ddp_losses)):
         torch.testing.assert_close(
@@ -581,7 +531,11 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k
         func_args = (config_class, config_dict, *test_args)
 
         results_file = tempfile.mktemp(suffix=".json")
-        port = _find_free_port()
+        # port binding
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            port = s.getsockname()[1]
+        
         try:
             mp.spawn(
                 _fsdp_global_wrapper,

From bec4d23a6a5d7de5afa3ba634cf049d866a6faea Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 07:22:38 +0000
Subject: [PATCH 09/33] restore test_modeling_utils

---
 tests/utils/test_modeling_utils.py | 40 ------------------------------
 1 file changed, 40 deletions(-)

diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py
index d8b1cebb77f5..7430e1ce6fb0 100644
--- a/tests/utils/test_modeling_utils.py
+++ b/tests/utils/test_modeling_utils.py
@@ -127,7 +127,6 @@
         T5ForConditionalGeneration,
     )
     from transformers.conversion_mapping import MergeModulelist, WeightConverter, get_model_conversion_mapping
-    from transformers.distributed import DistributedConfig
     from transformers.modeling_utils import (
         FLASH_ATTN_KERNEL_FALLBACK,
         _find_disjoint,
@@ -437,45 +436,6 @@ def test_get_total_byte_count_does_not_require_process_group(self):
         self.assertIn(torch.device("cpu"), total_byte_count)
         self.assertGreater(total_byte_count[torch.device("cpu")], 0)
 
-    def test_model_from_pretrained_fsdp_distributes_before_loading(self):
-        model = GPT2LMHeadModel(GPT2Config(n_layer=1, n_head=2, n_embd=8, n_positions=8, n_ctx=8, vocab_size=32))
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            model.save_pretrained(tmp_dir)
-            call_order = []
-
-            fake_mesh = mock.Mock()
-            fake_mesh.mesh_dim_names = ("fsdp",)
-            fake_mesh.ndim = 1
-            fake_mesh.device_type = "cpu"
-
-            def fake_distribute_model(model, distributed_config, device_mesh):
-                call_order.append("distribute")
-                self.assertIs(device_mesh, fake_mesh)
-                return model
-
-            def fake_load_pretrained_model(model, state_dict, checkpoint_files, load_config, expected_keys=None):
-                call_order.append("load")
-                self.assertIs(load_config.device_mesh, fake_mesh)
-                return mock.Mock(), None
-
-            with (
-                patch("transformers.modeling_utils.init_device_mesh", return_value=fake_mesh),
-                patch(
-                    "transformers.modeling_utils.distribute_model_distributed",
-                    side_effect=fake_distribute_model,
-                ),
-                patch.object(GPT2LMHeadModel, "_load_pretrained_model", side_effect=fake_load_pretrained_model),
-                patch.object(
-                    GPT2LMHeadModel,
-                    "_finalize_model_loading",
-                    side_effect=lambda model, load_config, loading_info: loading_info,
-                ),
-            ):
-                GPT2LMHeadModel.from_pretrained(tmp_dir, distributed_config=DistributedConfig(fsdp_size=2))
-
-        self.assertEqual(call_order, ["distribute", "load"])
-
     def test_hub_retry(self):
         @hub_retry(max_attempts=2)
         def test_func():

From 8d3d3298a2dac4372e773b40c94eec23391ad89f Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 07:25:13 +0000
Subject: [PATCH 10/33] linting

---
 tests/test_fsdp_mixin.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 84c63ea032da..9bfd5707dd36 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -462,7 +462,7 @@ def _test_fsdp2_plan_vs_ddp_impl(rank, config_class, config_dict, tie_word_embed
 
 class FSDPTesterMixin(ABC):
     fsdp_nproc_per_node: int = 2
-    #TODO(3outeille): do we put the CONSTANTS in the mixin class ?
+    # TODO(3outeille): do we put the CONSTANTS in the mixin class ?
 
     @property
     @abstractmethod
@@ -485,7 +485,10 @@ def _skip_if_fsdp_distributed_not_enabled(self):
 
         config = self.model_tester.get_config()
         # Only top-N models are tested, set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests.
-        if FSDP_DISTRIBUTED_TEST_MODEL_TYPES is not None and config.model_type not in FSDP_DISTRIBUTED_TEST_MODEL_TYPES:
+        if (
+            FSDP_DISTRIBUTED_TEST_MODEL_TYPES is not None
+            and config.model_type not in FSDP_DISTRIBUTED_TEST_MODEL_TYPES
+        ):
             self.skipTest(
                 f"FSDP distributed tests are not enabled for model_type={config.model_type!r} "
                 f"(enabled: {sorted(FSDP_DISTRIBUTED_TEST_MODEL_TYPES)}). Set FSDP_DISTRIBUTED_TEST_MODEL_TYPES = None to run all tests."
@@ -493,7 +496,7 @@ def _skip_if_fsdp_distributed_not_enabled(self):
 
     def _create_model_on_meta(self, config):
         """Instantiate a model on the meta device (no memory allocated)."""
-        auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM] #TODO(3outeille): why AutoModelForSeq2SeqLM ?
+        auto_classes = [AutoModelForCausalLM, AutoModelForSeq2SeqLM]  # TODO(3outeille): why AutoModelForSeq2SeqLM ?
         for auto_cls in auto_classes:
             try:
                 with torch.device("meta"):
@@ -535,7 +538,7 @@ def _run_fsdp2_distributed_test(self, test_name, test_impl, *test_args, **test_k
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             s.bind(("", 0))
             port = s.getsockname()[1]
-        
+
         try:
             mp.spawn(
                 _fsdp_global_wrapper,

From 6316ee1e644347b872d91901160e38ea39c3006f Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 08:38:37 +0000
Subject: [PATCH 11/33] start trim down stuff

---
 src/transformers/distributed/__init__.py      |   2 -
 .../distributed/sharding_utils.py             | 130 +-----------------
 src/transformers/distributed/utils.py         | 106 +-------------
 src/transformers/modeling_utils.py            |  23 +---
 tests/test_fsdp_mixin.py                      |  63 ++++++---
 5 files changed, 55 insertions(+), 269 deletions(-)

diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py
index dff6acb66217..a79c0b58dcd0 100644
--- a/src/transformers/distributed/__init__.py
+++ b/src/transformers/distributed/__init__.py
@@ -22,7 +22,6 @@
     "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"],
     "utils": [
         "distribute_model",
-        "gather_full_state_dict",
         "init_device_mesh",
         "load_optimizer_distributed",
         "save_model_checkpoint_distributed",
@@ -38,7 +37,6 @@
     from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan
     from .utils import (
         distribute_model,
-        gather_full_state_dict,
         init_device_mesh,
         load_optimizer_distributed,
         save_model_checkpoint_distributed,
diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py
index fb0d30368c2a..365713113cc9 100644
--- a/src/transformers/distributed/sharding_utils.py
+++ b/src/transformers/distributed/sharding_utils.py
@@ -25,10 +25,9 @@
 
 if is_torch_available():
     import torch
-    from torch.distributed._functional_collectives import wait_tensor
-    from torch.distributed.tensor import DTensor, Replicate
+    from torch.distributed.tensor import DTensor
     from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
-    from torch.distributed.tensor.placement_types import Shard, _StridedShard
+    from torch.distributed.tensor.placement_types import Shard
 
     # torch < 2.10 names as an underscore before `local_shard_size_and_offset`: alias it the non-underscored version
     if not hasattr(Shard, "local_shard_size_and_offset") and hasattr(Shard, "_local_shard_size_and_offset"):
@@ -331,128 +330,3 @@ def _dtensor_from_local_like(local_tensor: torch.Tensor, ref: DTensor) -> DTenso
         shape=ref.shape,
         stride=tuple(ref.stride()),
     )
-
-
-def _replicate_dtensor(tensor: DTensor) -> DTensor:
-    """All-gather a DTensor to fully Replicate, handling _StridedShard."""
-    mesh = tensor.device_mesh
-    placements = tensor.placements
-    replicate_all = tuple(Replicate() for _ in range(mesh.ndim))
-
-    if not any(isinstance(p, _StridedShard) for p in placements):
-        return tensor.redistribute(placements=replicate_all)
-
-    with torch.no_grad():
-        local = tensor._local_tensor
-        for i in reversed(range(mesh.ndim)):
-            p = placements[i]
-            if p.is_replicate():
-                continue
-            logical_shape = list(tensor.shape)
-            for j, pj in enumerate(placements[:i]):
-                if not pj.is_replicate():
-                    size, _ = Shard.local_shard_size_and_offset(
-                        logical_shape[pj.dim], mesh.size(j), mesh.get_local_rank(j)
-                    )
-                    logical_shape[pj.dim] = size
-            local = p._to_replicate_tensor(local, mesh, i, logical_shape)
-            local = wait_tensor(local)
-        return DTensor.from_local(local, mesh, replicate_all, run_check=False)
-
-
-def _find_strided_shard_placement_from_fused_params(placements):
-    for i, p in enumerate(placements):
-        if not isinstance(p, _StridedShard):
-            continue
-        has_partner_on_same_dim = any(
-            j != i and getattr(other, "dim", None) == p.dim for j, other in enumerate(placements)
-        )
-        if not has_partner_on_same_dim:
-            return p
-    return None
-
-
-def _split_fused_dtensor(dt, dim, n_pieces):
-    local_pieces = list(dt._local_tensor.chunk(n_pieces, dim=dim))
-    if len(local_pieces) != n_pieces:
-        raise RuntimeError(
-            f"Cannot split DTensor of shape {tuple(dt.shape)} into {n_pieces} pieces along dim {dim}: "
-            f"got {len(local_pieces)} chunks instead."
-        )
-    new_placements = tuple(
-        Shard(p.dim) if (isinstance(p, _StridedShard) and p.dim == dim) else p for p in dt.placements
-    )
-    return [
-        DTensor.from_local(lp.contiguous(), dt.device_mesh, new_placements, run_check=False) for lp in local_pieces
-    ]
-
-
-def _merge_unfused_dtensors(pieces, dim, target_placements):
-    local_cat = torch.cat([p._local_tensor for p in pieces], dim=dim).contiguous()
-    return DTensor.from_local(local_cat, pieces[0].device_mesh, target_placements, run_check=False)
-
-
-def get_fusion_metadata(optimizer_state_dict):
-    """Inspect the optimizer state dict and return metadata describing every
-    fused param that needs to be split for DCP."""
-    optimizer_state = optimizer_state_dict.get("state", {})
-    fusion_metadata: dict[str, dict] = {}
-
-    for param_name, optimizer_fields in optimizer_state.items():
-        dtensor = next((v for v in optimizer_fields.values() if isinstance(v, DTensor)), None)
-        if dtensor is None:
-            continue
-        strided_shard_placement = _find_strided_shard_placement_from_fused_params(dtensor.placements)
-        if strided_shard_placement is None:
-            continue
-        fusion_metadata[param_name] = {
-            "chunk_dim": strided_shard_placement.dim,
-            "placements": tuple(dtensor.placements),
-            "unfused_keys": [f"{param_name}.{i}" for i in range(strided_shard_placement.split_factor)],
-        }
-
-    return fusion_metadata
-
-
-def unfuse_optimizer_state(optimizer_state_dict, fusion_metadata):
-    """Apply `fusion_metadata` to split each fused param into plain-Shard pieces in place."""
-    optimizer_state = optimizer_state_dict.get("state", {})
-
-    for param_name, metadata in fusion_metadata.items():
-        optimizer_fields = optimizer_state[param_name]
-
-        for unfused_key in metadata["unfused_keys"]:
-            optimizer_state[unfused_key] = {}
-
-        for optim_param_name, value in optimizer_fields.items():
-            if isinstance(value, DTensor):
-                chunks = _split_fused_dtensor(value, metadata["chunk_dim"], len(metadata["unfused_keys"]))
-            else:
-                chunks = [value] * len(metadata["unfused_keys"])
-            for unfused_key, chunk in zip(metadata["unfused_keys"], chunks):
-                optimizer_state[unfused_key][optim_param_name] = chunk
-
-        del optimizer_state[param_name]
-
-
-def fuse_optimizer_state(optimizer_state_dict, fusion_metadata):
-    """Fuse the optimizer state dict back in place using the fusion info."""
-    optimizer_state = optimizer_state_dict.get("state", {})
-
-    for param_name, metadata in fusion_metadata.items():
-        first_piece = optimizer_state[metadata["unfused_keys"][0]]
-
-        merged_fields = {}
-        for optim_param_name, value in first_piece.items():
-            chunks = [optimizer_state[unfused_key][optim_param_name] for unfused_key in metadata["unfused_keys"]]
-            if isinstance(value, DTensor):
-                merged_fields[optim_param_name] = _merge_unfused_dtensors(
-                    chunks, metadata["chunk_dim"], metadata["placements"]
-                )
-            else:
-                merged_fields[optim_param_name] = value
-
-        optimizer_state[param_name] = merged_fields
-
-        for unfused_key in metadata["unfused_keys"]:
-            del optimizer_state[unfused_key]
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index a9b571aa85fc..011733002808 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -16,18 +16,8 @@
 import os
 from typing import TYPE_CHECKING
 
-from ..utils import is_torch_available, is_torch_greater_or_equal, logging
+from ..utils import is_torch_available, is_torch_greater_or_equal
 from .fsdp import apply_fully_sharded_data_parallel
-from .sharding_utils import (
-    _find_strided_shard_placement_from_fused_params,
-    _replicate_dtensor,
-    fuse_optimizer_state,
-    get_fusion_metadata,
-    unfuse_optimizer_state,
-)
-
-
-logger = logging.get_logger(__name__)
 
 
 if TYPE_CHECKING:
@@ -37,13 +27,13 @@
 
 if is_torch_available():
     import torch
+    #TODO(3outeille): guarding?
     import torch.distributed.checkpoint as dcp
     from torch.distributed.checkpoint.state_dict import (
         get_model_state_dict,
         get_optimizer_state_dict,
         set_optimizer_state_dict,
     )
-    from torch.distributed.tensor import DTensor
 
     if is_torch_greater_or_equal("2.7"):
         from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter
@@ -154,31 +144,6 @@ def distribute_model(model, distributed_config: DistributedConfig, device_mesh)
     return model
 
 
-def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
-    """Gather all sharded params to full plain tensors for saving.
-
-    Handles FSDP unshard and TP DTensor gather.
-    Streams one parameter at a time to avoid holding all full tensors on GPU.
-    Only rank 0 accumulates the result; other ranks return ``{}``.
-    """
-    is_rank0 = torch.distributed.get_rank() == 0
-    state_dict = get_model_state_dict(model)
-
-    result = {}
-    for key, tensor in state_dict.items():
-        if not isinstance(tensor, DTensor):
-            if is_rank0:
-                result[key] = tensor.detach().to(device="cpu", copy=True).contiguous()
-            continue
-
-        with torch.no_grad():
-            full = _replicate_dtensor(tensor).to_local()
-            if is_rank0:
-                result[key] = full.detach().to(device="cpu", copy=True).contiguous()
-            del full
-    return result
-
-
 def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
     """Save model parameters as standard HF-format sharded safetensors using
     DCP + HuggingFaceStorageWriter with consolidation enabled.
@@ -188,22 +153,11 @@ def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
     and emits HF-compatible `model-*-of-N.safetensors` (+ index) at
     `<checkpoint_dir>/`. The result is a directory `from_pretrained` reads
     through its normal path — no special flag needed at load time.
-
-    DTensors carrying an uncomposed `_StridedShard` placement (e.g. fused
-    gate||up MoE weights) are replicated to a full tensor on every rank
-    before the save, otherwise DCP cannot encode that placement.
     """
     if not is_torch_greater_or_equal("2.7"):
         raise OSError("Distributed checkpoint saving requires `torch>=2.7`.")
 
     state_dict = get_model_state_dict(model)
-    for key, value in list(state_dict.items()):
-        if (
-            isinstance(value, DTensor)
-            and _find_strided_shard_placement_from_fused_params(value.placements) is not None
-        ):
-            state_dict[key] = _replicate_dtensor(value)
-
     dcp.save(
         state_dict,
         storage_writer=HuggingFaceStorageWriter(
@@ -217,66 +171,14 @@ def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
     _distributed_barrier()
 
 
-def has_mixed_tensor_and_dtensor_params(params) -> bool:
-    has_dtensor = False
-    has_tensor = False
-    for param in params:
-        if isinstance(param, DTensor):
-            has_dtensor = True
-        elif isinstance(param, torch.Tensor):
-            has_tensor = True
-
-        if has_dtensor and has_tensor:
-            return True
-    return False
-
-
-def maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer) -> None:
-    """
-    When get_optimizer_state_dict() or set_optimizer_state_dict() runs on an optimizer with no state yet,
-    PyTorch first materializes that state by doing a no-op step() with zero gradients. If an optimizer
-    group mixes regular tensors and DTensors, the batched foreach/fused optimizer kernels cannot process
-    that mixed group, so we turn those kernels off for such groups before distributed optimizer save/
-    load.
-    """
-    for i, param_group in enumerate(optimizer.param_groups):
-        if has_mixed_tensor_and_dtensor_params(param_group.get("params", ())):
-            logger.warning_once(
-                f"Param group {i} mixes regular tensors and DTensors; disabling foreach/fused "
-                "optimizer kernels for that group so distributed optimizer save/load can materialize state."
-            )
-            param_group["foreach"] = False
-            if "fused" in param_group:
-                param_group["fused"] = False
-
-
 def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
-    """Save optimizer state via DCP.
-
-    Params whose DTensors carry a lonely `_StridedShard` placement (e.g. Mixtral
-    `gate_up_proj`) are locally split into plain-`Shard` pieces at the boundary
-    so DCP only ever sees DTensors it can encode as one contiguous chunk per
-    rank.
-    """
-    maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer)
+    """Save optimizer state via DCP."""
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
-    fusion_metadata = get_fusion_metadata(optimizer_state_dict)
-    unfuse_optimizer_state(optimizer_state_dict, fusion_metadata)
     dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
 
 
 def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
-    """Load optimizer state via DCP.
-
-    Symmetric to `save_optimizer_distributed`: build the unfused template, let DCP fill
-    it, then merge fused params back to their original `_StridedShard` form
-    before handing the state_dict back to the optimizer.
-    """
-    maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer)
+    """Load optimizer state via DCP."""
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
-    fusion_metadata = get_fusion_metadata(optimizer_state_dict)
-    unfuse_optimizer_state(optimizer_state_dict, fusion_metadata)
     dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
-    fuse_optimizer_state(optimizer_state_dict, fusion_metadata)
     set_optimizer_state_dict(model, optimizer, optimizer_state_dict)
-    maybe_disable_foreach_and_fused_for_mixed_dtensor_groups(optimizer)
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 1f738b8e1fa2..25688b5b3843 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -53,8 +53,6 @@
 )
 from .distributed import DistributedConfig
 from .distributed.utils import (
-    _distributed_barrier,
-    gather_full_state_dict,
     init_device_mesh,
     save_model_checkpoint_distributed,
 )
@@ -3543,15 +3541,13 @@ def save_pretrained(
             save_model_checkpoint_distributed(self, save_directory)
             return
 
-        # Get the model state_dict (handles FSDP unshard + TP gather in one call)
-        used_distributed_gather = False
         if state_dict is None:
             if getattr(self, "device_mesh", None) is not None:
-                # Pass self (not model_to_save) so device_mesh/tp_size/tp_plan are available
-                state_dict = gather_full_state_dict(self)
-                used_distributed_gather = True
-            else:
-                state_dict = model_to_save.state_dict()
+                #TODO(3outeille): check again
+                raise ValueError(
+                    "Saving an FSDP-sharded model requires save_pretrained(..., distributed_checkpoint=True)."
+                )
+            state_dict = model_to_save.state_dict()
 
         # if any model parameters are offloaded, we need to know it for later
         is_offloaded = False
@@ -3578,7 +3574,7 @@ def save_pretrained(
                     del state_dict[ignore_key]
 
         # If model was sharded with legacy TP, gather full tensors for saving
-        if self._tp_size is not None and not used_distributed_gather:
+        if self._tp_size is not None:
             state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size)
 
         # Remove tied weights as safetensors do not handle them
@@ -3686,13 +3682,6 @@ def save_pretrained(
                 create_pr=create_pr,
             )
 
-        # `gather_full_state_dict` concentrates the full state on rank 0 only;
-        # other ranks then loop over an empty shard list and would race ahead
-        # of rank 0's safetensors writes. Barrier so any subsequent
-        # `from_pretrained` on this path sees the consolidated files.
-        if used_distributed_gather:
-            _distributed_barrier()
-
     @wraps(PushToHubMixin.push_to_hub)
     def push_to_hub(self, *args, **kwargs):
         tags = self.model_tags if self.model_tags is not None else []
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 9bfd5707dd36..bfae6d67cc9b 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -51,7 +51,6 @@
     from transformers.distributed import DistributedConfig
     from transformers.distributed.fsdp import _resolve_tied_embed_lm_head_plan, expand_fsdp_plan
     from transformers.distributed.utils import (
-        gather_full_state_dict,
         load_optimizer_distributed,
         save_optimizer_distributed,
     )
@@ -239,6 +238,22 @@ def _load_training_state(model, optimizer, training_state_dir):
             accelerator_module.set_rng_state(rng["accel"])
 
 
+def _state_dict_from_distributed_checkpoint(model, dtype, rank):
+    """Rank 0 loads consolidated weights written via save_pretrained(distributed_checkpoint=True)."""
+    #TODO(3outeille): do we need to gather the state dict in tests?
+    with _distributed_tmpdir(rank) as tmpdir:
+        model.save_pretrained(tmpdir, is_main_process=(rank == 0), distributed_checkpoint=True)
+        dist.barrier()
+        if rank == 0:
+            cpu_model = AutoModelForCausalLM.from_pretrained(tmpdir, torch_dtype=dtype)
+            state_dict = {k: v.detach().cpu() for k, v in cpu_model.state_dict().items()}
+            del cpu_model
+        else:
+            state_dict = {}
+        dist.barrier()
+    return state_dict
+
+
 def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config, lr):
     """Save model+optimizer, scramble RNG, reload and restore training state."""
     rank = dist.get_rank()
@@ -246,7 +261,7 @@ def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config,
         model_dir = os.path.join(tmpdir, "model")
         training_state_dir = os.path.join(tmpdir, "training_state")
 
-        pre_model.save_pretrained(model_dir, is_main_process=(rank == 0))
+        pre_model.save_pretrained(model_dir, is_main_process=(rank == 0), distributed_checkpoint=True)
         _save_training_state(pre_model, pre_optimizer, training_state_dir)
         dist.barrier()
 
@@ -277,7 +292,7 @@ def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
     if dist.get_rank() != 0:
         state_dict = {}
     else:
-        # Only rank 0 returns data to match gather_full_state_dict semantics.
+        # Only rank 0 returns data for cross-rank comparison.
         state_dict = {k: v.clone().detach().cpu() for k, v in ddp_model.module.state_dict().items()}
 
     del optimizer, ddp_model, model
@@ -321,7 +336,7 @@ def train_fsdp2(
     return (
         pre_ckpt_losses + post_ckpt_losses,
         pre_ckpt_grad_norms + post_ckpt_grad_norms,
-        gather_full_state_dict(resumed_model),
+        _state_dict_from_distributed_checkpoint(resumed_model, dtype, rank),
     )
 
 
@@ -331,38 +346,46 @@ def train_fsdp2(
 
 
 def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
-    """Save FSDP2 model via save_pretrained, load via from_pretrained, compare state dicts."""
+    """Save FSDP2 model via distributed_checkpoint, load via from_pretrained, compare logits."""
     init_test_logger()
 
     config = config_class.from_dict(config_dict)
     distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
+    device = _get_rank_device(rank)
 
     with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
         _set_determinism(SEED)
         model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
         dist.barrier()
 
-        state_dict_before = gather_full_state_dict(model)
+        generator = torch.Generator(device=device)
+        generator.manual_seed(SEED)
+        input_ids = torch.randint(0, config.vocab_size, (BATCH_SIZE, SEQ_LEN), device=device, generator=generator)
+
+        model.eval()
+        with torch.no_grad():
+            logits_before = model(input_ids=input_ids, use_cache=False).logits
 
         with _distributed_tmpdir(rank) as tmpdir:
-            model.save_pretrained(tmpdir, is_main_process=(rank == 0))
+            model.save_pretrained(tmpdir, is_main_process=(rank == 0), distributed_checkpoint=True)
             dist.barrier()
             new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
             dist.barrier()
 
-        state_dict_after = gather_full_state_dict(new_model)
-        for key in state_dict_before:
-            assert key in state_dict_after, f"After save/load: Key {key} missing after load"
-            torch.testing.assert_close(
-                state_dict_before[key],
-                state_dict_after[key],
-                rtol=0,
-                atol=0,
-                msg=f"After save/load: Weight mismatch for {key}",
-            )
+        new_model.eval()
+        with torch.no_grad():
+            logits_after = new_model(input_ids=input_ids, use_cache=False).logits
+
+        torch.testing.assert_close(
+            logits_before,
+            logits_after,
+            rtol=0,
+            atol=0,
+            msg="After save/load: logits mismatch",
+        )
 
         if rank == 0:
-            logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
+            logger.debug("FSDP2 save/load test passed: logits match after distributed checkpoint round-trip.")
 
 
 def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings):
@@ -574,13 +597,13 @@ def test_fsdp2_sharding_structure(self, label):
             label == "tied",
         )
 
-    @require_torch_greater_or_equal("2.6")
+    @require_torch_greater_or_equal("2.7")
     @is_fsdp_test
     def test_fsdp2_save_load(self):
         self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl)
 
     @parameterized.expand(["untied", "tied"])
-    @require_torch_greater_or_equal("2.6")
+    @require_torch_greater_or_equal("2.7")
     @is_fsdp_test
     def test_fsdp2_plan_vs_ddp(self, label):
         self._run_fsdp2_distributed_test(

From 6e9004ec3e254eca0674a326418de775a10b590b Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 1 Jul 2026 08:41:41 +0000
Subject: [PATCH 12/33] fix

---
 src/transformers/distributed/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index 011733002808..0d0b815d867e 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -1,4 +1,4 @@
-# Copyright 2025 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 68df491a1f7eeaa60be64f6940106bce063eddd6 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Thu, 2 Jul 2026 04:38:33 +0000
Subject: [PATCH 13/33] breaking: cleaning modeling_utils.py

---
 src/transformers/distributed/__init__.py |   2 +
 src/transformers/distributed/utils.py    |  10 +
 src/transformers/modeling_utils.py       | 245 ++++++++---------------
 tests/test_fsdp_mixin.py                 |  63 ++----
 4 files changed, 121 insertions(+), 199 deletions(-)

diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py
index a79c0b58dcd0..dff6acb66217 100644
--- a/src/transformers/distributed/__init__.py
+++ b/src/transformers/distributed/__init__.py
@@ -22,6 +22,7 @@
     "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module", "verify_fsdp_plan"],
     "utils": [
         "distribute_model",
+        "gather_full_state_dict",
         "init_device_mesh",
         "load_optimizer_distributed",
         "save_model_checkpoint_distributed",
@@ -37,6 +38,7 @@
     from .fsdp import is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan
     from .utils import (
         distribute_model,
+        gather_full_state_dict,
         init_device_mesh,
         load_optimizer_distributed,
         save_model_checkpoint_distributed,
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index 0d0b815d867e..b60d299f919e 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -30,6 +30,7 @@
     #TODO(3outeille): guarding?
     import torch.distributed.checkpoint as dcp
     from torch.distributed.checkpoint.state_dict import (
+        StateDictOptions,
         get_model_state_dict,
         get_optimizer_state_dict,
         set_optimizer_state_dict,
@@ -144,6 +145,15 @@ def distribute_model(model, distributed_config: DistributedConfig, device_mesh)
     return model
 
 
+def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
+    """Gather FSDP-sharded params to full plain CPU tensors.
+
+    Only rank 0 receives the state dict; other ranks return ``{}``.
+    """
+    options = StateDictOptions(full_state_dict=True, cpu_offload=True)
+    return get_model_state_dict(model, options=options)
+
+
 def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
     """Save model parameters as standard HF-format sharded safetensors using
     DCP + HuggingFaceStorageWriter with consolidation enabled.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 25688b5b3843..42cb3d974b0e 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -52,13 +52,6 @@
     revert_weight_conversion,
 )
 from .distributed import DistributedConfig
-from .distributed.utils import (
-    init_device_mesh,
-    save_model_checkpoint_distributed,
-)
-from .distributed.utils import (
-    distribute_model as distribute_model_distributed,
-)
 from .dynamic_module_utils import custom_object_save
 from .generation import CompileConfig, GenerationConfig
 from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled, is_fsdp_enabled
@@ -216,12 +209,6 @@ def _get_torch_distributed_world_size() -> int:
     return torch.distributed.get_world_size()
 
 
-def _get_torch_distributed_rank() -> int:
-    if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_rank"):
-        return 0
-    return torch.distributed.get_rank()
-
-
 def is_local_dist_rank_0():
     return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
 
@@ -3377,7 +3364,6 @@ def save_pretrained(
         token: str | bool | None = None,
         save_peft_format: bool = True,
         save_original_format: bool = True,
-        distributed_checkpoint: bool = False,
         **kwargs,
     ):
         """
@@ -3453,13 +3439,9 @@ def save_pretrained(
             return
 
         os.makedirs(save_directory, exist_ok=True)
-        save_on_this_rank = is_main_process
-        if _is_torch_distributed_initialized() and getattr(self, "device_mesh", None) is not None:
-            # DTensor gather materializes a full plain state dict only on global rank 0.
-            save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
         save_directory_path = os.fspath(save_directory)
 
-        if push_to_hub and save_on_this_rank:
+        if push_to_hub:
             commit_message = kwargs.pop("commit_message", None)
             repo_id = kwargs.pop("repo_id", save_directory_path.split(os.path.sep)[-1])
             create_pr = kwargs.pop("create_pr", False)
@@ -3484,69 +3466,45 @@ def save_pretrained(
 
         # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
         # loaded from the Hub.
-        if self.is_remote_code() and save_on_this_rank:
+        if self.is_remote_code():
             custom_object_save(self, save_directory, config=self.config)
 
-        # Don't persist distributed_config in saved config — it's runtime-only
-        # (otherwise AutoConfig absorbs it on reload, preventing from_pretrained from seeing it as a kwarg).
-        # Keep a runtime copy around because TP/FSDP save helpers still rely on it after config serialization.
-        distributed_config = getattr(model_to_save.config, "distributed_config", None)
-        if distributed_config is not None:
-            del model_to_save.config.distributed_config
-
         # Save the config
-        try:
-            if save_on_this_rank:
-                if not _hf_peft_config_loaded:
-                    model_to_save.config.save_pretrained(save_directory)
-                if self.can_generate():
-                    model_to_save.generation_config.save_pretrained(save_directory)
+        if is_main_process:
+            if not _hf_peft_config_loaded:
+                model_to_save.config.save_pretrained(save_directory)
+            if self.can_generate():
+                model_to_save.generation_config.save_pretrained(save_directory)
 
-                if _hf_peft_config_loaded:
+            if _hf_peft_config_loaded:
+                logger.info(
+                    "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved."
+                )
+                state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict)
+
+                if save_peft_format:
                     logger.info(
-                        "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved."
+                        "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`."
                     )
-                    state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict)
-
-                    if save_peft_format:
-                        logger.info(
-                            "To match the expected format of the PEFT library, all keys of the state dict of adapters will be prepended with `base_model.model`."
-                        )
-                        peft_state_dict = {}
-                        for key, value in state_dict.items():
-                            peft_state_dict[f"base_model.model.{key}"] = value
-                        state_dict = peft_state_dict
+                    peft_state_dict = {}
+                    for key, value in state_dict.items():
+                        peft_state_dict[f"base_model.model.{key}"] = value
+                    state_dict = peft_state_dict
 
-                    active_adapter = self.active_adapters()
+                active_adapter = self.active_adapters()
 
-                    if len(active_adapter) > 1:
-                        raise ValueError(
-                            "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one "
-                            "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`"
-                        )
-                    active_adapter = active_adapter[0]
+                if len(active_adapter) > 1:
+                    raise ValueError(
+                        "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one "
+                        "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`"
+                    )
+                active_adapter = active_adapter[0]
 
-                    current_peft_config = self.peft_config[active_adapter]
-                    current_peft_config.save_pretrained(save_directory)
-        finally:
-            if distributed_config is not None:
-                model_to_save.config.distributed_config = distributed_config
-
-        if distributed_checkpoint:
-            if _is_torch_distributed_initialized() and getattr(self, "device_mesh", None) is None:
-                raise ValueError(
-                    "save_pretrained(distributed_checkpoint=True) requires the model to have been "
-                    "initialized with a distributed_config (device_mesh is None)."
-                )
-            save_model_checkpoint_distributed(self, save_directory)
-            return
+                current_peft_config = self.peft_config[active_adapter]
+                current_peft_config.save_pretrained(save_directory)
 
+        # Get the model state_dict
         if state_dict is None:
-            if getattr(self, "device_mesh", None) is not None:
-                #TODO(3outeille): check again
-                raise ValueError(
-                    "Saving an FSDP-sharded model requires save_pretrained(..., distributed_checkpoint=True)."
-                )
             state_dict = model_to_save.state_dict()
 
         # if any model parameters are offloaded, we need to know it for later
@@ -3573,7 +3531,7 @@ def save_pretrained(
                 if ignore_key in state_dict:
                     del state_dict[ignore_key]
 
-        # If model was sharded with legacy TP, gather full tensors for saving
+        # If model was sharded with TP, gather full tensors for saving
         if self._tp_size is not None:
             state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size)
 
@@ -3618,55 +3576,54 @@ def save_pretrained(
                 filename.startswith(weights_no_suffix)
                 and os.path.isfile(full_filename)
                 and filename not in state_dict_split.filename_to_tensors
-                and save_on_this_rank
+                and is_main_process
                 and reg.fullmatch(filename_no_suffix) is not None
             ):
                 os.remove(full_filename)
 
         # Save the model
-        if save_on_this_rank:
-            for shard_file, tensor_names in logging.tqdm(
-                state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
-            ):
-                filename = os.path.join(save_directory, shard_file)
-                shard_state_dict = {}
-                for tensor_name in tensor_names:
-                    # Get the tensor, and remove it from state_dict to avoid keeping the ref
-                    tensor = state_dict.pop(tensor_name)
-
-                    # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
-                    # but it would otherwise not be contained in the saved shard if we were to simply move the file
-                    # or something
-                    if is_offloaded and tensor.device.type == "meta":
-                        tensor = load_offloaded_parameter(model_to_save, tensor_name)
-
-                    # only do contiguous after it's permuted correctly in case of TP
-                    shard_state_dict[tensor_name] = tensor.contiguous()
-
-                # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
-                # so it's not possible for now....
-                # Write the shard to disk
-                safe_save_file(shard_state_dict, filename, metadata=metadata)
-                # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
-                del shard_state_dict
-
-            if index is None:
-                path_to_weights = os.path.join(save_directory, weights_name)
-                logger.info(f"Model weights saved in {path_to_weights}")
-            else:
-                save_index_file = SAFE_WEIGHTS_INDEX_NAME
-                save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
-                # Save the index as well
-                with open(save_index_file, "w", encoding="utf-8") as f:
-                    content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-                    f.write(content)
-                logger.info(
-                    f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
-                    f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
-                    f"index located at {save_index_file}."
-                )
+        for shard_file, tensor_names in logging.tqdm(
+            state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
+        ):
+            filename = os.path.join(save_directory, shard_file)
+            shard_state_dict = {}
+            for tensor_name in tensor_names:
+                # Get the tensor, and remove it from state_dict to avoid keeping the ref
+                tensor = state_dict.pop(tensor_name)
+
+                # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
+                # but it would otherwise not be contained in the saved shard if we were to simply move the file
+                # or something
+                if is_offloaded and tensor.device.type == "meta":
+                    tensor = load_offloaded_parameter(model_to_save, tensor_name)
+
+                # only do contiguous after it's permuted correctly in case of TP
+                shard_state_dict[tensor_name] = tensor.contiguous()
+
+            # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
+            # so it's not possible for now....
+            # Write the shard to disk
+            safe_save_file(shard_state_dict, filename, metadata=metadata)
+            # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
+            del shard_state_dict
+
+        if index is None:
+            path_to_weights = os.path.join(save_directory, weights_name)
+            logger.info(f"Model weights saved in {path_to_weights}")
+        else:
+            save_index_file = SAFE_WEIGHTS_INDEX_NAME
+            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
+            # Save the index as well
+            with open(save_index_file, "w", encoding="utf-8") as f:
+                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+                f.write(content)
+            logger.info(
+                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
+                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
+                f"index located at {save_index_file}."
+            )
 
-        if push_to_hub and save_on_this_rank:
+        if push_to_hub:
             # Eventually create an empty model card
             model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
 
@@ -4188,17 +4145,8 @@ def from_pretrained(
         kernel_config = kwargs.pop("kernel_config", None)
         key_mapping = kwargs.pop("key_mapping", None)
 
-        if distributed_config is not None:
-            if device_map is not None:
-                raise ValueError(
-                    "`distributed_config` and `device_map` are mutually exclusive. "
-                    "`distributed_config` handles device placement natively via torch.distributed."
-                )
-            # NOTE(3outeille): support quantization (fp4/fp8) with distributed training later
-            if quantization_config is not None:
-                raise ValueError(
-                    "Quantization is not currently supported with distributed training. Please disable quantization or distributed_config."
-                )
+        if distributed_config is not None and tp_plan is None:
+            tp_plan = "auto"
 
         # Not used anymore -- remove them from the kwargs
         for name in ["mirror", "_fast_init", "low_cpu_mem_usage", "from_tf", "from_flax", "offload_state_dict"]:
@@ -4229,26 +4177,17 @@ def from_pretrained(
                 "`state_dict` cannot be passed together with a model name or a `gguf_file`. Use one of the two loading strategies."
             )
 
-        if distributed_config is not None:
-            device_mesh = init_device_mesh(distributed_config)
-            # When using any distributed setup, we simply set the device_map here to the current rank so that we correctly
-            # load the params on the right device
-            if device_mesh.device_type == "cpu":
-                device_map = {"": torch.device("cpu")}
-            else:
-                device_map = {"": torch.device(device_mesh.device_type, int(os.environ["LOCAL_RANK"]))}
-        else:
-            if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")):
-                logger.info(
-                    "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. "
-                    "If your plan is to load the model on each device, you should set device_map={"
-                    ": PartialState().process_index} where PartialState comes from accelerate library"
-                )
+        if device_map == "auto" and int(os.environ.get("WORLD_SIZE", "0")):
+            logger.info(
+                "You've set device_map=`auto` while triggering a distributed run with torchrun. This might lead to unexpected behavior. "
+                "If your plan is to load the model on each device, you should set device_map={"
+                ": PartialState().process_index} where PartialState comes from accelerate library"
+            )
 
-            if tp_plan is not None or tp_size is not None:  # TP warnings, and setup
-                device_map, device_mesh, tp_size = initialize_tensor_parallelism(
-                    tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map
-                )
+        if tp_plan is not None or tp_size is not None:  # TP warnings, and setup
+            device_map, device_mesh, tp_size = initialize_tensor_parallelism(
+                tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map
+            )
 
         if gguf_file is not None and not is_accelerate_available():
             raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")
@@ -4261,8 +4200,7 @@ def from_pretrained(
             download_kwargs_with_commit,
             **adapter_kwargs,
         )
-        if distributed_config is None:
-            device_map = check_and_set_device_map(device_map)  # warn, error and fix the device map
+        device_map = check_and_set_device_map(device_map)  # warn, error and fix the device map
 
         user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class}
         if from_pipeline is not None:
@@ -4402,16 +4340,11 @@ def from_pretrained(
         # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively
         weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer)
 
-        if distributed_config is not None:
-            # Tie weights before sharding so `apply_fully_shard_data_parallel` /
-            # `apply_tensor_parallel` see the shared-parameter graph and can route tied
-            # entries (e.g. `lm_head` -> `embed_tokens`) correctly. `_finalize_model_loading`
-            # re-runs `tie_weights` after the checkpoint is loaded to handle missing-key edge cases.
-            model.tie_weights()
-            model = distribute_model_distributed(model, distributed_config, device_mesh)
-        elif _torch_distributed_available and device_mesh is not None:  # legacy TP hooks: no weights yet
+        if _torch_distributed_available and device_mesh is not None:  # add hooks to nn.Modules: no weights
             model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size)
-        elif device_map is not None:
+
+        # Prepare the full device map
+        if device_map is not None:
             device_map = _get_device_map(model, device_map, max_memory, hf_quantizer)
 
         # Finalize model weight initialization
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index bfae6d67cc9b..9bfd5707dd36 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -51,6 +51,7 @@
     from transformers.distributed import DistributedConfig
     from transformers.distributed.fsdp import _resolve_tied_embed_lm_head_plan, expand_fsdp_plan
     from transformers.distributed.utils import (
+        gather_full_state_dict,
         load_optimizer_distributed,
         save_optimizer_distributed,
     )
@@ -238,22 +239,6 @@ def _load_training_state(model, optimizer, training_state_dir):
             accelerator_module.set_rng_state(rng["accel"])
 
 
-def _state_dict_from_distributed_checkpoint(model, dtype, rank):
-    """Rank 0 loads consolidated weights written via save_pretrained(distributed_checkpoint=True)."""
-    #TODO(3outeille): do we need to gather the state dict in tests?
-    with _distributed_tmpdir(rank) as tmpdir:
-        model.save_pretrained(tmpdir, is_main_process=(rank == 0), distributed_checkpoint=True)
-        dist.barrier()
-        if rank == 0:
-            cpu_model = AutoModelForCausalLM.from_pretrained(tmpdir, torch_dtype=dtype)
-            state_dict = {k: v.detach().cpu() for k, v in cpu_model.state_dict().items()}
-            del cpu_model
-        else:
-            state_dict = {}
-        dist.barrier()
-    return state_dict
-
-
 def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config, lr):
     """Save model+optimizer, scramble RNG, reload and restore training state."""
     rank = dist.get_rank()
@@ -261,7 +246,7 @@ def _checkpoint_and_resume(pre_model, pre_optimizer, dtype, distributed_config,
         model_dir = os.path.join(tmpdir, "model")
         training_state_dir = os.path.join(tmpdir, "training_state")
 
-        pre_model.save_pretrained(model_dir, is_main_process=(rank == 0), distributed_checkpoint=True)
+        pre_model.save_pretrained(model_dir, is_main_process=(rank == 0))
         _save_training_state(pre_model, pre_optimizer, training_state_dir)
         dist.barrier()
 
@@ -292,7 +277,7 @@ def train_ddp(rank, batches, lr, device, dtype, init_model_dir):
     if dist.get_rank() != 0:
         state_dict = {}
     else:
-        # Only rank 0 returns data for cross-rank comparison.
+        # Only rank 0 returns data to match gather_full_state_dict semantics.
         state_dict = {k: v.clone().detach().cpu() for k, v in ddp_model.module.state_dict().items()}
 
     del optimizer, ddp_model, model
@@ -336,7 +321,7 @@ def train_fsdp2(
     return (
         pre_ckpt_losses + post_ckpt_losses,
         pre_ckpt_grad_norms + post_ckpt_grad_norms,
-        _state_dict_from_distributed_checkpoint(resumed_model, dtype, rank),
+        gather_full_state_dict(resumed_model),
     )
 
 
@@ -346,46 +331,38 @@ def train_fsdp2(
 
 
 def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
-    """Save FSDP2 model via distributed_checkpoint, load via from_pretrained, compare logits."""
+    """Save FSDP2 model via save_pretrained, load via from_pretrained, compare state dicts."""
     init_test_logger()
 
     config = config_class.from_dict(config_dict)
     distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
-    device = _get_rank_device(rank)
 
     with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
         _set_determinism(SEED)
         model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
         dist.barrier()
 
-        generator = torch.Generator(device=device)
-        generator.manual_seed(SEED)
-        input_ids = torch.randint(0, config.vocab_size, (BATCH_SIZE, SEQ_LEN), device=device, generator=generator)
-
-        model.eval()
-        with torch.no_grad():
-            logits_before = model(input_ids=input_ids, use_cache=False).logits
+        state_dict_before = gather_full_state_dict(model)
 
         with _distributed_tmpdir(rank) as tmpdir:
-            model.save_pretrained(tmpdir, is_main_process=(rank == 0), distributed_checkpoint=True)
+            model.save_pretrained(tmpdir, is_main_process=(rank == 0))
             dist.barrier()
             new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
             dist.barrier()
 
-        new_model.eval()
-        with torch.no_grad():
-            logits_after = new_model(input_ids=input_ids, use_cache=False).logits
-
-        torch.testing.assert_close(
-            logits_before,
-            logits_after,
-            rtol=0,
-            atol=0,
-            msg="After save/load: logits mismatch",
-        )
+        state_dict_after = gather_full_state_dict(new_model)
+        for key in state_dict_before:
+            assert key in state_dict_after, f"After save/load: Key {key} missing after load"
+            torch.testing.assert_close(
+                state_dict_before[key],
+                state_dict_after[key],
+                rtol=0,
+                atol=0,
+                msg=f"After save/load: Weight mismatch for {key}",
+            )
 
         if rank == 0:
-            logger.debug("FSDP2 save/load test passed: logits match after distributed checkpoint round-trip.")
+            logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
 
 
 def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings):
@@ -597,13 +574,13 @@ def test_fsdp2_sharding_structure(self, label):
             label == "tied",
         )
 
-    @require_torch_greater_or_equal("2.7")
+    @require_torch_greater_or_equal("2.6")
     @is_fsdp_test
     def test_fsdp2_save_load(self):
         self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl)
 
     @parameterized.expand(["untied", "tied"])
-    @require_torch_greater_or_equal("2.7")
+    @require_torch_greater_or_equal("2.6")
     @is_fsdp_test
     def test_fsdp2_plan_vs_ddp(self, label):
         self._run_fsdp2_distributed_test(

From 16b0b29102b66d6ffbca618ff3b2358cb61e99be Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Thu, 2 Jul 2026 06:21:42 +0000
Subject: [PATCH 14/33] load path with fsdp (dtensor) and tp (old tp) is linked

---
 src/transformers/distributed/__init__.py      |  4 +-
 .../distributed/configuration_utils.py        | 35 +++++++++++---
 src/transformers/distributed/fsdp.py          |  2 +-
 src/transformers/distributed/utils.py         | 46 +++++++++++--------
 .../integrations/tensor_parallel.py           |  6 +--
 src/transformers/modeling_utils.py            | 37 ++++++++++-----
 .../models/mixtral/configuration_mixtral.py   |  5 ++
 .../models/mixtral/modeling_mixtral.py        |  1 +
 .../models/mixtral/modular_mixtral.py         |  1 +
 9 files changed, 94 insertions(+), 43 deletions(-)

diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py
index dff6acb66217..9bb1db2d6853 100644
--- a/src/transformers/distributed/__init__.py
+++ b/src/transformers/distributed/__init__.py
@@ -23,7 +23,7 @@
     "utils": [
         "distribute_model",
         "gather_full_state_dict",
-        "init_device_mesh",
+        "initialize_fully_sharded_data_parallelism",
         "load_optimizer_distributed",
         "save_model_checkpoint_distributed",
         "save_optimizer_distributed",
@@ -39,7 +39,7 @@
     from .utils import (
         distribute_model,
         gather_full_state_dict,
-        init_device_mesh,
+        initialize_fully_sharded_data_parallelism,
         load_optimizer_distributed,
         save_model_checkpoint_distributed,
         save_optimizer_distributed,
diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py
index 9307398f6d37..b3ae3f830fa0 100644
--- a/src/transformers/distributed/configuration_utils.py
+++ b/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":
diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py
index 38baed939b60..86b29c776dfc 100644
--- a/src/transformers/distributed/fsdp.py
+++ b/src/transformers/distributed/fsdp.py
@@ -201,7 +201,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:
     """
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index b60d299f919e..65298eeb7a74 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -17,7 +17,8 @@
 from typing import TYPE_CHECKING
 
 from ..utils import is_torch_available, is_torch_greater_or_equal
-from .fsdp import apply_fully_sharded_data_parallel
+from .fsdp import apply_fully_sharded_data_parallelism
+from ..integrations.tensor_parallel import apply_tensor_parallelism
 
 
 if TYPE_CHECKING:
@@ -91,7 +92,7 @@ def _distributed_barrier():
         torch.distributed.barrier()
 
 
-def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed.device_mesh.DeviceMesh:
+def initialize_fully_sharded_data_parallelism(distributed_config: DistributedConfig):
     if not is_torch_greater_or_equal("2.5"):
         raise OSError("Distributed training with DistributedConfig requires `torch>=2.5`.")
 
@@ -103,22 +104,22 @@ def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed
 
     world_size = torch.distributed.get_world_size()
     if device_type != "cpu":
-        getattr(torch, device_type).set_device(int(os.environ.get("LOCAL_RANK", 0)))
+        local_rank = int(os.environ.get("LOCAL_RANK", 0))
+        getattr(torch, device_type).set_device(local_rank)
+        device_map = torch.device(device_type, local_rank)
+    else:
+        device_map = torch.device(device_type)
 
-    tp_size = distributed_config.tp_size
     fsdp_size = distributed_config.fsdp_size
 
-    assert world_size == tp_size * fsdp_size, (
-        f"world_size ({world_size}) must be equal to tp_size ({tp_size}) * fsdp_size ({fsdp_size})"
+    assert world_size == fsdp_size, (
+        f"world_size ({world_size}) must be equal to fsdp_size ({fsdp_size})"
     )
 
     dims, names = [], []
     if fsdp_size > 1:
         dims.append(fsdp_size)
         names.append("fsdp")
-    if tp_size > 1:
-        dims.append(tp_size)
-        names.append("tp")
 
     # Build the N-dimensional device mesh
     mesh = torch.distributed.init_device_mesh(device_type, tuple(dims), mesh_dim_names=tuple(names))
@@ -126,22 +127,29 @@ def init_device_mesh(distributed_config: DistributedConfig) -> torch.distributed
     if len(dims) > 1:
         mesh._flatten("_".join(names))
 
-    return mesh
+    return device_map, mesh
 
 
-def distribute_model(model, distributed_config: DistributedConfig, device_mesh) -> nn.Module:
-    """Apply TP and/or FSDP2 to `model` based on the mesh dims in `device_mesh`."""
+def distribute_model(
+    model,
+    distributed_config: DistributedConfig,
+    device_mesh,
+) -> nn.Module:
+    """Apply TP or FSDP2 to `model` based on ``distributed_config`` (mutually exclusive for now)."""
     model.config.distributed_config = distributed_config
     model.device_mesh = device_mesh
-    mesh_dim_names = device_mesh.mesh_dim_names or ()
-    if "tp" in mesh_dim_names:
-        from .tensor_parallel import apply_tensor_parallel
 
-        tp_mesh = device_mesh["tp"] if device_mesh.ndim > 1 else device_mesh
-        model = apply_tensor_parallel(model, tp_mesh, distributed_config.tp_plan)
-    if "fsdp" in mesh_dim_names:
+    if distributed_config.tp_size > 1:
+        model = apply_tensor_parallelism(
+            model,
+            distributed_config.tp_plan,
+            distributed_config,
+            device_mesh,
+        )
+    elif distributed_config.fsdp_size > 1:
         fsdp_mesh = device_mesh["fsdp"] if device_mesh.ndim > 1 else device_mesh
-        model = apply_fully_sharded_data_parallel(model, fsdp_mesh)
+        model = apply_fully_sharded_data_parallelism(model, fsdp_mesh)
+
     return model
 
 
diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py
index 1a9c9650b888..b8b98edb2868 100644
--- a/src/transformers/integrations/tensor_parallel.py
+++ b/src/transformers/integrations/tensor_parallel.py
@@ -1608,9 +1608,9 @@ def verify_tp_plan(expected_keys: list[str], tp_plan: dict[str, str] | None):
         logger.warning(f"The following layers were not sharded: {', '.join(unsharded_layers)}")
 
 
-def distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size):
-    """Distribute a model according to the TP plan."""
-    model._tp_size = tp_size
+def apply_tensor_parallelism(model, tp_plan, distributed_config, device_mesh):
+    """Apply tensor parallelism to a model according to the TP plan."""
+    model._tp_size = distributed_config.tp_size
     model._device_mesh = device_mesh
     if distributed_config is not None:
         if isinstance(distributed_config, dict):
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 42cb3d974b0e..43148dc2deb6 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -78,7 +78,6 @@
 from .integrations.tensor_parallel import (
     ALL_PARALLEL_STYLES,
     _get_parameter_tp_plan,
-    distribute_model,
     gather_state_dict_for_save,
     initialize_tensor_parallelism,
     shard_and_distribute_module,
@@ -135,7 +134,7 @@
 from .utils.loading_report import LoadStateDictInfo, log_state_dict_report
 from .utils.output_capturing import _CAN_RECORD_REGISTRY, OutputRecorder
 from .utils.quantization_config import QuantizationMethod
-
+from .distributed.utils import distribute_model, initialize_fully_sharded_data_parallelism
 
 if is_accelerate_available():
     from accelerate.hooks import add_hook_to_module
@@ -4135,19 +4134,17 @@ def from_pretrained(
         adapter_name = kwargs.pop("adapter_name", "default")
         generation_config = kwargs.pop("generation_config", None)
         gguf_file = kwargs.pop("gguf_file", None)
-        tp_plan = kwargs.pop("tp_plan", None)
-        tp_size = kwargs.pop("tp_size", None)
         distributed_config: DistributedConfig = kwargs.pop("distributed_config", None)
         device_mesh = kwargs.pop("device_mesh", None)
+        tp_size = None
+        use_fsdp_distributed = False
+        use_tp_distributed = False
         trust_remote_code = kwargs.pop("trust_remote_code", None)
         allow_all_kernels = kwargs.pop("allow_all_kernels", False)
         use_kernels = kwargs.pop("use_kernels", False)
         kernel_config = kwargs.pop("kernel_config", None)
         key_mapping = kwargs.pop("key_mapping", None)
 
-        if distributed_config is not None and tp_plan is None:
-            tp_plan = "auto"
-
         # Not used anymore -- remove them from the kwargs
         for name in ["mirror", "_fast_init", "low_cpu_mem_usage", "from_tf", "from_flax", "offload_state_dict"]:
             _ = kwargs.pop(name, None)
@@ -4184,10 +4181,26 @@ def from_pretrained(
                 ": PartialState().process_index} where PartialState comes from accelerate library"
             )
 
-        if tp_plan is not None or tp_size is not None:  # TP warnings, and setup
-            device_map, device_mesh, tp_size = initialize_tensor_parallelism(
-                tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map
-            )
+        if distributed_config is not None:
+            if isinstance(distributed_config, dict):
+                distributed_config = DistributedConfig.from_dict(distributed_config)
+
+            if distributed_config.tp_size > 1:
+                if distributed_config.tp_plan is None:
+                    distributed_config.tp_plan = "auto"
+                device_map, device_mesh, tp_size = initialize_tensor_parallelism(
+                    distributed_config.tp_plan,
+                    tp_size=distributed_config.tp_size,
+                    device_mesh=device_mesh,
+                    device_map=device_map,
+                )
+            elif distributed_config.fsdp_size > 1:
+                device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config)
+
+            distributed_config.validate()
+            config.distributed_config = distributed_config
+            # use_fsdp_distributed = distributed_config is not None and distributed_config.fsdp_size > 1
+            # use_tp_distributed = distributed_config is not None and distributed_config.tp_size > 1
 
         if gguf_file is not None and not is_accelerate_available():
             raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")
@@ -4341,7 +4354,7 @@ def from_pretrained(
         weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer)
 
         if _torch_distributed_available and device_mesh is not None:  # add hooks to nn.Modules: no weights
-            model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size)
+            model = distribute_model(model, distributed_config, device_mesh)
 
         # Prepare the full device map
         if device_map is not None:
diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py
index 707e98098357..16f0367d248f 100644
--- a/src/transformers/models/mixtral/configuration_mixtral.py
+++ b/src/transformers/models/mixtral/configuration_mixtral.py
@@ -62,6 +62,11 @@ class MixtralConfig(PreTrainedConfig):
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
         "layers.*.mlp.experts": "moe_tp_experts",
     }
+    base_model_fsdp_plan = {
+        "embed_tokens": "free_full_weight",
+        "layers.*": "free_full_weight",
+        "norm": "keep_full_weight",
+    }
     attribute_map = {"num_experts": "num_local_experts"}
 
     vocab_size: int = 32000
diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py
index 991851dbadd3..ecc0afcfb1fc 100644
--- a/src/transformers/models/mixtral/modeling_mixtral.py
+++ b/src/transformers/models/mixtral/modeling_mixtral.py
@@ -582,6 +582,7 @@ class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py
index 139e580fbca7..438005523f7e 100644
--- a/src/transformers/models/mixtral/modular_mixtral.py
+++ b/src/transformers/models/mixtral/modular_mixtral.py
@@ -335,6 +335,7 @@ def forward(
 
 class MixtralForCausalLM(MistralForCausalLM):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)

From e976a44c701c99aec5f73c29e1259acda84ea8bf Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Thu, 2 Jul 2026 06:25:25 +0000
Subject: [PATCH 15/33] linting

---
 src/transformers/distributed/utils.py | 9 ++++-----
 src/transformers/modeling_utils.py    | 8 ++------
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index 65298eeb7a74..87eea8cfbb0f 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -16,9 +16,9 @@
 import os
 from typing import TYPE_CHECKING
 
+from ..integrations.tensor_parallel import apply_tensor_parallelism
 from ..utils import is_torch_available, is_torch_greater_or_equal
 from .fsdp import apply_fully_sharded_data_parallelism
-from ..integrations.tensor_parallel import apply_tensor_parallelism
 
 
 if TYPE_CHECKING:
@@ -28,7 +28,8 @@
 
 if is_torch_available():
     import torch
-    #TODO(3outeille): guarding?
+
+    # TODO(3outeille): guarding?
     import torch.distributed.checkpoint as dcp
     from torch.distributed.checkpoint.state_dict import (
         StateDictOptions,
@@ -112,9 +113,7 @@ def initialize_fully_sharded_data_parallelism(distributed_config: DistributedCon
 
     fsdp_size = distributed_config.fsdp_size
 
-    assert world_size == fsdp_size, (
-        f"world_size ({world_size}) must be equal to fsdp_size ({fsdp_size})"
-    )
+    assert world_size == fsdp_size, f"world_size ({world_size}) must be equal to fsdp_size ({fsdp_size})"
 
     dims, names = [], []
     if fsdp_size > 1:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 43148dc2deb6..e6f2e9aba38c 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -52,6 +52,7 @@
     revert_weight_conversion,
 )
 from .distributed import DistributedConfig
+from .distributed.utils import distribute_model, initialize_fully_sharded_data_parallelism
 from .dynamic_module_utils import custom_object_save
 from .generation import CompileConfig, GenerationConfig
 from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled, is_fsdp_enabled
@@ -134,7 +135,7 @@
 from .utils.loading_report import LoadStateDictInfo, log_state_dict_report
 from .utils.output_capturing import _CAN_RECORD_REGISTRY, OutputRecorder
 from .utils.quantization_config import QuantizationMethod
-from .distributed.utils import distribute_model, initialize_fully_sharded_data_parallelism
+
 
 if is_accelerate_available():
     from accelerate.hooks import add_hook_to_module
@@ -4136,9 +4137,6 @@ def from_pretrained(
         gguf_file = kwargs.pop("gguf_file", None)
         distributed_config: DistributedConfig = kwargs.pop("distributed_config", None)
         device_mesh = kwargs.pop("device_mesh", None)
-        tp_size = None
-        use_fsdp_distributed = False
-        use_tp_distributed = False
         trust_remote_code = kwargs.pop("trust_remote_code", None)
         allow_all_kernels = kwargs.pop("allow_all_kernels", False)
         use_kernels = kwargs.pop("use_kernels", False)
@@ -4199,8 +4197,6 @@ def from_pretrained(
 
             distributed_config.validate()
             config.distributed_config = distributed_config
-            # use_fsdp_distributed = distributed_config is not None and distributed_config.fsdp_size > 1
-            # use_tp_distributed = distributed_config is not None and distributed_config.tp_size > 1
 
         if gguf_file is not None and not is_accelerate_available():
             raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")

From a2fb155fe70d339eac0adeade38bada479db67b6 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Thu, 2 Jul 2026 09:06:38 +0000
Subject: [PATCH 16/33] add saving

---
 src/transformers/configuration_utils.py |   4 +-
 src/transformers/distributed/fsdp.py    |   6 +-
 src/transformers/distributed/utils.py   |  11 +-
 src/transformers/modeling_utils.py      | 176 +++++++++++++++++-------
 tests/test_fsdp_mixin.py                |  48 ++++++-
 5 files changed, 184 insertions(+), 61 deletions(-)

diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py
index 01d6b3e156d5..8d97fa7b37b9 100755
--- a/src/transformers/configuration_utils.py
+++ b/src/transformers/configuration_utils.py
@@ -1039,9 +1039,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]
@@ -1188,6 +1185,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)
 
diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py
index 86b29c776dfc..74b4ec80fdca 100644
--- a/src/transformers/distributed/fsdp.py
+++ b/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
 
@@ -210,8 +210,8 @@ def apply_fully_sharded_data_parallelism(
     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:
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index 87eea8cfbb0f..e6318724cf85 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -97,8 +97,8 @@ def initialize_fully_sharded_data_parallelism(distributed_config: DistributedCon
     if not is_torch_greater_or_equal("2.5"):
         raise OSError("Distributed training with DistributedConfig requires `torch>=2.5`.")
 
-    if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.6"):
-        raise OSError("FSDP2 requires `torch>=2.6`.")
+    if distributed_config.fsdp_size > 1 and not is_torch_greater_or_equal("2.7"):
+        raise OSError("FSDP2 requires `torch>=2.7`.")
 
     device_type = torch._C._get_accelerator().type
     _ensure_torch_distributed(device_type)
@@ -155,10 +155,13 @@ def distribute_model(
 def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
     """Gather FSDP-sharded params to full plain CPU tensors.
 
-    Only rank 0 receives the state dict; other ranks return ``{}``.
+    Only rank 0 accumulates the result; other ranks return ``{}``.
     """
     options = StateDictOptions(full_state_dict=True, cpu_offload=True)
-    return get_model_state_dict(model, options=options)
+    full_state_dict = get_model_state_dict(model, options=options)
+    if torch.distributed.get_rank() == 0:
+        return full_state_dict
+    return {}
 
 
 def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index e6f2e9aba38c..8c2d4438fbd2 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -52,7 +52,14 @@
     revert_weight_conversion,
 )
 from .distributed import DistributedConfig
-from .distributed.utils import distribute_model, initialize_fully_sharded_data_parallelism
+from .distributed.fsdp import is_fsdp_managed_module
+from .distributed.utils import (
+    _distributed_barrier,
+    distribute_model,
+    gather_full_state_dict,
+    initialize_fully_sharded_data_parallelism,
+    save_model_checkpoint_distributed,
+)
 from .dynamic_module_utils import custom_object_save
 from .generation import CompileConfig, GenerationConfig
 from .integrations import PeftAdapterMixin, deepspeed_config, hub_kernels, is_deepspeed_zero3_enabled, is_fsdp_enabled
@@ -130,6 +137,7 @@
     is_huggingface_hub_greater_or_equal,
     is_sagemaker_mp_enabled,
     is_torch_cuda_available,
+    is_torch_greater_or_equal,
     is_tracing,
 )
 from .utils.loading_report import LoadStateDictInfo, log_state_dict_report
@@ -213,6 +221,12 @@ def is_local_dist_rank_0():
     return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
 
 
+def _get_torch_distributed_rank() -> int:
+    if not _is_torch_distributed_initialized():
+        return 0
+    return torch.distributed.get_rank()
+
+
 @contextmanager
 def set_quantized_state():
     global _is_quantized
@@ -3364,6 +3378,7 @@ def save_pretrained(
         token: str | bool | None = None,
         save_peft_format: bool = True,
         save_original_format: bool = True,
+        distributed_checkpoint: bool = False,
         **kwargs,
     ):
         """
@@ -3387,7 +3402,7 @@ def save_pretrained(
                 namespace).
             max_shard_size (`int` or `str`, *optional*, defaults to `"50GB"`):
                 The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
-                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).
+                lower than this size. If expressed as a string, needs be digits followed by a unit (like `"5MB"`).
 
                 <Tip warning={true}>
 
@@ -3409,6 +3424,11 @@ def save_pretrained(
                 For backward compatibility with the previous versions of `transformers` you can save the checkpoint with
                 its reverse mapping. The reverse mapping needs to exists even if the model was loaded from a None legacy
                 checkpoint.
+            distributed_checkpoint (`bool`, *optional*, defaults to `False`):
+                When saving an FSDP-wrapped model, use the distributed checkpoint (DCP) path instead of gathering weights
+                to CPU first. Every rank must call this method; rank 0 writes the consolidated Hugging Face safetensors.
+                When `False`, FSDP weights are gathered to CPU on rank 0 via `gather_full_state_dict` before writing.
+                Native FSDP requires `torch>=2.7`.
             kwargs (`dict[str, Any]`, *optional*):
                 Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
         """
@@ -3455,6 +3475,23 @@ def save_pretrained(
 
         # Only save the model itself if we are using distributed training
         model_to_save = unwrap_model(self)
+        is_fsdp_model = is_fsdp_managed_module(model_to_save)
+
+        if distributed_checkpoint:
+            if not is_torch_greater_or_equal("2.7"):
+                raise OSError(
+                    "save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7."
+                )
+            if not is_fsdp_model:
+                raise ValueError(
+                    "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models."
+                )
+
+        # Collectives run on every rank; checkpoint files are written on global rank 0 only.
+        save_on_this_rank = is_main_process
+        if _is_torch_distributed_initialized():
+            save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
+        
         # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
         # we currently don't use this setting automatically, but may start to use with v5
         dtype = model_to_save.dtype
@@ -3466,11 +3503,11 @@ def save_pretrained(
 
         # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
         # loaded from the Hub.
-        if self.is_remote_code():
+        if save_on_this_rank and self.is_remote_code():
             custom_object_save(self, save_directory, config=self.config)
 
         # Save the config
-        if is_main_process:
+        if save_on_this_rank:
             if not _hf_peft_config_loaded:
                 model_to_save.config.save_pretrained(save_directory)
             if self.can_generate():
@@ -3503,9 +3540,41 @@ def save_pretrained(
                 current_peft_config = self.peft_config[active_adapter]
                 current_peft_config.save_pretrained(save_directory)
 
-        # Get the model state_dict
+        # --- FSDP path: DCP save (large models, torch >= 2.7) ---
+        if distributed_checkpoint:
+            if getattr(model_to_save, "device_mesh", None) is None:
+                raise ValueError(
+                    "save_pretrained(..., distributed_checkpoint=True) requires the model to have been "
+                    "initialized with a distributed_config (device_mesh is None)."
+                )
+            save_model_checkpoint_distributed(model_to_save, save_directory)
+            if push_to_hub and save_on_this_rank:
+                model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
+                model_card.save(os.path.join(save_directory, "README.md"))
+                self._upload_modified_files(
+                    save_directory,
+                    repo_id,
+                    files_timestamps,
+                    commit_message=commit_message,
+                    token=token,
+                    create_pr=create_pr,
+                )
+            return
+
+        
+        # Get the state dict
+        used_distributed_gather = False
         if state_dict is None:
-            state_dict = model_to_save.state_dict()
+            if is_fsdp_model:
+                if not _is_torch_distributed_initialized():
+                    raise ValueError(
+                        "Saving an FSDP-wrapped model requires torch.distributed to be initialized. "
+                        "Call save_pretrained from every rank after init_process_group."
+                    )
+                state_dict = gather_full_state_dict(model_to_save)
+                used_distributed_gather = True
+            else:
+                state_dict = model_to_save.state_dict()
 
         # if any model parameters are offloaded, we need to know it for later
         is_offloaded = False
@@ -3534,6 +3603,9 @@ def save_pretrained(
         # If model was sharded with TP, gather full tensors for saving
         if self._tp_size is not None:
             state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size)
+            used_distributed_gather = True
+            if not save_on_this_rank:
+                state_dict = {}
 
         # Remove tied weights as safetensors do not handle them
         state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save)
@@ -3576,54 +3648,55 @@ def save_pretrained(
                 filename.startswith(weights_no_suffix)
                 and os.path.isfile(full_filename)
                 and filename not in state_dict_split.filename_to_tensors
-                and is_main_process
+                and save_on_this_rank
                 and reg.fullmatch(filename_no_suffix) is not None
             ):
                 os.remove(full_filename)
 
         # Save the model
-        for shard_file, tensor_names in logging.tqdm(
-            state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
-        ):
-            filename = os.path.join(save_directory, shard_file)
-            shard_state_dict = {}
-            for tensor_name in tensor_names:
-                # Get the tensor, and remove it from state_dict to avoid keeping the ref
-                tensor = state_dict.pop(tensor_name)
-
-                # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
-                # but it would otherwise not be contained in the saved shard if we were to simply move the file
-                # or something
-                if is_offloaded and tensor.device.type == "meta":
-                    tensor = load_offloaded_parameter(model_to_save, tensor_name)
-
-                # only do contiguous after it's permuted correctly in case of TP
-                shard_state_dict[tensor_name] = tensor.contiguous()
-
-            # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
-            # so it's not possible for now....
-            # Write the shard to disk
-            safe_save_file(shard_state_dict, filename, metadata=metadata)
-            # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
-            del shard_state_dict
-
-        if index is None:
-            path_to_weights = os.path.join(save_directory, weights_name)
-            logger.info(f"Model weights saved in {path_to_weights}")
-        else:
-            save_index_file = SAFE_WEIGHTS_INDEX_NAME
-            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
-            # Save the index as well
-            with open(save_index_file, "w", encoding="utf-8") as f:
-                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-                f.write(content)
-            logger.info(
-                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
-                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
-                f"index located at {save_index_file}."
-            )
+        if save_on_this_rank:
+            for shard_file, tensor_names in logging.tqdm(
+                state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
+            ):
+                filename = os.path.join(save_directory, shard_file)
+                shard_state_dict = {}
+                for tensor_name in tensor_names:
+                    # Get the tensor, and remove it from state_dict to avoid keeping the ref
+                    tensor = state_dict.pop(tensor_name)
+
+                    # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
+                    # but it would otherwise not be contained in the saved shard if we were to simply move the file
+                    # or something
+                    if is_offloaded and tensor.device.type == "meta":
+                        tensor = load_offloaded_parameter(model_to_save, tensor_name)
+
+                    # only do contiguous after it's permuted correctly in case of TP
+                    shard_state_dict[tensor_name] = tensor.contiguous()
+
+                # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
+                # so it's not possible for now....
+                # Write the shard to disk
+                safe_save_file(shard_state_dict, filename, metadata=metadata)
+                # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
+                del shard_state_dict
+
+            if index is None:
+                path_to_weights = os.path.join(save_directory, weights_name)
+                logger.info(f"Model weights saved in {path_to_weights}")
+            else:
+                save_index_file = SAFE_WEIGHTS_INDEX_NAME
+                save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
+                # Save the index as well
+                with open(save_index_file, "w", encoding="utf-8") as f:
+                    content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+                    f.write(content)
+                logger.info(
+                    f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
+                    f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
+                    f"index located at {save_index_file}."
+                )
 
-        if push_to_hub:
+        if push_to_hub and save_on_this_rank:
             # Eventually create an empty model card
             model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
 
@@ -3639,6 +3712,13 @@ def save_pretrained(
                 create_pr=create_pr,
             )
 
+        # `gather_full_state_dict` concentrates the full state on rank 0 only;
+        # other ranks then loop over an empty shard list and would race ahead
+        # of rank 0's safetensors writes. Barrier so any subsequent
+        # `from_pretrained` on this path sees the consolidated files.
+        if used_distributed_gather:
+            _distributed_barrier()
+
     @wraps(PushToHubMixin.push_to_hub)
     def push_to_hub(self, *args, **kwargs):
         tags = self.model_tags if self.model_tags is not None else []
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 9bfd5707dd36..c0c147c28953 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -365,6 +365,43 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
             logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
 
 
+def _test_fsdp2_save_load_dcp_impl(rank, config_class, config_dict):
+    """Save FSDP2 model via save_pretrained(distributed_checkpoint=True), reload, compare state dicts."""
+    init_test_logger()
+
+    config = config_class.from_dict(config_dict)
+    distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
+
+    with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir:
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(init_dir, distributed_config=distributed_config)
+        dist.barrier()
+
+        state_dict_before = gather_full_state_dict(model)
+
+        with _distributed_tmpdir(rank) as tmpdir:
+            model.save_pretrained(tmpdir, is_main_process=(rank == 0), distributed_checkpoint=True)
+            dist.barrier()
+            new_model = AutoModelForCausalLM.from_pretrained(tmpdir, distributed_config=distributed_config)
+            dist.barrier()
+
+        state_dict_after = gather_full_state_dict(new_model)
+        for key in state_dict_before:
+            assert key in state_dict_after, f"After DCP save/load: Key {key} missing after load"
+            torch.testing.assert_close(
+                state_dict_before[key],
+                state_dict_after[key],
+                rtol=0,
+                atol=0,
+                msg=f"After DCP save/load: Weight mismatch for {key}",
+            )
+
+        if rank == 0:
+            logger.debug(
+                f"FSDP2 DCP save/load test passed: all {len(state_dict_before)} parameters match exactly."
+            )
+
+
 def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings):
     """Verify that apply_fully_sharded_data_parallel wraps exactly the right modules."""
     init_test_logger()
@@ -565,7 +602,7 @@ def test_fsdp_plan_declared(self):
         self.assertTrue(model._fsdp_plan, f"No _fsdp_plan declared for {type(model).__name__}")
 
     @parameterized.expand(["untied", "tied"])
-    @require_torch_greater_or_equal("2.6")
+    @require_torch_greater_or_equal("2.7")
     @is_fsdp_test
     def test_fsdp2_sharding_structure(self, label):
         self._run_fsdp2_distributed_test(
@@ -574,13 +611,18 @@ def test_fsdp2_sharding_structure(self, label):
             label == "tied",
         )
 
-    @require_torch_greater_or_equal("2.6")
+    @require_torch_greater_or_equal("2.7")
     @is_fsdp_test
     def test_fsdp2_save_load(self):
         self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl)
 
+    @require_torch_greater_or_equal("2.7")
+    @is_fsdp_test
+    def test_fsdp2_save_load_dcp(self):
+        self._run_fsdp2_distributed_test("test_fsdp2_save_load_dcp", _test_fsdp2_save_load_dcp_impl)
+
     @parameterized.expand(["untied", "tied"])
-    @require_torch_greater_or_equal("2.6")
+    @require_torch_greater_or_equal("2.7")
     @is_fsdp_test
     def test_fsdp2_plan_vs_ddp(self, label):
         self._run_fsdp2_distributed_test(

From 5f52f1965030c690837b5f4458f19e95ec370757 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Thu, 2 Jul 2026 09:14:03 +0000
Subject: [PATCH 17/33] styling

---
 src/transformers/modeling_utils.py | 7 ++-----
 tests/test_fsdp_mixin.py           | 4 +---
 2 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 8c2d4438fbd2..0e0c85550e2b 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3479,9 +3479,7 @@ def save_pretrained(
 
         if distributed_checkpoint:
             if not is_torch_greater_or_equal("2.7"):
-                raise OSError(
-                    "save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7."
-                )
+                raise OSError("save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7.")
             if not is_fsdp_model:
                 raise ValueError(
                     "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models."
@@ -3491,7 +3489,7 @@ def save_pretrained(
         save_on_this_rank = is_main_process
         if _is_torch_distributed_initialized():
             save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
-        
+
         # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
         # we currently don't use this setting automatically, but may start to use with v5
         dtype = model_to_save.dtype
@@ -3561,7 +3559,6 @@ def save_pretrained(
                 )
             return
 
-        
         # Get the state dict
         used_distributed_gather = False
         if state_dict is None:
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index c0c147c28953..b0f1491a0bcf 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -397,9 +397,7 @@ def _test_fsdp2_save_load_dcp_impl(rank, config_class, config_dict):
             )
 
         if rank == 0:
-            logger.debug(
-                f"FSDP2 DCP save/load test passed: all {len(state_dict_before)} parameters match exactly."
-            )
+            logger.debug(f"FSDP2 DCP save/load test passed: all {len(state_dict_before)} parameters match exactly.")
 
 
 def _test_fsdp2_sharding_structure_impl(rank, config_class, config_dict, tie_word_embeddings):

From 7f5430172555ad75f4494d2e382632a7eb25a678 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 01:28:45 +0000
Subject: [PATCH 18/33] fix tp ci

---
 src/transformers/modeling_utils.py                |  4 +++-
 tests/generation/test_continuous_batching.py      | 15 ++++++++++++---
 .../glm_moe_dsa/test_modeling_glm_moe_dsa.py      |  4 +++-
 tests/models/gpt_oss/test_modeling_gpt_oss.py     |  3 ++-
 tests/test_fsdp_mixin.py                          |  1 +
 tests/test_tensor_parallel_mixin.py               | 15 +++++++++++----
 6 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 0e0c85550e2b..d37e78569ee5 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -4273,7 +4273,6 @@ def from_pretrained(
                 device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config)
 
             distributed_config.validate()
-            config.distributed_config = distributed_config
 
         if gguf_file is not None and not is_accelerate_available():
             raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")
@@ -4317,6 +4316,9 @@ def from_pretrained(
             model_kwargs = kwargs
             commit_hash = getattr(config, "_commit_hash", commit_hash)
 
+        if distributed_config is not None:
+            config.distributed_config = distributed_config
+
         download_kwargs_with_commit["commit_hash"] = commit_hash
 
         # Because some composite configs call super().__init__ before instantiating the sub-configs, we need this call
diff --git a/tests/generation/test_continuous_batching.py b/tests/generation/test_continuous_batching.py
index 44d132886a45..4ad084f66b66 100644
--- a/tests/generation/test_continuous_batching.py
+++ b/tests/generation/test_continuous_batching.py
@@ -1739,13 +1739,14 @@ def _tp_continuous_batching_worker(
     use_cuda_graph: bool,
     use_async_batching: bool,
 ) -> None:
-    """Loads `model_id` with `tp_plan="auto"`, checks three TP-specific paths in the same process: (a) direct
+    """Loads `model_id` with `DistributedConfig(tp_size=...)`, checks three TP-specific paths in the same process: (a) direct
     broadcasts via `DistributedHelper`, (b) per-rank parity of CB-generated tokens via `dist.all_gather_object`, and
     (c) reproducibility across two CB runs sharing the same seed. Rank 0 owns all the assertions; the other ranks
     only need to participate in the collectives."""
     import torch
     import torch.distributed as dist
 
+    from transformers.distributed import DistributedConfig
     from transformers.generation.continuous_batching.distributed import DistributedHelper
 
     tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
@@ -1753,7 +1754,10 @@ def _tp_continuous_batching_worker(
         tokenizer.pad_token = tokenizer.eos_token
 
     model = AutoModelForCausalLM.from_pretrained(
-        model_id, attn_implementation=attn_implementation, tp_plan="auto", dtype=torch.float32
+        model_id,
+        attn_implementation=attn_implementation,
+        distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])),
+        dtype=torch.float32,
     ).eval()
 
     # Direct broadcast tests: only rank 0's value should propagate to every TP rank
@@ -1831,6 +1835,8 @@ def _tp_cancellation_worker(
 
     import torch
 
+    from transformers.distributed import DistributedConfig
+
     cb_config = ContinuousBatchingConfig(use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching)
 
     tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
@@ -1838,7 +1844,10 @@ def _tp_cancellation_worker(
         tokenizer.pad_token = tokenizer.eos_token
 
     model = AutoModelForCausalLM.from_pretrained(
-        model_id, attn_implementation=attn_implementation, tp_plan="auto", dtype=torch.float32
+        model_id,
+        attn_implementation=attn_implementation,
+        distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])),
+        dtype=torch.float32,
     ).eval()
 
     chat = [{"role": "user", "content": "Tell me a long story about a robot exploring the galaxy."}]
diff --git a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py
index aa87525eede2..45f760fc2f2d 100644
--- a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py
+++ b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Testing suite for the PyTorch GlmMoeDsa model."""
 
+import os
 import unittest
 
 import torch
@@ -27,6 +28,7 @@
     is_torch_available,
     set_seed,
 )
+from transformers.distributed import DistributedConfig
 from transformers.testing_utils import (
     require_torch,
     require_torch_accelerator,
@@ -204,7 +206,7 @@ def test_glm_moe_dsa_fp8_inference(self):
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
             quantization_config=quantization_config,
-            tp_plan="auto",
+            distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])),
             attn_implementation="eager",
         )
 
diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py
index 5b3c13f8c25a..1dcd911bdbe4 100644
--- a/tests/models/gpt_oss/test_modeling_gpt_oss.py
+++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py
@@ -159,6 +159,7 @@ def distributed_worker(quantized, model_size, kernels, attn_impl, mode):
     import os
 
     from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers.distributed import DistributedConfig
     from transformers.testing_utils import torch_device
 
     def generate_config_key(quantized, model, kernels, attn_impl, mode):
@@ -179,7 +180,7 @@ def generate_config_key(quantized, model, kernels, attn_impl, mode):
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         dtype="auto",
-        tp_plan="auto",  # distributed inference
+        distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])),
         use_kernels=kernels,
     ).to(torch_device)
     model.set_attn_implementation(attn_impl)
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index b0f1491a0bcf..3d0faef77f7f 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -73,6 +73,7 @@
 FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {
     "llama",
     "qwen3_moe",
+    "mixtral"
 }
 
 
diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py
index a40649ad883c..7f18a7678ba6 100644
--- a/tests/test_tensor_parallel_mixin.py
+++ b/tests/test_tensor_parallel_mixin.py
@@ -118,7 +118,10 @@ def _load_tp_and_reference_models(model_path, model_class):
     Returns:
         tuple: (model_tp, model_ref, device)
     """
-    model_tp = model_class.from_pretrained(model_path, tp_plan="auto")
+    model_tp = model_class.from_pretrained(
+        model_path,
+        distributed_config=DistributedConfig(tp_size=dist.get_world_size()),
+    )
     dist.barrier()
 
     device = model_tp.device
@@ -146,7 +149,7 @@ def _verify_tp_sharding(rank, model_tp, model_ref):
             # Verify sharding is correct
             for dim in range(param.ndim):
                 if param.size(dim) != param_full.size(dim):
-                    param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True)
+                    param_plan = _get_parameter_tp_plan(name, model_tp._tp_plan, is_weight=True)
                     if param_plan in ("packed_colwise",):
                         expected_size = param_full.size(dim) // world_size
                         assert param.size(dim) == expected_size, (
@@ -227,7 +230,7 @@ def _test_tp_backward_impl(rank, model_path, model_class, atol, rtol):
             if grad.shape != grad_tp.shape:
                 for dim in range(grad.ndim):
                     if grad.size(dim) != grad_tp.size(dim):
-                        param_plan = _get_parameter_tp_plan(name, model_tp.tp_plan, is_weight=True)
+                        param_plan = _get_parameter_tp_plan(name, model_tp._tp_plan, is_weight=True)
                         if param_plan in ("packed_colwise",):
                             # interleaved slicing
                             grad = get_packed_grad_shard(grad, world_size, rank, dim)
@@ -297,7 +300,11 @@ def _test_tp_generation_quantized_impl(_rank, model_path, model_class, max_new_t
 
     quantization_config = TorchAoConfig(Float8WeightOnlyConfig())
 
-    model_tp = model_class.from_pretrained(model_path, tp_plan="auto", quantization_config=quantization_config)
+    model_tp = model_class.from_pretrained(
+        model_path,
+        distributed_config=DistributedConfig(tp_size=dist.get_world_size()),
+        quantization_config=quantization_config,
+    )
     dist.barrier()
 
     device = model_tp.device

From 99f79acc993ebe696a2728cc8110277f37c9d8e7 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 01:38:47 +0000
Subject: [PATCH 19/33] add fsdp to ci

---
 .circleci/create_circleci_config.py | 14 ++++++++++++--
 conftest.py                         |  1 +
 docs/source/en/pr_checks.md         |  1 +
 docs/source/ro/pr_checks.md         |  1 +
 utils/tests_fetcher.py              |  1 +
 5 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/.circleci/create_circleci_config.py b/.circleci/create_circleci_config.py
index de0646af2dbf..d004a4c37cf0 100644
--- a/.circleci/create_circleci_config.py
+++ b/.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):
diff --git a/conftest.py b/conftest.py
index 3970d3ffee9b..899c38acebd3 100644
--- a/conftest.py
+++ b/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)
diff --git a/docs/source/en/pr_checks.md b/docs/source/en/pr_checks.md
index a17d84700989..704546f24f61 100644
--- a/docs/source/en/pr_checks.md
+++ b/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
 
diff --git a/docs/source/ro/pr_checks.md b/docs/source/ro/pr_checks.md
index bceb75d73e60..e452ed4b08ac 100644
--- a/docs/source/ro/pr_checks.md
+++ b/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
 
diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py
index a138ef2eaacb..095df71293f7 100644
--- a/utils/tests_fetcher.py
+++ b/utils/tests_fetcher.py
@@ -1099,6 +1099,7 @@ def parse_commit_message(commit_message: str) -> dict[str, bool]:
     "tests_non_model": r"tests/[^/]*?/test_.*\.py",
     "tests_training_ci": r"tests/models/.*/test_modeling_.*",
     "tests_tensor_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/tensor_parallel(?:/test_tensor_parallel\.py)?)",
+    "tests_fsdp_ci": r"tests/models/.*/test_modeling_.*",
 }
 
 

From b45149068fac2f0de17692ca94d9b645d94169b3 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 01:40:54 +0000
Subject: [PATCH 20/33] linting

---
 tests/test_fsdp_mixin.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index 3d0faef77f7f..f937c7634c6d 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -70,11 +70,7 @@
 DDP_FSDP_ATOL = 1e-5
 
 # Set to None to run distributed FSDP tests for every model with a plan.
-FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {
-    "llama",
-    "qwen3_moe",
-    "mixtral"
-}
+FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {"llama", "qwen3_moe", "mixtral"}
 
 
 # =============================================================================

From 54ff4d1d3eaf43d7ed597699d67de00e1b85d08e Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 02:14:17 +0000
Subject: [PATCH 21/33] pick one model only for this PR

---
 .../models/cohere2_moe/configuration_cohere2_moe.py   |  5 +++++
 .../models/cohere2_moe/modeling_cohere2_moe.py        |  1 +
 .../models/cohere2_moe/modular_cohere2_moe.py         |  2 ++
 src/transformers/models/llama/configuration_llama.py  |  5 -----
 src/transformers/models/llama/modeling_llama.py       |  1 -
 .../models/mixtral/configuration_mixtral.py           | 11 -----------
 src/transformers/models/mixtral/modeling_mixtral.py   |  1 -
 src/transformers/models/mixtral/modular_mixtral.py    |  1 -
 .../models/qwen3_moe/configuration_qwen3_moe.py       |  5 -----
 .../models/qwen3_moe/modeling_qwen3_moe.py            |  1 -
 .../models/qwen3_moe/modular_qwen3_moe.py             |  2 --
 tests/test_fsdp_mixin.py                              |  2 +-
 12 files changed, 9 insertions(+), 28 deletions(-)

diff --git a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
index 32ef6fdf66ad..83af4aa21894 100644
--- a/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
+++ b/src/transformers/models/cohere2_moe/configuration_cohere2_moe.py
@@ -90,6 +90,11 @@ class Cohere2MoeConfig(PreTrainedConfig):
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
         "layers.*.mlp.experts": "moe_tp_experts",
     }
+    base_model_fsdp_plan = {
+        "embed_tokens": "free_full_weight",
+        "layers.*": "free_full_weight",
+        "norm": "keep_full_weight",
+    }
 
     vocab_size: int = 256000
     hidden_size: int = 8192
diff --git a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py
index e0290b835948..bdd0ed5b596f 100644
--- a/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py
+++ b/src/transformers/models/cohere2_moe/modeling_cohere2_moe.py
@@ -583,6 +583,7 @@ class Cohere2MoeForCausalLM(Cohere2MoePreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/cohere2_moe/modular_cohere2_moe.py b/src/transformers/models/cohere2_moe/modular_cohere2_moe.py
index b5d4c3b49f57..57c76cb0805a 100644
--- a/src/transformers/models/cohere2_moe/modular_cohere2_moe.py
+++ b/src/transformers/models/cohere2_moe/modular_cohere2_moe.py
@@ -279,6 +279,8 @@ def forward(
 
 
 class Cohere2MoeForCausalLM(Cohere2ForCausalLM):
+    _fsdp_plan = {"lm_head": "keep_full_weight"}
+
     def forward(
         self,
         input_ids: torch.LongTensor | None = None,
diff --git a/src/transformers/models/llama/configuration_llama.py b/src/transformers/models/llama/configuration_llama.py
index 9d84c0c16d3c..6960a6970592 100644
--- a/src/transformers/models/llama/configuration_llama.py
+++ b/src/transformers/models/llama/configuration_llama.py
@@ -60,11 +60,6 @@ class LlamaConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
-    base_model_fsdp_plan = {
-        "embed_tokens": "free_full_weight",
-        "layers.*": "free_full_weight",
-        "norm": "keep_full_weight",
-    }
 
     vocab_size: int = 32000
     hidden_size: int = 4096
diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py
index d56a47c3651c..9d659c7c6f08 100644
--- a/src/transformers/models/llama/modeling_llama.py
+++ b/src/transformers/models/llama/modeling_llama.py
@@ -430,7 +430,6 @@ class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
-    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py
index 16f0367d248f..240f24411031 100644
--- a/src/transformers/models/mixtral/configuration_mixtral.py
+++ b/src/transformers/models/mixtral/configuration_mixtral.py
@@ -56,17 +56,6 @@ class MixtralConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
-    base_model_ep_plan = {
-        "layers.*.mlp.gate": "ep_router",
-        "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
-        "layers.*.mlp.experts.down_proj": "grouped_gemm",
-        "layers.*.mlp.experts": "moe_tp_experts",
-    }
-    base_model_fsdp_plan = {
-        "embed_tokens": "free_full_weight",
-        "layers.*": "free_full_weight",
-        "norm": "keep_full_weight",
-    }
     attribute_map = {"num_experts": "num_local_experts"}
 
     vocab_size: int = 32000
diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py
index ecc0afcfb1fc..991851dbadd3 100644
--- a/src/transformers/models/mixtral/modeling_mixtral.py
+++ b/src/transformers/models/mixtral/modeling_mixtral.py
@@ -582,7 +582,6 @@ class MixtralForCausalLM(MixtralPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
-    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py
index 438005523f7e..139e580fbca7 100644
--- a/src/transformers/models/mixtral/modular_mixtral.py
+++ b/src/transformers/models/mixtral/modular_mixtral.py
@@ -335,7 +335,6 @@ def forward(
 
 class MixtralForCausalLM(MistralForCausalLM):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
index ca040eedbec4..9a7d4b4c8b5b 100644
--- a/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/configuration_qwen3_moe.py
@@ -81,11 +81,6 @@ class Qwen3MoeConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
-    base_model_fsdp_plan = {
-        "embed_tokens": "free_full_weight",
-        "layers.*": "free_full_weight",
-        "norm": "keep_full_weight",
-    }
 
     vocab_size: int = 151936
     hidden_size: int = 2048
diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
index 7f4c7404356d..ddf84fc575b7 100644
--- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py
@@ -611,7 +611,6 @@ class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
     _tp_plan = {"lm_head": "colwise_gather_output"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
-    _fsdp_plan = {"lm_head": "keep_full_weight"}
 
     def __init__(self, config):
         super().__init__(config)
diff --git a/src/transformers/models/qwen3_moe/modular_qwen3_moe.py b/src/transformers/models/qwen3_moe/modular_qwen3_moe.py
index 6453f1754a6a..0fd5b451959c 100644
--- a/src/transformers/models/qwen3_moe/modular_qwen3_moe.py
+++ b/src/transformers/models/qwen3_moe/modular_qwen3_moe.py
@@ -95,8 +95,6 @@ class Qwen3MoeModel(MixtralModel):
 
 
 class Qwen3MoeForCausalLM(MixtralForCausalLM):
-    _fsdp_plan = {"lm_head": "keep_full_weight"}
-
     def __init__(self, config):
         super().__init__(config)
         self.model = Qwen3MoeModel(config)
diff --git a/tests/test_fsdp_mixin.py b/tests/test_fsdp_mixin.py
index f937c7634c6d..2a966328fc22 100644
--- a/tests/test_fsdp_mixin.py
+++ b/tests/test_fsdp_mixin.py
@@ -70,7 +70,7 @@
 DDP_FSDP_ATOL = 1e-5
 
 # Set to None to run distributed FSDP tests for every model with a plan.
-FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {"llama", "qwen3_moe", "mixtral"}
+FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {"cohere2_moe"}
 
 
 # =============================================================================

From 339953990fc917fbc257f65a28ea2129b9654cc9 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 02:17:41 +0000
Subject: [PATCH 22/33] restore

---
 src/transformers/models/mixtral/configuration_mixtral.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/transformers/models/mixtral/configuration_mixtral.py b/src/transformers/models/mixtral/configuration_mixtral.py
index 240f24411031..707e98098357 100644
--- a/src/transformers/models/mixtral/configuration_mixtral.py
+++ b/src/transformers/models/mixtral/configuration_mixtral.py
@@ -56,6 +56,12 @@ class MixtralConfig(PreTrainedConfig):
         "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
         "norm": (["hidden_states"], ["hidden_states"]),
     }
+    base_model_ep_plan = {
+        "layers.*.mlp.gate": "ep_router",
+        "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+        "layers.*.mlp.experts.down_proj": "grouped_gemm",
+        "layers.*.mlp.experts": "moe_tp_experts",
+    }
     attribute_map = {"num_experts": "num_local_experts"}
 
     vocab_size: int = 32000

From 11cf79c23aeb795942b73643bbe598b8b891e638 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 02:19:48 +0000
Subject: [PATCH 23/33] trigger fsdp ci

---
 utils/tests_fetcher.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py
index 095df71293f7..f09480a59875 100644
--- a/utils/tests_fetcher.py
+++ b/utils/tests_fetcher.py
@@ -1099,7 +1099,7 @@ def parse_commit_message(commit_message: str) -> dict[str, bool]:
     "tests_non_model": r"tests/[^/]*?/test_.*\.py",
     "tests_training_ci": r"tests/models/.*/test_modeling_.*",
     "tests_tensor_parallel_ci": r"(tests/models/.*/test_modeling_.*|tests/tensor_parallel(?:/test_tensor_parallel\.py)?)",
-    "tests_fsdp_ci": r"tests/models/.*/test_modeling_.*",
+    "tests_fsdp_ci": r"(tests/models/.*/test_modeling_.*|tests/test_fsdp_mixin\.py)",
 }
 
 

From 5b7ac3e55c9cca38c56253a38510461ed24a8fa8 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 02:38:07 +0000
Subject: [PATCH 24/33] doc cleaning +  tp_size remove

---
 .../generation/continuous_batching/distributed.py          | 3 ++-
 src/transformers/integrations/tensor_parallel.py           | 4 +---
 src/transformers/modeling_utils.py                         | 7 ++++++-
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/transformers/generation/continuous_batching/distributed.py b/src/transformers/generation/continuous_batching/distributed.py
index fd4bd78a2338..86871e8040bd 100644
--- a/src/transformers/generation/continuous_batching/distributed.py
+++ b/src/transformers/generation/continuous_batching/distributed.py
@@ -161,7 +161,8 @@ def maybe_warn_nccl_graph_mixing(self) -> None:
         if tp_on and graph_mixing_not_disabled:
             logger.warning(
                 "NCCL_GRAPH_MIXING_SUPPORT was not set to '0' before init_process_group: performance will be harmed. "
-                "Construct your `ContinuousBatchingConfig(...)` BEFORE calling `from_pretrained(tp_plan='auto')`, or "
+                "Construct your `ContinuousBatchingConfig(...)` BEFORE calling "
+                "`from_pretrained(distributed_config=DistributedConfig(tp_size=...))`, or "
                 "set NCCL_GRAPH_MIXING_SUPPORT=0 in the launch environment."
             )
 
diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py
index b8b98edb2868..1a31b1110d30 100644
--- a/src/transformers/integrations/tensor_parallel.py
+++ b/src/transformers/integrations/tensor_parallel.py
@@ -108,7 +108,6 @@ def initialize_tensor_parallelism(
             tp_device = torch.device(device_type)
             device_map = device_type or {}
 
-        tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size()
         device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,))
     else:
         if device_mesh.ndim > 1:
@@ -118,10 +117,9 @@ def initialize_tensor_parallelism(
                     "Please provide a valid `device_mesh`."
                 )
             device_mesh = device_mesh["tp"]
-        tp_size = device_mesh.size()
         device_map = torch.device(f"{device_mesh.device_type}:{int(os.environ['LOCAL_RANK'])}")
 
-    return device_map, device_mesh, tp_size
+    return device_map, device_mesh
 
 
 def replace_layer_number_by_wildcard(name: str) -> str:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index d37e78569ee5..01da0d308a02 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -4112,6 +4112,11 @@ def from_pretrained(
             max_memory (`Dict`, *optional*):
                 A dictionary device identifier to maximum memory if using `device_map`. Will default to the maximum memory available for each
                 GPU and the available CPU RAM if unset.
+            distributed_config ([`~transformers.distributed.configuration_utils.DistributedConfig`], *optional*):
+                Configuration for native distributed loading with tensor parallelism or FSDP2. Pass
+                `DistributedConfig(tp_size=N)` for tensor parallelism, or
+                `DistributedConfig(fsdp_size=N)` for FSDP2. Requires `torchrun` and an initialized
+                process group when `tp_size > 1` or `fsdp_size > 1`. Mutually exclusive with `device_map`.
             tp_plan (`Optional[Union[dict, str]]`, *optional*):
                 A torch tensor parallel plan, see [here](https://pytorch.org/tutorials/intermediate/TP_tutorial.html). Use `tp_plan="auto"` to
                 use the predefined plan based on the model. If it's a dict, then it should match between module names and desired layout.
@@ -4263,7 +4268,7 @@ def from_pretrained(
             if distributed_config.tp_size > 1:
                 if distributed_config.tp_plan is None:
                     distributed_config.tp_plan = "auto"
-                device_map, device_mesh, tp_size = initialize_tensor_parallelism(
+                device_map, device_mesh = initialize_tensor_parallelism(
                     distributed_config.tp_plan,
                     tp_size=distributed_config.tp_size,
                     device_mesh=device_mesh,

From 06b0c394affe2d887594e695af55ebf1b211b760 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 02:44:52 +0000
Subject: [PATCH 25/33] fix tp ci for ep

---
 tests/models/deepseek_v4/test_modeling_deepseek_v4.py | 8 ++++++--
 tests/test_tensor_parallel_mixin.py                   | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py
index 3cb4205d7567..23aca5b1c4ed 100644
--- a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py
+++ b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py
@@ -468,7 +468,9 @@ def main() -> int:
         dtype="auto",
         attn_implementation="eager",
         experts_implementation=LOADTIME_DISPATCH,
-        distributed_config=DistributedConfig(enable_expert_parallel=True),
+        distributed_config=DistributedConfig(
+            tp_size=int(os.environ["WORLD_SIZE"]), enable_expert_parallel=True
+        ),
     )
     model.eval()
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -545,7 +547,9 @@ def main() -> int:
         dtype="auto",
         attn_implementation="eager",
         experts_implementation=LOADTIME_DISPATCH,
-        distributed_config=DistributedConfig(enable_expert_parallel=True),
+        distributed_config=DistributedConfig(
+            tp_size=int(os.environ["WORLD_SIZE"]), enable_expert_parallel=True
+        ),
     )
     model.eval()
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
diff --git a/tests/test_tensor_parallel_mixin.py b/tests/test_tensor_parallel_mixin.py
index 7f18a7678ba6..42b23d9ce2c9 100644
--- a/tests/test_tensor_parallel_mixin.py
+++ b/tests/test_tensor_parallel_mixin.py
@@ -356,7 +356,7 @@ def _load_ep_and_reference_models(model_path, model_class):
     """Load EP model and non-EP reference model for comparison."""
     model_ep = model_class.from_pretrained(
         model_path,
-        distributed_config=DistributedConfig(enable_expert_parallel=True),
+        distributed_config=DistributedConfig(tp_size=dist.get_world_size(), enable_expert_parallel=True),
     )
     dist.barrier()
 

From 7a94d77aab77898ff895bdb263b820a1438a5bce Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 03:00:23 +0000
Subject: [PATCH 26/33] edit doc

---
 docs/source/en/expert_parallelism.md      | 7 ++++++-
 docs/source/en/model_doc/minimax_m3_vl.md | 6 ++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/docs/source/en/expert_parallelism.md b/docs/source/en/expert_parallelism.md
index 61d298bb5cef..63256d518309 100644
--- a/docs/source/en/expert_parallelism.md
+++ b/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",
diff --git a/docs/source/en/model_doc/minimax_m3_vl.md b/docs/source/en/model_doc/minimax_m3_vl.md
index 5a73737b7e67..56a07cfab807 100644
--- a/docs/source/en/model_doc/minimax_m3_vl.md
+++ b/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()

From f3c742b730214d37c8d34938d275bb62e2656223 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 03:17:26 +0000
Subject: [PATCH 27/33] move distributed function to utils + guarding

---
 src/transformers/distributed/utils.py         | 50 +++++++++++++++----
 src/transformers/modeling_utils.py            | 28 ++---------
 .../test_trainer_distributed_fsdp.py          |  4 +-
 3 files changed, 46 insertions(+), 36 deletions(-)

diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index e6318724cf85..4f10a630895b 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -29,17 +29,41 @@
 if is_torch_available():
     import torch
 
-    # TODO(3outeille): guarding?
-    import torch.distributed.checkpoint as dcp
-    from torch.distributed.checkpoint.state_dict import (
-        StateDictOptions,
-        get_model_state_dict,
-        get_optimizer_state_dict,
-        set_optimizer_state_dict,
-    )
+    _torch_distributed_available = torch.distributed.is_available()
 
     if is_torch_greater_or_equal("2.7"):
+        import torch.distributed.checkpoint as dcp
+        from torch.distributed.checkpoint.state_dict import (
+            StateDictOptions,
+            get_model_state_dict,
+            get_optimizer_state_dict,
+            set_optimizer_state_dict,
+        )
         from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter
+else:
+    _torch_distributed_available = False
+
+
+def _is_torch_distributed_initialized() -> bool:
+    if not _torch_distributed_available:
+        return False
+    return hasattr(torch.distributed, "is_initialized") and torch.distributed.is_initialized()
+
+
+def _get_torch_distributed_rank() -> int:
+    if not _is_torch_distributed_initialized():
+        return 0
+    return torch.distributed.get_rank()
+
+
+def _get_torch_distributed_world_size() -> int:
+    if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_world_size"):
+        return 1
+    return torch.distributed.get_world_size()
+
+
+def is_local_dist_rank_0() -> bool:
+    return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
 
 
 def _ensure_torch_distributed(device_type: str):
@@ -84,7 +108,7 @@ def _distributed_barrier():
     `device_id`; with it, the call is a no-op compared to plain `barrier()`. Safe to call
     when torch.distributed has not been initialized — returns immediately.
     """
-    if not torch.distributed.is_initialized():
+    if not _is_torch_distributed_initialized():
         return
     device_type = torch._C._get_accelerator().type
     if device_type != "cpu":
@@ -157,9 +181,11 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
 
     Only rank 0 accumulates the result; other ranks return ``{}``.
     """
+    if not is_torch_greater_or_equal("2.7"):
+        raise OSError("Distributed state dict gathering requires `torch>=2.7`.")
     options = StateDictOptions(full_state_dict=True, cpu_offload=True)
     full_state_dict = get_model_state_dict(model, options=options)
-    if torch.distributed.get_rank() == 0:
+    if _get_torch_distributed_rank() == 0:
         return full_state_dict
     return {}
 
@@ -193,12 +219,16 @@ def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
 
 def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
     """Save optimizer state via DCP."""
+    if not is_torch_greater_or_equal("2.7"):
+        raise OSError("Distributed optimizer saving requires `torch>=2.7`.")
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
     dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
 
 
 def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
     """Load optimizer state via DCP."""
+    if not is_torch_greater_or_equal("2.7"):
+        raise OSError("Distributed optimizer loading requires `torch>=2.7`.")
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
     dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
     set_optimizer_state_dict(model, optimizer, optimizer_state_dict)
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 01da0d308a02..d281a41f3ed4 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -55,9 +55,13 @@
 from .distributed.fsdp import is_fsdp_managed_module
 from .distributed.utils import (
     _distributed_barrier,
+    _get_torch_distributed_rank,
+    _get_torch_distributed_world_size,
+    _is_torch_distributed_initialized,
     distribute_model,
     gather_full_state_dict,
     initialize_fully_sharded_data_parallelism,
+    is_local_dist_rank_0,
     save_model_checkpoint_distributed,
 )
 from .dynamic_module_utils import custom_object_save
@@ -203,30 +207,6 @@ def is_quantized(self) -> bool:
         return self.hf_quantizer is not None
 
 
-def _is_torch_distributed_initialized() -> bool:
-    return (
-        _torch_distributed_available
-        and hasattr(torch.distributed, "is_initialized")
-        and torch.distributed.is_initialized()
-    )
-
-
-def _get_torch_distributed_world_size() -> int:
-    if not _is_torch_distributed_initialized() or not hasattr(torch.distributed, "get_world_size"):
-        return 1
-    return torch.distributed.get_world_size()
-
-
-def is_local_dist_rank_0():
-    return _is_torch_distributed_initialized() and int(os.environ.get("LOCAL_RANK", "-1")) == 0
-
-
-def _get_torch_distributed_rank() -> int:
-    if not _is_torch_distributed_initialized():
-        return 0
-    return torch.distributed.get_rank()
-
-
 @contextmanager
 def set_quantized_state():
     global _is_quantized
diff --git a/tests/trainer/distributed/test_trainer_distributed_fsdp.py b/tests/trainer/distributed/test_trainer_distributed_fsdp.py
index 250ef854dcfc..aaaafe18d820 100644
--- a/tests/trainer/distributed/test_trainer_distributed_fsdp.py
+++ b/tests/trainer/distributed/test_trainer_distributed_fsdp.py
@@ -144,7 +144,7 @@ def test_move_missing_keys_fsdp_non_rank0_moves_meta_to_cpu(self):
 
         with (
             patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True),
-            patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False),
+            patch("transformers.distributed.utils.is_local_dist_rank_0", return_value=False),
         ):
             model._move_missing_keys_from_meta_to_device(
                 missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None
@@ -161,7 +161,7 @@ def test_fsdp_non_rank0_end_to_end_no_reinit(self):
 
         with (
             patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True),
-            patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False),
+            patch("transformers.distributed.utils.is_local_dist_rank_0", return_value=False),
         ):
             model._move_missing_keys_from_meta_to_device(
                 missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None

From 37df13ee7bd6c37773eb394d53ea83933fbc11ff Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 03:26:57 +0000
Subject: [PATCH 28/33] linting

---
 src/transformers/distributed/utils.py | 26 ++++++++++----------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index 4f10a630895b..d9b14b8217f9 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -30,19 +30,19 @@
     import torch
 
     _torch_distributed_available = torch.distributed.is_available()
-
-    if is_torch_greater_or_equal("2.7"):
-        import torch.distributed.checkpoint as dcp
-        from torch.distributed.checkpoint.state_dict import (
-            StateDictOptions,
-            get_model_state_dict,
-            get_optimizer_state_dict,
-            set_optimizer_state_dict,
-        )
-        from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter
 else:
     _torch_distributed_available = False
 
+if is_torch_available() and is_torch_greater_or_equal("2.7"):
+    import torch.distributed.checkpoint as dcp
+    from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageWriter
+    from torch.distributed.checkpoint.state_dict import (
+        StateDictOptions,
+        get_model_state_dict,
+        get_optimizer_state_dict,
+        set_optimizer_state_dict,
+    )
+
 
 def _is_torch_distributed_initialized() -> bool:
     if not _torch_distributed_available:
@@ -181,8 +181,6 @@ def gather_full_state_dict(model) -> dict[str, torch.Tensor]:
 
     Only rank 0 accumulates the result; other ranks return ``{}``.
     """
-    if not is_torch_greater_or_equal("2.7"):
-        raise OSError("Distributed state dict gathering requires `torch>=2.7`.")
     options = StateDictOptions(full_state_dict=True, cpu_offload=True)
     full_state_dict = get_model_state_dict(model, options=options)
     if _get_torch_distributed_rank() == 0:
@@ -219,16 +217,12 @@ def save_model_checkpoint_distributed(model, checkpoint_dir: str) -> None:
 
 def save_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
     """Save optimizer state via DCP."""
-    if not is_torch_greater_or_equal("2.7"):
-        raise OSError("Distributed optimizer saving requires `torch>=2.7`.")
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
     dcp.save({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
 
 
 def load_optimizer_distributed(model, optimizer, checkpoint_dir: str) -> None:
     """Load optimizer state via DCP."""
-    if not is_torch_greater_or_equal("2.7"):
-        raise OSError("Distributed optimizer loading requires `torch>=2.7`.")
     optimizer_state_dict = get_optimizer_state_dict(model, optimizer)
     dcp.load({"optimizer": optimizer_state_dict}, checkpoint_id=checkpoint_dir)
     set_optimizer_state_dict(model, optimizer, optimizer_state_dict)

From 91260ba4af88f88bc69496e1b820cc74914fdbe3 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Fri, 3 Jul 2026 05:51:11 +0000
Subject: [PATCH 29/33] add DistributedMixin

---
 src/transformers/distributed/__init__.py |   2 +
 src/transformers/distributed/mixin.py    | 213 +++++++++++++++++++++++
 src/transformers/distributed/utils.py    |   2 +-
 src/transformers/modeling_utils.py       | 172 +++---------------
 4 files changed, 237 insertions(+), 152 deletions(-)
 create mode 100644 src/transformers/distributed/mixin.py

diff --git a/src/transformers/distributed/__init__.py b/src/transformers/distributed/__init__.py
index 9bb1db2d6853..eb770a5f8399 100644
--- a/src/transformers/distributed/__init__.py
+++ b/src/transformers/distributed/__init__.py
@@ -20,6 +20,7 @@
 _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",
@@ -36,6 +37,7 @@
         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,
diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
new file mode 100644
index 000000000000..5ce0b68f7884
--- /dev/null
+++ b/src/transformers/distributed/mixin.py
@@ -0,0 +1,213 @@
+# 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.
+# 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 os
+import re
+import warnings
+from typing import TYPE_CHECKING
+
+from ..integrations.tensor_parallel import (
+    ALL_PARALLEL_STYLES,
+    gather_state_dict_for_save,
+    initialize_tensor_parallelism,
+)
+from ..utils import is_torch_available, is_torch_greater_or_equal
+from .configuration_utils import DistributedConfig
+from .fsdp import is_fsdp_managed_module
+from .utils import (
+    _get_torch_distributed_rank,
+    _is_torch_distributed_initialized,
+    distribute_model,
+    gather_full_state_dict,
+    initialize_fully_sharded_data_parallelism,
+    save_model_checkpoint_distributed,
+)
+
+
+if TYPE_CHECKING:
+    import torch.nn as nn
+
+if is_torch_available():
+    import torch
+
+    _torch_distributed_available = torch.distributed.is_available()
+else:
+    _torch_distributed_available = False
+
+
+class DistributedMixin:
+    """Distributed save/load hooks for [`PreTrainedModel`].
+
+    Stateless heavy lifting stays in `transformers.distributed.*` and
+    `integrations.tensor_parallel`. This mixin owns orchestration and instance state.
+    """
+
+    _device_mesh = None
+    _tp_plan: dict[str, str] | None = None
+    _tp_size = None
+    _pp_plan: dict[str, tuple[str, str]] | None = None
+    _fsdp_plan: dict[str, str] | None = None
+
+    @property
+    def tp_plan(self) -> dict[str, str]:
+        """The full tp plan for the model's modules."""
+        if hasattr(self.config, "distributed_config") and self.config.distributed_config.enable_expert_parallel:
+            if not self._ep_plan:
+                raise ValueError(
+                    f"Expert parallelism was requested (`enable_expert_parallel=True`), but "
+                    f"`{self.__class__.__name__}` does not define an expert-parallel plan. Add a "
+                    f"`base_model_ep_plan` to its config, or disable expert parallelism."
+                )
+            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
+
+    @tp_plan.setter
+    def tp_plan(self, plan: dict[str, str] | None):
+        if plan is None:
+            self._tp_plan = {}
+            return
+        if not isinstance(plan, dict):
+            raise ValueError("Can only set a dictionary as `tp_plan`")
+
+        for layer_pattern, parallel_style in plan.items():
+            if parallel_style not in ALL_PARALLEL_STYLES:
+                raise ValueError(
+                    f"Unsupported tensor parallel style '{parallel_style}' for layer '{layer_pattern}'. "
+                    f"Supported styles are {list(ALL_PARALLEL_STYLES.keys())}"
+                )
+
+        model_param_names = [name for name, _ in self.named_parameters()]
+        for layer_pattern in plan.keys():
+            regex_pattern = layer_pattern.replace("*", r"\d+")
+            pattern_matched = False
+            for param_name in model_param_names:
+                if re.match(regex_pattern, param_name):
+                    pattern_matched = True
+                    break
+            if not pattern_matched:
+                warnings.warn(
+                    f"Layer pattern '{layer_pattern}' does not match any parameters in the model. This rule may not "
+                    "be applied during tensor parallelization, or may lead to dimension mismatches"
+                )
+
+        self._tp_plan = plan
+
+    @pp_plan.setter
+    def pp_plan(self, plan: dict[str, tuple[str, str]] | None):
+        if plan is None:
+            self._pp_plan = {}
+            return
+        if not isinstance(plan, dict):
+            raise ValueError("Can only set a dictionary as `pp_plan`")
+
+        self._pp_plan = plan
+
+    @classmethod
+    def prepare_distribute_model(
+        cls,
+        distributed_config: DistributedConfig | dict | None,
+        *,
+        device_mesh=None,
+        device_map=None,
+    ) -> tuple[DistributedConfig | None, object, object]:
+        """Parse ``distributed_config``, init TP/FSDP mesh, and validate."""
+        if distributed_config is None:
+            return None, device_map, device_mesh
+
+        if isinstance(distributed_config, dict):
+            distributed_config = DistributedConfig.from_dict(distributed_config)
+
+        if distributed_config.tp_size > 1:
+            if distributed_config.tp_plan is None:
+                distributed_config.tp_plan = "auto"
+            device_map, device_mesh = initialize_tensor_parallelism(
+                distributed_config.tp_plan,
+                tp_size=distributed_config.tp_size,
+                device_mesh=device_mesh,
+                device_map=device_map,
+            )
+        elif distributed_config.fsdp_size > 1:
+            device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config)
+
+        distributed_config.validate()
+        return distributed_config, device_map, device_mesh
+
+    @classmethod
+    def maybe_distribute_model(
+        cls,
+        model: nn.Module,
+        distributed_config: DistributedConfig | None,
+        device_mesh,
+    ):
+        """Apply TP or FSDP2 after model init, before weight loading."""
+        if _torch_distributed_available and device_mesh is not None:
+            return distribute_model(model, distributed_config, device_mesh)
+        return model
+
+    def should_save_on_this_rank(self, is_main_process: bool) -> bool:
+        """Return whether this rank should write checkpoint files."""
+        save_on_this_rank = is_main_process
+        if _is_torch_distributed_initialized():
+            save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
+        return save_on_this_rank
+
+    def _ensure_dcp_save_supported(self, model_to_save) -> None:
+        """Raise if DCP save prerequisites are not met (torch>=2.7, FSDP-wrapped model)."""
+        if not is_torch_greater_or_equal("2.7"):
+            raise OSError("save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7.")
+        if not is_fsdp_managed_module(model_to_save):
+            raise ValueError(
+                "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models."
+            )
+
+    def save_distributed_checkpoint(self, model_to_save, save_directory: str | os.PathLike) -> None:
+        """Save an FSDP-wrapped model via DCP."""
+        if getattr(model_to_save, "_device_mesh", None) is None:
+            raise ValueError(
+                "save_pretrained(..., distributed_checkpoint=True) requires the model to have been "
+                "initialized with a distributed_config (_device_mesh is None)."
+            )
+        save_model_checkpoint_distributed(model_to_save, save_directory)
+
+    def _gather_tp_state_dict_for_save(
+        self,
+        local_state_dict: dict,
+        *,
+        is_checkpoint_writer: bool = True,
+    ) -> tuple[dict, bool]:
+        """All-gather TP-sharded weights for checkpoint writing."""
+        full_state_dict = gather_state_dict_for_save(
+            local_state_dict, self._tp_plan, self._device_mesh, self._tp_size
+        )
+        if not is_checkpoint_writer:
+            full_state_dict = {}
+        return full_state_dict, True
+
+    def _gather_fsdp_state_dict_for_save(self, model_to_save) -> tuple[dict, bool]:
+        """Gather FSDP-sharded weights to full CPU tensors on rank 0."""
+        if not _is_torch_distributed_initialized():
+            raise ValueError(
+                "Saving an FSDP-wrapped model requires torch.distributed to be initialized. "
+                "Call save_pretrained from every rank after init_process_group."
+            )
+        return gather_full_state_dict(model_to_save), True
diff --git a/src/transformers/distributed/utils.py b/src/transformers/distributed/utils.py
index d9b14b8217f9..3ed5eab1a16d 100644
--- a/src/transformers/distributed/utils.py
+++ b/src/transformers/distributed/utils.py
@@ -160,7 +160,7 @@ def distribute_model(
 ) -> nn.Module:
     """Apply TP or FSDP2 to `model` based on ``distributed_config`` (mutually exclusive for now)."""
     model.config.distributed_config = distributed_config
-    model.device_mesh = device_mesh
+    model._device_mesh = device_mesh
 
     if distributed_config.tp_size > 1:
         model = apply_tensor_parallelism(
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index d281a41f3ed4..1a54532e6216 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -53,16 +53,12 @@
 )
 from .distributed import DistributedConfig
 from .distributed.fsdp import is_fsdp_managed_module
+from .distributed.mixin import DistributedMixin
 from .distributed.utils import (
     _distributed_barrier,
-    _get_torch_distributed_rank,
     _get_torch_distributed_world_size,
     _is_torch_distributed_initialized,
-    distribute_model,
-    gather_full_state_dict,
-    initialize_fully_sharded_data_parallelism,
     is_local_dist_rank_0,
-    save_model_checkpoint_distributed,
 )
 from .dynamic_module_utils import custom_object_save
 from .generation import CompileConfig, GenerationConfig
@@ -88,10 +84,7 @@
 from .integrations.sdpa_attention import sdpa_attention_forward
 from .integrations.sdpa_paged import sdpa_attention_paged_forward
 from .integrations.tensor_parallel import (
-    ALL_PARALLEL_STYLES,
     _get_parameter_tp_plan,
-    gather_state_dict_for_save,
-    initialize_tensor_parallelism,
     shard_and_distribute_module,
     verify_tp_plan,
 )
@@ -141,7 +134,6 @@
     is_huggingface_hub_greater_or_equal,
     is_sagemaker_mp_enabled,
     is_torch_cuda_available,
-    is_torch_greater_or_equal,
     is_tracing,
 )
 from .utils.loading_report import LoadStateDictInfo, log_state_dict_report
@@ -1177,7 +1169,9 @@ def set_output_embeddings(self, new_embeddings):
             self.lm_head = new_embeddings
 
 
-class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToHubMixin, PeftAdapterMixin):
+class PreTrainedModel(
+    nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToHubMixin, PeftAdapterMixin, DistributedMixin
+):
     r"""
     Base class for all models.
 
@@ -1238,22 +1232,6 @@ class PreTrainedModel(nn.Module, EmbeddingAccessMixin, ModuleUtilsMixin, PushToH
     # Model's compatible flash kernels (e.g., "kernels-community/flash-mla") defaulting to the first in the list
     _compatible_flash_implementations: list[str] | None = None
 
-    # Tensor-parallelism-related properties
-    # A tensor parallel plan of the form `{"model.layer.mlp.param": "colwise"}` to be applied to the model when TP is enabled.
-    # For top-level models, this attribute is currently defined in respective model code. For base models, this attribute comes
-    # from `config.base_model_tp_plan` during `post_init`.
-    _tp_plan: dict[str, str] = None
-    # Tensor parallel degree to which model is sharded to
-    _tp_size = None
-    # A pipeline parallel plan specifying the layers which may not be present on all ranks when PP is enabled. For top-level
-    # 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
     _can_compile_fullgraph: bool = False
@@ -1455,75 +1433,6 @@ def post_init(self):
         self.init_weights()
         self._backward_compatibility_gradient_checkpointing()
 
-    @property
-    def tp_plan(self) -> dict[str, str]:
-        """
-        The full tp plan for the model's modules
-        """
-        if hasattr(self.config, "distributed_config") and self.config.distributed_config.enable_expert_parallel:
-            if not self._ep_plan:
-                raise ValueError(
-                    f"Expert parallelism was requested (`enable_expert_parallel=True`), but "
-                    f"`{self.__class__.__name__}` does not define an expert-parallel plan. Add a "
-                    f"`base_model_ep_plan` to its config, or disable expert parallelism."
-                )
-            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
-
-    @tp_plan.setter
-    def tp_plan(self, plan: dict[str, str] | None):
-        if plan is None:
-            self._tp_plan = {}
-            return
-        if not isinstance(plan, dict):
-            raise ValueError("Can only set a dictionary as `tp_plan`")
-
-        # Ensure the styles are all valid
-        for layer_pattern, parallel_style in plan.items():
-            if parallel_style not in ALL_PARALLEL_STYLES:
-                raise ValueError(
-                    f"Unsupported tensor parallel style '{parallel_style}' for layer '{layer_pattern}'. "
-                    f"Supported styles are {list(ALL_PARALLEL_STYLES.keys())}"
-                )
-
-        # Validate that the layer patterns match existing model structure. We check this by getting all parameter
-        # names and seeing if any match the patterns
-        model_param_names = [name for name, _ in self.named_parameters()]
-        for layer_pattern in plan.keys():
-            # Convert pattern to regex (replace * with .*)
-            regex_pattern = layer_pattern.replace("*", r"\d+")
-            pattern_matched = False
-            for param_name in model_param_names:
-                if re.match(regex_pattern, param_name):
-                    pattern_matched = True
-                    break
-            if not pattern_matched:
-                warnings.warn(
-                    f"Layer pattern '{layer_pattern}' does not match any parameters in the model. This rule may not "
-                    "be applied during tensor parallelization, or may lead to dimension mismatches"
-                )
-
-        # Set the plan
-        self._tp_plan = plan
-
-    @pp_plan.setter
-    def pp_plan(self, plan: dict[str, tuple[str, str]] | None):
-        if plan is None:
-            self._pp_plan = {}
-            return
-        if not isinstance(plan, dict):
-            raise ValueError("Can only set a dictionary as `pp_plan`")
-
-        self._pp_plan = plan
-
     def dequantize(self, dtype=None):
         """
         Potentially dequantize the model in case it has been quantized by a quantization method that support
@@ -3455,20 +3364,10 @@ def save_pretrained(
 
         # Only save the model itself if we are using distributed training
         model_to_save = unwrap_model(self)
-        is_fsdp_model = is_fsdp_managed_module(model_to_save)
-
         if distributed_checkpoint:
-            if not is_torch_greater_or_equal("2.7"):
-                raise OSError("save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7.")
-            if not is_fsdp_model:
-                raise ValueError(
-                    "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models."
-                )
+            self._ensure_dcp_save_supported(model_to_save)
 
-        # Collectives run on every rank; checkpoint files are written on global rank 0 only.
-        save_on_this_rank = is_main_process
-        if _is_torch_distributed_initialized():
-            save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
+        save_on_this_rank = self.should_save_on_this_rank(is_main_process)
 
         # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
         # we currently don't use this setting automatically, but may start to use with v5
@@ -3520,12 +3419,7 @@ def save_pretrained(
 
         # --- FSDP path: DCP save (large models, torch >= 2.7) ---
         if distributed_checkpoint:
-            if getattr(model_to_save, "device_mesh", None) is None:
-                raise ValueError(
-                    "save_pretrained(..., distributed_checkpoint=True) requires the model to have been "
-                    "initialized with a distributed_config (device_mesh is None)."
-                )
-            save_model_checkpoint_distributed(model_to_save, save_directory)
+            self.save_distributed_checkpoint(model_to_save, save_directory)
             if push_to_hub and save_on_this_rank:
                 model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
                 model_card.save(os.path.join(save_directory, "README.md"))
@@ -3539,19 +3433,16 @@ def save_pretrained(
                 )
             return
 
-        # Get the state dict
-        used_distributed_gather = False
+        needs_dist_barrier = False
         if state_dict is None:
-            if is_fsdp_model:
-                if not _is_torch_distributed_initialized():
-                    raise ValueError(
-                        "Saving an FSDP-wrapped model requires torch.distributed to be initialized. "
-                        "Call save_pretrained from every rank after init_process_group."
-                    )
-                state_dict = gather_full_state_dict(model_to_save)
-                used_distributed_gather = True
-            else:
-                state_dict = model_to_save.state_dict()
+            state_dict = model_to_save.state_dict()
+
+        if self._tp_size is not None:
+            state_dict, needs_dist_barrier = self._gather_tp_state_dict_for_save(
+                state_dict, is_checkpoint_writer=save_on_this_rank
+            )
+        elif is_fsdp_managed_module(model_to_save):
+            state_dict, needs_dist_barrier = self._gather_fsdp_state_dict_for_save(model_to_save)
 
         # if any model parameters are offloaded, we need to know it for later
         is_offloaded = False
@@ -3577,13 +3468,6 @@ def save_pretrained(
                 if ignore_key in state_dict:
                     del state_dict[ignore_key]
 
-        # If model was sharded with TP, gather full tensors for saving
-        if self._tp_size is not None:
-            state_dict = gather_state_dict_for_save(state_dict, self._tp_plan, self._device_mesh, self._tp_size)
-            used_distributed_gather = True
-            if not save_on_this_rank:
-                state_dict = {}
-
         # Remove tied weights as safetensors do not handle them
         state_dict = remove_tied_weights_from_state_dict(state_dict, model_to_save)
 
@@ -3693,7 +3577,7 @@ def save_pretrained(
         # other ranks then loop over an empty shard list and would race ahead
         # of rank 0's safetensors writes. Barrier so any subsequent
         # `from_pretrained` on this path sees the consolidated files.
-        if used_distributed_gather:
+        if needs_dist_barrier:
             _distributed_barrier()
 
     @wraps(PushToHubMixin.push_to_hub)
@@ -4242,22 +4126,9 @@ def from_pretrained(
             )
 
         if distributed_config is not None:
-            if isinstance(distributed_config, dict):
-                distributed_config = DistributedConfig.from_dict(distributed_config)
-
-            if distributed_config.tp_size > 1:
-                if distributed_config.tp_plan is None:
-                    distributed_config.tp_plan = "auto"
-                device_map, device_mesh = initialize_tensor_parallelism(
-                    distributed_config.tp_plan,
-                    tp_size=distributed_config.tp_size,
-                    device_mesh=device_mesh,
-                    device_map=device_map,
-                )
-            elif distributed_config.fsdp_size > 1:
-                device_map, device_mesh = initialize_fully_sharded_data_parallelism(distributed_config)
-
-            distributed_config.validate()
+            distributed_config, device_map, device_mesh = cls.prepare_distribute_model(
+                distributed_config, device_mesh=device_mesh, device_map=device_map
+            )
 
         if gguf_file is not None and not is_accelerate_available():
             raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.")
@@ -4413,8 +4284,7 @@ def from_pretrained(
         # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively
         weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer)
 
-        if _torch_distributed_available and device_mesh is not None:  # add hooks to nn.Modules: no weights
-            model = distribute_model(model, distributed_config, device_mesh)
+        model = cls.maybe_distribute_model(model, distributed_config, device_mesh)
 
         # Prepare the full device map
         if device_map is not None:

From 22215b7c78367697972f80f082cd61a54c8c903a Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 8 Jul 2026 04:01:47 +0000
Subject: [PATCH 30/33] linting

---
 src/transformers/distributed/mixin.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
index 5ce0b68f7884..36b60f2e9e1f 100644
--- a/src/transformers/distributed/mixin.py
+++ b/src/transformers/distributed/mixin.py
@@ -57,7 +57,7 @@ class DistributedMixin:
     _device_mesh = None
     _tp_plan: dict[str, str] | None = None
     _tp_size = None
-    _pp_plan: dict[str, tuple[str, str]] | None = None
+    _pp_plan: dict[str, tuple[str, str]] = None
     _fsdp_plan: dict[str, str] | None = None
 
     @property
@@ -196,9 +196,7 @@ def _gather_tp_state_dict_for_save(
         is_checkpoint_writer: bool = True,
     ) -> tuple[dict, bool]:
         """All-gather TP-sharded weights for checkpoint writing."""
-        full_state_dict = gather_state_dict_for_save(
-            local_state_dict, self._tp_plan, self._device_mesh, self._tp_size
-        )
+        full_state_dict = gather_state_dict_for_save(local_state_dict, self._tp_plan, self._device_mesh, self._tp_size)
         if not is_checkpoint_writer:
             full_state_dict = {}
         return full_state_dict, True

From 18d709c65abf14db8fb5ed9c751f0d8b42668809 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 8 Jul 2026 04:26:47 +0000
Subject: [PATCH 31/33] refactor

---
 src/transformers/distributed/mixin.py | 32 ++++++++++++++++++++++-----
 src/transformers/modeling_utils.py    | 32 +++++++++++++--------------
 2 files changed, 43 insertions(+), 21 deletions(-)

diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
index 36b60f2e9e1f..88dc33c13de1 100644
--- a/src/transformers/distributed/mixin.py
+++ b/src/transformers/distributed/mixin.py
@@ -24,6 +24,7 @@
     initialize_tensor_parallelism,
 )
 from ..utils import is_torch_available, is_torch_greater_or_equal
+from ..utils.hub import create_and_tag_model_card
 from .configuration_utils import DistributedConfig
 from .fsdp import is_fsdp_managed_module
 from .utils import (
@@ -171,17 +172,26 @@ def should_save_on_this_rank(self, is_main_process: bool) -> bool:
             save_on_this_rank = save_on_this_rank and _get_torch_distributed_rank() == 0
         return save_on_this_rank
 
-    def _ensure_dcp_save_supported(self, model_to_save) -> None:
-        """Raise if DCP save prerequisites are not met (torch>=2.7, FSDP-wrapped model)."""
+    def save_distributed_checkpoint(
+        self,
+        model_to_save,
+        save_directory: str | os.PathLike,
+        *,
+        push_to_hub: bool = False,
+        save_on_this_rank: bool = True,
+        repo_id: str | None = None,
+        files_timestamps: dict | None = None,
+        commit_message: str | None = None,
+        token: str | bool | None = None,
+        create_pr: bool = False,
+    ) -> None:
+        """Save an FSDP-wrapped model via DCP and optionally push to the Hub."""
         if not is_torch_greater_or_equal("2.7"):
             raise OSError("save_pretrained(..., distributed_checkpoint=True) requires torch>=2.7.")
         if not is_fsdp_managed_module(model_to_save):
             raise ValueError(
                 "save_pretrained(..., distributed_checkpoint=True) is only supported for FSDP-wrapped models."
             )
-
-    def save_distributed_checkpoint(self, model_to_save, save_directory: str | os.PathLike) -> None:
-        """Save an FSDP-wrapped model via DCP."""
         if getattr(model_to_save, "_device_mesh", None) is None:
             raise ValueError(
                 "save_pretrained(..., distributed_checkpoint=True) requires the model to have been "
@@ -189,6 +199,18 @@ def save_distributed_checkpoint(self, model_to_save, save_directory: str | os.Pa
             )
         save_model_checkpoint_distributed(model_to_save, save_directory)
 
+        if push_to_hub and save_on_this_rank:
+            model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
+            model_card.save(os.path.join(save_directory, "README.md"))
+            self._upload_modified_files(
+                save_directory,
+                repo_id,
+                files_timestamps,
+                commit_message=commit_message,
+                token=token,
+                create_pr=create_pr,
+            )
+
     def _gather_tp_state_dict_for_save(
         self,
         local_state_dict: dict,
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 265c42419a94..af13fc81ca1a 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3377,9 +3377,6 @@ def save_pretrained(
 
         # Only save the model itself if we are using distributed training
         model_to_save = unwrap_model(self)
-        if distributed_checkpoint:
-            self._ensure_dcp_save_supported(model_to_save)
-
         save_on_this_rank = self.should_save_on_this_rank(is_main_process)
 
         # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
@@ -3430,20 +3427,23 @@ def save_pretrained(
                 current_peft_config = self.peft_config[active_adapter]
                 current_peft_config.save_pretrained(save_directory)
 
-        # --- FSDP path: DCP save (large models, torch >= 2.7) ---
         if distributed_checkpoint:
-            self.save_distributed_checkpoint(model_to_save, save_directory)
-            if push_to_hub and save_on_this_rank:
-                model_card = create_and_tag_model_card(repo_id, self.model_tags, token=token)
-                model_card.save(os.path.join(save_directory, "README.md"))
-                self._upload_modified_files(
-                    save_directory,
-                    repo_id,
-                    files_timestamps,
-                    commit_message=commit_message,
-                    token=token,
-                    create_pr=create_pr,
-                )
+            hub_kwargs = {}
+            if push_to_hub:
+                hub_kwargs = {
+                    "repo_id": repo_id,
+                    "files_timestamps": files_timestamps,
+                    "commit_message": commit_message,
+                    "create_pr": create_pr,
+                }
+            self.save_distributed_checkpoint(
+                model_to_save,
+                save_directory,
+                push_to_hub=push_to_hub,
+                save_on_this_rank=save_on_this_rank,
+                token=token,
+                **hub_kwargs,
+            )
             return
 
         needs_dist_barrier = False

From 566ef6e48f510cd3121ab7e099725a3fb813ce46 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 8 Jul 2026 05:19:12 +0000
Subject: [PATCH 32/33] abstract to mixin

---
 src/transformers/distributed/mixin.py | 42 ++++++++++++++-
 src/transformers/modeling_utils.py    | 76 +++++++++++----------------
 2 files changed, 71 insertions(+), 47 deletions(-)

diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
index 88dc33c13de1..33b7f1a36fa8 100644
--- a/src/transformers/distributed/mixin.py
+++ b/src/transformers/distributed/mixin.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import json
 import os
 import re
 import warnings
@@ -23,7 +24,8 @@
     gather_state_dict_for_save,
     initialize_tensor_parallelism,
 )
-from ..utils import is_torch_available, is_torch_greater_or_equal
+from ..modeling_utils import _add_variant
+from ..utils import SAFE_WEIGHTS_INDEX_NAME, is_torch_available, is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
 from .configuration_utils import DistributedConfig
 from .fsdp import is_fsdp_managed_module
@@ -37,6 +39,9 @@
 )
 
 
+logger = logging.get_logger(__name__)
+
+
 if TYPE_CHECKING:
     import torch.nn as nn
 
@@ -211,6 +216,41 @@ def save_distributed_checkpoint(
                 create_pr=create_pr,
             )
 
+    def save_gathered_checkpoint_index(
+        self,
+        save_directory: str | os.PathLike,
+        *,
+        state_dict_split,
+        weight_map: dict | None,
+        weights_name: str,
+        variant: str | None,
+        max_shard_size: int | str,
+    ) -> None:
+        """Write sharded safetensors index after a gathered state-dict save."""
+        # Save index if sharded
+        index = None
+        if state_dict_split.is_sharded:
+            index = {
+                "metadata": {"total_parameters": self.num_parameters(), **state_dict_split.metadata},
+                "weight_map": weight_map,
+            }
+
+        if index is None:
+            path_to_weights = os.path.join(save_directory, weights_name)
+            logger.info(f"Model weights saved in {path_to_weights}")
+        else:
+            save_index_file = SAFE_WEIGHTS_INDEX_NAME
+            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
+            # Save the index as well
+            with open(save_index_file, "w", encoding="utf-8") as f:
+                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+                f.write(content)
+            logger.info(
+                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
+                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
+                f"index located at {save_index_file}."
+            )
+
     def _gather_tp_state_dict_for_save(
         self,
         local_state_dict: dict,
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index af13fc81ca1a..05531bca6f50 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -16,7 +16,6 @@
 import copy
 import functools
 import inspect
-import json
 import os
 import re
 import sys
@@ -3554,53 +3553,38 @@ def save_pretrained(
                     # only do contiguous after it's permuted correctly in case of TP
                     shard_state_dict[tensor_name] = tensor.contiguous()
 
-            # As explained above, for offloaded scenarios, weight format could not be reverted before due to meta weights,
-            # so do it now after they were loaded onto cpu. For one-weight-to-many operations, it may be an issue, but usually the shards
-            # contain all the necessary params, except if we are quite unlucky on the sharding. The failure surface is (very few models
-            # with one-weight-to-many + offloading to disk + unlucky sharding), so it will almost never happen
-            if is_offloaded and save_original_format and not _hf_peft_config_loaded:
-                try:
-                    shard_state_dict = revert_weight_conversion(model_to_save, shard_state_dict)
-                    # Save the weight_map, since some names etc may have changed due to conversion compared to initial `state_dict_split`
-                    if state_dict_split.is_sharded:
-                        weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys())  # ty: ignore[unresolved-attribute]
-                except Exception:
-                    raise RuntimeError(
-                        "We could not revert some weight conversions because of offlading, and several weights needed for a single "
-                        "conversion operation living in different shard files. Try reducing `max_shard_size` a bit, or worst case "
-                        "set `save_original_format=False`."
-                    )
+                # As explained above, for offloaded scenarios, weight format could not be reverted before due to meta weights,
+                # so do it now after they were loaded onto cpu. For one-weight-to-many operations, it may be an issue, but usually the shards
+                # contain all the necessary params, except if we are quite unlucky on the sharding. The failure surface is (very few models
+                # with one-weight-to-many + offloading to disk + unlucky sharding), so it will almost never happen
+                if is_offloaded and save_original_format and not _hf_peft_config_loaded:
+                    try:
+                        shard_state_dict = revert_weight_conversion(model_to_save, shard_state_dict)
+                        # Save the weight_map, since some names etc may have changed due to conversion compared to initial `state_dict_split`
+                        if state_dict_split.is_sharded:
+                            weight_map.update({k: os.path.basename(shard_file)} for k in shard_state_dict.keys())  # ty: ignore[unresolved-attribute]
+                    except Exception:
+                        raise RuntimeError(
+                            "We could not revert some weight conversions because of offlading, and several weights needed for a single "
+                            "conversion operation living in different shard files. Try reducing `max_shard_size` a bit, or worst case "
+                            "set `save_original_format=False`."
+                        )
 
-            # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
-            # so it's not possible for now....
-            # Write the shard to disk
-            safe_save_file(shard_state_dict, filename, metadata=metadata)
-            # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
-            del shard_state_dict
-
-            # Save index if sharded
-            index = None
-            if state_dict_split.is_sharded:
-                index = {
-                    "metadata": {"total_parameters": self.num_parameters(), **state_dict_split.metadata},
-                    "weight_map": weight_map,
-                }
+                # TODO: it would be very nice to do the writing concurrently, but safetensors never releases the GIL,
+                # so it's not possible for now....
+                # Write the shard to disk
+                safe_save_file(shard_state_dict, filename, metadata=metadata)
+                # Cleanup the data before next loop (important with offloading, so we don't blowup cpu RAM)
+                del shard_state_dict
 
-            if index is None:
-                path_to_weights = os.path.join(save_directory, weights_name)
-                logger.info(f"Model weights saved in {path_to_weights}")
-            else:
-                save_index_file = SAFE_WEIGHTS_INDEX_NAME
-                save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
-                # Save the index as well
-                with open(save_index_file, "w", encoding="utf-8") as f:
-                    content = json.dumps(index, indent=2, sort_keys=True) + "\n"
-                    f.write(content)
-                logger.info(
-                    f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
-                    f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
-                    f"index located at {save_index_file}."
-                )
+            self.save_gathered_checkpoint_index(
+                save_directory,
+                state_dict_split=state_dict_split,
+                weight_map=weight_map,
+                weights_name=weights_name,
+                variant=variant,
+                max_shard_size=max_shard_size,
+            )
 
         if push_to_hub and save_on_this_rank:
             # Eventually create an empty model card

From b632db38cc765c6cb57a62a02e906753474d7f14 Mon Sep 17 00:00:00 2001
From: 3outeille <ferdinand.mom@epita.fr>
Date: Wed, 8 Jul 2026 05:21:23 +0000
Subject: [PATCH 33/33] typo

---
 src/transformers/distributed/mixin.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py
index 33b7f1a36fa8..c41f70466d28 100644
--- a/src/transformers/distributed/mixin.py
+++ b/src/transformers/distributed/mixin.py
@@ -24,7 +24,6 @@
     gather_state_dict_for_save,
     initialize_tensor_parallelism,
 )
-from ..modeling_utils import _add_variant
 from ..utils import SAFE_WEIGHTS_INDEX_NAME, is_torch_available, is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
 from .configuration_utils import DistributedConfig
@@ -239,6 +238,8 @@ def save_gathered_checkpoint_index(
             path_to_weights = os.path.join(save_directory, weights_name)
             logger.info(f"Model weights saved in {path_to_weights}")
         else:
+            from ..modeling_utils import _add_variant
+
             save_index_file = SAFE_WEIGHTS_INDEX_NAME
             save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
             # Save the index as well