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

[qgallouedec]

Stacked on #48205 (which carries #48208); the diff to review is this branch against fix-ep-training-gradients.

Expert parallelism shards only the experts. Everything else and its optimizer state is replicated on every EP rank, so the model size you can train is bounded by what one rank can hold of the dense part plus its Adam state.
DistributedConfig(tp_size=E, fsdp_size=D, enable_expert_parallel=True) now builds a 2-D (fsdp, tp) mesh and the Trainer trains it end to end. Until now tp_size > 1 with fsdp_size > 1 raised FSDP+TP is not supported yet.

Correctness

Trainer-level parity (fp32, 8 GPUs, 6 steps, max_grad_norm=1.0 active, same total batch of 8): per-step loss and gradient norm of every arm against a single-process run of the same tiny MoE (16 experts).

arm	max abs diff in loss	max abs diff in grad norm	saved weights vs single-process save
EP only, tp_size=8	0.0	1.2e-7	117 tensors, max rel diff 2.2e-6
2-D, tp_size=4, fsdp_size=2	4.8e-7	1.2e-7	117 tensors, max rel diff 2.4e-6
2-D, tp_size=2, fsdp_size=4	4.8e-7	1.2e-7	117 tensors, max rel diff 1.6e-6

This exercises the batch split across fsdp, FSDP2's reduction of the expert gradients, the mesh-aware norm and clipping, the per-parameter optimizer, the loss normalization across ranks, and save_model. Driver and comparison scripts in the details ⬇️.

fp32 gradient certification against a single-GPU reference (real Qwen3-30B-A3B weights, first N layers): every parameter's gradient, relative max-abs difference.

layers	EP only (tp_size=4)	2-D (tp_size=2, fsdp_size=2)
2	25/25 params, max 2.0e-6	25/25 params, max 5.2e-5
4	47/47 params, max 2.7e-6	47/47 params, max 1.7e-2 (one expert tensor; the loss differs by 1.2e-5 relative, consistent with a single routing flip; every other parameter is below 6e-3)

Loss curves (OLMoE-1B-7B, bf16 full fine-tuning, tulu-3 data, same total batch of 8 x 1024 tokens):

bf16 runs with different reduction orders drift by a few percent per step (peak memory: 64.5 GB single GPU, 12.5 GB EP, 9.5 GB and 8.7 GB for the two 2-D layouts); the fp32 tables above are the exact check.

Throughput and memory

Qwen3-30B-A3B full fine-tuning, bf16, 8xH100, sequence length 2048, per-device batch 1, sdpa, AdamW:

configuration	tokens/s/GPU	peak memory/GPU
tp_size=8	3485	38.6 GB
tp_size=4, fsdp_size=2	2900	34.2 GB
tp_size=2, fsdp_size=4	2830	32.3 GB

The 2-D configurations pay FSDP2's all-gather/reduce-scatter of the experts across fsdp. Splitting the optimizer param groups by mesh (instead of stepping per parameter as #48208 did) is worth 189 -> 41 ms per step at tp_size=4, fsdp_size=2 on this model, which is where most of the +15-25% over the umbrella branch's numbers comes from.

