[Umbrella] Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh

[qgallouedec]

Important

Draft to unblock experimentation (fine-tuning large MoEs in TRL without replicating dense weights on every EP rank). The Trainer-side fallbacks are deliberately blunt; we'll migrate to whatever native design transformers lands.

• Today: tp_size=N, fsdp_size=M with both > 1 raises FSDP+TP is not supported yet.
• But EP shards only the experts — everything else (and its optimizer state) replicates on every EP rank. Big MoEs need both.
• This draft: 2-D (fsdp, tp) mesh, EP plan on tp, fully_shard on fsdp, trains end-to-end with Trainer.

Try it

pip install git+https://github.com/huggingface/transformers@ep-fsdp-2d-mesh
pip install git+https://github.com/huggingface/accelerate@fsdp2-ep-integration
pip install git+https://github.com/huggingface/peft@fsdp2-dtensor-fixes

Validated on 64 H100 with TRL's SFTTrainer:

• GLM-5.2 753B, LoRA: loss 3.3 → 2.3
• GLM-4.5-Air 110B, full fine-tuning: 3.35 → 2.13, full 206 GiB save
• Runnable examples: Add expert-parallelism example: SFT of 100B-753B MoE models trl#6869

How

• Build the 2-D mesh, record _tp_size/_fsdp_size so Trainer hands accelerate a ParallelismConfig (otherwise accelerate wraps the DTensor model in DDP, which raises).
• Trainer fallbacks: all-reduce expert gradients over fsdp (FSDP2 only reduces what it shards), a grad-norm/clip that spans parameters on different meshes, per-parameter AdamW when fused/foreach can't span the set.

Correctness

Validation surfaced two pre-existing EP bugs, split out as #48205:

• uninitialized torch._grouped_mm sentinel rows → NaN on step 2
• router hook drops non-local score gradients → wrong gradients for every non-expert parameter

fp32 gradient certification vs a single-GPU reference, after the fixes:

• EP: 179/179 params within float epsilon (max 2.7e-5; before: 3/179)
• 2-D mesh: 179/179 (5.6e-5) · real Qwen3-30B weights: 47/47 (2.7e-6)

Loss parity, throughput and memory tables

OLMoE-1B-7B FFT: EP-only tracks the single-GPU control step by step; the 2-D run has 2× the batch and agrees in trend.

Qwen3-30B-A3B full fine-tuning, 8×H100, seq 2048 (both converge, grad norms finite):

config	tok/s/GPU	peak mem/GPU
EP only, ep=8	3,089	40.9 GB
FSDP2 × EP, ep=4 × dp=2	2,379	34.2 GB

OLMoE-1B-7B full FT, 8×H100: ep=8 → 10,052 tok/s at 10.4 GB; ep=4×dp=2 → 7,143 at 8.9 GB; ep=2×dp=4 → 7,023 at 8.7 GB. LoRA at small scale fits everywhere and EP-only is faster; the 2-D mesh is for full FT and beyond.

Minimal repro (min_fsdp_ep.py)

# torchrun --nproc_per_node 4 min_fsdp_ep.py
import torch
from torch.utils.data import Dataset

from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from transformers.distributed import DistributedConfig


class RandomTokens(Dataset):
    def __len__(self):
        return 512

    def __getitem__(self, i):
        g = torch.Generator().manual_seed(i)
        ids = torch.randint(0, 50000, (512,), generator=g)
        return {"input_ids": ids, "labels": ids.clone()}


model = AutoModelForCausalLM.from_pretrained(
    "allenai/OLMoE-1B-7B-0924",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2, enable_expert_parallel=True),
)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="/tmp/fsdp-ep",
        per_device_train_batch_size=1,
        max_steps=8,
        logging_steps=1,
        report_to=[],
        save_strategy="no",
    ),
    train_dataset=RandomTokens(),
)
trainer.train()

Known limitations

1. tp_size doubles as the EP size (a dedicated ep name would be clearer).
2. LoRA needs two PEFT fixes: Fix LoRA on FSDP2-sharded models: shape inference and module hooks peft#3578.
3. Pure FSDP2 from from_pretrained (tp_size=1) still fails in accelerate (pre-existing, tracked in Should Trainer support models FSDP2-sharded at load time (DistributedConfig(fsdp_size=N))? #48210).
4. Checkpoint saving works (every rank gathers, main writes, HF format). Resume does not: the Trainer's resume path raw-copies plain tensors into DTensor params and doesn't convert per-expert → fused; it needs to reuse the from_pretrained sharded-load path.
5. The proper design likely puts expert DTensors on the full 2-D mesh so FSDP2 handles reduction itself (torchtitan-style).

What this branch carries

What	Standalone
NaN + missing gradients	#48205
EP reachable through Trainer	#48208
load-scaled pg timeout (rank skew vs 10-min NCCL watchdog)	#48228
don't device-move sharded-at-load models in Trainer	this branch
HF_SHARD_PREFETCH: 4–6× cold multi-node loads	#48227
token-dispatch + trunk-data-parallel EP (env-gated, fp32-certified, 31–59× samples/s at 64 GPUs)	follow-up PR
HF_DEBUG_INIT_SWEEP logging	this branch
fix empty FSDP shards in sharded loading (crash + silent group hang; froze 357B runs at ep=8 × fsdp=8)	#48237

[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.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 32661785504:1
Result: failure | Jobs: 2 | Tests: 23 | Failures: 1 | Duration: 1m 55s

Code quality check failed: test jobs were skipped. Fix the code quality issues and push again to run tests.