FSDP2 native support in transformers

[3outeille]

• TODO:
• fsdp => faire comme tp en mode fsdp_plan manual qui devient l'auto par défaut

---

This PR introduces first-class FSDP2 (Fully Sharded Data Parallel v2) support directly in Transformers, bypassing the need for Accelerate's FSDP wrapper. It covers the full lifecycle: model distribution, training, checkpointing, and CI testing across dozens of models.

A standalone script for usage and expected throughput will be available at https://github.com/huggingface/distributed-training-cookbook

TODO: add example train_fsdp.py in the PR

The correctness has been tested to match the Torchtitan implementation (cf https://github.com/huggingface/torchtitan/blob/sanity-check-fsdp/torchtitan/experiments/debug_fsdp/README.md)

NOTE: Transformers modeling do take more memory for some reasons. To investigate later

1. Native FSDP2 Integration (src/transformers/integrations/fsdp.py)

The core addition is a new FSDP integration module that provides:

• initialize_fsdp() -- Sets up the DeviceMesh and process group for FSDP2 (requires PyTorch >= 2.5). Handles automatic backend detection (NCCL, GLOO, XCCL, etc.) and device assignment.

• apply_fsdp2() with two modes:

  • Auto mode ({"mode": "auto"}) -- Automatically discovers transformer block classes (DecoderLayer, EncoderLayer, etc.), shards input embeddings, all transformer blocks, and groups the final norm + output head together. Supports optional cpu_offload and mixed_precision policies.
  • Manual mode ({"mode": "manual", "modules": {...}}) -- Lets users specify exactly which modules to shard, with per-module options like "free_full_weight", "keep_full_weight", "cpu_offload", and "mixed_precision".

• Smart block detection (get_transformer_block_classes()) -- Finds transformer block classes by name pattern and filters out nested blocks (e.g., MoeBlock inside DecoderLayer) to only FSDP-wrap the outermost ones. This enables MoE model support.

• Tied weight handling -- Properly handles weight tying (e.g., lm_head.weight == embed_tokens.weight) by grouping tied modules and re-tying after fully_shard replaces parameters with DTensors.

2. FSDP2-Aware Save/Load via DCP + Safetensors

• save_fsdp_model() -- Saves FSDP2 model weights using PyTorch's Distributed Checkpoint (DCP) with HuggingFaceStorageWriter, enabling parallel distributed save with automatic consolidation into standard HF-compatible safetensors files.

• save_pretrained() integration -- PreTrainedModel.save_pretrained() now detects FSDP2 models (_is_fsdp_managed_module) and automatically routes to the DCP save path.

• from_pretrained() integration -- Accepts new fsdp_plan and fsdp_device_mesh kwargs. After loading weights, it applies FSDP2 distribution via distribute_fsdp_model().

3. Comprehensive FSDP Test Suite (tests/test_fsdp_mixin.py)

A new FSDPTesterMixin class is added to the standard test infrastructure, automatically inherited by all CausalLMModelTest classes. It includes 7 batched subtests per model, all run on CPU through gloo backend + mp.spawn:

Test	What it validates
sharding_structure_untied/tied	Correct FSDP wrapping targets match expectations
auto_plan_vs_ddp (untied/tied)	FSDP2 auto mode produces identical losses, grad norms, and final weights as DDP
manual_plan_vs_ddp (untied/tied)	FSDP2 manual mode matches DDP
save_load	Save via save_pretrained + reload via from_pretrained produces bit-exact weights

The tests also validate checkpoint resumability: train for N/2 steps, save checkpoint (model via DCP+safetensors, optimizer+RNG via distcp), load into a fresh model, continue training, and verify the full trace matches an uninterrupted DDP run.

4. CI Test for Broad Model Coverage

Two bash scripts run the FSDP mixin tests across many models in parallel:

• Dense models: Tests ~10 active models (GPT-2, Qwen3, Phi, Llama, ModernBERT-decoder, OLMo3, Phi3, Mistral, LFM2, Qwen3.5) out of 40 total, ranked by HuggingFace Hub downloads.

• MoE models: Tests ~10 active MoE models (GPT-OSS, GLM-MoE-DSA, Qwen3-MoE, GLM4-MoE-Lite, Qwen3.5-MoE, DeepSeek-V2, Qwen3-Next, Mixtral, Qwen2-MoE, PhiMoE) out of 24 total.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks, left a couple of comments !

[SunMarc]

let's create a new folder called tests/training and put those there instead. It will be better I think

[SunMarc]

even the fsdp folder we can move that there

[SunMarc]

maybe it could make sense to have a nice dataclass for fsdp_plan instead of a dict ?

[ArthurZucker]

Yep, I think the idea was to stay simple like the pp / tp plan. But for fsdp we might want more control.

[SunMarc]

It would be nice to check how it integrates well with trainer. We can pass fsdp and fsdp_config in training_args and we would have to do the mapping to create the correct fsdp_plan + we need to call apply_fsdp2 instead of prepare() on the model.

[cursor]

Unconditional torch imports break non-torch environments

High Severity

Top-level import torch, import torch.distributed, import torch.multiprocessing, and from torch.distributed.device_mesh import init_device_mesh are added unconditionally at module level. This will cause an ImportError for anyone importing from testing_utils in an environment without PyTorch (or with an older PyTorch lacking device_mesh). The existing file already uses is_torch_available() guards for other torch imports — these new ones need the same treatment.

 

[cursor]

Duplicate FSDPTesterMixin import in causal_lm_tester

Low Severity

from .test_fsdp_mixin import FSDPTesterMixin appears twice — once at line 32 and again at line 43. One of these is redundant and was likely left in by mistake during development.

Additional Locations (1)

• tests/causal_lm_tester.py#L31-L32

 

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[ArthurZucker]

Will look at the tests next time but my main comment:

You are not integrated with the core-model-loading API noo?

My understanding is that you should shard the weights exactly the same way we do for TP, maybe push further as you want to shard ALL layers, then when running the forward, each process materialize the weights locally (GP0-8 end up with the same full tensor) and discard it keeping only its slice.

We need 2 points:

1. TENSOR_PARALLEL_LAYERS integration in core_model_loading needs to support fsdp. This is what's gonna be responsible for loading the weights
2. distribute_module, which needs tto happen before the load, and is responsible for attaching tthe appropriate hooks for fsdp2.
   THis also needs to be explained like the above comment: what is fsdp? -> 1 sharding plan 2. hook plan.

Now if you have to rely on DTensor, you might need to change set_param? or you just apply dtensor conversion post all the loading.

The most important is to test say MIxtral with the dynamic weight loader. The way I see it you'll load all weights on all layers, then shard (discard) some of them.

[ArthurZucker]

can't we re-use the func we defined in tensor_parallel ?

[ArthurZucker]

pretty sure you can just check if the layer is GradientCheckpointingLayer 😉

transformers/src/transformers/modeling_layers.py

Line 34 in 3491f11

class GradientCheckpointingLayer(nn.Module):

we make all blocks inherit from this!

[ArthurZucker]

    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

this looks like something we can define in metadata / take from PP no?

[ArthurZucker]

can be taken from:

    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

transformers/src/transformers/models/llama/modeling_llama.py

Line 432 in 3491f11

_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

No?

[ArthurZucker]

can you document what this does?
not entirely sure we need it, but if we do, probably something TP plan or PP plan or Dtype plan are gonna be using ?

[ArthurZucker]

Yep, I think the idea was to stay simple like the pp / tp plan. But for fsdp we might want more control.

[kashif]

Based on my CP experiments:

1. Accept fsdp_plan="auto" string shorthand: Currently fsdp_plan={"mode": "auto"} works but fsdp_plan="auto" raises a confusing error. Since tp_plan="auto" is accepted as a string, fsdp_plan should be consistent:

  # _parse_fsdp_plan_mode could normalize strings:
  if isinstance(fsdp_plan, str):
      fsdp_plan = {"mode": fsdp_plan}

2. Document CP + FSDP2 device mesh pattern: When using Context Parallelism with native FSDP2, the fsdp_device_mesh must be the flattened dp_cp mesh, not just the dp submesh otherwise FSDP2 doesn't shard parameters across CP ranks and you OOM. This is non-obvious:

  # Users need to do this:
  world_mesh = init_device_mesh("cuda", (dp_size, cp_size), mesh_dim_names=("dp", "cp"))
  fsdp_mesh = world_mesh["dp", "cp"]._flatten(mesh_dim_name="dp_cp")
  model = AutoModelForCausalLM.from_pretrained(..., fsdp_device_mesh=fsdp_mesh, fsdp_plan={"mode": "auto"})

The PR's docstring or the 3D_parallel.py example should show this pattern.

3. Add a CP example using native FSDP2: The current 3D_parallel.py uses FSDP1 (FullyShardedDataParallel with NO_SHARD). An example showing native FSDP2 (fsdp_plan) + CP would be good to have.

4. Note about CP rotation method: The default CP rotation (allgather) requires allgather_into_tensor_coalesced which may not be available in all PyTorch builds. The alltoall rotation works universally. Worth noting in the docs or defaulting to alltoall.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: clap, deit

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=44083&sha=37dcc1