At scale (installed stack: this branch + accelerate/peft/trl from main, 64 H100, trl.SFTTrainer, huggingface/trl#6869's script unchanged): GLM-4.5-Air 110B full fine-tuning with tp_size=32, fsdp_size=2, 20 steps at ~3 s/step steady state, loss 3.9 -> 1.2, then save_model writes the 200 GB checkpoint (5 safetensors shards) through the collective gather.

Known limitations

• tp_size doubles as the EP size.
• Resume from a checkpoint is not supported for models sharded at load time; the Trainer refuses to write optimizer checkpoints for them.
• accelerate builds its own (dp_shard, tp) mesh from the ParallelismConfig next to the model's; it only reads sizes and ranks from it, but that is a second set of communicators per rank.

Trainer-level parity driver (single process vs EP vs 2-D through the Trainer, same total batch)

parity_train.py

"""Per-step loss / grad-norm parity of EP and 2-D (fsdp, tp) training through the Trainer against a
single-process control on the same total batch.

    python parity_train.py --mode single ...                       (1 GPU)
    torchrun --nproc_per_node 8 parity_train.py --mode ep ...      (tp=8)
    torchrun --nproc_per_node 8 parity_train.py --mode 2d --ep 4   (tp=4 x fsdp=2)

`--total-bs` is the batch every mode sees per step: the control and EP take all of it on each
process, the 2-D run splits it across the `fsdp` ranks (per-device batch = total / dp).
"""

import argparse
import json
import os

import torch
from datasets import Dataset, load_dataset

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


ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--mode", choices=["single", "ep", "2d"], required=True)
ap.add_argument("--ep", type=int, default=None, help="tp/ep size for --mode 2d")
ap.add_argument("--steps", type=int, default=5)
ap.add_argument("--total-bs", type=int, default=8)
ap.add_argument("--seq-len", type=int, default=64)
ap.add_argument("--dtype", default="float32")
ap.add_argument("--real-data", default=None, help="HF dataset id with a `messages` column; synthetic ids otherwise")
ap.add_argument("--lr", type=float, default=1e-4)
ap.add_argument("--gc", action="store_true")
ap.add_argument("--out", required=True)
ap.add_argument("--save-dir", default=None, help="call trainer.save_model here after training")
args = ap.parse_args()

world = int(os.environ.get("WORLD_SIZE", "1"))
rank = int(os.environ.get("RANK", "0"))
dtype = getattr(torch, args.dtype)
dp = 1
kwargs = {}
if args.mode == "ep":
    kwargs["distributed_config"] = DistributedConfig(tp_size=world, fsdp_size=1, enable_expert_parallel=True)
elif args.mode == "2d":
    dp = world // args.ep
    kwargs["distributed_config"] = DistributedConfig(tp_size=args.ep, fsdp_size=dp, enable_expert_parallel=True)
model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype, **kwargs)
if args.mode == "single":
    model.cuda()

n_rows = args.steps * args.total_bs
if args.real_data:
    tok = AutoTokenizer.from_pretrained(args.model)
    raw = load_dataset(args.real_data, split=f"train[:{n_rows}]")
    if tok.chat_template is not None:
        texts = [tok.apply_chat_template(m, tokenize=False) for m in raw["messages"]]
    else:  # base model: plain turns
        texts = ["\n".join(t["content"] for t in m) for m in raw["messages"]]
    enc = tok(texts, max_length=args.seq_len, truncation=True, padding="max_length", return_tensors="pt")
    ids = enc["input_ids"]
    labels = ids.clone()
    labels[enc["attention_mask"] == 0] = -100
    ds = Dataset.from_dict({"input_ids": ids.tolist(), "attention_mask": enc["attention_mask"].tolist(), "labels": labels.tolist()})
else:
    g = torch.Generator().manual_seed(1234)
    ids = torch.randint(0, model.config.vocab_size, (n_rows, args.seq_len), generator=g)
    ds = Dataset.from_dict({"input_ids": ids.tolist(), "labels": ids.tolist()})

targs = TrainingArguments(
    output_dir=f"/tmp/parity_{args.mode}_{args.ep}_{rank}",
    per_device_train_batch_size=args.total_bs // dp,
    max_steps=args.steps,
    logging_steps=1,
    report_to=[],
    max_grad_norm=1.0,
    save_strategy="no",
    seed=0,
    learning_rate=args.lr,
    lr_scheduler_type="constant",
    warmup_steps=0,
    optim="adamw_torch",
    bf16=dtype == torch.bfloat16,
    gradient_checkpointing=args.gc,
    dataloader_drop_last=True,
)
trainer = Trainer(model=model, args=targs, train_dataset=ds)
trainer.train()
if args.save_dir:
    trainer.save_model(args.save_dir)

