[Feature] Decouple expert parallelism from FSDP sharding (dp2ep layout)

examples/v1/config/sft_glm5p2.py

@@ -108,10 +108,15 @@ def _get_float8_config() -> Float8Config | None:
 else:
     raise ValueError(f"Unsupported OPTIMIZER={optimizer!r}. Use adamw or muon.")
 lr_cfg = LRConfig(lr_type=os.environ.get("LR_TYPE", "cosine"), warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0")))
+hsdp_sharding_size = os.environ.get("HSDP_SHARDING_SIZE")
 fsdp_cfg = FSDPConfig(
     cpu_offload=_get_bool_env("CPU_OFFLOAD", False),
     ep_size=ep_size,
     torch_compile=_get_bool_env("TORCH_COMPILE", False),
+    # DECOUPLE_EP_FSDP=1: dp2ep layout — dense params sharded over the full FSDP mesh instead of
+    # being replicated `ep_size` times; experts sharded `dp_shard / ep_size` ways on top of EP.
+    decouple_ep_fsdp=_get_bool_env("DECOUPLE_EP_FSDP", False),
+    hsdp_sharding_size=int(hsdp_sharding_size) if hsdp_sharding_size else None,
 )
 
 trainer = TrainerConfig(

tests/model/run_decoupled_ep_fsdp_ckpt.py

@@ -0,0 +1,151 @@
+"""DCP round trip and HF export checks for the decoupled EP/FSDP layout (DESIGN.md §6 L3).
+
+For every mode:
+
+1. ``from_hf`` the tiny checkpoint and immediately ``save_hf`` → must be bit-identical to the source;
+2. train 5 steps, ``save_dcp``, train 5 more steps, ``save_hf`` (the "continuous" run);
+3. build a fresh engine, ``load_dcp`` the step-5 checkpoint, train steps 5-9 → losses must continue
+   the continuous curve, and its ``save_hf`` must match the continuous export;
+4. the step-10 HF exports of all modes are compared against the ``ep=1`` baseline export.
+
+Optionally the step-5 DCP checkpoint of one layout is loaded into another layout
+(``--cross-load SRC:DST``) to check DCP resharding across the switch.
+
+The pytest gate for these checks lives in ``tests/engine/test_decoupled_ep_fsdp_train_engine.py``;
+this script writes the JSON behind the markdown report.
+
+    torchrun --nproc-per-node 8 tests/model/run_decoupled_ep_fsdp_ckpt.py \\
+        --modes A:ep=1 C:ep=8,decouple=1 C4:ep=4,decouple=1 H41:ep=4,decouple=1,hsdp=4 --out /tmp/l3
+"""
+
+import argparse
+import json
+import os
+import shutil
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from safetensors.torch import save_file
+
+from xtuner._testing.decoupled_ep_fsdp import (
+    build_engine,
+    build_hf_checkpoint,
+    compare_hf,
+    load_hf_dir,
+    max_rel_diff,
+    parse_mode,
+    release,
+    train,
+)
+from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--modes", nargs="+", required=True)
+    parser.add_argument("--cross-load", nargs="*", default=[], help="SRC:DST mode-name pairs")
+    parser.add_argument("--seq-len", type=int, default=2048)
+    parser.add_argument("--hf-dir", type=str, default=None)
+    parser.add_argument("--out", type=str, required=True)
+    args = parser.parse_args()
+
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    out = Path(args.out)
+    hf_dir = Path(args.hf_dir) if args.hf_dir else out / "tiny_qwen3_moe_hf"
+    if rank == 0:
+        out.mkdir(parents=True, exist_ok=True)
+        if not (hf_dir / "config.json").exists():
+            build_hf_checkpoint(hf_dir, 0, "tiny")
+    dist.barrier()
+    vocab_size = Qwen3MoEConfig.from_hf(hf_dir).vocab_size
+
+    # Source checkpoint in bf16, the dtype `save_hf` writes.
+    src_bf16 = out / "source_bf16"
+    if rank == 0:
+        src_bf16.mkdir(exist_ok=True)
+        tensors = {k: v.to(torch.bfloat16).contiguous() for k, v in load_hf_dir(hf_dir).items()}
+        save_file(tensors, str(src_bf16 / "model.safetensors"))
+    dist.barrier()
+
+    report: dict[str, Any] = {"modes": {}, "cross_hf_vs_baseline": {}, "cross_load": {}}
+    modes = [parse_mode(spec) for spec in args.modes]
+    for mode in modes:
+        name = mode["name"]
+        mode_dir = out / name
+        if rank == 0:
+            shutil.rmtree(mode_dir, ignore_errors=True)
+            mode_dir.mkdir(parents=True)
+        dist.barrier()
+        if rank == 0:
+            print(f"===== {mode}", flush=True)
+
+        engine = build_engine(mode, hf_dir, "ckpt_main")
+        engine.from_hf(hf_path=hf_dir, strict=True)
+        engine.save_hf(str(mode_dir / "hf_step0"))
+        dist.barrier()
+        entry: dict[str, Any] = {"mode": mode}
+        if rank == 0:
+            entry["hf_step0_vs_source"] = compare_hf(src_bf16, mode_dir / "hf_step0")
+
+        losses_first = train(engine, range(0, 5), vocab_size, args.seq_len)
+        engine.save_dcp(mode_dir / "dcp_step5")
+        dist.barrier()
+        losses_cont = train(engine, range(5, 10), vocab_size, args.seq_len)
+        engine.save_hf(str(mode_dir / "hf_step10"))
+        dist.barrier()
+        release(engine)
+
+        resumed = build_engine(mode, hf_dir, "ckpt_resume")
+        resumed.load_dcp(mode_dir / "dcp_step5")
+        losses_resumed = train(resumed, range(5, 10), vocab_size, args.seq_len)
+        resumed.save_hf(str(mode_dir / "hf_step10_resumed"))
+        dist.barrier()
+        release(resumed)
+
+        entry["losses_steps_0_4"] = losses_first
+        entry["losses_steps_5_9_continuous"] = losses_cont
+        entry["losses_steps_5_9_resumed"] = losses_resumed
+        entry["resume_max_rel_loss_diff"] = max_rel_diff(losses_resumed, losses_cont)
+        if rank == 0:
+            entry["hf_step10_resumed_vs_continuous"] = compare_hf(
+                mode_dir / "hf_step10", mode_dir / "hf_step10_resumed"
+            )
+            print(json.dumps({k: v for k, v in entry.items() if k != "mode"}, indent=1), flush=True)
+        report["modes"][name] = entry
+
+    if rank == 0:
+        baseline = modes[0]["name"]
+        for mode in modes[1:]:
+            report["cross_hf_vs_baseline"][mode["name"]] = compare_hf(
+                out / baseline / "hf_step10", out / mode["name"] / "hf_step10"
+            )
+
+    for pair in args.cross_load:
+        src, dst = pair.split(":")
+        dst_mode = next(m for m in modes if m["name"] == dst)
+        if rank == 0:
+            print(f"===== cross load {src} -> {dst}", flush=True)
+        engine = build_engine(dst_mode, hf_dir, f"ckpt_cross_{src}")
+        engine.load_dcp(out / src / "dcp_step5")
+        losses = train(engine, range(5, 10), vocab_size, args.seq_len)
+        release(engine)
+        ref = report["modes"][src]["losses_steps_5_9_continuous"]
+        report["cross_load"][pair] = {
+            "losses_steps_5_9": losses,
+            "max_rel_loss_diff_vs_src_continuous": max_rel_diff(losses, ref),
+        }
+
+    if rank == 0:
+        (out / "l3.json").write_text(json.dumps(report, indent=2))
+        print(json.dumps({k: report[k] for k in ("cross_hf_vs_baseline", "cross_load")}, indent=1), flush=True)
+        print(f"wrote {out / 'l3.json'}")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

tests/model/run_decoupled_ep_fsdp_numerics.py

@@ -0,0 +1,82 @@
+"""Multi-GPU numerical-equivalence runs for the EP/FSDP decoupling (DESIGN.md §6 L1 / L2).
+
+Every mode trains the same tiny Qwen3-MoE (random HF checkpoint built once from a fixed seed)
+on the same token stream, so loss / grad-norm curves of different layouts can be compared. The
+pytest gate for the same comparison lives in ``tests/engine/test_decoupled_ep_fsdp_train_engine.py``;
+this script exists for the larger ``--model-size medium`` memory / step-time runs, ``--fp8`` and
+``--dispatcher deepep``, and for producing the JSON behind the markdown reports
+(``summarize_decoupled_ep_fsdp_numerics.py``).
+
+Example (L1):
+
+    torchrun --nproc-per-node 8 tests/model/run_decoupled_ep_fsdp_numerics.py \\
+        --modes A:ep=1 B:ep=8 C:ep=8,decouple=1 --steps 50 --out reports/l1.json
+
+Mode spec: ``<name>:key=value[,key=value...]`` with keys ``ep`` (int), ``decouple`` (0/1) and
+``hsdp`` (hsdp_sharding_size, int).
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from xtuner._testing.decoupled_ep_fsdp import MODEL_SIZES, build_hf_checkpoint, parse_mode, run_mode
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--modes", nargs="+", required=True)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--seq-len", type=int, default=2048)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--model-size", choices=tuple(MODEL_SIZES), default="tiny")
+    parser.add_argument("--dispatcher", choices=("all2all", "deepep"), default="all2all")
+    parser.add_argument("--fp8", action="store_true", help="tile-wise float8 for linear and grouped linear")
+    parser.add_argument("--hf-dir", type=str, default=None, help="where to build / reuse the tiny HF checkpoint")
+    parser.add_argument("--out", type=str, required=True)
+    args = parser.parse_args()
+
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+    hf_dir = Path(args.hf_dir or Path(args.out).with_suffix("")).with_name(f"{args.model_size}_qwen3_moe_hf")
+    if rank == 0 and not (hf_dir / "config.json").exists():
+        hf_dir.mkdir(parents=True, exist_ok=True)
+        build_hf_checkpoint(hf_dir, args.seed, args.model_size)
+    dist.barrier()
+
+    results = []
+    for spec in args.modes:
+        mode = parse_mode(spec)
+        if rank == 0:
+            print(f"===== running mode {mode}", flush=True)
+        results.append(run_mode(mode, hf_dir, args.steps, args.seq_len, args.lr, args.dispatcher, args.fp8))
+        dist.barrier()
+        if rank == 0:
+            last = results[-1]
+            print(
+                f"[{mode['name']}] loss[0]={last['losses'][0]:.6f} loss[-1]={last['losses'][-1]:.6f} "
+                f"grad_norm[0]={last['grad_norms'][0]:.6f} step_time={last['step_time_s'] * 1000:.1f}ms "
+                f"mem={last['memory']}",
+                flush=True,
+            )
+
+    if rank == 0:
+        out = Path(args.out)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(
+            json.dumps({"args": vars(args), "world_size": dist.get_world_size(), "results": results}, indent=2)
+        )
+        print(f"wrote {out}")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

tests/model/summarize_decoupled_ep_fsdp_numerics.py

@@ -0,0 +1,90 @@
+"""Turn the JSON written by ``run_decoupled_ep_fsdp_numerics.py`` into markdown tables."""
+
+import argparse
+import json
+from pathlib import Path
+
+
+def rel(a: float, b: float) -> float:
+    return abs(a - b) / max(abs(b), 1e-12)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("json_path")
+    parser.add_argument("--ref", default=None, help="reference mode name (default: first mode)")
+    args = parser.parse_args()
+    data = json.loads(Path(args.json_path).read_text())
+    results = {r["mode"]["name"]: r for r in data["results"]}
+    names = list(results)
+    ref = args.ref or names[0]
+    steps = len(results[ref]["losses"])
+
+    print("## Modes\n")
+    print("| mode | ep | decouple | hsdp_sharding_size |")
+    print("|---|---|---|---|")
+    for n in names:
+        m = results[n]["mode"]
+        print(f"| {n} | {m['ep']} | {m['decouple']} | {m['hsdp']} |")
+
+    print("\n## Loss curve (reduced_llm_loss)\n")
+    print("| step | " + " | ".join(names) + " | " + " | ".join(f"rel({n} vs {ref})" for n in names if n != ref) + " |")
+    print("|---|" + "---|" * (2 * len(names) - 1))
+    for s in list(range(0, steps, 5)) + [steps - 1]:
+        row = [f"{results[n]['losses'][s]:.6f}" for n in names]
+        diffs = [f"{rel(results[n]['losses'][s], results[ref]['losses'][s]):.2e}" for n in names if n != ref]
+        print(f"| {s} | " + " | ".join(row) + " | " + " | ".join(diffs) + " |")
+    print("\n| pair | max rel loss diff over all steps | mean rel loss diff |")
+    print("|---|---|---|")
+    for n in names:
+        if n == ref:
+            continue
+        d = [rel(a, b) for a, b in zip(results[n]["losses"], results[ref]["losses"])]
+        print(f"| {n} vs {ref} | {max(d):.2e} | {sum(d) / len(d):.2e} |")
+
+    print("\n## Total grad norm (pre-clip)\n")
+    print("| step | " + " | ".join(names) + " | " + " | ".join(f"rel({n} vs {ref})" for n in names if n != ref) + " |")
+    print("|---|" + "---|" * (2 * len(names) - 1))
+    for s in list(range(0, steps, 5)) + [steps - 1]:
+        row = [f"{results[n]['grad_norms'][s]:.6f}" for n in names]
+        diffs = [f"{rel(results[n]['grad_norms'][s], results[ref]['grad_norms'][s]):.2e}" for n in names if n != ref]
+        print(f"| {s} | " + " | ".join(row) + " | " + " | ".join(diffs) + " |")
+
+    print("\n## Per-parameter grad norms (max / mean relative diff vs reference)\n")
+    print("| step | pair | dense params | expert params | worst param |")
+    print("|---|---|---|---|---|")
+    for s in results[ref]["param_grad_norms"]:
+        ref_norms = results[ref]["param_grad_norms"][s]
+        for n in names:
+            if n == ref:
+                continue
+            norms = results[n]["param_grad_norms"][s]
+            dense, expert, worst = [], [], (0.0, "")
+            for k, v in ref_norms.items():
+                d = rel(norms[k], v)
+                (expert if ".experts" in k else dense).append(d)
+                if d > worst[0]:
+                    worst = (d, k)
+            print(
+                f"| {s} | {n} vs {ref} | max {max(dense):.2e} / mean {sum(dense) / len(dense):.2e} "
+                f"| max {max(expert):.2e} / mean {sum(expert) / len(expert):.2e} | {worst[1]} ({worst[0]:.2e}) |"
+            )
+
+    print("\n## Per-rank memory (MiB) and step time\n")
+    keys = [
+        "dense_param_mib",
+        "expert_param_mib",
+        "allocated_after_load_mib",
+        "allocated_after_step0_mib",
+        "peak_allocated_mib",
+        "peak_reserved_mib",
+    ]
+    print("| mode | " + " | ".join(keys) + " | step time (ms) |")
+    print("|---|" + "---|" * (len(keys) + 1))
+    for n in names:
+        mem = results[n]["memory"]
+        print(f"| {n} | " + " | ".join(f"{mem[k]:.1f}" for k in keys) + f" | {results[n]['step_time_s'] * 1000:.1f} |")
+
+
+if __name__ == "__main__":
+    main()