if rank == 0:
    hist = [
        {"step": h["step"], "loss": h["loss"], "grad_norm": h["grad_norm"]}
        for h in trainer.state.log_history
        if "loss" in h
    ]
    res = {
        "mode": args.mode, "ep": args.ep if args.mode == "2d" else (world if args.mode == "ep" else 1),
        "dp": dp, "world": world, "model": args.model, "dtype": args.dtype, "total_bs": args.total_bs,
        "seq_len": args.seq_len, "real_data": args.real_data,
        "peak_mem_gb": torch.cuda.max_memory_allocated() / 2**30, "history": hist,
    }
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    json.dump(res, open(args.out, "w"), indent=2)
    print("### PARITY", json.dumps({k: v for k, v in res.items() if k != "history"}), flush=True)
    for h in hist:
        print(f"### step {h['step']}: loss {h['loss']:.6f} grad_norm {h['grad_norm']:.6f}", flush=True)
if world > 1:
    torch.distributed.barrier()

compare_parity.py

"""Print per-step loss / grad-norm deltas of every arm against the single-process control."""
import json
import sys

runs = {p: json.load(open(p)) for p in sys.argv[1:]}
ctrl = next(r for r in runs.values() if r["mode"] == "single")
print(f"control: {ctrl['model']} {ctrl['dtype']} total_bs={ctrl['total_bs']} seq={ctrl['seq_len']}")
print(f"{'arm':<22} {'step':>4} {'loss':>10} {'d_loss':>10} {'grad_norm':>10} {'d_gnorm':>10}")
for r in runs.values():
    name = r["mode"] if r["mode"] != "2d" else f"2d ep={r['ep']} dp={r['dp']}"
    for h0, h in zip(ctrl["history"], r["history"]):
        print(f"{name:<22} {h['step']:>4} {h['loss']:>10.6f} {h['loss'] - h0['loss']:>+10.2e} {h['grad_norm']:>10.6f} {h['grad_norm'] - h0['grad_norm']:>+10.2e}")
    worst_l = max(abs(h["loss"] - h0["loss"]) for h0, h in zip(ctrl["history"], r["history"]))
    worst_g = max(abs(h["grad_norm"] - h0["grad_norm"]) for h0, h in zip(ctrl["history"], r["history"]))
    print(f"### {name}: max |d_loss| {worst_l:.2e}  max |d_grad_norm| {worst_g:.2e}  peak_mem {r['peak_mem_gb']:.1f} GB")

compare_saves.py

"""Compare every tensor of two saved checkpoints (safetensors dirs)."""
import glob
import os
import sys

import torch
from safetensors.torch import load_file

def load(d):
    out = {}
    for f in sorted(glob.glob(os.path.join(d, "*.safetensors"))):
        out.update(load_file(f))
    return out

ref, other = load(sys.argv[1]), load(sys.argv[2])
missing = set(ref) ^ set(other)
worst = 0.0
for k in ref:
    if k in other:
        a, b = ref[k].float(), other[k].float()
        worst = max(worst, ((a - b).abs().max() / (a.abs().max() + 1e-12)).item())
print(f"### SAVE_COMPARE {sys.argv[2]}: {len(ref)} tensors, {len(missing)} key mismatches, max rel diff {worst:.3e}")

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

[qgallouedec]

Merge-order note!
#48201 makes sentinel_mask None when expert parallelism is off, and the two per-matmul masked_fill(sentinel_mask, 0.0) sites this PR adds in moe.py need the same if self.is_expert_parallel: gate once both are on main.

Found by merging both into #48204: the non-EP forward crashes otherwise. Whichever lands second should add the gate.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33819692765:1
Result: success | Jobs: 16 | Tests: 185,507 | Failures: 0 | Duration: 15h 49m