From 1f919bec8a50c2555069225b99c9349f8e45fd52 Mon Sep 17 00:00:00 2001 From: daixu Date: Fri, 21 Aug 2026 22:46:05 +0000 Subject: [PATCH 01/14] [Test] Add fake-PG L0 placement tests pinning the legacy MoE EP/FSDP layout Single-process fake ProcessGroup tests asserting the mesh ranks and DTensor placements that MoE.fully_shard produces today for 8/16/64 ranks, plus the Phase-0 baseline snapshot and the decision log for the EP/FSDP decoupling work. --- tests/model/test_decoupled_ep_fsdp_mesh.py | 233 +++++++++++++++++++++ 1 file changed, 233 insertions(+) create mode 100644 tests/model/test_decoupled_ep_fsdp_mesh.py diff --git a/tests/model/test_decoupled_ep_fsdp_mesh.py b/tests/model/test_decoupled_ep_fsdp_mesh.py new file mode 100644 index 0000000000..8fd9d3d67c --- /dev/null +++ b/tests/model/test_decoupled_ep_fsdp_mesh.py @@ -0,0 +1,233 @@ +"""L0 placement tests for the MoE EP/FSDP mesh layouts. + +These tests run in a single process on top of a fake ``ProcessGroup`` so that +8/16/64-rank layouts can be asserted without any multi-process launch. The +model stays on the ``meta`` device; only the mesh bookkeeping and the DTensor +placements produced by ``MoE.fully_shard`` are inspected. +""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch +from pydantic import ValidationError +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, _mesh_resources +from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor.placement_types import _StridedShard +from torch.testing._internal.distributed.fake_pg import FakeStore + +from xtuner.v1.config import FSDPConfig +from xtuner.v1.model.base import BaseModel +from xtuner.v1.model.moe.moe import MoE, MoEConfig +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.router import NoAuxRouterConfig + + +PREFIX = "l0mesh" +HIDDEN = 128 +N_EXPERTS = 64 +MOE_INTER = 128 + +DENSE_PARAMS = ( + "embed_tokens.weight", + "layers.0.mlp.gate_proj.weight", + "layers.1.self_attn.q_proj.weight", + "layers.1.gate.weight", + "layers.1.shared_experts.gate_proj.weight", + "layers.1.input_layernorm.weight", + "norm.weight", + "lm_head.weight", +) +EXPERT_PARAMS = ( + "layers.1.experts.fused_w1w3.weight", + "layers.1.experts.fused_w2.weight", +) + + +def _reset_mesh_resources() -> None: + _mesh_resources.mesh_stack.clear() + _mesh_resources.child_to_root_mapping.clear() + _mesh_resources.root_to_flatten_mapping.clear() + _mesh_resources.flatten_name_to_root_dims.clear() + + +@contextmanager +def _fake_world(world_size: int, rank: int) -> Iterator[None]: + dist.init_process_group("fake", store=FakeStore(), rank=rank, world_size=world_size) + try: + yield + finally: + dist.destroy_process_group() + _reset_mesh_resources() + + +# `MoE.__init__` unconditionally creates CUDA streams and dispatcher buffers, so a +# CUDA context is required even though the fake ProcessGroup never communicates. +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="MoE construction requires a CUDA device") + + +@pytest.fixture(autouse=True) +def _keep_model_on_meta(monkeypatch: pytest.MonkeyPatch) -> None: + # Placements are fully determined before materialization; skipping + # `_to_empty_meta` keeps the test free of device allocations. + monkeypatch.setattr(BaseModel, "_to_empty_meta", lambda self: None) + + +def _build_model(ep_size: int) -> MoE: + config = MoEConfig( + vocab_size=1024, + max_position_embeddings=256, + pad_token_id=0, + eos_token_id=0, + num_hidden_layers=2, + hidden_size=HIDDEN, + intermediate_size=256, + rms_norm_eps=1e-6, + rope_theta=1e6, + hidden_act="silu", + attention=MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=32), + tie_word_embeddings=False, + n_routed_experts=N_EXPERTS, + n_shared_experts=1, + num_experts_per_tok=2, + first_k_dense_replace=1, + hidden_factor=1.0, + moe_intermediate_size=MOE_INTER, + router=NoAuxRouterConfig( + scoring_func="sigmoid", + router_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, + ), + compile_cfg=False, + ep_size=ep_size, + dispatcher="all2all", + mesh_prefix=PREFIX, + ) + with torch.device("meta"): + return MoE(config) + + +def _shard_model(ep_size: int, **fsdp_kwargs) -> MoE: + model = _build_model(ep_size) + fsdp_config = FSDPConfig(ep_size=ep_size, torch_compile=False, **fsdp_kwargs) + return model.fully_shard(fsdp_config) + + +def _named_params(model: MoE) -> dict[str, DTensor]: + params: dict[str, DTensor] = {} + for name, param in model.named_parameters(): + params[name.replace("_checkpoint_wrapped_module.", "")] = param + return params + + +def _ranks(mesh: DeviceMesh) -> list: + return mesh.mesh.tolist() + + +def _legacy_ep_ranks(world_size: int, ep_size: int, rank: int) -> list[int]: + start = (rank // ep_size) * ep_size + return list(range(start, start + ep_size)) + + +def _legacy_fsdp_ranks(world_size: int, ep_size: int, rank: int) -> list[int]: + return [rank % ep_size + k * ep_size for k in range(world_size // ep_size)] + + +def _assert_placements(param: DTensor, mesh_dim_names: tuple[str, ...], placements: tuple) -> None: + assert isinstance(param, DTensor) + assert param.device_mesh.mesh_dim_names == mesh_dim_names + assert len(param.placements) == len(placements) + for actual, expected in zip(param.placements, placements): + assert type(actual) is type(expected), (actual, expected) + if isinstance(expected, _StridedShard): + assert actual.dim == expected.dim and actual.split_factor == expected.split_factor + elif isinstance(expected, Shard): + assert actual.dim == expected.dim + + +class TestLegacyMoEMeshLayout: + """Layout produced by the original EP-orthogonal-to-FSDP path. + + ``decouple_ep_fsdp`` is off by default, so these expectations are the + regression contract for the legacy path: ``root = (fsdp = world / ep, ep)`` + with ``ep`` on the innermost (contiguous-rank) dimension. + """ + + @pytest.mark.parametrize( + "world_size,ep_size,rank", + [(8, 8, 0), (8, 8, 5), (8, 4, 6), (8, 2, 3), (16, 8, 9), (64, 8, 37)], + ) + def test_mesh_ranks(self, world_size: int, ep_size: int, rank: int) -> None: + with _fake_world(world_size, rank): + model = _shard_model(ep_size) + + assert model.hsdp_mesh is None + assert model._world_mesh is not None + assert tuple(model._world_mesh.shape) == (world_size // ep_size, ep_size) + assert model._world_mesh.mesh_dim_names == (f"{PREFIX}.fsdp", f"{PREFIX}.ep") + + assert model.fsdp_mesh is not None and model.ep_mesh is not None + assert model.fsdp_mesh.mesh_dim_names == (f"{PREFIX}.fsdp",) + assert _ranks(model.fsdp_mesh) == _legacy_fsdp_ranks(world_size, ep_size, rank) + assert model.ep_mesh.mesh_dim_names == (f"{PREFIX}.ep",) + assert _ranks(model.ep_mesh) == _legacy_ep_ranks(world_size, ep_size, rank) + assert dist.get_process_group_ranks(model.ep_mesh.get_group()) == _legacy_ep_ranks( + world_size, ep_size, rank + ) + + @pytest.mark.parametrize("world_size,ep_size,rank", [(8, 8, 0), (16, 8, 9), (64, 8, 37)]) + def test_param_placements(self, world_size: int, ep_size: int, rank: int) -> None: + fsdp_size = world_size // ep_size + with _fake_world(world_size, rank): + model = _shard_model(ep_size) + params = _named_params(model) + mesh_dim_names = (f"{PREFIX}.fsdp", f"{PREFIX}.ep") + + for name in DENSE_PARAMS: + param = params[name] + # Dense params are replicated over ep and only sharded world/ep ways. + _assert_placements(param, mesh_dim_names, (Shard(0), Replicate())) + assert param.to_local().shape[0] == param.shape[0] // fsdp_size, name + + for name in EXPERT_PARAMS: + param = params[name] + _assert_placements(param, mesh_dim_names, (_StridedShard(0, split_factor=ep_size), Shard(0))) + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + def test_ep1_is_plain_fsdp(self) -> None: + world_size, rank = 8, 3 + with _fake_world(world_size, rank): + model = _shard_model(ep_size=1) + params = _named_params(model) + + assert model.ep_mesh is not None and model.ep_mesh.size() == 1 + assert model.fsdp_mesh is not None and _ranks(model.fsdp_mesh) == list(range(world_size)) + for name in DENSE_PARAMS + EXPERT_PARAMS: + param = params[name] + _assert_placements(param, (f"{PREFIX}.fsdp",), (Shard(0),)) + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + def test_hsdp_ep1_layout(self) -> None: + world_size, shard_size, rank = 16, 8, 11 + with _fake_world(world_size, rank): + model = _shard_model(ep_size=1, hsdp_sharding_size=shard_size) + params = _named_params(model) + + assert model.hsdp_mesh is not None + assert tuple(model.hsdp_mesh.shape) == (world_size // shard_size, shard_size) + assert model.hsdp_mesh.mesh_dim_names == (f"{PREFIX}.hsdp_replicate", f"{PREFIX}.hsdp_shard") + assert model.fsdp_mesh is not None + assert _ranks(model.fsdp_mesh) == list(range((rank // shard_size) * shard_size, (rank // shard_size + 1) * shard_size)) + for name in DENSE_PARAMS + EXPERT_PARAMS: + param = params[name] + _assert_placements(param, model.hsdp_mesh.mesh_dim_names, (Replicate(), Shard(0))) + assert param.to_local().shape[0] == param.shape[0] // shard_size, name + + def test_hsdp_with_ep_is_rejected(self) -> None: + with pytest.raises(ValidationError, match="HSDP requires expert parallel size to be 1"): + FSDPConfig(ep_size=8, hsdp_sharding_size=8) + From da6612689d9be52f80a40be136fa170ab9c0471e Mon Sep 17 00:00:00 2001 From: daixu Date: Fri, 21 Aug 2026 22:59:46 +0000 Subject: [PATCH 02/14] [Feature] Decouple expert parallel from FSDP sharding (dp2ep layout) Add FSDPConfig.decouple_ep_fsdp. When enabled, MoE builds a single (replicate, efsdp, ep) root mesh with efsdp = dp_shard / ep, shards routed experts with their own FSDP group on efsdp and every other parameter over the full dp_shard mesh (no more ep-fold replication of dense params), and scales expert gradients by 1/ep instead of all-reducing dense gradients over EP. The legacy layout is untouched when the flag is off. Includes fake-PG L0 placement tests for 8/16/64 ranks, the 8-GPU numerics script used for the L1 report, and reports/L1.md. --- tests/model/run_decoupled_ep_fsdp_numerics.py | 260 ++++++++++++++++++ .../summarize_decoupled_ep_fsdp_numerics.py | 90 ++++++ tests/model/test_decoupled_ep_fsdp_mesh.py | 154 ++++++++++- xtuner/v1/config/fsdp.py | 16 +- xtuner/v1/model/moe/moe.py | 200 +++++++++++++- 5 files changed, 713 insertions(+), 7 deletions(-) create mode 100644 tests/model/run_decoupled_ep_fsdp_numerics.py create mode 100644 tests/model/summarize_decoupled_ep_fsdp_numerics.py diff --git a/tests/model/run_decoupled_ep_fsdp_numerics.py b/tests/model/run_decoupled_ep_fsdp_numerics.py new file mode 100644 index 0000000000..0b5e026292 --- /dev/null +++ b/tests/model/run_decoupled_ep_fsdp_numerics.py @@ -0,0 +1,260 @@ +"""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. + +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: ``:key=value[,key=value...]`` with keys ``ep`` (int), ``decouple`` (0/1) and +``hsdp`` (hsdp_sharding_size, int). +""" + +import argparse +import gc +import json +import os +import time +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor + +from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.base import ModelItem +from xtuner.v1.model.moe.moe import SequenceContext +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig +from xtuner.v1.utils.device import get_device +from xtuner.v1.utils.dtensor import cal_total_norm + + +DEVICE = get_device() +GRAD_NORM_STEPS = (0, 1, 25) + + +def parse_mode(spec: str) -> dict[str, Any]: + name, _, kv = spec.partition(":") + mode: dict[str, Any] = {"name": name, "ep": 1, "decouple": False, "hsdp": None} + for item in filter(None, kv.split(",")): + key, value = item.split("=") + if key == "ep": + mode["ep"] = int(value) + elif key == "decouple": + mode["decouple"] = bool(int(value)) + elif key == "hsdp": + mode["hsdp"] = int(value) + else: + raise ValueError(f"unknown mode key {key}") + return mode + + +MODEL_SIZES: dict[str, dict[str, int]] = { + # ~108M params: numerics (fast, 50 steps in < 10 s) + "tiny": { + "vocab_size": 4096, + "hidden_size": 512, + "intermediate_size": 1024, + "moe_intermediate_size": 256, + "num_hidden_layers": 4, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "head_dim": 64, + "num_experts": 64, + "num_experts_per_tok": 4, + }, + # ~3.4B params (3.2B in routed experts): step-time / memory comparison + "medium": { + "vocab_size": 32768, + "hidden_size": 2048, + "intermediate_size": 4096, + "moe_intermediate_size": 1024, + "num_hidden_layers": 8, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 128, + "num_experts": 64, + "num_experts_per_tok": 4, + }, +} + + +def build_hf_checkpoint(path: Path, seed: int, model_size: str) -> None: + config = Qwen3MoeConfig( + **MODEL_SIZES[model_size], + max_position_embeddings=8192, + norm_topk_prob=True, + tie_word_embeddings=False, + bos_token_id=1, + eos_token_id=2, + rope_theta=1000000.0, + use_sliding_window=False, + max_window_layers=MODEL_SIZES[model_size]["num_hidden_layers"], + ) + torch.manual_seed(seed) + model = Qwen3MoeForCausalLM(config) + model.save_pretrained(path, safe_serialization=True) + + +def make_batch(step: int, rank: int, vocab_size: int, seq_len: int) -> ModelItem: + generator = torch.Generator().manual_seed(100_000 + step * 1024 + rank) + tokens = torch.randint(0, vocab_size, (1, seq_len + 1), generator=generator) + input_ids = tokens[:, :-1] + labels = tokens[:, 1:].to(DEVICE) + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) + loss_cfg = CELossConfig() + loss_ctx = loss_cfg.build(data={"shifted_labels": labels}, sp_mesh=None) + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + return ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) + + +def clean_name(name: str) -> str: + return name.replace("_checkpoint_wrapped_module.", "") + + +def param_memory(model: torch.nn.Module) -> dict[str, float]: + expert_bytes = 0 + dense_bytes = 0 + for name, param in model.named_parameters(): + local = param.to_local() if isinstance(param, DTensor) else param + nbytes = local.numel() * local.element_size() + if ".experts" in name: + expert_bytes += nbytes + else: + dense_bytes += nbytes + return {"expert_param_mib": expert_bytes / 2**20, "dense_param_mib": dense_bytes / 2**20} + + +def per_param_grad_norms(model: torch.nn.Module) -> dict[str, float]: + norms: dict[str, float] = {} + for name, param in model.named_parameters(): + if param.grad is None: + continue + if isinstance(param.grad, DTensor): + norm = cal_total_norm([param.grad], norm_type=2.0, foreach=True, dtype=torch.float32) + else: + norm = torch.linalg.vector_norm(param.grad, 2.0, dtype=torch.float32) + norms[clean_name(name)] = float(norm) + return norms + + +def run_mode( + mode: dict[str, Any], hf_path: Path, steps: int, seq_len: int, lr: float, dispatcher: str +) -> dict[str, Any]: + rank = dist.get_rank() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + base_allocated = torch.cuda.memory_allocated() + + model_cfg = Qwen3MoEConfig.from_hf(hf_path) + model_cfg.ep_size = mode["ep"] + model_cfg.dispatcher = dispatcher + model_cfg.compile_cfg = False + model_cfg.mesh_prefix = f"numerics_{mode['name']}" + fsdp_cfg = FSDPConfig( + ep_size=mode["ep"], + decouple_ep_fsdp=mode["decouple"], + hsdp_sharding_size=mode["hsdp"], + torch_compile=False, + ) + optim_cfg = AdamWConfig(lr=lr) + + engine = TrainEngine(model_cfg=model_cfg, optim_cfg=optim_cfg, fsdp_cfg=fsdp_cfg) + engine.from_hf(hf_path=hf_path, strict=True) + torch.cuda.synchronize() + result: dict[str, Any] = { + "mode": mode, + "memory": param_memory(engine.model), + "losses": [], + "grad_norms": [], + "param_grad_norms": {}, + "step_time_s": None, + } + result["memory"]["allocated_after_load_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 + + step_times = [] + for step in range(steps): + torch.cuda.synchronize() + t0 = time.perf_counter() + batch = make_batch(step, rank, model_cfg.vocab_size, seq_len) + info = engine.train_step([batch]) + grad_norm = engine.clip_grad_norm() + if step in GRAD_NORM_STEPS: + result["param_grad_norms"][str(step)] = per_param_grad_norms(engine.model) + engine.step_optimizer(grad_norm) + torch.cuda.synchronize() + step_times.append(time.perf_counter() - t0) + result["losses"].append(float(info["logs_info"]["reduced_llm_loss"])) + result["grad_norms"].append(float(grad_norm)) + if step == 0: + result["memory"]["allocated_after_step0_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 + + warm = step_times[5:] if len(step_times) > 10 else step_times + result["step_time_s"] = sum(warm) / len(warm) + result["memory"]["peak_allocated_mib"] = (torch.cuda.max_memory_allocated() - base_allocated) / 2**20 + result["memory"]["peak_reserved_mib"] = torch.cuda.max_memory_reserved() / 2**20 + + del engine + gc.collect() + torch.cuda.empty_cache() + return result + + +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("--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)) + 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() diff --git a/tests/model/summarize_decoupled_ep_fsdp_numerics.py b/tests/model/summarize_decoupled_ep_fsdp_numerics.py new file mode 100644 index 0000000000..516f5cf97c --- /dev/null +++ b/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() diff --git a/tests/model/test_decoupled_ep_fsdp_mesh.py b/tests/model/test_decoupled_ep_fsdp_mesh.py index 8fd9d3d67c..978ca26ae3 100644 --- a/tests/model/test_decoupled_ep_fsdp_mesh.py +++ b/tests/model/test_decoupled_ep_fsdp_mesh.py @@ -11,8 +11,8 @@ import pytest import torch -from pydantic import ValidationError import torch.distributed as dist +from pydantic import ValidationError from torch.distributed.device_mesh import DeviceMesh, _mesh_resources from torch.distributed.tensor import DTensor, Replicate, Shard from torch.distributed.tensor.placement_types import _StridedShard @@ -221,7 +221,9 @@ def test_hsdp_ep1_layout(self) -> None: assert tuple(model.hsdp_mesh.shape) == (world_size // shard_size, shard_size) assert model.hsdp_mesh.mesh_dim_names == (f"{PREFIX}.hsdp_replicate", f"{PREFIX}.hsdp_shard") assert model.fsdp_mesh is not None - assert _ranks(model.fsdp_mesh) == list(range((rank // shard_size) * shard_size, (rank // shard_size + 1) * shard_size)) + assert _ranks(model.fsdp_mesh) == list( + range((rank // shard_size) * shard_size, (rank // shard_size + 1) * shard_size) + ) for name in DENSE_PARAMS + EXPERT_PARAMS: param = params[name] _assert_placements(param, model.hsdp_mesh.mesh_dim_names, (Replicate(), Shard(0))) @@ -231,3 +233,151 @@ def test_hsdp_with_ep_is_rejected(self) -> None: with pytest.raises(ValidationError, match="HSDP requires expert parallel size to be 1"): FSDPConfig(ep_size=8, hsdp_sharding_size=8) + +def _decoupled_ranks(world_size: int, ep_size: int, dp_shard: int, rank: int) -> dict[str, list]: + replicate = world_size // dp_shard + efsdp = dp_shard // ep_size + block_start = (rank // dp_shard) * dp_shard + ep_start = (rank // ep_size) * ep_size + efsdp_in_block = [(rank % ep_size) + k * ep_size for k in range(efsdp)] + return { + "ep": list(range(ep_start, ep_start + ep_size)), + "efsdp": [block_start + r for r in efsdp_in_block], + "dp_shard": list(range(block_start, block_start + dp_shard)), + "hsdp": [[r * dp_shard + j for j in range(dp_shard)] for r in range(replicate)], + "expert_hsdp": [[r * dp_shard + j for j in efsdp_in_block] for r in range(replicate)], + } + + +class TestDecoupledMoEMeshShapes: + """Mesh bookkeeping of the decoupled (dp2ep) path: ``root = (replicate, efsdp, ep)``.""" + + @pytest.mark.parametrize( + "world_size,ep_size,rank", + [(8, 8, 5), (8, 4, 6), (16, 8, 9), (64, 8, 37), (64, 32, 50)], + ) + def test_mesh_without_hsdp(self, world_size: int, ep_size: int, rank: int) -> None: + expected = _decoupled_ranks(world_size, ep_size, world_size, rank) + with _fake_world(world_size, rank): + model = _build_model(ep_size) + model._init_device_mesh(FSDPConfig(ep_size=ep_size, torch_compile=False, decouple_ep_fsdp=True)) + + root = model._world_mesh + assert root is not None + assert tuple(root.shape) == (1, world_size // ep_size, ep_size) + assert root.mesh_dim_names == (f"{PREFIX}.replicate", f"{PREFIX}.efsdp", f"{PREFIX}.ep") + assert _mesh_resources.get_root_mesh(model.ep_mesh) is root + assert _mesh_resources.get_root_mesh(model.fsdp_mesh) is root + assert _mesh_resources.get_root_mesh(model.expert_fsdp_mesh) is root + + assert model.hsdp_mesh is None + assert model.ep_mesh.mesh_dim_names == (f"{PREFIX}.ep",) + assert _ranks(model.ep_mesh) == expected["ep"] + assert dist.get_process_group_ranks(model.ep_mesh.get_group()) == expected["ep"] + assert model.expert_fsdp_mesh.mesh_dim_names == (f"{PREFIX}.efsdp",) + assert _ranks(model.expert_fsdp_mesh) == expected["efsdp"] + assert model.fsdp_mesh.mesh_dim_names == (f"{PREFIX}.dp_shard",) + assert _ranks(model.fsdp_mesh) == expected["dp_shard"] + assert dist.get_process_group_ranks(model.fsdp_mesh.get_group()) == expected["dp_shard"] + + @pytest.mark.parametrize( + "world_size,ep_size,dp_shard,rank", + [(16, 8, 8, 11), (64, 8, 32, 37), (64, 8, 8, 3)], + ) + def test_mesh_with_hsdp(self, world_size: int, ep_size: int, dp_shard: int, rank: int) -> None: + expected = _decoupled_ranks(world_size, ep_size, dp_shard, rank) + replicate = world_size // dp_shard + with _fake_world(world_size, rank): + model = _build_model(ep_size) + model._init_device_mesh( + FSDPConfig(ep_size=ep_size, torch_compile=False, decouple_ep_fsdp=True, hsdp_sharding_size=dp_shard) + ) + + root = model._world_mesh + assert root is not None + assert tuple(root.shape) == (replicate, dp_shard // ep_size, ep_size) + assert _ranks(model.ep_mesh) == expected["ep"] + assert _ranks(model.fsdp_mesh) == expected["dp_shard"] + + assert model.hsdp_mesh is not None + assert model.hsdp_mesh.mesh_dim_names == (f"{PREFIX}.replicate", f"{PREFIX}.dp_shard") + assert _ranks(model.hsdp_mesh) == expected["hsdp"] + assert model.expert_fsdp_mesh.mesh_dim_names == (f"{PREFIX}.replicate", f"{PREFIX}.efsdp") + assert _ranks(model.expert_fsdp_mesh) == expected["expert_hsdp"] + assert _mesh_resources.get_root_mesh(model.hsdp_mesh) is root + assert _mesh_resources.get_root_mesh(model.expert_fsdp_mesh) is root + + def test_config_requires_ep_to_divide_shard_size(self) -> None: + FSDPConfig(ep_size=8, hsdp_sharding_size=8, decouple_ep_fsdp=True) + FSDPConfig(ep_size=4, hsdp_sharding_size=8, decouple_ep_fsdp=True) + with pytest.raises(ValidationError, match="divisible by `ep_size`"): + FSDPConfig(ep_size=8, hsdp_sharding_size=4, decouple_ep_fsdp=True) + with pytest.raises(ValidationError, match="divisible by `ep_size`"): + FSDPConfig(ep_size=3, hsdp_sharding_size=8, decouple_ep_fsdp=True) + + +class TestDecoupledMoEParamPlacements: + """DTensor placements after the two-level ``fully_shard`` of the decoupled path.""" + + @pytest.mark.parametrize("world_size,ep_size,rank", [(8, 8, 0), (16, 8, 9), (64, 8, 37)]) + def test_param_placements(self, world_size: int, ep_size: int, rank: int) -> None: + efsdp = world_size // ep_size + with _fake_world(world_size, rank): + model = _shard_model(ep_size, decouple_ep_fsdp=True) + params = _named_params(model) + + for name in DENSE_PARAMS: + param = params[name] + # Dense params: plain FSDP over the full dp_shard mesh, no EP replication. + _assert_placements(param, (f"{PREFIX}.dp_shard",), (Shard(0),)) + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + for name in EXPERT_PARAMS: + param = params[name] + _assert_placements( + param, + (f"{PREFIX}.efsdp", f"{PREFIX}.ep"), + (_StridedShard(0, split_factor=ep_size), Shard(0)), + ) + assert param.device_mesh.size(0) == efsdp + assert param.to_local().shape[0] == param.shape[0] // world_size, name + + from torch.distributed.fsdp import FSDPModule + + layer = model.layers["1"] + assert isinstance(layer.experts, FSDPModule) + assert isinstance(layer, FSDPModule) + assert not isinstance(model.layers["0"].mlp, FSDPModule) + + @pytest.mark.parametrize("world_size,ep_size,dp_shard,rank", [(16, 8, 8, 11), (64, 8, 32, 37)]) + def test_param_placements_with_hsdp(self, world_size: int, ep_size: int, dp_shard: int, rank: int) -> None: + efsdp = dp_shard // ep_size + with _fake_world(world_size, rank): + model = _shard_model(ep_size, decouple_ep_fsdp=True, hsdp_sharding_size=dp_shard) + params = _named_params(model) + + for name in DENSE_PARAMS: + param = params[name] + _assert_placements(param, (f"{PREFIX}.replicate", f"{PREFIX}.dp_shard"), (Replicate(), Shard(0))) + assert param.to_local().shape[0] == param.shape[0] // dp_shard, name + + for name in EXPERT_PARAMS: + param = params[name] + _assert_placements( + param, + (f"{PREFIX}.replicate", f"{PREFIX}.efsdp", f"{PREFIX}.ep"), + (Replicate(), _StridedShard(0, split_factor=ep_size), Shard(0)), + ) + assert param.device_mesh.size(1) == efsdp + assert param.to_local().shape[0] == param.shape[0] // dp_shard, name + + def test_legacy_layout_is_untouched_by_default(self) -> None: + # Same construction as the decoupled tests, but with the flag left at its default. + world_size, ep_size, rank = 16, 8, 9 + with _fake_world(world_size, rank): + model = _shard_model(ep_size) + params = _named_params(model) + assert model.expert_fsdp_mesh is None + _assert_placements( + params["layers.1.self_attn.q_proj.weight"], (f"{PREFIX}.fsdp", f"{PREFIX}.ep"), (Shard(0), Replicate()) + ) diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py index 7335d6d7a5..85d32887c5 100644 --- a/xtuner/v1/config/fsdp.py +++ b/xtuner/v1/config/fsdp.py @@ -41,10 +41,24 @@ class FSDPConfig(BaseModel): hsdp_sharding_size: Annotated[ Optional[int], Parameter(help="Sharding size for HSDP (Hybrid Sharding Data Parallel)") ] = None + # Decoupled EP/FSDP ("dp2ep") layout. When enabled, `ep_size` is a sub-dimension of the + # FSDP shard dimension instead of being orthogonal to it: routed experts are sharded + # `dp_shard / ep_size` ways on top of EP, while every other parameter is sharded over the + # full `dp_shard` (= `hsdp_sharding_size` or world size) without being replicated across EP + # ranks. Requires `dp_shard % ep_size == 0`. When disabled, the legacy layout is untouched. + decouple_ep_fsdp: Annotated[ + bool, Parameter(help="Decouple expert parallel from FSDP: shard dense params over the full FSDP mesh") + ] = False def model_post_init(self, __context: Any) -> None: if self.hsdp_sharding_size is not None: - assert self.ep_size == 1, "Currently, HSDP requires expert parallel size to be 1" + if self.decouple_ep_fsdp: + assert self.hsdp_sharding_size % self.ep_size == 0, ( + "`decouple_ep_fsdp` requires `hsdp_sharding_size` to be divisible by `ep_size`, " + f"got hsdp_sharding_size={self.hsdp_sharding_size}, ep_size={self.ep_size}" + ) + else: + assert self.ep_size == 1, "Currently, HSDP requires expert parallel size to be 1" @field_serializer("param_dtype", "reduce_dtype") def serialize_param_dtype(self, value: torch.dtype) -> str: diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 15dc0d628e..44c61ef974 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -201,6 +201,9 @@ class MoE(BaseModel): ep_mesh: DeviceMesh | None = None expert_tp_mesh: DeviceMesh | None = None ep_tp_mesh: DeviceMesh | None = None + # Only set on the decoupled EP/FSDP path: mesh used to `fully_shard` routed experts, + # i.e. `(efsdp,)` or `(replicate, efsdp)` with `efsdp = dp_shard / ep_size`. + expert_fsdp_mesh: DeviceMesh | None = None dense_decoder_layer_cls = DenseDecoderLayer moe_decoder_layer_cls = MoEDecoderLayer mtp_layer_cls = MTPLayer @@ -1276,6 +1279,10 @@ def fully_shard( self._init_device_mesh(fsdp_config) if self.config.float8_cfg is not None: + if self.fsdp_config.decouple_ep_fsdp: + # Experts and dense params are sharded on meshes of different sizes; the fp8 + # padding / tile-wise reduce meshes need per-class shard sizes (PR-3). + raise NotImplementedError("float8 training with `decouple_ep_fsdp` is not supported yet") # As we modify the shape of the model's parameters, # we need to reinitialize the load spec mapping. Float8Handler.pad_for_fsdp(self, cast(DeviceMesh, self.fsdp_mesh), callback_after_pad=self._init_load_spec) @@ -1296,9 +1303,13 @@ def fully_shard( param.requires_grad = False tp_enabled = self.expert_tp_mesh is not None and self.expert_tp_mesh.size() > 1 - if self.ep_mesh.size() > 1 or tp_enabled: + decoupled = self.fsdp_config.decouple_ep_fsdp + if not decoupled and (self.ep_mesh.size() > 1 or tp_enabled): # 中文注释：不开 EP 但开启 expert TP 时，非 expert 参数仍是 TP rank 间的逻辑副本， # 需要显式放到 Replicate DTensor 上，后续梯度才会跨 expert TP 平均。 + # + # On the decoupled path dense params are NOT replicated over EP: they stay plain + # tensors and get sharded over the full `dp_shard` mesh by FSDP below. self._replicate_other_params(self) # Although rotary_emb was already constructed in __init__, it was built on the meta device. @@ -1313,6 +1324,18 @@ def fully_shard( for layer_idx, layer in tqdm(self.layers.items(), desc="[FSDP Sharding]"): layer_idx = int(layer_idx) + if decoupled: + # Routed experts get their own FSDP group on `expert_fsdp_mesh` (efsdp) before the + # layer is wrapped on the dense mesh; FSDP then excludes them from the layer group. + self._fully_shard_expert_blocks( + layer, + mp_policy=mp_policy, + reshard_after_forward=( + False + if layer_idx >= len(self.layers) - 1 and self.mtp_block is None + else self.fsdp_config.reshard_after_forward + ), + ) if self._should_recompute( layer_idx=layer_idx, mtp_idx=None, @@ -1340,7 +1363,17 @@ def fully_shard( list(self.layers.values())[:-1], list(self.layers.values())[1:], ): - layer_cur.set_modules_to_forward_prefetch([layer_next]) # type: ignore + if decoupled: + # Issue the expert all-gather of the current layer together with the next layer's + # dense all-gather so it overlaps with attention instead of stalling the MoE block. + layer_cur.set_modules_to_forward_prefetch( # type: ignore + [*self._expert_blocks(layer_cur), layer_next] + ) + else: + layer_cur.set_modules_to_forward_prefetch([layer_next]) # type: ignore + if decoupled: + last_layer = list(self.layers.values())[-1] + last_layer.set_modules_to_forward_prefetch(self._expert_blocks(last_layer)) # type: ignore self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, @@ -1379,6 +1412,10 @@ def fully_shard( self.mtp_block.layers[mtp_idx] = mtp_layer reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1 + if decoupled: + self._fully_shard_expert_blocks( + mtp_layer, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward + ) self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, mp_policy=mp_policy, @@ -1387,14 +1424,27 @@ def fully_shard( module=mtp_layer, ) if mtp_idx == 0: - layer_next.set_modules_to_forward_prefetch([mtp_layer]) # type: ignore + if decoupled: + layer_next.set_modules_to_forward_prefetch( # type: ignore + [*self._expert_blocks(layer_next), mtp_layer] + ) + else: + layer_next.set_modules_to_forward_prefetch([mtp_layer]) # type: ignore if self.config.mtp_config is not None and self.config.mtp_config.num_layers > 0: for prev_mtp_layer, next_mtp_layer in zip( list(self.mtp_block.layers)[:-1], list(self.mtp_block.layers)[1:], ): - prev_mtp_layer.set_modules_to_forward_prefetch([next_mtp_layer]) # type: ignore + if decoupled: + prev_mtp_layer.set_modules_to_forward_prefetch( # type: ignore + [*self._expert_blocks(prev_mtp_layer), next_mtp_layer] + ) + else: + prev_mtp_layer.set_modules_to_forward_prefetch([next_mtp_layer]) # type: ignore + if decoupled: + last_mtp_layer = list(self.mtp_block.layers)[-1] + last_mtp_layer.set_modules_to_forward_prefetch(self._expert_blocks(last_mtp_layer)) # type: ignore self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, @@ -1426,6 +1476,10 @@ def need_update_bias(self) -> bool: @torch.no_grad # type: ignore def scale_and_reduce_grad(self): + if self.fsdp_config is not None and self.fsdp_config.decouple_ep_fsdp: + self._scale_and_reduce_grad_decoupled() + return + # Bucket gradients that need a cross-rank reduction by their target process # group. Each bucket is reduced with a single coalesced NCCL all_reduce # instead of one launch per parameter, which used to dominate latency for @@ -1518,6 +1572,10 @@ def cal_grad_norm(self, grads: list[DTensor], dtype=torch.float32): def _init_device_mesh(self, fsdp_config: FSDPConfig): self.fsdp_config = fsdp_config + if self.fsdp_config.decouple_ep_fsdp: + self._init_decoupled_device_mesh(fsdp_config) + return + device = DEVICE world_size = dist.get_world_size() expert_tp_size = self.config.expert_tp_size if self.config.expert_tp_size > 1 else 1 @@ -1608,6 +1666,140 @@ def _init_device_mesh(self, fsdp_config: FSDPConfig): ) self.fsdp_mesh = self.hsdp_mesh[f"{self.config.mesh_prefix}.hsdp_shard"] + def _scale_and_reduce_grad_decoupled(self) -> None: + # Invariant (DESIGN §4.4): every gradient ends up as the mean over the full data-parallel + # world, independent of ep / efsdp. + # * dense params are FSDP-sharded over `dp_shard` (+ `replicate` with HSDP); FSDP's + # reduce-scatter / all-reduce already average them -> nothing to do; + # * routed experts are reduce-scattered over `efsdp = dp_shard / ep` only, while each + # expert already saw the tokens of its whole EP group -> divide by `ep`; + # * DTensors replicated on every mesh dim (fp32 params kept out of FSDP via + # `fp32_keys_pattern`) are averaged manually, as on the legacy path. + ep_size = self.ep_mesh.size() if self.ep_mesh is not None else 1 + grads_by_group: dict[dist.ProcessGroup, list[torch.Tensor]] = {} + + for name, param in self.trainable_parameters(): + if param.grad is None: + continue + + if ep_size > 1 and ".experts" in name: + param.grad.div_(ep_size) # type: ignore + continue + + if not isinstance(param, DTensor): + continue + + if any(isinstance(placement, Shard) for placement in param.placements): + # FSDP-managed (`_StridedShard` is a `Shard` subclass). + continue + + replicate_dim_names = tuple( + param.device_mesh.mesh_dim_names[i] for i, p in enumerate(param.placements) if isinstance(p, Replicate) + ) + if not replicate_dim_names: + continue + + if len(replicate_dim_names) > 1: + flat_mesh = param.device_mesh[replicate_dim_names]._flatten() + else: + flat_mesh = param.device_mesh[replicate_dim_names[0]] + + grad = param.grad.to_local() if isinstance(param.grad, DTensor) else param.grad + grad.div_(flat_mesh.size()) # type: ignore + grads_by_group.setdefault(flat_mesh.get_group(), []).append(grad) # type: ignore + + for group, grads in grads_by_group.items(): + with dist._coalescing_manager(group=group): + for grad in grads: + dist.all_reduce(grad, ReduceOp.SUM, group=group) + + @staticmethod + def _expert_blocks(module: nn.Module) -> list[nn.Module]: + return [submodule for submodule in module.modules() if isinstance(submodule, MoEBlock)] + + def _fully_shard_expert_blocks( + self, + module: nn.Module, + mp_policy: MixedPrecisionPolicy, + reshard_after_forward: bool, + ) -> None: + assert self.expert_fsdp_mesh is not None + assert self.fsdp_config is not None + for expert_block in self._expert_blocks(module): + self._fully_shard( + mesh=self.expert_fsdp_mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, + module=expert_block, + ) + + def _init_decoupled_device_mesh(self, fsdp_config: FSDPConfig) -> None: + # Decoupled ("dp2ep") layout: EP is a sub-dimension of the FSDP shard dimension. + # + # root = (replicate, efsdp, ep) replicate = world / dp_shard (1 without HSDP) + # efsdp = dp_shard / ep + # dense params: FSDP over flatten(efsdp, ep) = dp_shard (+ replicate for HSDP) + # expert params: Shard(0) over ep (built in GroupedLinear) + FSDP over efsdp + # (+ replicate for HSDP) + # + # `ep` stays the innermost dimension so the EP group is still made of contiguous + # (intra-node) ranks, exactly like the legacy `(fsdp, ep)` root mesh. Every sub-mesh is + # derived from this single root so FSDP accepts the EP-sharded DTensor expert params. + # The `efsdp` dimension is kept even when its size is 1: the FSDP wrapping of the + # experts is what applies the mixed-precision policy to them. + if self.config.expert_tp_size > 1: + raise NotImplementedError("`decouple_ep_fsdp` with ExpertTP is not supported") + + device = DEVICE + world_size = dist.get_world_size() + ep_size = fsdp_config.ep_size + dp_shard = fsdp_config.hsdp_sharding_size if fsdp_config.hsdp_sharding_size is not None else world_size + assert world_size % dp_shard == 0, ( + f"world_size ({world_size}) must be divisible by hsdp_sharding_size ({dp_shard})" + ) + assert dp_shard % ep_size == 0, ( + f"`decouple_ep_fsdp` requires the FSDP shard size ({dp_shard}) to be divisible by ep_size ({ep_size})" + ) + replicate_size = world_size // dp_shard + efsdp_size = dp_shard // ep_size + + prefix = self.config.mesh_prefix + replicate_name = f"{prefix}.replicate" + efsdp_name = f"{prefix}.efsdp" + ep_name = f"{prefix}.ep" + dp_shard_name = f"{prefix}.dp_shard" + + root_mesh = init_device_mesh( + device, + (replicate_size, efsdp_size, ep_size), + mesh_dim_names=(replicate_name, efsdp_name, ep_name), + ) + self._world_mesh = root_mesh + + # Same requirement as the legacy path: the `ep_mesh` created in `__init__` must map to the + # `ep` dimension of this root mesh (see the comment in `_init_device_mesh`). + new_ep_mesh = root_mesh[ep_name] + if self.ep_mesh is not None: + assert new_ep_mesh.mesh_dim_names == self.ep_mesh.mesh_dim_names, ( + f"FSDP enabled, it requires the name of new created `ep_mesh`: {new_ep_mesh.mesh_dim_names}" + f"equals to the origin one: {self.ep_mesh.mesh_dim_names}" + ) + assert torch.equal(self.ep_mesh.mesh, new_ep_mesh.mesh), ( + "FSDP enabled, it requires the `ep_size` of model config equals to the `ep_size` of FSDPConfig." + ) + else: + self.ep_mesh = new_ep_mesh + + # `fsdp_mesh` keeps its meaning of "the 1D shard group of dense parameters". + self.fsdp_mesh = root_mesh[efsdp_name, ep_name]._flatten(dp_shard_name) + if replicate_size > 1: + self.hsdp_mesh = root_mesh[replicate_name, dp_shard_name] + self.expert_fsdp_mesh = root_mesh[replicate_name, efsdp_name] + else: + self.hsdp_mesh = None + self.expert_fsdp_mesh = root_mesh[efsdp_name] + def _replicate_other_params(self, model: nn.Module): def traverse(module: nn.Module) -> None: if isinstance(module, MoEBlock): From f101de11d50261edaee7b62da8b978b620738128 Mon Sep 17 00:00:00 2001 From: daixu Date: Fri, 21 Aug 2026 23:30:02 +0000 Subject: [PATCH 03/14] [Feature] Float8, HF/DCP and RL weight-sync support for the decoupled EP/FSDP layout - Float8Handler: per-parameter-class fp8 padding chunk counts and tile-wise reduce meshes (routed experts follow the efsdp shard stride, dense params the flattened dp_shard); legacy single-mesh path unchanged. - BaseModel._fsdp_foreach_allgather picks the FSDP gather group per LoadSpec so EP-local expert slices are still reconstructed for RL weight sync. - examples/v1/config/sft_glm5p2.py: DECOUPLE_EP_FSDP / HSDP_SHARDING_SIZE env switches. - L0 fake-PG tests for the fp8 meshes, the L3 DCP/HF checkpoint script and reports/L3.md (bit-exact HF export, DCP resume, cross-layout DCP reshard, fp8 numerics within the fp8 run-to-run noise floor). --- examples/v1/config/sft_glm5p2.py | 5 + tests/model/run_decoupled_ep_fsdp_ckpt.py | 222 ++++++++++++++++++ tests/model/run_decoupled_ep_fsdp_numerics.py | 23 +- tests/model/test_decoupled_ep_fsdp_mesh.py | 72 ++++++ xtuner/v1/float8/float8_handler.py | 121 +++++++++- xtuner/v1/float8/fsdp_utils.py | 6 +- xtuner/v1/model/base.py | 29 ++- xtuner/v1/model/moe/moe.py | 20 +- 8 files changed, 473 insertions(+), 25 deletions(-) create mode 100644 tests/model/run_decoupled_ep_fsdp_ckpt.py diff --git a/examples/v1/config/sft_glm5p2.py b/examples/v1/config/sft_glm5p2.py index 2b991b92c0..947f21a342 100644 --- a/examples/v1/config/sft_glm5p2.py +++ b/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( diff --git a/tests/model/run_decoupled_ep_fsdp_ckpt.py b/tests/model/run_decoupled_ep_fsdp_ckpt.py new file mode 100644 index 0000000000..06f505c66a --- /dev/null +++ b/tests/model/run_decoupled_ep_fsdp_ckpt.py @@ -0,0 +1,222 @@ +"""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. + + 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 gc +import json +import os +import shutil +import sys +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from safetensors.torch import load_file + + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from run_decoupled_ep_fsdp_numerics import build_hf_checkpoint, make_batch, parse_mode # noqa: E402 + +from xtuner.v1.config import AdamWConfig, FSDPConfig # noqa: E402 +from xtuner.v1.engine.train_engine import TrainEngine # noqa: E402 +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig # noqa: E402 + + +def build_engine(mode: dict[str, Any], hf_path: Path, tag: str) -> TrainEngine: + model_cfg = Qwen3MoEConfig.from_hf(hf_path) + model_cfg.ep_size = mode["ep"] + model_cfg.dispatcher = "all2all" + model_cfg.compile_cfg = False + model_cfg.mesh_prefix = f"ckpt_{mode['name']}_{tag}" + fsdp_cfg = FSDPConfig( + ep_size=mode["ep"], + decouple_ep_fsdp=mode["decouple"], + hsdp_sharding_size=mode["hsdp"], + torch_compile=False, + ) + return TrainEngine(model_cfg=model_cfg, optim_cfg=AdamWConfig(lr=1e-4), fsdp_cfg=fsdp_cfg) + + +def train(engine: TrainEngine, steps: range, vocab_size: int, seq_len: int) -> list[float]: + losses = [] + for step in steps: + info = engine.train_step([make_batch(step, dist.get_rank(), vocab_size, seq_len)]) + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + losses.append(float(info["logs_info"]["reduced_llm_loss"])) + return losses + + +def load_hf_dir(path: Path) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + for file in sorted(path.glob("*.safetensors")): + tensors.update(load_file(str(file))) + return tensors + + +def compare_hf(lhs: Path, rhs: Path) -> dict[str, Any]: + a = load_hf_dir(lhs) + b = load_hf_dir(rhs) + missing = sorted(set(a) - set(b)) + extra = sorted(set(b) - set(a)) + max_abs = 0.0 + max_rel = 0.0 + mismatched = 0 + total = 0 + worst = "" + for key in sorted(set(a) & set(b)): + x = a[key].to(torch.float32) + y = b[key].to(torch.float32) + if x.shape != y.shape: + return {"error": f"shape mismatch for {key}: {tuple(x.shape)} vs {tuple(y.shape)}"} + diff = (x - y).abs() + cur = float(diff.max()) + if cur > max_abs: + max_abs, worst = cur, key + max_rel = max(max_rel, float((diff / x.abs().clamp_min(1e-6)).max())) + mismatched += int((diff > 0).sum()) + total += diff.numel() + return { + "keys": len(set(a) & set(b)), + "missing_in_rhs": missing[:5], + "extra_in_rhs": extra[:5], + "max_abs_diff": max_abs, + "max_rel_diff": max_rel, + "mismatched_elements": mismatched, + "total_elements": total, + "worst_key": worst, + } + + +def release(engine: TrainEngine) -> None: + del engine + gc.collect() + torch.cuda.empty_cache() + + +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) + from safetensors.torch import save_file + + 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, "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, "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( + abs(a - b) / abs(b) for a, b in zip(losses_resumed, losses_cont, strict=True) + ) + 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"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(abs(a - b) / abs(b) for a, b in zip(losses, ref, strict=True)), + } + + 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() diff --git a/tests/model/run_decoupled_ep_fsdp_numerics.py b/tests/model/run_decoupled_ep_fsdp_numerics.py index 0b5e026292..ed819277c9 100644 --- a/tests/model/run_decoupled_ep_fsdp_numerics.py +++ b/tests/model/run_decoupled_ep_fsdp_numerics.py @@ -27,6 +27,7 @@ from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM from xtuner.v1.config import AdamWConfig, FSDPConfig from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.float8.config import Float8Config, ScalingGranularity from xtuner.v1.loss.ce_loss import CELossConfig from xtuner.v1.model.base import ModelItem from xtuner.v1.model.moe.moe import SequenceContext @@ -145,7 +146,7 @@ def per_param_grad_norms(model: torch.nn.Module) -> dict[str, float]: def run_mode( - mode: dict[str, Any], hf_path: Path, steps: int, seq_len: int, lr: float, dispatcher: str + mode: dict[str, Any], hf_path: Path, steps: int, seq_len: int, lr: float, dispatcher: str, fp8: bool ) -> dict[str, Any]: rank = dist.get_rank() torch.cuda.reset_peak_memory_stats() @@ -157,12 +158,16 @@ def run_mode( model_cfg.dispatcher = dispatcher model_cfg.compile_cfg = False model_cfg.mesh_prefix = f"numerics_{mode['name']}" - fsdp_cfg = FSDPConfig( - ep_size=mode["ep"], - decouple_ep_fsdp=mode["decouple"], - hsdp_sharding_size=mode["hsdp"], - torch_compile=False, - ) + if fp8: + model_cfg.float8_cfg = Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ) + fsdp_kwargs: dict[str, Any] = {"ep_size": mode["ep"], "hsdp_sharding_size": mode["hsdp"], "torch_compile": False} + if mode["decouple"]: + # Only pass the flag when set so the script also runs against checkouts that predate it. + fsdp_kwargs["decouple_ep_fsdp"] = True + fsdp_cfg = FSDPConfig(**fsdp_kwargs) optim_cfg = AdamWConfig(lr=lr) engine = TrainEngine(model_cfg=model_cfg, optim_cfg=optim_cfg, fsdp_cfg=fsdp_cfg) @@ -197,6 +202,7 @@ def run_mode( warm = step_times[5:] if len(step_times) > 10 else step_times result["step_time_s"] = sum(warm) / len(warm) + result["step_times_s"] = step_times result["memory"]["peak_allocated_mib"] = (torch.cuda.max_memory_allocated() - base_allocated) / 2**20 result["memory"]["peak_reserved_mib"] = torch.cuda.max_memory_reserved() / 2**20 @@ -215,6 +221,7 @@ def main() -> None: 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() @@ -234,7 +241,7 @@ def main() -> None: 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)) + 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] diff --git a/tests/model/test_decoupled_ep_fsdp_mesh.py b/tests/model/test_decoupled_ep_fsdp_mesh.py index 978ca26ae3..06161848f6 100644 --- a/tests/model/test_decoupled_ep_fsdp_mesh.py +++ b/tests/model/test_decoupled_ep_fsdp_mesh.py @@ -381,3 +381,75 @@ def test_legacy_layout_is_untouched_by_default(self) -> None: _assert_placements( params["layers.1.self_attn.q_proj.weight"], (f"{PREFIX}.fsdp", f"{PREFIX}.ep"), (Shard(0), Replicate()) ) + + +class TestDecoupledFloat8Meshes: + """Tile-wise fp8 bookkeeping on the decoupled path: per-class FSDP chunk counts and reduce meshes.""" + + def test_reduce_meshes_follow_each_class_shard_stride(self) -> None: + from xtuner.v1.float8.config import Float8Config, ScalingGranularity + + world_size, ep_size, rank = 64, 8, 37 + with _fake_world(world_size, rank): + config = _build_model(ep_size).config.model_copy( + update={ + "hidden_size": 256, + "n_routed_experts": 16, + "moe_intermediate_size": 128, + "attention": MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=64), + "float8_cfg": Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ), + } + ) + with torch.device("meta"): + model = MoE(config) + model.fully_shard(FSDPConfig(ep_size=ep_size, torch_compile=False, decouple_ep_fsdp=True)) + params = _named_params(model) + handler = model.float8_handler + assert handler is not None + + # dense q_proj (256, 256) is sharded 64 ways -> 4 local rows -> absmax reduced over 32 contiguous ranks + dense_local = tuple(params["layers.1.self_attn.q_proj.weight"].to_local().shape) + assert dense_local == (4, 256) + assert _ranks(handler.tilewise_reduce_mesh_mapping[dense_local]) == list(range(32, 64)) + # experts: 16 experts x 256 rows = 4096 rows sharded (efsdp=8) x (ep=8) ways -> 64 local rows + # -> absmax reduced over 2 ranks that are `ep` apart (neighbours on the efsdp dim) + expert_local = tuple(params["layers.1.experts.fused_w1w3.weight"].to_local().shape) + assert expert_local == (64, 256) + assert handler.expert_tilewise_reduce_mesh_mapping is not None + assert _ranks(handler.expert_tilewise_reduce_mesh_mapping[expert_local]) == [37, 45] + assert dense_local not in handler.expert_tilewise_reduce_mesh_mapping + assert expert_local not in handler.tilewise_reduce_mesh_mapping + + # "n*128+64" straddling blocks: dense neighbours are adjacent ranks, expert neighbours are `ep` apart + assert _ranks(handler.tilewise_reduce_mesh_devided_64) == [36, 37] + assert _ranks(handler.expert_tilewise_reduce_mesh_devided_64) == [37, 45] + + def test_legacy_float8_meshes_unchanged(self) -> None: + from xtuner.v1.float8.config import Float8Config, ScalingGranularity + + world_size, ep_size, rank = 64, 8, 37 + with _fake_world(world_size, rank): + config = _build_model(ep_size).config.model_copy( + update={ + "hidden_size": 256, + "n_routed_experts": 16, + "moe_intermediate_size": 128, + "attention": MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=64), + "float8_cfg": Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ), + } + ) + with torch.device("meta"): + model = MoE(config) + model.fully_shard(FSDPConfig(ep_size=ep_size, torch_compile=False)) + handler = model.float8_handler + assert handler is not None + assert handler.expert_tilewise_reduce_mesh_mapping is None + assert handler.expert_tilewise_reduce_mesh_devided_64 is None + # legacy: every class is sharded on the same (fsdp=8, stride ep=8) mesh + assert _ranks(handler.tilewise_reduce_mesh_devided_64) == [37, 45] diff --git a/xtuner/v1/float8/float8_handler.py b/xtuner/v1/float8/float8_handler.py index c6ecd2c66a..510483995a 100644 --- a/xtuner/v1/float8/float8_handler.py +++ b/xtuner/v1/float8/float8_handler.py @@ -42,6 +42,11 @@ class Float8Handler: fsdp_mesh: Optional[DeviceMesh] = None tilewise_reduce_mesh_devided_64: Optional[DeviceMesh] = None tilewise_reduce_mesh_mapping: Dict[Tuple[int, int], DeviceMesh] = {} + # Decoupled EP/FSDP layout only: routed experts are FSDP-sharded on `expert_fsdp_mesh` + # (efsdp, strided by ep) instead of `fsdp_mesh`, so they need their own reduce meshes. + expert_fsdp_mesh: Optional[DeviceMesh] = None + expert_tilewise_reduce_mesh_devided_64: Optional[DeviceMesh] = None + expert_tilewise_reduce_mesh_mapping: Optional[Dict[Tuple[int, int], DeviceMesh]] = None def __init__( self, @@ -113,7 +118,12 @@ def get_shard_size_on_dim(tensor: torch.Tensor | DTensor, dim: int) -> int: return RuntimeLayout.from_dtensor(tensor).shard_size(dim) @staticmethod - def pad_for_fsdp(model: nn.Module, fsdp_mesh: DeviceMesh, callback_after_pad: Callable | None = None): + def pad_for_fsdp( + model: nn.Module, + fsdp_mesh: DeviceMesh, + callback_after_pad: Callable | None = None, + expert_fsdp_mesh: DeviceMesh | None = None, + ): from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear from xtuner.v1.float8.float8_linear_tensor_wise import TensorWiseFloat8Linear from xtuner.v1.float8.float8_linear_tile_wise import TileWiseFloat8Linear @@ -131,20 +141,102 @@ def pad_for_fsdp(model: nn.Module, fsdp_mesh: DeviceMesh, callback_after_pad: Ca else: tensor_size = module.weight.size() parallel_size = 1 - padded_out_features = Float8Handler.get_num_features_after_pad(tensor_size, 0, fsdp_mesh.size(-1)) + num_fsdp_chunks = fsdp_mesh.size(-1) + if expert_fsdp_mesh is not None and isinstance(module, TileWiseFloat8GroupedLinear): + # Decoupled EP/FSDP: routed experts are sharded `efsdp` ways, not `dp_shard` ways. + num_fsdp_chunks = expert_fsdp_mesh.size(-1) + padded_out_features = Float8Handler.get_num_features_after_pad(tensor_size, 0, num_fsdp_chunks) padded_out_features *= parallel_size module.pad_for_fsdp(padded_out_features=padded_out_features) if callback_after_pad is not None: callback_after_pad() - def build_reduce_mesh(self, model: nn.Module, fsdp_mesh: DeviceMesh): + def build_reduce_mesh(self, model: nn.Module, fsdp_mesh: DeviceMesh, expert_fsdp_mesh: DeviceMesh | None = None): self.fsdp_mesh = fsdp_mesh + self.expert_fsdp_mesh = expert_fsdp_mesh if self.is_tilewise_fp8: + if expert_fsdp_mesh is not None: + self._build_decoupled_reduce_meshes(model, fsdp_mesh, expert_fsdp_mesh) + return if fsdp_mesh.size(-1) >= 2: self._build_reduce_mesh_devided_64(fsdp_mesh) self._build_reduce_mesh_mapping(model, fsdp_mesh) + def _build_decoupled_reduce_meshes( + self, model: nn.Module, fsdp_mesh: DeviceMesh, expert_fsdp_mesh: DeviceMesh + ) -> None: + # Decoupled EP/FSDP layout (root mesh `(replicate, efsdp, ep)`): + # * dense fp8 weights are sharded over `dp_shard` = flatten(efsdp, ep): consecutive FSDP + # shards live on consecutive ranks (stride 1); + # * routed-expert fp8 weights are sharded over `efsdp`, whose ranks are `ep` apart. + # Each class gets its own "reduce max" meshes, built with the matching rank stride. + from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear + from xtuner.v1.float8.float8_linear_tile_wise import TileWiseFloat8Linear + + world_size = dist.get_world_size() + dense_shard_size = fsdp_mesh.size(-1) + expert_shard_size = expert_fsdp_mesh.size(-1) + expert_stride = world_size // expert_fsdp_mesh.size() # == ep_size + + self.tilewise_reduce_mesh_devided_64 = ( + self._build_strided_reduce_mesh(2, 1) if dense_shard_size >= 2 and dense_shard_size % 2 == 0 else None + ) + self.expert_tilewise_reduce_mesh_devided_64 = ( + self._build_strided_reduce_mesh(2, expert_stride) + if expert_shard_size >= 2 and expert_shard_size % 2 == 0 + else None + ) + self.tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8Linear,), dense_shard_size, 1 + ) + self.expert_tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8GroupedLinear,), expert_shard_size, expert_stride + ) + + @staticmethod + def _build_strided_reduce_mesh(num_ranks: int, stride: int) -> DeviceMesh: + # Groups of `num_ranks` ranks spaced `stride` apart, e.g. (r, r + stride, ...). + world_size = dist.get_world_size() + assert world_size % (num_ranks * stride) == 0, (world_size, num_ranks, stride) + return init_device_mesh( + "cuda", + (world_size // (num_ranks * stride), num_ranks, stride), + mesh_dim_names=("_", "tilewise_reduce", "ep_or_tp"), + )["tilewise_reduce"] + + def _build_strided_reduce_mesh_mapping( + self, model: nn.Module, module_types: Tuple[type, ...], shard_size: int, stride: int + ) -> Dict[Tuple[int, int], DeviceMesh]: + SHARD_DIM = 0 + mapping: Dict[Tuple[int, int], DeviceMesh] = {} + for module in model.modules(): + if not isinstance(module, module_types): + continue + assert isinstance(module.weight, DTensor), ( + "`build_reduce_mesh` should be called after apply fully_shard to the model." + ) + local_shape = module.weight._local_tensor.shape + if local_shape[SHARD_DIM] >= 128: + assert local_shape[SHARD_DIM] % 128 in (0, 64), ( + f"Currently only local_shape[SHARD_DIM] % 128 == 0 or " + f"local_shape[SHARD_DIM] % 128 == 64 is supported, got {local_shape}. Please contact us." + ) + continue + assert 128 % local_shape[SHARD_DIM] == 0, ( + f"Currently only local_shape[SHARD_DIM] % 128 == 0 is supported, got {local_shape}. Please contact us." + ) + reduce_world_size = 128 // local_shape[SHARD_DIM] + if local_shape in mapping: + assert mapping[local_shape].size() == reduce_world_size + continue + assert shard_size >= reduce_world_size and shard_size % reduce_world_size == 0, ( + f"Expect FSDP shard size >= reduce_world_size and shard size % reduce_world_size == 0, " + f"got shard size = {shard_size}, reduce_world_size = {reduce_world_size}. Please contact us." + ) + mapping[local_shape] = self._build_strided_reduce_mesh(reduce_world_size, stride) + return mapping + def _build_reduce_mesh_devided_64(self, fsdp_mesh: DeviceMesh): # 为了支持 moe 参数被 fsdp 和 ep 切成 dout = n * 128 + 64 (n >= 1) 的情况 # fsdp rank 0 的后 64 个 dim 要跟 fsdp rank 1 的前 64 个 dim 共同组成一个 block @@ -221,9 +313,26 @@ def precompute_float8_dynamic_scale_for_fsdp(self, model: Union[nn.Module, List[ for m in models: if self.is_tilewise_fp8: - precompute_tilewise_float8_scale_for_fsdp( - m, self.tilewise_reduce_mesh_mapping, self.tilewise_reduce_mesh_devided_64 - ) + if self.expert_tilewise_reduce_mesh_mapping is None: + precompute_tilewise_float8_scale_for_fsdp( + m, self.tilewise_reduce_mesh_mapping, self.tilewise_reduce_mesh_devided_64 + ) + else: + from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear + from xtuner.v1.float8.float8_linear_tile_wise import TileWiseFloat8Linear + + precompute_tilewise_float8_scale_for_fsdp( + m, + self.tilewise_reduce_mesh_mapping, + self.tilewise_reduce_mesh_devided_64, + module_types=(TileWiseFloat8Linear,), + ) + precompute_tilewise_float8_scale_for_fsdp( + m, + self.expert_tilewise_reduce_mesh_mapping, + self.expert_tilewise_reduce_mesh_devided_64, + module_types=(TileWiseFloat8GroupedLinear,), + ) if self.is_tensorwise_fp8: assert self.fsdp_mesh is not None, "FSDP mesh must be set for tensorwise float8 training." precompute_tensorwise_float8_scale_for_fsdp(m, self.fsdp_mesh) diff --git a/xtuner/v1/float8/fsdp_utils.py b/xtuner/v1/float8/fsdp_utils.py index 8444672fcb..75b8b8372c 100644 --- a/xtuner/v1/float8/fsdp_utils.py +++ b/xtuner/v1/float8/fsdp_utils.py @@ -124,13 +124,17 @@ def precompute_tilewise_float8_scale_for_fsdp( module: nn.Module, reduce_mesh_mapping: Dict[Tuple[int, int], DeviceMesh], # absmax need to be reduced in this group reduce_mesh_devided_64: Optional[DeviceMesh] = None, # All params share the same reduce mesh + module_types: Optional[Tuple[type, ...]] = None, # restrict to these fp8 module classes (default: all) ) -> None: from xtuner.v1.float8 import TileWiseFloat8GroupedLinear, TileWiseFloat8Linear + if module_types is None: + module_types = (TileWiseFloat8Linear, TileWiseFloat8GroupedLinear) + weights: List[WeightWithDynamicTilewiseFloat8CastTensor] = [] for m in module.modules(): if ( - isinstance(m, (TileWiseFloat8Linear, TileWiseFloat8GroupedLinear)) + isinstance(m, module_types) and isinstance(m.weight, DTensor) and isinstance(m.weight._local_tensor, WeightWithDynamicTilewiseFloat8CastTensor) ): diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 50e4ee59dc..39352cbec9 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -564,6 +564,9 @@ class BaseModel(nn.Module): load_spec_mapping: dict[str, LoadSpec] = {} fsdp_mesh: DeviceMesh | None = None hsdp_mesh: DeviceMesh | None = None + # Only set by MoE models on the decoupled EP/FSDP path: the mesh routed experts are + # `fully_shard`ed on, `(efsdp,)` or `(replicate, efsdp)` with `efsdp = dp_shard / ep_size`. + expert_fsdp_mesh: DeviceMesh | None = None fsdp_config: FSDPConfig | None = None config: XTunerBaseModelConfig @@ -1009,7 +1012,7 @@ def float8_handler(self): self._float8_handler = self.config.float8_cfg.build() if self.fsdp_mesh is not None: - self._float8_handler.build_reduce_mesh(self, self.fsdp_mesh) + self._float8_handler.build_reduce_mesh(self, self.fsdp_mesh, expert_fsdp_mesh=self.expert_fsdp_mesh) return self._float8_handler @torch.no_grad() @@ -1917,9 +1920,31 @@ def _fsdp_foreach_allgather( return tensor_list fsdp_group = self.fsdp_mesh.get_group() - save_plan_list = [load_spec.plan_hf_save(gather_process_group=fsdp_group) for load_spec in load_spec_list] + expert_fsdp_group: dist.ProcessGroup | None = None + if self.expert_fsdp_mesh is not None: + # Decoupled EP/FSDP: routed experts are FSDP-sharded on the innermost (`efsdp`) dim of + # `expert_fsdp_mesh`, so their FSDP-only gather must use that group instead. + expert_fsdp_group = self.expert_fsdp_mesh.get_group(self.expert_fsdp_mesh.ndim - 1) + save_plan_list = [ + load_spec.plan_hf_save( + gather_process_group=self._fsdp_gather_group(load_spec, fsdp_group, expert_fsdp_group) + ) + for load_spec in load_spec_list + ] return unshard_tensors_for_hf_save(list(tensor_list), save_plan_list) + def _fsdp_gather_group( + self, + load_spec: LoadSpec, + fsdp_group: dist.ProcessGroup, + expert_fsdp_group: dist.ProcessGroup | None, + ) -> dist.ProcessGroup: + if expert_fsdp_group is None: + return fsdp_group + if any(self._is_same_process_group(shard.group, expert_fsdp_group) for shard in load_spec.shards): + return expert_fsdp_group + return fsdp_group + @staticmethod def _is_same_process_group(left: dist.ProcessGroup, right: dist.ProcessGroup) -> bool: if left is right: diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 44c61ef974..3a48d05568 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -201,9 +201,6 @@ class MoE(BaseModel): ep_mesh: DeviceMesh | None = None expert_tp_mesh: DeviceMesh | None = None ep_tp_mesh: DeviceMesh | None = None - # Only set on the decoupled EP/FSDP path: mesh used to `fully_shard` routed experts, - # i.e. `(efsdp,)` or `(replicate, efsdp)` with `efsdp = dp_shard / ep_size`. - expert_fsdp_mesh: DeviceMesh | None = None dense_decoder_layer_cls = DenseDecoderLayer moe_decoder_layer_cls = MoEDecoderLayer mtp_layer_cls = MTPLayer @@ -1279,13 +1276,20 @@ def fully_shard( self._init_device_mesh(fsdp_config) if self.config.float8_cfg is not None: - if self.fsdp_config.decouple_ep_fsdp: - # Experts and dense params are sharded on meshes of different sizes; the fp8 - # padding / tile-wise reduce meshes need per-class shard sizes (PR-3). - raise NotImplementedError("float8 training with `decouple_ep_fsdp` is not supported yet") # As we modify the shape of the model's parameters, # we need to reinitialize the load spec mapping. - Float8Handler.pad_for_fsdp(self, cast(DeviceMesh, self.fsdp_mesh), callback_after_pad=self._init_load_spec) + if self.fsdp_config.decouple_ep_fsdp: + # Routed experts are padded for `efsdp` FSDP chunks, everything else for `dp_shard`. + Float8Handler.pad_for_fsdp( + self, + cast(DeviceMesh, self.fsdp_mesh), + callback_after_pad=self._init_load_spec, + expert_fsdp_mesh=self.expert_fsdp_mesh, + ) + else: + Float8Handler.pad_for_fsdp( + self, cast(DeviceMesh, self.fsdp_mesh), callback_after_pad=self._init_load_spec + ) # Just for narrowing the type of self.fsdp_mesh and self.ep_mesh assert self.fsdp_mesh is not None From d521615e99c71e5919dfe0e088f01637419a819d Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 3 Sep 2026 03:42:45 +0000 Subject: [PATCH 04/14] [Fix] Gather RL layer-wise weights with the submodule that owns them Root cause: `WeightIterator.iter_layer_batches` (IPC + Turbomind) took the parameters and `LoadSpec`s of a compose model's language tower but ran the FSDP-only all-gather through the outer compose model. The compose model is wrapped on the world mesh, has no `expert_fsdp_mesh`, and never builds a `load_spec_mapping` of its own, so on the decoupled EP/FSDP layout the language tower's `efsdp` expert shards (and, with HSDP, its `dp_shard` dense shards) were treated as shards to preserve and never gathered: Turbomind received rank-local fragments. The vision / projector / non-layer language parameters additionally hit the empty `load_spec_mapping` of the compose model. Fix: resolve the owning submodule (language_model / vision_tower / multi_modal_projector, or the model itself) for every parameter and run `_fsdp_foreach_allgather`, `_to_float8` and the `load_spec_mapping` lookup on that owner. Plain models are unaffected. Test plan: `tests/rl/test_weight_iterator.py::TestLayerBatchesGatherWithParamOwner` builds a compose model whose language tower is EP-sharded and FSDP-sharded on `efsdp` (with and without an HSDP replicate dim) on fake process groups and asserts every streamed tensor is complete (EP slice kept; `efsdp`, `dp_shard` and world shards gathered). It fails before the fix and passes after; the existing weight-iterator, load-spec and L0 mesh tests still pass. --- tests/rl/test_weight_iterator.py | 194 ++++++++++++++++++ xtuner/v1/rl/weight_update/weight_iterator.py | 31 ++- 2 files changed, 219 insertions(+), 6 deletions(-) diff --git a/tests/rl/test_weight_iterator.py b/tests/rl/test_weight_iterator.py index ec7a3f7884..83bdb19fe5 100644 --- a/tests/rl/test_weight_iterator.py +++ b/tests/rl/test_weight_iterator.py @@ -10,6 +10,7 @@ from xtuner.v1.float8.fsdp_utils import WeightWithDynamicTilewiseFloat8CastTensor from xtuner.v1.float8.triton_kernels.per_block_quant_gemm import per_block_quant_torch from xtuner.v1.model.base import BaseModel, HFSaveCfg, XTunerBaseModelConfig +from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.rl.weight_update.data import RolloutBackend, RolloutWeightUpdateInfo, RolloutWeightUpdateTarget from xtuner.v1.rl.weight_update.weight_iterator import WeightIterator from xtuner.v1.utils import get_device @@ -231,3 +232,196 @@ def gather_train_ep_shards( state_dict[f"expert_{index}"].cpu(), torch.tensor([index], dtype=torch.bfloat16), ) + + +def _bf16_param(values: list[int]) -> nn.Parameter: + return nn.Parameter(torch.tensor(values, device=get_device(), dtype=torch.bfloat16)) + + +class _DecoupledLanguageTower(BaseModel): + """Language tower on the decoupled EP/FSDP layout as seen from rank 0. + + ``layers.0.experts.weight`` is EP-sharded and then FSDP-sharded on ``efsdp``; the dense + parameters are FSDP-sharded on ``dp_shard``. Every local block holds ``[0, 1]``. + """ + + def __init__(self, groups: dict[str, dist.ProcessGroup], expert_fsdp_ndim: int) -> None: + super().__init__(XTunerBaseModelConfig()) + layer = nn.Module() + layer.attn = nn.Module() + layer.attn.weight = _bf16_param([0, 1]) + layer.experts = nn.Module() + layer.experts.weight = _bf16_param([0, 1]) + self.layers = nn.ModuleDict({"0": layer}) + self.norm = nn.Module() + self.norm.weight = _bf16_param([0, 1]) + self.fsdp_mesh = SimpleNamespace(get_group=lambda: groups["dp_shard"]) + self.expert_fsdp_mesh = SimpleNamespace(ndim=expert_fsdp_ndim, get_group=lambda dim=None: groups["efsdp"]) + dense_shards = [ShardDescriptor(dim=0, group=groups["dp_shard"])] + self.load_spec_mapping = { + "layers.0.attn.weight": LoadSpec( + name="layers.0.attn.weight", + global_hf_keys=["attn"], + global_shape=(8,), + shards=dense_shards, + local_shape=(2,), + ), + "layers.0.experts.weight": LoadSpec( + name="layers.0.experts.weight", + global_hf_keys=[f"expert_{index}" for index in range(8)], + global_shape=(8,), + fused_dim=0, + shards=[ShardDescriptor(dim=0, group=groups["ep"]), ShardDescriptor(dim=0, group=groups["efsdp"])], + local_shape=(2,), + ), + "norm.weight": LoadSpec( + name="norm.weight", + global_hf_keys=["norm"], + global_shape=(8,), + shards=dense_shards, + local_shape=(2,), + ), + } + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + +class _WorldShardedTower(BaseModel): + """Vision tower / projector stand-in: one parameter FSDP-sharded on the world mesh.""" + + def __init__(self, block_name: str, world_group: dist.ProcessGroup, world_size: int) -> None: + super().__init__(XTunerBaseModelConfig()) + block = nn.Module() + block.weight = _bf16_param(list(range(8 // world_size))) + setattr(self, block_name, block) + self.fsdp_mesh = SimpleNamespace(get_group=lambda: world_group) + self.load_spec_mapping = { + f"{block_name}.weight": LoadSpec( + name=f"{block_name}.weight", + global_hf_keys=[block_name], + global_shape=(8,), + shards=[ShardDescriptor(dim=0, group=world_group)], + local_shape=(8 // world_size,), + ) + } + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + +class _ComposeModel(BaseModel): + """Compose model stand-in: like ``BaseComposeModel.fully_shard`` it is wrapped on the world + mesh and has neither an ``expert_fsdp_mesh`` nor a ``load_spec_mapping`` of its own.""" + + def __init__( + self, + language_model: BaseModel, + vision_tower: BaseModel, + projector: BaseModel, + world_group: dist.ProcessGroup, + ) -> None: + super().__init__(XTunerBaseModelConfig()) + # `iter_layer_batches` only inspects `isinstance(model.config, BaseComposeConfig)`. + self.config = BaseComposeConfig.model_construct() + self.language_model = language_model + self.vision_tower = vision_tower + self.multi_modal_projector = projector + self.fsdp_mesh = SimpleNamespace(get_group=lambda: world_group) + + +class TestLayerBatchesGatherWithParamOwner: + """``iter_layer_batches`` must gather every parameter with the module that owns it. + + Regression test for compose MoE models on the decoupled EP/FSDP layout with the IPC/Turbomind + transport: the outer compose model is wrapped on the world mesh, so gathering with it kept the + language tower's ``efsdp`` (and HSDP ``dp_shard``) shards local and streamed rank-local + fragments instead of complete weights. + """ + + @pytest.mark.parametrize( + ("world_size", "expert_fsdp_ndim"), + [ + pytest.param(4, 1, id="ep2-efsdp2"), + pytest.param(8, 2, id="hsdp-replicate2-ep2-efsdp2"), + ], + ) + def test_compose_layer_batches_are_gathered_per_owner( + self, + monkeypatch: pytest.MonkeyPatch, + world_size: int, + expert_fsdp_ndim: int, + ) -> None: + groups = {name: dist.ProcessGroup(dist.HashStore(), 0, 1) for name in ("world", "dp_shard", "ep", "efsdp")} + group_ranks = { + groups["world"]: list(range(world_size)), + groups["dp_shard"]: [0, 1, 2, 3], + groups["ep"]: [0, 1], + groups["efsdp"]: [0, 2], + } + monkeypatch.setattr( + dist, + "get_world_size", + lambda group=None: world_size if group is None else len(group_ranks[group]), + ) + monkeypatch.setattr(dist, "get_rank", lambda group=None: 0) + monkeypatch.setattr(dist, "get_process_group_ranks", lambda group: group_ranks[group]) + + gathered_groups: list[dist.ProcessGroup] = [] + + def gather_contiguous_blocks( + tensor_list: list[torch.Tensor], + group: dist.ProcessGroup, + ) -> list[list[torch.Tensor]]: + # Rank ``r`` of ``group`` holds the r-th contiguous block, each block the size of ours. + gathered_groups.append(group) + return [ + [tensor + rank * tensor.numel() for rank in range(len(group_ranks[group]))] for tensor in tensor_list + ] + + monkeypatch.setattr(load_spec_module, "foreach_all_gather", gather_contiguous_blocks) + + language_model = _DecoupledLanguageTower(groups, expert_fsdp_ndim) + vision_tower = _WorldShardedTower("patch", groups["world"], world_size) + projector = _WorldShardedTower("proj", groups["world"], world_size) + model = _ComposeModel(language_model, vision_tower, projector, groups["world"]) + rollout_info = RolloutWeightUpdateInfo( + rollout_config=cast(Any, SimpleNamespace()), + weight_update_targets=( + RolloutWeightUpdateTarget( + endpoint_rank=0, + update_ranks=(0,), + inference_engine_ranks=(0,), + server_url="http://rollout", + lifecycle_state="active", + ), + ), + train_rank=0, + transport_type="ipc", + backend="turbomind", + ) + iterator = WeightIterator( + config=SimpleNamespace(model_cfg=model.config), + engine=SimpleNamespace(model=model), + rollout_info=rollout_info, + global_hf_keys_mapping_cache={}, + ) + + state_dict = { + name: tensor.cpu() for batch in iterator.iter_layer_batches() for name, tensor in batch.state_dict.items() + } + + full = torch.arange(8, dtype=torch.bfloat16) + expected = { + # The EP shard stays local (experts 0-3 of this EP rank); only the `efsdp` shard is gathered. + "model.language_model.layers.0.mlp.experts.weight": full[:4], + "model.language_model.layers.0.attn.weight": full, + "model.language_model.norm.weight": full, + "model.vision_tower.patch.weight": full, + "model.multi_modal_projector.proj.weight": full, + } + assert set(state_dict) == set(expected) + for name, tensor in expected.items(): + torch.testing.assert_close(state_dict[name], tensor, rtol=0, atol=0) + assert {groups["efsdp"], groups["dp_shard"], groups["world"]} <= set(gathered_groups) + assert groups["ep"] not in gathered_groups diff --git a/xtuner/v1/rl/weight_update/weight_iterator.py b/xtuner/v1/rl/weight_update/weight_iterator.py index 02da41bc47..c908d761d7 100644 --- a/xtuner/v1/rl/weight_update/weight_iterator.py +++ b/xtuner/v1/rl/weight_update/weight_iterator.py @@ -139,12 +139,16 @@ def iter_layer_batches(self): else: dtype = torch.bfloat16 - def get_params(tensor_list, name_list, save_dtype): + def get_params(owner, tensor_list, name_list, save_dtype): + # Gather with the module that owns the parameters, not with the outer model. A + # compose model is wrapped on the world mesh, while its language tower may be + # sharded on its own meshes (EP, `dp_shard`, `efsdp`, HSDP shard); only the owner + # knows which shards are FSDP shards to gather and which EP shards to keep local. _tensor_list, _spec_list = list(zip(*tensor_list)) - fsdp_unshard_tensor_list = model._fsdp_foreach_allgather(_tensor_list, _spec_list) + fsdp_unshard_tensor_list = owner._fsdp_foreach_allgather(_tensor_list, _spec_list) if save_dtype == torch.float8_e4m3fn: runtime_is_float8_list = [is_float8_weight(tensor) for tensor in _tensor_list] - fsdp_unshard_tensor_list, name_list = model._to_float8( + fsdp_unshard_tensor_list, name_list = owner._to_float8( fsdp_unshard_tensor_list, name_list, runtime_is_float8_list, @@ -191,7 +195,7 @@ def get_params(tensor_list, name_list, save_dtype): name = name.replace(".gate.", ".mlp.gate.") name_list.append(name) tensor_list.append((local_tensor, load_spec)) - fsdp_unshard_tensor_list, name_list = get_params(tensor_list, name_list, dtype) + fsdp_unshard_tensor_list, name_list = get_params(language_model, tensor_list, name_list, dtype) state_dict = dict(zip(name_list, fsdp_unshard_tensor_list)) yield WeightUpdateBatch(state_dict) @@ -200,7 +204,10 @@ def get_params(tensor_list, name_list, save_dtype): continue local_tensor = param._local_tensor if isinstance(param, DTensor) else param local_tensor = local_tensor.bfloat16() - load_spec = model.load_spec_mapping.get(name) + owner, owner_name = self._param_owner(model, name) + load_spec = owner.load_spec_mapping.get(owner._clean_param_name(owner_name)) + if load_spec is None: + raise ValueError(f"Internal Error. Parameter {name} not found in load_spec_mapping.") if isinstance(model.config, BaseComposeConfig): if "vision_tower." in name: @@ -220,7 +227,7 @@ def get_params(tensor_list, name_list, save_dtype): name = "model.embed_tokens.weight" tensor_list = [(local_tensor, load_spec)] name_list = [name] - fsdp_unshard_tensor_list, name_list = get_params(tensor_list, name_list, dtype) + fsdp_unshard_tensor_list, name_list = get_params(owner, tensor_list, name_list, dtype) state_dict = dict(zip(name_list, fsdp_unshard_tensor_list)) yield WeightUpdateBatch(state_dict) @@ -228,3 +235,15 @@ def get_params(tensor_list, name_list, save_dtype): yield WeightUpdateBatch({}, finished=True) DEVICE_MODULE.empty_cache() + + @staticmethod + def _param_owner(model: Any, name: str) -> tuple[Any, str]: + # Map a `model.state_dict()` key to the submodule that owns the parameter and to the + # key relative to that submodule, which is how each owner's `load_spec_mapping` is keyed. + # Compose models never build a `load_spec_mapping` of their own; their children do. + if isinstance(model.config, BaseComposeConfig): + for submodule in ("language_model", "vision_tower", "multi_modal_projector"): + prefix = f"{submodule}." + if name.startswith(prefix): + return getattr(model, submodule), name[len(prefix) :] + return model, name From 7df5a7ddad0b2ab73f4b3263660ffb6f07a6e80a Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 3 Sep 2026 06:36:59 +0000 Subject: [PATCH 05/14] [Docs] Add the decoupled EP/FSDP design note for maintainers --- docs/design/decouple_ep_fsdp.md | 177 ++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 docs/design/decouple_ep_fsdp.md diff --git a/docs/design/decouple_ep_fsdp.md b/docs/design/decouple_ep_fsdp.md new file mode 100644 index 0000000000..241ba7960a --- /dev/null +++ b/docs/design/decouple_ep_fsdp.md @@ -0,0 +1,177 @@ +# MoE EP 与 FSDP 解耦（dp2ep 布局） + +Expert Parallelism（EP）与 FSDP 解耦，是把 EP 从"与 FSDP 正交的 mesh 维度"改成"FSDP shard 维度的子维度"（torchtitan 称之为 dp2ep）。解耦后 dense 参数在完整的 data-parallel shard 域上分片、不再沿 EP 复制；routed expert 保持按 EP 切分专家，再只沿剩余的 `efsdp` 维做 FSDP。配置入口是 `FSDPConfig.decouple_ep_fsdp`，默认关闭，关闭时旧路径不受影响。 + +```mermaid +flowchart LR + R["root mesh (replicate, efsdp, ep)"] --> EP["ep_mesh = root[ep]
token dispatch / expert 归属"] + R --> DS["fsdp_mesh = flatten(efsdp, ep) = dp_shard
dense 参数的 FSDP 分片域"] + R --> EF["expert_fsdp_mesh = root[efsdp]
routed expert 在 EP 之上的 FSDP 分片域"] + R -. "replicate > 1 (HSDP)" .-> HS["hsdp_mesh = root[replicate, dp_shard]
expert_fsdp_mesh = root[replicate, efsdp]"] +``` + +## 1. 问题：EP 把 FSDP 分片度锁死 + +旧路径的 MoE 训练 mesh 是二维 `(fsdp, ep)`，且 `fsdp = world_size / ep_size`。两类参数的 data-parallel 语义不同，却共用这一张 mesh： + +| 参数 | 旧路径 placement | 后果 | +|---|---|---| +| routed expert | `Shard(0)` on `ep`，再 FSDP over `fsdp` | 正确，且与新路径数学同构 | +| dense（attention / router / shared expert / embedding / lm_head） | `Replicate()` on `ep`（`_replicate_other_params`），再 FSDP over `fsdp` | 参数、梯度、优化器状态在 EP 维复制 `ep` 份；每步一次额外的 coalesced all-reduce 补梯度 | +| HSDP | `assert ep_size == 1` | HSDP 与 EP 不能共存 | + +EP 越大，dense 侧越亏。以 GLM-5.2-30B 单机 8 卡为例，EP=8 时每卡保存一份完整的 dense 参数与 fp32 优化器状态，legacy 路径 EP=8 的 reserved 显存顶到约 132 GiB 的 allocator 上限（allocated 峰值 120 GiB），每步触发 allocator 释放重试，步时退化到 EP=4 的 7 倍（见 §5）。 + +## 2. 目标布局 + +```text +dp_shard = hsdp_sharding_size or world_size # dense 参数的 FSDP 分片数 +replicate = world_size / dp_shard # HSDP 副本数，无 HSDP 时为 1 +efsdp = dp_shard / ep # routed expert 在 EP 之上的 FSDP 分片数 + +root mesh = (replicate, efsdp, ep) # ep 保持最内维：EP group 仍是连续 rank +dense : FSDP over flatten(efsdp, ep) = dp_shard (+ replicate 上 HSDP) +expert : Shard(0) over ep + FSDP over efsdp (+ replicate 上 HSDP) +``` + +约束从 `fsdp = world / ep`（死）变成 `ep | dp_shard`（活）。routed expert 的数学与旧路径完全相同：旧路径在 `world/ep` 个 rank 上 FSDP，新路径在 `efsdp = dp_shard/ep` 个 rank 上 FSDP，无 HSDP 时二者就是同一组 rank。新旧差异全部集中在 dense 侧。 + +单机 8 卡的三个典型拓扑： + +| 配置 | root mesh | dense 每卡持有 | routed expert 每卡持有 | +|---|---|---|---| +| EP=8（旧） | `(fsdp=1, ep=8)` | 100%（8 份复制） | 1/8 的 expert | +| EP=8，解耦 | `(1, efsdp=1, ep=8)` | 1/8 | 1/8 的 expert（与旧路径相同） | +| EP=4（旧） | `(fsdp=2, ep=4)` | 1/2（4 份复制） | 1/4 的 expert，再在 rank r 与 r+4 之间 FSDP 对半 | +| EP=4，解耦 | `(1, efsdp=2, ep=4)` | 1/8 | 同上（rank r 与 r+4 构成 efsdp group） | +| 16 卡，EP=8，`hsdp_sharding_size=8` | `(replicate=2, efsdp=1, ep=8)` | 节点内 1/8，跨节点复制 | 节点内 EP=8，跨节点复制并由 HSDP 归约 | + +"EP=8 且 FSDP=8"不再是矛盾配置：对 dense 而言 FSDP shard 数是 8；对 expert 而言 `efsdp = 1`，因为 expert 已被 EP 切成 8 份。 + +## 3. 实现 + +### 3.1 配置与校验 + +`FSDPConfig.decouple_ep_fsdp: bool = False`。旧断言"HSDP 要求 `ep_size == 1`"只在开关关闭时保留；开关打开时改为要求 `hsdp_sharding_size % ep_size == 0`，即必须存在整数 `efsdp`。`MoE._init_decoupled_device_mesh` 再次校验 `world_size % dp_shard == 0` 与 `dp_shard % ep_size == 0`。ExpertTP（`expert_tp_size > 1`）与解耦同时开启时抛 `NotImplementedError`，ETP 的 `(Shard, InterleavedShard)` 二维 placement 需要单独设计 mesh 维。 + +### 3.2 单一 root mesh + +`MoE._init_decoupled_device_mesh` 一次性建立 `(replicate, efsdp, ep)` root，所有子 mesh 从它派生： + +- `ep_mesh = root[ep]`：dispatcher / DeepEP 拿到的 EP group 不变，仍是节点内连续 rank； +- `fsdp_mesh = root[efsdp, ep]._flatten("dp_shard")`：语义仍是"dense 参数的一维 shard group"，`_fsdp_foreach_allgather`、`Float8Handler.build_reduce_mesh` 等既有消费者无需改动； +- `hsdp_mesh = root[replicate, dp_shard]`（仅 `replicate > 1`）：dense 的 HSDP mesh，与旧路径契约一致； +- `expert_fsdp_mesh = root[efsdp]` 或 `root[replicate, efsdp]`：新增属性，只在解耦路径上非空，dense 模型恒为 `None`。 + +所有子 mesh 同源很重要：FSDP2 要求已是 DTensor 的参数与 FSDP mesh 共享同一个 root，expert 参数最终要同时带 EP shard 和 `efsdp` 上的 `_StridedShard`。`efsdp = 1` 时该维保留而不是省略——对 expert 做的那次 `fully_shard` 是 mixed precision policy 生效的地方（torchtitan #1324 的同一注释）。 + +### 3.3 两级 `fully_shard` + +解耦路径跳过 `_replicate_other_params`：dense 参数进入 `fully_shard` 时是普通 tensor，由 FSDP 在完整 `dp_shard` 上分片。每个 decoder layer 的包装顺序变为： + +1. `_fully_shard_expert_blocks`：找出该层的所有 `MoEBlock`（只含 routed expert 的 grouped linear 与激活），在 `expert_fsdp_mesh` 上先做一次 `fully_shard`； +2. 再在 dense 的 `fsdp_mesh` / `hsdp_mesh` 上包装整个 layer。FSDP2 会保留内层的 expert wrapper，把 expert 参数从外层 param group 中排除。 + +MTP layer 使用同样的两级包装。forward prefetch 列表同时包含下一层与其 expert block，让 expert 的 all-gather 与 attention 重叠，而不是串行等在 MoE block 前。 + +### 3.4 梯度归约 + +不变式：每个参数的最终梯度等于全 data-parallel world 上的均值，与 `ep` / `efsdp` 取值无关。`MoE._scale_and_reduce_grad_decoupled` 只做三件事： + +| 参数类 | 谁完成归约 | 解耦路径的额外动作 | +|---|---|---| +| dense FSDP 参数（任一 `Shard` placement） | FSDP 在 `dp_shard` 上 reduce-scatter；HSDP 再跨 `replicate` all-reduce | 无（旧路径的手工跨 EP all-reduce 删除） | +| routed expert | FSDP 在 `efsdp` 上 reduce-scatter，但每个 expert 已看过整个 EP group 路由来的 token | `grad.div_(ep)`，与 NeMo AutoModel 的 `ep_ratio = dp_shard / efsdp` 相同 | +| FSDP ignored、各维全 `Replicate` 的 fp32 参数（`fp32_keys_pattern` 命中的 dense 参数） | 无人归约 | 沿所有 `Replicate` 维 flatten 后做一次 coalesced all-reduce 求均值 | + +这也是解耦路径不能复用旧梯度逻辑的原因：旧 dense 是 EP replicate，需要手工跨 EP 归约；新 dense 已由 FSDP 在完整 `dp_shard` 上归约，再做一次就是重复归约。注意第三行与 HSDP 无关：无 HSDP 时 world mesh 仍是三维 `(1, efsdp, ep)`，ignored 参数在三个维度上都是 `Replicate`，仍需这次 all-reduce。 + +### 3.5 FP8 + +tile-wise FP8 需要知道每个本地权重会被多少个 FSDP rank 切分（决定 padding），以及 scale 的 reduce-max 应跨哪些 rank。解耦后 dense 与 expert 的答案不同： + +- dense linear 按 `dp_shard` 的 chunk 数 padding，reduce mesh 的 rank stride 为 1； +- grouped-expert linear 按 `efsdp` 的 chunk 数 padding，reduce mesh 的 rank stride 为 `ep`； +- `Float8Handler.pad_for_fsdp` 接受可选的 `expert_fsdp_mesh`，`build_reduce_mesh` 为两类参数各建一套 tile-wise reduce mesh，scale 预计算按类各执行一次。 + +旧路径 `expert_fsdp_mesh is None`，走原有的单 mesh 代码。tensor-wise FP8 只作用于 dense linear，shard group 仍是 `fsdp_mesh`，无需改动。 + +### 3.6 HF 权重与 DCP + +`LoadSpec.from_tensor` 直接从 DTensor placement 记录 shard 历史（EP 的 `Shard(0)`、`efsdp` 上的 `_StridedShard`、HSDP 的 `Replicate` 不产生 shard）。同步 HF save 按 shard 描述逆向 unshard 后再走模型已有的 HF key / fused-expert 映射；DCP 按 DTensor placement 保存并在载入时 reshard。因此没有为任何模型重写 `state_dict_adapter`，工作量在于让新 mesh 产出正确的 `LoadSpec`。 + +### 3.7 RL 权重同步 + +Turbomind 的 layer-wise 更新只 gather FSDP shard、保留 EP-local 的 expert 切片。`BaseModel._fsdp_foreach_allgather` 按每个 `LoadSpec` 选择 gather group：spec 含 `efsdp` 上的 shard 时用 `expert_fsdp_mesh` 的 group，否则用 `fsdp_mesh` 的 group。 + +compose 模型（如 Qwen3.5-VL MoE）曾有一个 owner 选择错误：`WeightIterator.iter_layer_batches` 从 language tower 取参数与 `LoadSpec`，却用外层 compose 模型的 mesh 做 gather。compose 模型本身在 world mesh 上包装、没有 `expert_fsdp_mesh`，也不维护自己的 `load_spec_mapping`，于是 language tower 的 `efsdp` shard（HSDP 时连 `dp_shard` shard）被当作"应保留"的 shard 而不 gather，Turbomind 收到的是 rank-local 碎片。现在 gather 由参数所属的子模块（language / vision / projector）执行，并从该子模块的 `load_spec_mapping` 取 spec；`tests/rl/test_weight_iterator.py::TestLayerBatchesGatherWithParamOwner` 用假 process group 固定了 `efsdp = 2` 与 HSDP 两种拓扑下每个 tensor 的完整形状与取值。 + +## 4. 与参考实现的对照 + +| 项目 | mesh | expert 的 FSDP 维 | 梯度修正 | HSDP + EP | +|---|---|---|---|---| +| torchtitan（#1324 起） | `(pp, dp_replicate, dp_shard_mod_ep, dp_shard_in_ep, cp, tp)` | `dp_shard_mod_ep = dp_shard·cp / ep`，派生 | FSDP gradient divide factor（#1551） | 首日支持 | +| NeMo AutoModel | `device_mesh(dp_replicate, dp_shard_cp)` + `moe_mesh(ep, ep_shard)` | `ep_shard`，派生 | 显式 `grad.div_(ep_ratio)` | 支持 | +| Megatron-Core（Megatron-FSDP） | expert 参数在 expert data-parallel group 上分片 | `dp·cp / ep`，派生 | 框架内部 | 支持 | +| XTuner 本方案 | `(replicate, efsdp, ep)` | `efsdp = dp_shard / ep`，派生 | 显式 `grad.div_(ep)` | mesh 与数学路径支持，多机未实测 | + +三个参考实现都把 expert 的 FSDP 维当作派生量而不是独立配置，本方案与之一致。实现上采用 torchtitan #1324 的"两次 `fully_shard`"法，对 torch 版本要求最低；`shard_placement_fn` 单次包装、expert dim-1 分片（#1561）留作后续优化。 + +## 5. 验证结论与边界 + +| 项目 | 证据 | 边界 | +|---|---|---| +| mesh / placement / 两级 FSDP | `tests/model/test_decoupled_ep_fsdp_mesh.py`，fake-PG 下 8/16/64 rank 共 29 例；旧布局在开关关闭时逐项固定 | 需要一张空闲 GPU 建 CUDA context，不通信 | +| BF16 训练语义（L1） | tiny Qwen3-MoE 50 步：解耦 EP8 vs 旧 EP8 loss 最大相对差 `1.9e-5`（旧 EP8 vs EP1 为 `2.5e-5`）；3.4B 模型每卡 dense 参数 836 → 105 MiB，峰值 allocated 11861 → 8203 MiB，步时 170 → 158 ms | 单机；reduction order 不同，不是逐 bit 相同 | +| `efsdp > 1`、HSDP + EP（L2） | 单机缩比拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)` 均在噪声内；16/64 rank 形状由 L0 固定 | 真正的双节点 collective 未跑 | +| HF / DCP（L3） | 刚 `from_hf` 后导出 807 个 key 全部 bit-exact；DCP 同布局 resume 与跨布局 reshard 后 loss 连续 | 未覆盖 async save、FP8 checkpoint、GLM/MTP 导出、world size 变化 | +| FP8 | 3.4B 与 GLM-5.2 tile-wise FP8 训练 20 步稳定，差异量级接近同配置重复噪声 | FP8 + HSDP、FP8 + checkpoint 未测 | +| GLM-5.2-30B，8×H200 目标配方 | EP4：1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB；EP8：11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB | EP8 的加速来自离开 allocator 上限（legacy 每步 2–3 次 alloc retry，解耦为 0），不是 collective 本身快 7 倍 | +| RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | +| legacy 兼容 | 开关默认关闭，旧分支执行逻辑未改动，L0 旧布局回归通过 | 未做 base-vs-branch 端到端逐 tensor 对照 | + +详细数据见 `reports/L1.md`、`reports/L2.md`、`reports/L3.md`、`reports/GLM52.md`（不随代码 PR 合入）。 + +## 6. 已知限制与后续工作 + +- `fp32_keys_pattern` 命中 routed expert 时（当前内置 pattern 均不命中），该 expert 只被 `ignored_params` 排除而不会沿 `efsdp` 分片，梯度路径只做 `/ep`，`efsdp > 1` 或 `replicate > 1` 时副本会分叉。短期应显式拒绝；完整方案见 §7.2，与"expert 不做 FSDP"模式是同一实现。 +- expert 的分类在 FSDP 包装时按 `MoEBlock` 类型，在梯度缩放时按参数名含 `.experts`，不是同一事实来源；应在包装时记录 expert 参数身份，梯度、FP8、save / RL 复用。 +- 新拓扑约束目前用 `assert`，应改为 Pydantic validator + `ValueError`。 +- 两次不同 mesh 的 `fully_shard` 与 `DeviceMesh._flatten(name)` 只在 torch 2.8 / 2.9 验证；`pyproject.toml` 声明 `torch>=2.6`，需要补 CI 或为该开关声明更高的最低版本。 +- L1–L3 是人工实验脚本（写 JSON、打印统计），不是会失败的 pytest gate。 +- ExpertTP 与解耦互斥；expert dim-1 分片、`shard_placement_fn` 单次包装留作优化。 + +## 7. 开放问题（请架构师决策） + +### 7.1 是否保留 HSDP 与 EP 的共存 + +保留的成本很低。`replicate` 维由同一个 `init_device_mesh` 调用产生，HSDP 专属代码只有 `hsdp_mesh` / `expert_fsdp_mesh` 的二维选择；`replicate` 上的 all-reduce 由 FSDP2 完成。梯度路径没有任何 HSDP 分支：§3.4 三行规则中，前两行与 HSDP 无关，第三行在无 HSDP 时同样需要（world mesh 仍是三维）。去掉 HSDP 能省的只是几行 mesh 选择和 L2 的两个验证拓扑，梯度处理不会因此变简单。 + +保留的收益是多机拓扑：EP 留在节点内时，`hsdp_sharding_size = 8` 让 dense 的 all-gather / reduce-scatter 也留在 NVLink 域内，跨机只剩一次梯度 all-reduce；这是 torchtitan、AutoModel 都支持的组合，也是 DESIGN 中 64 卡场景的前提。目前 XTuner 对 HSDP 使用确实不多（compose 模型的 `init_world_mesh` 还标着 TODO），且本方案的 HSDP + EP 只在单机缩比拓扑验证过。 + +建议：保留 mesh 与数学路径，在文档中标为"experimental，多机未验证"；若维护者希望收窄支持面，只需在 `FSDPConfig` 中对 `decouple_ep_fsdp and hsdp_sharding_size is not None` 抛错，待双节点验证后再放开，代码不必删除。 + +### 7.2 `efsdp` 是否作为独立配置暴露 + +三个参考实现都不暴露：expert 的 data-parallel 语义要求 `ep × efsdp = dp_shard`，即每个 token 在整个 world 上恰好被计算一次，`efsdp` 因此只能是派生量。Megatron-Core 只暴露 `--expert-model-parallel-size`，expert 的 FSDP 分片域是 `dp·cp / ep`；torchtitan、AutoModel 同理。 + +接入 MoonEP / UltraEP 这类动态冗余 expert 的通信库时，需要的不是另一个 `efsdp` 数值，而是一种新模式："expert 不做 FSDP"。从公开资料看，这两个库在每层 MoE 前按实时负载把热 expert 复制到其它 EP rank 并预取权重，要求 EP-local 的 expert 权重在 MoE 计算窗口之外也是完整可寻址的；而 FSDP 分片的 expert 只在 all-gather 窗口内完整，且 FSDP 单元是整个 `MoEBlock`。当 `ep == dp_shard` 时（如单机 EP=8）`efsdp = 1` 已自动满足；当 `ep < dp_shard` 时，"不做 FSDP"意味着 expert 在 `efsdp` 维上复制，这是 placement 的改变而不是分片数的改变： + +```text +expert_fsdp = True （现状）: Shard(0) on ep + _StridedShard on efsdp FSDP reduce-scatter，再 /ep +expert_fsdp = False（新增）: Shard(0) on ep + Replicate on (replicate, efsdp) 手工 all-reduce 求均值，再 /ep +``` + +后者的梯度处理正是 §3.4 第三行已经存在的"沿 Replicate 维 all-reduce"分支，也正是修复 §6 第一条（fp32 pattern 命中 expert）所需的实现，两件事可以一次做完。代价是 expert 参数与优化器状态在 `efsdp` 维复制、失去 FSDP wrapper 提供的 mixed precision（需手工 cast）。 + +建议：`efsdp` 保持派生；不新增整数配置。等 MoonEP / UltraEP 接入启动时，在 `FSDPConfig` 增加 `expert_fsdp: bool = True`（或 `expert_shard_mode: "fsdp" | "replicate"`）。root mesh 不变，`ep_mesh` 仍是最内层连续 rank 的 group，dispatcher 的替换与本方案正交。 + +## 8. 合入建议 + +按 torchtitan 的演进顺序拆分： + +1. 基础解耦：配置开关、root mesh、两级 `fully_shard`、梯度归约、L0 fake-PG 测试； +2. 生态：FP8 per-class padding / reduce mesh、`LoadSpec` 驱动的 HF / DCP、RL per-spec gather 与 compose owner 修复； +3. 文档与后续：本文、§6 的限制项。 + +`decouple_ep_fsdp` 默认 `False`，两个代码 PR 都不改动旧路径的执行逻辑，出问题时一键回退。 From 9f5eefa920b7f1a77ebdfa08d138b2c11ba532d3 Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 3 Sep 2026 06:59:42 +0000 Subject: [PATCH 06/14] [Docs] Expand the decoupled EP/FSDP design note: code map, evidence links, limitation details --- docs/design/decouple_ep_fsdp.md | 257 ++++++++++++++++++++++++++++---- 1 file changed, 229 insertions(+), 28 deletions(-) diff --git a/docs/design/decouple_ep_fsdp.md b/docs/design/decouple_ep_fsdp.md index 241ba7960a..4bfb9be056 100644 --- a/docs/design/decouple_ep_fsdp.md +++ b/docs/design/decouple_ep_fsdp.md @@ -50,33 +50,149 @@ expert : Shard(0) over ep + FSDP over efsdp (+ replicate 上 HS ## 3. 实现 +行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `5923dc1d`（2026-09-03）。改动的文件与位置： + +| 文件 | 行号 | 内容 | +|---|---|---| +| [`xtuner/v1/config/fsdp.py`](../../xtuner/v1/config/fsdp.py#L48-L65) | 48–65 | `decouple_ep_fsdp` 开关与校验 | +| [`xtuner/v1/model/moe/moe.py`](../../xtuner/v1/model/moe/moe.py#L1651-L1716) | 1651–1716 | `_init_decoupled_device_mesh`：root mesh 与子 mesh | +| | [1486–1492](../../xtuner/v1/model/moe/moe.py#L1486-L1492) | `_init_device_mesh` 分流到新路径 | +| | [1194–1201](../../xtuner/v1/model/moe/moe.py#L1194-L1201) | 解耦路径跳过 `_replicate_other_params` | +| | [1214–1225](../../xtuner/v1/model/moe/moe.py#L1214-L1225)、[1250–1260](../../xtuner/v1/model/moe/moe.py#L1250-L1260) | decoder layer 两级 `fully_shard` 与 forward prefetch | +| | [1328–1360](../../xtuner/v1/model/moe/moe.py#L1328-L1360) | MTP layer 的同样处理 | +| | [1631–1649](../../xtuner/v1/model/moe/moe.py#L1631-L1649) | `_expert_blocks` / `_fully_shard_expert_blocks` | +| | [1392–1396](../../xtuner/v1/model/moe/moe.py#L1392-L1396)、[1583–1629](../../xtuner/v1/model/moe/moe.py#L1583-L1629) | `scale_and_reduce_grad` 分流与 `_scale_and_reduce_grad_decoupled` | +| | [1165–1176](../../xtuner/v1/model/moe/moe.py#L1165-L1176) | FP8 padding 传入 `expert_fsdp_mesh` | +| [`xtuner/v1/model/base.py`](../../xtuner/v1/model/base.py#L567-L569) | 567–569 | 新增 `expert_fsdp_mesh` 属性（dense 模型恒为 `None`） | +| | [1015](../../xtuner/v1/model/base.py#L1015) | float8 handler 建 reduce mesh 时传入 expert mesh | +| | [1916–1947](../../xtuner/v1/model/base.py#L1916-L1947) | `_fsdp_foreach_allgather` / `_fsdp_gather_group`：RL 按 spec 选 gather group | +| [`xtuner/v1/float8/float8_handler.py`](../../xtuner/v1/float8/float8_handler.py#L121-L161) | 45–49、121–161 | `pad_for_fsdp` / `build_reduce_mesh` 的 expert 分支 | +| | [166–239](../../xtuner/v1/float8/float8_handler.py#L166-L239) | `_build_decoupled_reduce_meshes` 与 strided reduce mesh helper | +| | [316–335](../../xtuner/v1/float8/float8_handler.py#L316-L335) | tile-wise scale 预计算按类各执行一次 | +| [`xtuner/v1/float8/fsdp_utils.py`](../../xtuner/v1/float8/fsdp_utils.py#L127-L137) | 127–137 | `precompute_tilewise_float8_scale_for_fsdp` 增加 `module_types` | +| [`xtuner/v1/rl/weight_update/weight_iterator.py`](../../xtuner/v1/rl/weight_update/weight_iterator.py#L140-L149) | 140–149、196–208、228、236–247 | layer-wise 权重同步按参数 owner gather | +| [`examples/v1/config/sft_glm5p2.py`](../../examples/v1/config/sft_glm5p2.py#L107-L115) | 107–115 | `DECOUPLE_EP_FSDP` / `HSDP_SHARDING_SIZE` 环境变量 | +| [`tests/model/test_decoupled_ep_fsdp_mesh.py`](../../tests/model/test_decoupled_ep_fsdp_mesh.py) | 全文 | L0 fake-PG placement 测试（29 例） | +| [`tests/rl/test_weight_iterator.py`](../../tests/rl/test_weight_iterator.py#L234-L418) | 234–418 | compose owner 回归测试 | +| [`tests/model/run_decoupled_ep_fsdp_numerics.py`](../../tests/model/run_decoupled_ep_fsdp_numerics.py)、[`run_decoupled_ep_fsdp_ckpt.py`](../../tests/model/run_decoupled_ep_fsdp_ckpt.py)、[`summarize_decoupled_ep_fsdp_numerics.py`](../../tests/model/summarize_decoupled_ep_fsdp_numerics.py) | 全文 | L1–L3 实验脚本 | + +相关但**未改动**的文件：[`xtuner/v1/utils/load_spec.py`](../../xtuner/v1/utils/load_spec.py)（`LoadSpec.from_tensor`、`plan_hf_save`、`_preserved_shard_indices`，已能表达新 placement）、[`xtuner/v1/module/decoder_layer/moe_decoder_layer.py`](../../xtuner/v1/module/decoder_layer/moe_decoder_layer.py)（`MoEBlock` 定义，是 expert FSDP 的包装单元）、`xtuner/v1/module/grouped_linear/moe_group_linear.py`（expert 仍在 `ep_mesh` 上 `Shard(0)`）。 + ### 3.1 配置与校验 -`FSDPConfig.decouple_ep_fsdp: bool = False`。旧断言"HSDP 要求 `ep_size == 1`"只在开关关闭时保留；开关打开时改为要求 `hsdp_sharding_size % ep_size == 0`，即必须存在整数 `efsdp`。`MoE._init_decoupled_device_mesh` 再次校验 `world_size % dp_shard == 0` 与 `dp_shard % ep_size == 0`。ExpertTP（`expert_tp_size > 1`）与解耦同时开启时抛 `NotImplementedError`，ETP 的 `(Shard, InterleavedShard)` 二维 placement 需要单独设计 mesh 维。 +位置：`xtuner/v1/config/fsdp.py` 48–65，`xtuner/v1/model/moe/moe.py` 1665–1678。 + +```python +# xtuner/v1/config/fsdp.py +decouple_ep_fsdp: bool = False # 默认关闭，旧路径不受影响 + +def model_post_init(self, __context): + if self.hsdp_sharding_size is not None: + if self.decouple_ep_fsdp: + assert self.hsdp_sharding_size % self.ep_size == 0 # 必须存在整数 efsdp + else: + assert self.ep_size == 1 # 旧断言只在旧路径保留 +``` + +`MoE._init_decoupled_device_mesh` 再次校验 `world_size % dp_shard == 0` 与 `dp_shard % ep_size == 0`。ExpertTP（`expert_tp_size > 1`）与解耦同时开启时抛 `NotImplementedError`，ETP 的 `(Shard, InterleavedShard)` 二维 placement 需要单独设计 mesh 维。 ### 3.2 单一 root mesh -`MoE._init_decoupled_device_mesh` 一次性建立 `(replicate, efsdp, ep)` root，所有子 mesh 从它派生： +位置：`xtuner/v1/model/moe/moe.py` 1651–1716（`_init_decoupled_device_mesh`），1486–1492（分流）。 + +```python +# xtuner/v1/model/moe/moe.py::_init_decoupled_device_mesh（节选） +world_size = dist.get_world_size() +dp_shard = fsdp_config.hsdp_sharding_size or world_size +replicate_size = world_size // dp_shard +efsdp_size = dp_shard // ep_size + +root_mesh = init_device_mesh( + device, (replicate_size, efsdp_size, ep_size), + mesh_dim_names=(f"{prefix}.replicate", f"{prefix}.efsdp", f"{prefix}.ep"), +) +self._world_mesh = root_mesh +self.ep_mesh = root_mesh[ep_name] # dispatcher 用，接口不变 +self.fsdp_mesh = root_mesh[efsdp_name, ep_name]._flatten(dp_shard_name) # dense 的一维 dp_shard +if replicate_size > 1: # HSDP + self.hsdp_mesh = root_mesh[replicate_name, dp_shard_name] + self.expert_fsdp_mesh = root_mesh[replicate_name, efsdp_name] +else: + self.hsdp_mesh = None + self.expert_fsdp_mesh = root_mesh[efsdp_name] +``` + +所有子 mesh 从同一个 root 派生： -- `ep_mesh = root[ep]`：dispatcher / DeepEP 拿到的 EP group 不变，仍是节点内连续 rank； +- `ep_mesh = root[ep]`：dispatcher / DeepEP 拿到的 EP group 不变，仍是节点内连续 rank。`MoE.__init__` 里先建的 `ep_mesh` 通过 PyTorch 的 mesh 相等性 hash 被重新挂到这个 root 上，因此 `ep` 维必须沿用旧名字 `{prefix}.ep`； - `fsdp_mesh = root[efsdp, ep]._flatten("dp_shard")`：语义仍是"dense 参数的一维 shard group"，`_fsdp_foreach_allgather`、`Float8Handler.build_reduce_mesh` 等既有消费者无需改动； - `hsdp_mesh = root[replicate, dp_shard]`（仅 `replicate > 1`）：dense 的 HSDP mesh，与旧路径契约一致； - `expert_fsdp_mesh = root[efsdp]` 或 `root[replicate, efsdp]`：新增属性，只在解耦路径上非空，dense 模型恒为 `None`。 -所有子 mesh 同源很重要：FSDP2 要求已是 DTensor 的参数与 FSDP mesh 共享同一个 root，expert 参数最终要同时带 EP shard 和 `efsdp` 上的 `_StridedShard`。`efsdp = 1` 时该维保留而不是省略——对 expert 做的那次 `fully_shard` 是 mixed precision policy 生效的地方（torchtitan #1324 的同一注释）。 +同源很重要：FSDP2 要求已是 DTensor 的参数与 FSDP mesh 共享同一个 root，expert 参数最终要同时带 EP shard 和 `efsdp` 上的 `_StridedShard`。`efsdp = 1` 时该维保留而不是省略——对 expert 做的那次 `fully_shard` 是 mixed precision policy 生效的地方（torchtitan #1324 的同一注释）。 ### 3.3 两级 `fully_shard` -解耦路径跳过 `_replicate_other_params`：dense 参数进入 `fully_shard` 时是普通 tensor，由 FSDP 在完整 `dp_shard` 上分片。每个 decoder layer 的包装顺序变为： - -1. `_fully_shard_expert_blocks`：找出该层的所有 `MoEBlock`（只含 routed expert 的 grouped linear 与激活），在 `expert_fsdp_mesh` 上先做一次 `fully_shard`； -2. 再在 dense 的 `fsdp_mesh` / `hsdp_mesh` 上包装整个 layer。FSDP2 会保留内层的 expert wrapper，把 expert 参数从外层 param group 中排除。 +位置：`xtuner/v1/model/moe/moe.py` 1194–1225（decoder layer），1250–1260（prefetch），1328–1360（MTP），1631–1649（helper）。 + +```python +# xtuner/v1/model/moe/moe.py::fully_shard（节选） +decoupled = self.fsdp_config.decouple_ep_fsdp +if not decoupled and (self.ep_mesh.size() > 1 or tp_enabled): + self._replicate_other_params(self) # 旧路径：dense 在 ep 维 Replicate；解耦路径跳过 + +for layer_idx, layer in self.layers.items(): + if decoupled: # 1) 先包 MoEBlock -> expert_fsdp_mesh + self._fully_shard_expert_blocks(layer, mp_policy=mp_policy, reshard_after_forward=...) + ... + self._fully_shard( # 2) 再包整层 -> dense 的 fsdp_mesh / hsdp_mesh + mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, module=layer, ... + ) + +for layer_cur, layer_next in zip(layers[:-1], layers[1:]): + if decoupled: # expert 的 all-gather 与下一层 dense 一起预取 + layer_cur.set_modules_to_forward_prefetch([*self._expert_blocks(layer_cur), layer_next]) + +@staticmethod +def _expert_blocks(module): + return [m for m in module.modules() if isinstance(m, MoEBlock)] + +def _fully_shard_expert_blocks(self, module, mp_policy, reshard_after_forward): + for expert_block in self._expert_blocks(module): + self._fully_shard(mesh=self.expert_fsdp_mesh, module=expert_block, mp_policy=mp_policy, ...) +``` -MTP layer 使用同样的两级包装。forward prefetch 列表同时包含下一层与其 expert block，让 expert 的 all-gather 与 attention 重叠，而不是串行等在 MoE block 前。 +解耦路径跳过 `_replicate_other_params`：dense 参数进入 `fully_shard` 时是普通 tensor，由 FSDP 在完整 `dp_shard` 上分片。每个 decoder layer 先把所有 `MoEBlock`（只含 routed expert 的 grouped linear 与激活）在 `expert_fsdp_mesh` 上 `fully_shard`，再在 dense mesh 上包装整层；FSDP2 会保留内层的 expert wrapper，把 expert 参数从外层 param group 中排除。MTP layer 使用同样的两级包装。forward prefetch 列表同时包含下一层与其 expert block，让 expert 的 all-gather 与 attention 重叠，而不是串行等在 MoE block 前。 ### 3.4 梯度归约 -不变式：每个参数的最终梯度等于全 data-parallel world 上的均值，与 `ep` / `efsdp` 取值无关。`MoE._scale_and_reduce_grad_decoupled` 只做三件事： +位置：`xtuner/v1/model/moe/moe.py` 1392–1396（分流），1583–1629（`_scale_and_reduce_grad_decoupled`）。 + +不变式：每个参数的最终梯度等于全 data-parallel world 上的均值，与 `ep` / `efsdp` 取值无关。 + +```python +# xtuner/v1/model/moe/moe.py::_scale_and_reduce_grad_decoupled（节选） +for name, param in self.trainable_parameters(): + if param.grad is None: + continue + if ep_size > 1 and ".experts" in name: + param.grad.div_(ep_size) # expert：FSDP 已在 efsdp 上求均值，剩余因子是 ep + continue + if not isinstance(param, DTensor): + continue + if any(isinstance(p, Shard) for p in param.placements): + continue # FSDP 管理的参数（含 _StridedShard）：FSDP 已归约 + replicate_dims = [names[i] for i, p in enumerate(param.placements) if isinstance(p, Replicate)] + flat_mesh = param.device_mesh[replicate_dims]._flatten() if len(replicate_dims) > 1 else param.device_mesh[replicate_dims[0]] + grad = param.grad.to_local() + grad.div_(flat_mesh.size()) # 全 Replicate 的 ignored fp32 参数：手工求均值 + grads_by_group.setdefault(flat_mesh.get_group(), []).append(grad) + +for group, grads in grads_by_group.items(): # 每个 group 一次 coalesced all-reduce + with dist._coalescing_manager(group=group): + for grad in grads: + dist.all_reduce(grad, ReduceOp.SUM, group=group) +``` | 参数类 | 谁完成归约 | 解耦路径的额外动作 | |---|---|---| @@ -88,23 +204,79 @@ MTP layer 使用同样的两级包装。forward prefetch 列表同时包含下 ### 3.5 FP8 +位置：`xtuner/v1/float8/float8_handler.py` 121–161、166–239、316–335；`xtuner/v1/float8/fsdp_utils.py` 127–137；`xtuner/v1/model/moe/moe.py` 1165–1176；`xtuner/v1/model/base.py` 1015。 + tile-wise FP8 需要知道每个本地权重会被多少个 FSDP rank 切分（决定 padding），以及 scale 的 reduce-max 应跨哪些 rank。解耦后 dense 与 expert 的答案不同： +```python +# xtuner/v1/float8/float8_handler.py::pad_for_fsdp（节选） +num_fsdp_chunks = fsdp_mesh.size(-1) # dense：dp_shard 份 +if expert_fsdp_mesh is not None and isinstance(module, TileWiseFloat8GroupedLinear): + num_fsdp_chunks = expert_fsdp_mesh.size(-1) # routed expert：efsdp 份 +padded_out_features = Float8Handler.get_num_features_after_pad(tensor_size, 0, num_fsdp_chunks) + +# ::_build_decoupled_reduce_meshes（节选） +dense_shard_size = fsdp_mesh.size(-1) # 连续 rank，stride 1 +expert_shard_size = expert_fsdp_mesh.size(-1) +expert_stride = world_size // expert_fsdp_mesh.size() # == ep_size +self.tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8Linear,), dense_shard_size, 1) +self.expert_tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mapping( + model, (TileWiseFloat8GroupedLinear,), expert_shard_size, expert_stride) +``` + - dense linear 按 `dp_shard` 的 chunk 数 padding，reduce mesh 的 rank stride 为 1； - grouped-expert linear 按 `efsdp` 的 chunk 数 padding，reduce mesh 的 rank stride 为 `ep`； -- `Float8Handler.pad_for_fsdp` 接受可选的 `expert_fsdp_mesh`，`build_reduce_mesh` 为两类参数各建一套 tile-wise reduce mesh，scale 预计算按类各执行一次。 +- scale 预计算（`precompute_tilewise_float8_scale_for_fsdp`）新增 `module_types` 参数，按 dense / expert 两类各执行一次。 旧路径 `expert_fsdp_mesh is None`，走原有的单 mesh 代码。tensor-wise FP8 只作用于 dense linear，shard group 仍是 `fsdp_mesh`，无需改动。 ### 3.6 HF 权重与 DCP -`LoadSpec.from_tensor` 直接从 DTensor placement 记录 shard 历史（EP 的 `Shard(0)`、`efsdp` 上的 `_StridedShard`、HSDP 的 `Replicate` 不产生 shard）。同步 HF save 按 shard 描述逆向 unshard 后再走模型已有的 HF key / fused-expert 映射；DCP 按 DTensor placement 保存并在载入时 reshard。因此没有为任何模型重写 `state_dict_adapter`，工作量在于让新 mesh 产出正确的 `LoadSpec`。 +位置：无代码改动。依赖 `xtuner/v1/utils/load_spec.py` 的 `LoadSpec.from_tensor`、`RuntimeLayout.from_dtensor`、`plan_hf_save`；验证脚本 `tests/model/run_decoupled_ep_fsdp_ckpt.py`。 + +`LoadSpec.from_tensor` 直接从 DTensor placement 记录 shard 历史，新 placement 无需特殊处理： + +```text +routed expert（解耦 + HSDP） + DTensor placement : (Replicate, _StridedShard(0), Shard(0)) on (replicate, efsdp, ep) + LoadSpec.shards : [ShardDescriptor(dim=0, group=ep), ShardDescriptor(dim=0, group=efsdp)] # Replicate 不产生 shard +dense（解耦 + HSDP） + DTensor placement : (Replicate, Shard(0)) on (replicate, dp_shard) + LoadSpec.shards : [ShardDescriptor(dim=0, group=dp_shard)] +``` + +同步 HF save 按 shard 描述逆向 unshard 后再走模型已有的 HF key / fused-expert 映射；DCP 按 DTensor placement 保存并在载入时 reshard。因此没有为任何模型重写 `state_dict_adapter`，工作量在于让新 mesh 产出正确的 `LoadSpec`。 ### 3.7 RL 权重同步 -Turbomind 的 layer-wise 更新只 gather FSDP shard、保留 EP-local 的 expert 切片。`BaseModel._fsdp_foreach_allgather` 按每个 `LoadSpec` 选择 gather group：spec 含 `efsdp` 上的 shard 时用 `expert_fsdp_mesh` 的 group，否则用 `fsdp_mesh` 的 group。 +位置：`xtuner/v1/model/base.py` 1916–1947；`xtuner/v1/rl/weight_update/weight_iterator.py` 140–149、196–208、236–247；测试 `tests/rl/test_weight_iterator.py` 234–418。 -compose 模型（如 Qwen3.5-VL MoE）曾有一个 owner 选择错误：`WeightIterator.iter_layer_batches` 从 language tower 取参数与 `LoadSpec`，却用外层 compose 模型的 mesh 做 gather。compose 模型本身在 world mesh 上包装、没有 `expert_fsdp_mesh`，也不维护自己的 `load_spec_mapping`，于是 language tower 的 `efsdp` shard（HSDP 时连 `dp_shard` shard）被当作"应保留"的 shard 而不 gather，Turbomind 收到的是 rank-local 碎片。现在 gather 由参数所属的子模块（language / vision / projector）执行，并从该子模块的 `load_spec_mapping` 取 spec；`tests/rl/test_weight_iterator.py::TestLayerBatchesGatherWithParamOwner` 用假 process group 固定了 `efsdp = 2` 与 HSDP 两种拓扑下每个 tensor 的完整形状与取值。 +Turbomind 的 layer-wise 更新只 gather FSDP shard、保留 EP-local 的 expert 切片。`BaseModel._fsdp_foreach_allgather` 按每个 `LoadSpec` 选择 gather group： + +```python +# xtuner/v1/model/base.py::_fsdp_gather_group +def _fsdp_gather_group(self, load_spec, fsdp_group, expert_fsdp_group): + if expert_fsdp_group is None: + return fsdp_group + if any(self._is_same_process_group(shard.group, expert_fsdp_group) for shard in load_spec.shards): + return expert_fsdp_group # spec 带 efsdp shard：只 gather efsdp，EP shard 保留 + return fsdp_group # dense：gather dp_shard +``` + +compose 模型（如 Qwen3.5-VL MoE）曾有一个 owner 选择错误：`WeightIterator.iter_layer_batches` 从 language tower 取参数与 `LoadSpec`，却用外层 compose 模型的 mesh 做 gather。compose 模型本身在 world mesh 上包装、没有 `expert_fsdp_mesh`，也不维护自己的 `load_spec_mapping`，于是 language tower 的 `efsdp` shard（HSDP 时连 `dp_shard` shard）被当作"应保留"的 shard 而不 gather，Turbomind 收到的是 rank-local 碎片。现在 gather 由参数所属的子模块执行，并从该子模块的 `load_spec_mapping` 取 spec： + +```python +# xtuner/v1/rl/weight_update/weight_iterator.py::_param_owner +def _param_owner(model, name): + if isinstance(model.config, BaseComposeConfig): + for submodule in ("language_model", "vision_tower", "multi_modal_projector"): + if name.startswith(f"{submodule}."): + return getattr(model, submodule), name[len(submodule) + 1 :] + return model, name # plain 模型：owner 就是自己，行为不变 +``` + +`tests/rl/test_weight_iterator.py::TestLayerBatchesGatherWithParamOwner` 用假 process group 固定了 `efsdp = 2` 与 HSDP 两种拓扑下每个 tensor 的完整形状与取值。 ## 4. 与参考实现的对照 @@ -119,23 +291,52 @@ compose 模型（如 Qwen3.5-VL MoE）曾有一个 owner 选择错误：`WeightI ## 5. 验证结论与边界 -| 项目 | 证据 | 边界 | -|---|---|---| -| mesh / placement / 两级 FSDP | `tests/model/test_decoupled_ep_fsdp_mesh.py`，fake-PG 下 8/16/64 rank 共 29 例；旧布局在开关关闭时逐项固定 | 需要一张空闲 GPU 建 CUDA context，不通信 | -| BF16 训练语义（L1） | tiny Qwen3-MoE 50 步：解耦 EP8 vs 旧 EP8 loss 最大相对差 `1.9e-5`（旧 EP8 vs EP1 为 `2.5e-5`）；3.4B 模型每卡 dense 参数 836 → 105 MiB，峰值 allocated 11861 → 8203 MiB，步时 170 → 158 ms | 单机；reduction order 不同，不是逐 bit 相同 | -| `efsdp > 1`、HSDP + EP（L2） | 单机缩比拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)` 均在噪声内；16/64 rank 形状由 L0 固定 | 真正的双节点 collective 未跑 | -| HF / DCP（L3） | 刚 `from_hf` 后导出 807 个 key 全部 bit-exact；DCP 同布局 resume 与跨布局 reshard 后 loss 连续 | 未覆盖 async save、FP8 checkpoint、GLM/MTP 导出、world size 变化 | -| FP8 | 3.4B 与 GLM-5.2 tile-wise FP8 训练 20 步稳定，差异量级接近同配置重复噪声 | FP8 + HSDP、FP8 + checkpoint 未测 | -| GLM-5.2-30B，8×H200 目标配方 | EP4：1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB；EP8：11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB | EP8 的加速来自离开 allocator 上限（legacy 每步 2–3 次 alloc retry，解耦为 0），不是 collective 本身快 7 倍 | -| RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | -| legacy 兼容 | 开关默认关闭，旧分支执行逻辑未改动，L0 旧布局回归通过 | 未做 base-vs-branch 端到端逐 tensor 对照 | +| 项目 | 证据 | 边界 | 链接 | +|---|---|---|---| +| mesh / placement / 两级 FSDP（L0） | fake-PG 下 8/16/64 rank 共 29 例；旧布局在开关关闭时逐项固定 | 需要一张空闲 GPU 建 CUDA context，不通信 | [test_decoupled_ep_fsdp_mesh.py](../../tests/model/test_decoupled_ep_fsdp_mesh.py)、[baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md) | +| BF16 训练语义（L1） | tiny Qwen3-MoE 50 步：解耦 EP8 vs 旧 EP8 loss 最大相对差 `1.9e-5`（旧 EP8 vs EP1 为 `2.5e-5`）；3.4B 模型每卡 dense 参数 836 → 105 MiB，峰值 allocated 11861 → 8203 MiB，步时 170 → 158 ms | 单机；reduction order 不同，不是逐 bit 相同 | [L1.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L1.md)、[run_decoupled_ep_fsdp_numerics.py](../../tests/model/run_decoupled_ep_fsdp_numerics.py) | +| `efsdp > 1`、HSDP + EP（L2） | 单机缩比拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)` 均在噪声内；16/64 rank 形状由 L0 固定 | 真正的双节点 collective 未跑 | [L2.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L2.md)、[decisions.md D10](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | +| HF / DCP（L3） | 刚 `from_hf` 后导出 807 个 key 全部 bit-exact；DCP 同布局 resume 与跨布局 reshard 后 loss 连续 | 未覆盖 async save、FP8 checkpoint、GLM/MTP 导出、world size 变化 | [L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[run_decoupled_ep_fsdp_ckpt.py](../../tests/model/run_decoupled_ep_fsdp_ckpt.py) | +| FP8 | 3.4B 与 GLM-5.2 tile-wise FP8 训练 20 步稳定，差异量级接近同配置重复噪声 | FP8 + HSDP、FP8 + checkpoint 未测 | [L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | +| GLM-5.2-30B，8×H200 目标配方 | EP4：1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB；EP8：11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB | EP8 的加速来自离开 allocator 上限（legacy 每步 2–3 次 alloc retry，解耦为 0），不是 collective 本身快 7 倍 | [GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | +| RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | [test_weight_iterator.py](../../tests/rl/test_weight_iterator.py#L234-L418) | +| legacy 兼容 | 开关默认关闭，旧分支执行逻辑未改动，L0 旧布局回归通过 | 未做 base-vs-branch 端到端逐 tensor 对照 | [baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md)、[decisions.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | -详细数据见 `reports/L1.md`、`reports/L2.md`、`reports/L3.md`、`reports/GLM52.md`（不随代码 PR 合入）。 +`reports/` 下的报告不随代码 PR 合入，链接指向 fork 分支 `feat/decouple-ep-fsdp` 上的固定地址；测试与脚本是仓库内相对链接。 ## 6. 已知限制与后续工作 -- `fp32_keys_pattern` 命中 routed expert 时（当前内置 pattern 均不命中），该 expert 只被 `ignored_params` 排除而不会沿 `efsdp` 分片，梯度路径只做 `/ep`，`efsdp > 1` 或 `replicate > 1` 时副本会分叉。短期应显式拒绝；完整方案见 §7.2，与"expert 不做 FSDP"模式是同一实现。 -- expert 的分类在 FSDP 包装时按 `MoEBlock` 类型，在梯度缩放时按参数名含 `.experts`，不是同一事实来源；应在包装时记录 expert 参数身份，梯度、FP8、save / RL 复用。 +### 6.1 `fp32_keys_pattern` 命中 routed expert 时梯度不归约 + +`HFSaveCfg.fp32_keys_pattern` 用 HF key 正则把少量参数保留在 fp32 并排除在 FSDP 之外，内置用途是 Qwen3.5 linear-attention 的 `A_log`、norm 这类 dense 小参数。`BaseModel._fully_shard` 对命中的参数分两种处理：普通 tensor 先 `distribute_tensor(Replicate on world_mesh)` 再放进 `ignored_params`；已经是 DTensor 的参数原样放进 `ignored_params`。routed expert 在 `build_grouped_linear` 里就已经是 `Shard(0)` on `ep` 的 DTensor，走的是第二种。 + +于是一旦某条 pattern 命中 expert，会同时发生两件事： + +1. expert 不再被 `_fully_shard_expert_blocks` 沿 `efsdp` 分片，每个 `efsdp` rank（HSDP 时每个 replica）各持一份完整的 EP-local expert； +2. `_scale_and_reduce_grad_decoupled` 看到参数名含 `.experts`，只做 `/ep` 就 `continue`，不会走第三行的 Replicate 维 all-reduce。 + +每个 `efsdp` rank 的 expert 梯度只含自己 EP group 的 token 贡献，optimizer 一步之后同一个 expert 在不同 `efsdp` rank 上的权重就不再相同，训练静默出错。`efsdp = replicate = 1`（如单机 EP=8）时没有第二份副本，问题不显现；当前内置 pattern 也都只匹配 dense，GLM-5.2 没有 pattern，所以现有实验没有触发。旧路径同样存在这个边界：ignored expert 不会沿 `fsdp` 分片，`fsdp > 1` 时一样分叉。 + +修法：短期在解耦路径检测到 pattern 命中 `MoEBlock` 参数时直接报错；完整方案是把这类 expert 表达成 `Replicate(replicate), Replicate(efsdp), Shard(0)(ep)`，先走 §3.4 第三行的 Replicate 维 all-reduce 求均值，再 `/ep`。这与 §7.2 的 `expert_fsdp = False` 模式是同一套实现。 + +### 6.2 "什么是 expert 参数"有两套判定标准 + +| 环节 | 判定依据 | 位置 | +|---|---|---| +| FSDP 包装 | 模块类型 `isinstance(m, MoEBlock)` | `MoE._expert_blocks` | +| 梯度缩放 | 参数名 `".experts" in name` | `MoE._scale_and_reduce_grad_decoupled`（旧路径 `scale_and_reduce_grad` 同样） | +| FP8 padding / reduce mesh | 模块类型 `TileWiseFloat8GroupedLinear` | `Float8Handler.pad_for_fsdp` / `_build_decoupled_reduce_meshes` | +| RL gather | `LoadSpec.shards` 中是否有 `efsdp` group 的 shard | `BaseModel._fsdp_gather_group` | + +今天它们恰好一致，因为 `MoEDecoderLayer` 把 `MoEBlock` 挂在 `experts` 属性下（`shared_experts` 不含 `.experts` 子串，不会误伤）。两个方向都可能失效： + +- 新模型把 `MoEBlock` 挂到别的属性名（如 `moe`）：expert 会被 `expert_fsdp_mesh` 分片，reduce-scatter 只在 `efsdp` 上求均值，但名字里没有 `.experts`，不做 `/ep`，梯度被放大 `ep` 倍，等价于 expert 的学习率乘以 `ep`，没有任何报错； +- 某个 dense 参数名恰好含 `.experts`（如 `experts_router`）：会被 `/ep`，梯度缩小 `ep` 倍。 + +修法：在 `_fully_shard_expert_blocks` 包装时记录 expert 参数的 id / FQN 集合，梯度缩放、FP8、save / RL 都查这一份，删除按名字匹配的分支。 + +### 6.3 其它 + - 新拓扑约束目前用 `assert`，应改为 Pydantic validator + `ValueError`。 - 两次不同 mesh 的 `fully_shard` 与 `DeviceMesh._flatten(name)` 只在 torch 2.8 / 2.9 验证；`pyproject.toml` 声明 `torch>=2.6`，需要补 CI 或为该开关声明更高的最低版本。 - L1–L3 是人工实验脚本（写 JSON、打印统计），不是会失败的 pytest gate。 @@ -162,7 +363,7 @@ expert_fsdp = True （现状）: Shard(0) on ep + _StridedShard on efsdp expert_fsdp = False（新增）: Shard(0) on ep + Replicate on (replicate, efsdp) 手工 all-reduce 求均值，再 /ep ``` -后者的梯度处理正是 §3.4 第三行已经存在的"沿 Replicate 维 all-reduce"分支，也正是修复 §6 第一条（fp32 pattern 命中 expert）所需的实现，两件事可以一次做完。代价是 expert 参数与优化器状态在 `efsdp` 维复制、失去 FSDP wrapper 提供的 mixed precision（需手工 cast）。 +后者的梯度处理正是 §3.4 第三行已经存在的"沿 Replicate 维 all-reduce"分支，也正是修复 §6.1（fp32 pattern 命中 expert）所需的实现，两件事可以一次做完。代价是 expert 参数与优化器状态在 `efsdp` 维复制、失去 FSDP wrapper 提供的 mixed precision（需手工 cast）。 建议：`efsdp` 保持派生；不新增整数配置。等 MoonEP / UltraEP 接入启动时，在 `FSDPConfig` 增加 `expert_fsdp: bool = True`（或 `expert_shard_mode: "fsdp" | "replicate"`）。root mesh 不变，`ep_mesh` 仍是最内层连续 rank 的 group，dispatcher 的替换与本方案正交。 From 23751f427c9309812dea33a166984f1d7026d171 Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 3 Sep 2026 07:13:21 +0000 Subject: [PATCH 07/14] [Docs] Drop the per-file code map from the decoupled EP/FSDP design note --- docs/design/decouple_ep_fsdp.md | 28 +--------------------------- 1 file changed, 1 insertion(+), 27 deletions(-) diff --git a/docs/design/decouple_ep_fsdp.md b/docs/design/decouple_ep_fsdp.md index 4bfb9be056..7d1a4476f2 100644 --- a/docs/design/decouple_ep_fsdp.md +++ b/docs/design/decouple_ep_fsdp.md @@ -50,33 +50,7 @@ expert : Shard(0) over ep + FSDP over efsdp (+ replicate 上 HS ## 3. 实现 -行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `5923dc1d`（2026-09-03）。改动的文件与位置： - -| 文件 | 行号 | 内容 | -|---|---|---| -| [`xtuner/v1/config/fsdp.py`](../../xtuner/v1/config/fsdp.py#L48-L65) | 48–65 | `decouple_ep_fsdp` 开关与校验 | -| [`xtuner/v1/model/moe/moe.py`](../../xtuner/v1/model/moe/moe.py#L1651-L1716) | 1651–1716 | `_init_decoupled_device_mesh`：root mesh 与子 mesh | -| | [1486–1492](../../xtuner/v1/model/moe/moe.py#L1486-L1492) | `_init_device_mesh` 分流到新路径 | -| | [1194–1201](../../xtuner/v1/model/moe/moe.py#L1194-L1201) | 解耦路径跳过 `_replicate_other_params` | -| | [1214–1225](../../xtuner/v1/model/moe/moe.py#L1214-L1225)、[1250–1260](../../xtuner/v1/model/moe/moe.py#L1250-L1260) | decoder layer 两级 `fully_shard` 与 forward prefetch | -| | [1328–1360](../../xtuner/v1/model/moe/moe.py#L1328-L1360) | MTP layer 的同样处理 | -| | [1631–1649](../../xtuner/v1/model/moe/moe.py#L1631-L1649) | `_expert_blocks` / `_fully_shard_expert_blocks` | -| | [1392–1396](../../xtuner/v1/model/moe/moe.py#L1392-L1396)、[1583–1629](../../xtuner/v1/model/moe/moe.py#L1583-L1629) | `scale_and_reduce_grad` 分流与 `_scale_and_reduce_grad_decoupled` | -| | [1165–1176](../../xtuner/v1/model/moe/moe.py#L1165-L1176) | FP8 padding 传入 `expert_fsdp_mesh` | -| [`xtuner/v1/model/base.py`](../../xtuner/v1/model/base.py#L567-L569) | 567–569 | 新增 `expert_fsdp_mesh` 属性（dense 模型恒为 `None`） | -| | [1015](../../xtuner/v1/model/base.py#L1015) | float8 handler 建 reduce mesh 时传入 expert mesh | -| | [1916–1947](../../xtuner/v1/model/base.py#L1916-L1947) | `_fsdp_foreach_allgather` / `_fsdp_gather_group`：RL 按 spec 选 gather group | -| [`xtuner/v1/float8/float8_handler.py`](../../xtuner/v1/float8/float8_handler.py#L121-L161) | 45–49、121–161 | `pad_for_fsdp` / `build_reduce_mesh` 的 expert 分支 | -| | [166–239](../../xtuner/v1/float8/float8_handler.py#L166-L239) | `_build_decoupled_reduce_meshes` 与 strided reduce mesh helper | -| | [316–335](../../xtuner/v1/float8/float8_handler.py#L316-L335) | tile-wise scale 预计算按类各执行一次 | -| [`xtuner/v1/float8/fsdp_utils.py`](../../xtuner/v1/float8/fsdp_utils.py#L127-L137) | 127–137 | `precompute_tilewise_float8_scale_for_fsdp` 增加 `module_types` | -| [`xtuner/v1/rl/weight_update/weight_iterator.py`](../../xtuner/v1/rl/weight_update/weight_iterator.py#L140-L149) | 140–149、196–208、228、236–247 | layer-wise 权重同步按参数 owner gather | -| [`examples/v1/config/sft_glm5p2.py`](../../examples/v1/config/sft_glm5p2.py#L107-L115) | 107–115 | `DECOUPLE_EP_FSDP` / `HSDP_SHARDING_SIZE` 环境变量 | -| [`tests/model/test_decoupled_ep_fsdp_mesh.py`](../../tests/model/test_decoupled_ep_fsdp_mesh.py) | 全文 | L0 fake-PG placement 测试（29 例） | -| [`tests/rl/test_weight_iterator.py`](../../tests/rl/test_weight_iterator.py#L234-L418) | 234–418 | compose owner 回归测试 | -| [`tests/model/run_decoupled_ep_fsdp_numerics.py`](../../tests/model/run_decoupled_ep_fsdp_numerics.py)、[`run_decoupled_ep_fsdp_ckpt.py`](../../tests/model/run_decoupled_ep_fsdp_ckpt.py)、[`summarize_decoupled_ep_fsdp_numerics.py`](../../tests/model/summarize_decoupled_ep_fsdp_numerics.py) | 全文 | L1–L3 实验脚本 | - -相关但**未改动**的文件：[`xtuner/v1/utils/load_spec.py`](../../xtuner/v1/utils/load_spec.py)（`LoadSpec.from_tensor`、`plan_hf_save`、`_preserved_shard_indices`，已能表达新 placement）、[`xtuner/v1/module/decoder_layer/moe_decoder_layer.py`](../../xtuner/v1/module/decoder_layer/moe_decoder_layer.py)（`MoEBlock` 定义，是 expert FSDP 的包装单元）、`xtuner/v1/module/grouped_linear/moe_group_linear.py`（expert 仍在 `ep_mesh` 上 `Shard(0)`）。 +各小节的"位置"行给出文件与行号，行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `5923dc1d`（2026-09-03）。 ### 3.1 配置与校验 From de4cc89c9fbece7ac38f59d13892b855fe6e6627 Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 17 Sep 2026 10:04:50 +0000 Subject: [PATCH 08/14] [Fix] Validate the decoupled EP/FSDP topology with a Pydantic validator - FSDPConfig: replace the `assert`s in `model_post_init` with a `@model_validator(mode="after")` that raises `ValueError` (surfaced as a pydantic `ValidationError`), so the checks survive `python -O`; also reject non-positive `ep_size` / `hsdp_sharding_size`. - MoE._init_decoupled_device_mesh: raise `ValueError` instead of asserting `world_size % dp_shard == 0` and `dp_shard % ep_size == 0`. - MoE._scale_and_reduce_grad_decoupled: narrow `mesh_dim_names` before indexing so the decoupled path type-checks. - L0 regression tests for the config validator and the runtime check. --- tests/model/test_decoupled_ep_fsdp_mesh.py | 16 +++++++++++++ xtuner/v1/config/fsdp.py | 26 +++++++++++++++------- xtuner/v1/model/moe/moe.py | 17 ++++++++------ 3 files changed, 44 insertions(+), 15 deletions(-) diff --git a/tests/model/test_decoupled_ep_fsdp_mesh.py b/tests/model/test_decoupled_ep_fsdp_mesh.py index 06161848f6..ae0c1d0480 100644 --- a/tests/model/test_decoupled_ep_fsdp_mesh.py +++ b/tests/model/test_decoupled_ep_fsdp_mesh.py @@ -315,6 +315,22 @@ def test_config_requires_ep_to_divide_shard_size(self) -> None: with pytest.raises(ValidationError, match="divisible by `ep_size`"): FSDPConfig(ep_size=3, hsdp_sharding_size=8, decouple_ep_fsdp=True) + def test_config_rejects_non_positive_sizes(self) -> None: + # The checks are pydantic validators raising `ValueError`, not `assert`s that `python -O` + # would drop; `ValidationError` is the `ValueError` subclass pydantic wraps them in. + with pytest.raises(ValidationError, match="`ep_size` must be a positive integer"): + FSDPConfig(ep_size=0) + with pytest.raises(ValidationError, match="`hsdp_sharding_size` must be a positive integer"): + FSDPConfig(hsdp_sharding_size=0, decouple_ep_fsdp=True) + assert issubclass(ValidationError, ValueError) + + def test_runtime_rejects_shard_size_not_dividing_world_size(self) -> None: + # `hsdp_sharding_size=3` passes the config validator (3 % ep_size == 0) but cannot tile + # an 8-rank world; the mesh construction must fail loudly instead of asserting. + with _fake_world(8, 0): + with pytest.raises(ValueError, match=r"world_size \(8\) must be divisible by hsdp_sharding_size \(3\)"): + _shard_model(1, hsdp_sharding_size=3, decouple_ep_fsdp=True) + class TestDecoupledMoEParamPlacements: """DTensor placements after the two-level ``fully_shard`` of the decoupled path.""" diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py index 85d32887c5..f0d023c057 100644 --- a/xtuner/v1/config/fsdp.py +++ b/xtuner/v1/config/fsdp.py @@ -1,8 +1,8 @@ -from typing import Any, Optional +from typing import Optional import torch from cyclopts import Parameter -from pydantic import BaseModel, ConfigDict, field_serializer, field_validator +from pydantic import BaseModel, ConfigDict, field_serializer, field_validator, model_validator from typing_extensions import Annotated @@ -50,15 +50,25 @@ class FSDPConfig(BaseModel): bool, Parameter(help="Decouple expert parallel from FSDP: shard dense params over the full FSDP mesh") ] = False - def model_post_init(self, __context: Any) -> None: - if self.hsdp_sharding_size is not None: - if self.decouple_ep_fsdp: - assert self.hsdp_sharding_size % self.ep_size == 0, ( + @model_validator(mode="after") + def _validate_ep_fsdp_topology(self) -> "FSDPConfig": + # Explicit `ValueError`s instead of `assert`: the topology checks must survive `python -O`, + # and pydantic reports them as a `ValidationError` when the config is built. + if self.ep_size < 1: + raise ValueError(f"`ep_size` must be a positive integer, got {self.ep_size}") + if self.hsdp_sharding_size is None: + return self + if self.hsdp_sharding_size < 1: + raise ValueError(f"`hsdp_sharding_size` must be a positive integer, got {self.hsdp_sharding_size}") + if self.decouple_ep_fsdp: + if self.hsdp_sharding_size % self.ep_size != 0: + raise ValueError( "`decouple_ep_fsdp` requires `hsdp_sharding_size` to be divisible by `ep_size`, " f"got hsdp_sharding_size={self.hsdp_sharding_size}, ep_size={self.ep_size}" ) - else: - assert self.ep_size == 1, "Currently, HSDP requires expert parallel size to be 1" + elif self.ep_size != 1: + raise ValueError("Currently, HSDP requires expert parallel size to be 1") + return self @field_serializer("param_dtype", "reduce_dtype") def serialize_param_dtype(self, value: torch.dtype) -> str: diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 3a48d05568..31ad44d967 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1697,8 +1697,10 @@ def _scale_and_reduce_grad_decoupled(self) -> None: # FSDP-managed (`_StridedShard` is a `Shard` subclass). continue + mesh_dim_names = param.device_mesh.mesh_dim_names + assert mesh_dim_names is not None # every mesh of this model is built with named dims replicate_dim_names = tuple( - param.device_mesh.mesh_dim_names[i] for i, p in enumerate(param.placements) if isinstance(p, Replicate) + mesh_dim_names[i] for i, p in enumerate(param.placements) if isinstance(p, Replicate) ) if not replicate_dim_names: continue @@ -1759,12 +1761,13 @@ def _init_decoupled_device_mesh(self, fsdp_config: FSDPConfig) -> None: world_size = dist.get_world_size() ep_size = fsdp_config.ep_size dp_shard = fsdp_config.hsdp_sharding_size if fsdp_config.hsdp_sharding_size is not None else world_size - assert world_size % dp_shard == 0, ( - f"world_size ({world_size}) must be divisible by hsdp_sharding_size ({dp_shard})" - ) - assert dp_shard % ep_size == 0, ( - f"`decouple_ep_fsdp` requires the FSDP shard size ({dp_shard}) to be divisible by ep_size ({ep_size})" - ) + # Explicit exceptions (not `assert`): the mesh shape below is only valid with these. + if world_size % dp_shard != 0: + raise ValueError(f"world_size ({world_size}) must be divisible by hsdp_sharding_size ({dp_shard})") + if dp_shard % ep_size != 0: + raise ValueError( + f"`decouple_ep_fsdp` requires the FSDP shard size ({dp_shard}) to be divisible by ep_size ({ep_size})" + ) replicate_size = world_size // dp_shard efsdp_size = dp_shard // ep_size From cc1336292106ee40ae911b655d8ac0775a3c7c91 Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 17 Sep 2026 10:04:50 +0000 Subject: [PATCH 09/14] [Test] Gate the decoupled EP/FSDP L1-L3 checks in pytest - xtuner/_testing/decoupled_ep_fsdp.py: the tiny Qwen3-MoE checkpoint, token stream, per-layout training run and HF-checkpoint comparison shared by the gates and the manual experiment scripts. - tests/engine/test_decoupled_ep_fsdp_train_engine.py: 8-GPU `DeterministicDDPTestCase` gates. L1 / L2 compare the decoupled layouts (efsdp == 1, efsdp > 1, HSDP + EP) with the legacy ones on loss curves, total grad norms, per-parameter grad norms at step 0 and per-rank parameter memory; L3 checks bit-exact HF export after `from_hf`, DCP resume, the HF export after resume and cross-layout DCP resharding. CPU tests pin the checkpoint comparison itself. - run_decoupled_ep_fsdp_numerics.py / run_decoupled_ep_fsdp_ckpt.py now import the shared helpers; their CLI and JSON output are unchanged. --- .../test_decoupled_ep_fsdp_train_engine.py | 324 ++++++++++++ tests/model/run_decoupled_ep_fsdp_ckpt.py | 113 +---- tests/model/run_decoupled_ep_fsdp_numerics.py | 197 +------- xtuner/_testing/decoupled_ep_fsdp.py | 461 ++++++++++++++++++ 4 files changed, 812 insertions(+), 283 deletions(-) create mode 100644 tests/engine/test_decoupled_ep_fsdp_train_engine.py create mode 100644 xtuner/_testing/decoupled_ep_fsdp.py diff --git a/tests/engine/test_decoupled_ep_fsdp_train_engine.py b/tests/engine/test_decoupled_ep_fsdp_train_engine.py new file mode 100644 index 0000000000..f3cbe177c3 --- /dev/null +++ b/tests/engine/test_decoupled_ep_fsdp_train_engine.py @@ -0,0 +1,324 @@ +"""GPU gates for the decoupled EP/FSDP layout (``FSDPConfig.decouple_ep_fsdp``). + +Every test trains the same tiny random Qwen3-MoE on the same token stream under several EP/FSDP +layouts and fails when the decoupled layout departs from the legacy one: + +* ``TestDecoupledEpFsdpNumerics`` (L1 / L2): loss curves, total gradient norms, per-parameter + gradient norms at step 0 (identical weights and data, so they catch a wrong expert-gradient scale + that AdamW would hide from the loss curve; after the first update bf16 weight differences flip + individual top-k routing decisions and router gradients legitimately differ by percents) and + per-rank parameter memory, for ``efsdp == 1``, ``efsdp > 1`` and HSDP + EP. +* ``TestDecoupledEpFsdpCheckpoint`` (L3): bit-exact HF export right after ``from_hf``, DCP resume, + HF export after resume and cross-layout DCP resharding. + +The observed noise floor between layouts is ~1e-5 on losses and ≤ 2e-3 on per-parameter grad +norms (see ``docs/design/decouple_ep_fsdp.md`` §5); the gates sit 1-2 orders of magnitude above +it and orders of magnitude below any layout bug seen so far. +""" + +import shutil +import tempfile +import unittest +from pathlib import Path +from typing import Any + +import pytest +import torch +import torch.distributed as dist +from safetensors.torch import save_file + +from xtuner._testing import DeterministicDDPTestCase +from xtuner._testing.decoupled_ep_fsdp import ( + LayoutMode, + assert_hf_dirs_close, + build_engine, + build_hf_checkpoint, + load_hf_dir, + max_rel_diff, + parse_mode, + release, + run_mode, + train, +) +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig + + +SEQ_LEN = 2048 +LR = 1e-4 +# Loss curves of two layouts: max relative difference over all steps (noise ~2e-5). +LOSS_RTOL = 1e-3 +# Total grad norms per step and per-parameter grad norms at step 0 (noise ≤ 2e-3 dense, ≤ 2e-5 expert). +GRAD_NORM_RTOL = 1e-2 +# Per-rank parameter memory vs. the value implied by the layout (FSDP padding slack). +MEMORY_RTOL = 0.25 +# HF exports: 1 bf16 ulp is 2**-7 relative; resumed vs. continuous differs by single ulps on a +# subset of elements, exports of different layouts by a few ulps. +HF_RESUME_RTOL, HF_RESUME_ATOL = 2**-6, 1e-3 +HF_LAYOUT_RTOL, HF_LAYOUT_ATOL = 2**-6, 2e-3 + + +class _SharedTmpDirMixin: + """Temp directory created by rank 0 and shared with every rank of the process group.""" + + def _make_shared_tmpdir(self) -> Path: + holder = [tempfile.mkdtemp(prefix="decoupled_ep_fsdp_") if dist.get_rank() == 0 else None] + dist.broadcast_object_list(holder, src=0) + assert holder[0] is not None + return Path(holder[0]) + + @staticmethod + def _cleanup_shared_tmpdir(path: Path) -> None: + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(path, ignore_errors=True) + + +def _build_tiny_checkpoint(root: Path) -> Path: + hf_dir = root / "tiny_qwen3_moe_hf" + if dist.get_rank() == 0: + build_hf_checkpoint(hf_dir, seed=0, model_size="tiny") + dist.barrier() + return hf_dir + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestDecoupledEpFsdpNumerics(_SharedTmpDirMixin, DeterministicDDPTestCase): + def test_ep8_decoupled_matches_legacy_ep8_and_ep1(self) -> None: + # L1: `efsdp == 1`; dense params go from replicated over EP to sharded over all 8 ranks. + results = self._run_layouts(("A:ep=1", "B:ep=8", "C:ep=8,decouple=1"), steps=20) + self._assert_training_matches(results["C"], results["B"]) + self._assert_training_matches(results["C"], results["A"]) + # 中文注释：每卡参数显存断言。记 D / E 为 dense / expert 参数总量： + # B（legacy，ep=8）：mesh 为 (fsdp=1, ep=8)，dense 在 ep 维整份复制 → 每卡 D；expert 按 ep 分 8 份 → E/8。 + # C（解耦，ep=8）：dense 在整个 dp_shard=8 上分片 → D/8；expert 按 ep×efsdp=8×1 分片 → E/8。 + # 所以 C 的 dense 应是 B 的 1/8，expert 相同。loss / 梯度对比区分不了"dense 真的分片了"和 + # "dense 仍在 EP 维复制"（两者数值完全一样），只有这条显存断言能证明解耦布局确实生效。 + self._assert_param_memory(results["C"], results["B"], dense_ratio=1 / 8, expert_ratio=1.0) + + def test_efsdp_and_hsdp_layouts_match_legacy_ep4(self) -> None: + # L2 on 8 ranks: `efsdp == 2` (C4), HSDP + EP with `efsdp == 1` (H41) and `efsdp == 2` (H22). + results = self._run_layouts( + ("B4:ep=4", "C4:ep=4,decouple=1", "H41:ep=4,decouple=1,hsdp=4", "H22:ep=2,decouple=1,hsdp=4"), + steps=10, + ) + for name in ("C4", "H41", "H22"): + self._assert_training_matches(results[name], results["B4"]) + # Per-rank shards. Legacy ep=4 on 8 ranks: dense over fsdp=2 (D/2), experts over ep=4 then + # fsdp=2 (E/8). Decoupled: dense over dp_shard (D/dp_shard), experts over ep * efsdp == + # dp_shard (E/dp_shard); dp_shard is 8 for C4 and 4 for the HSDP layouts H41 / H22. + # + # 中文注释：ratio = 被测布局的每卡参数量 / 参考布局 B4 的每卡参数量，按各布局的分片数推导： + # B4 （legacy，ep=4） ：mesh (fsdp=2, ep=4)；dense 在 ep 维复制、在 fsdp 维分 2 份 → D/2； + # expert 先按 ep 分 4 份、再按 fsdp 分 2 份 → E/8。 + # C4 （解耦，ep=4） ：dp_shard=8，efsdp=8/4=2；dense → D/8；expert 按 ep×efsdp=8 → E/8。 + # H41（解耦，ep=4，hsdp=4）：dp_shard=4，replicate=2，efsdp=1；dense → D/4；expert 按 4×1 → E/4。 + # H22（解耦，ep=2，hsdp=4）：dp_shard=4，replicate=2，efsdp=2；dense → D/4；expert 按 2×2 → E/4。 + # 于是 C4 vs B4：dense (D/8)/(D/2)=1/4，expert (E/8)/(E/8)=1； + # H41、H22 vs B4：dense (D/4)/(D/2)=1/2，expert (E/4)/(E/8)=2。 + # 写成 (1/8)/(1/2) 而不是 0.25，是为了让"每卡份额 / 参考份额"的推导过程留在代码里。 + self._assert_param_memory(results["C4"], results["B4"], dense_ratio=(1 / 8) / (1 / 2), expert_ratio=1.0) + self._assert_param_memory(results["H41"], results["B4"], dense_ratio=(1 / 4) / (1 / 2), expert_ratio=2.0) + self._assert_param_memory(results["H22"], results["B4"], dense_ratio=(1 / 4) / (1 / 2), expert_ratio=2.0) + + @property + def world_size(self) -> int: + return 8 + + def _run_layouts(self, specs: tuple[str, ...], steps: int) -> dict[str, dict[str, Any]]: + self.create_pg("cuda") + root = self._make_shared_tmpdir() + try: + hf_dir = _build_tiny_checkpoint(root) + results: dict[str, dict[str, Any]] = {} + for spec in specs: + mode = parse_mode(spec) + results[mode["name"]] = run_mode( + mode, + hf_dir, + steps=steps, + seq_len=SEQ_LEN, + lr=LR, + dispatcher="all2all", + fp8=False, + grad_norm_steps=(0,), + tag=self._testMethodName, + ) + dist.barrier() + return results + finally: + self._cleanup_shared_tmpdir(root) + + def _assert_training_matches(self, result: dict[str, Any], reference: dict[str, Any]) -> None: + name, ref_name = result["mode"]["name"], reference["mode"]["name"] + loss_diff = max_rel_diff(result["losses"], reference["losses"]) + self.assertLessEqual(loss_diff, LOSS_RTOL, f"{name} vs {ref_name}: loss curves differ by {loss_diff:.2e}") + grad_diff = max_rel_diff(result["grad_norms"], reference["grad_norms"]) + self.assertLessEqual(grad_diff, GRAD_NORM_RTOL, f"{name} vs {ref_name}: grad norms differ by {grad_diff:.2e}") + for step, ref_norms in reference["param_grad_norms"].items(): + norms = result["param_grad_norms"][step] + self.assertEqual(set(norms), set(ref_norms), f"{name} vs {ref_name}: parameter sets differ at step {step}") + worst = max(ref_norms, key=lambda key: abs(norms[key] - ref_norms[key]) / max(abs(ref_norms[key]), 1e-12)) + diff = abs(norms[worst] - ref_norms[worst]) / max(abs(ref_norms[worst]), 1e-12) + self.assertLessEqual( + diff, + GRAD_NORM_RTOL, + f"{name} vs {ref_name}: grad norm of {worst} differs by {diff:.2e} at step {step} " + f"({norms[worst]:.6g} vs {ref_norms[worst]:.6g})", + ) + + def _assert_param_memory( + self, result: dict[str, Any], reference: dict[str, Any], dense_ratio: float, expert_ratio: float + ) -> None: + # 中文注释：比较两个布局训练时记录的"每卡本地参数字节数"（`param_memory`：对每个参数取 + # DTensor 的本地分片求和，按名字里是否含 `.experts` 分成 expert / dense 两类）。 + # 断言 result 的值 ≈ reference 的值 × ratio，容差 MEMORY_RTOL（±25%）留给 FSDP 的 padding。 + # 期望比值由调用方按布局的分片数给出，见各调用点的推导。 + name, ref_name = result["mode"]["name"], reference["mode"]["name"] + for key, ratio in (("dense_param_mib", dense_ratio), ("expert_param_mib", expert_ratio)): + expected = reference["memory"][key] * ratio + actual = result["memory"][key] + self.assertAlmostEqual( + actual, + expected, + delta=expected * MEMORY_RTOL, + msg=f"{name} vs {ref_name}: {key} {actual:.1f} MiB, expected {expected:.1f} MiB per rank", + ) + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestDecoupledEpFsdpCheckpoint(_SharedTmpDirMixin, DeterministicDDPTestCase): + def test_hf_export_dcp_resume_and_cross_layout_reshard(self) -> None: + self.create_pg("cuda") + root = self._make_shared_tmpdir() + try: + hf_dir = _build_tiny_checkpoint(root) + vocab_size = Qwen3MoEConfig.from_hf(hf_dir).vocab_size + src_bf16 = root / "source_bf16" + if dist.get_rank() == 0: + # `save_hf` writes bf16; compare against the bf16-cast source. + src_bf16.mkdir() + 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() + + modes = [ + parse_mode(s) + for s in ("A:ep=1", "C:ep=8,decouple=1", "C4:ep=4,decouple=1", "H41:ep=4,decouple=1,hsdp=4") + ] + continuous: dict[str, list[float]] = {} + for mode in modes: + continuous[mode["name"]] = self._check_round_trip( + mode, hf_dir, src_bf16, root / mode["name"], vocab_size + ) + + # Step-10 exports of every layout vs. the ep=1 baseline (same data, same optimizer). + for mode in modes[1:]: + if dist.get_rank() == 0: + assert_hf_dirs_close( + root / "A" / "hf_step10", + root / mode["name"] / "hf_step10", + rtol=HF_LAYOUT_RTOL, + atol=HF_LAYOUT_ATOL, + label=f"hf_step10 {mode['name']} vs A", + ) + dist.barrier() + + # DCP resharding across layouts: the step-5 checkpoint of `src` continues under `dst`. + for src, dst in (("C", "A"), ("A", "C"), ("C4", "H41")): + dst_mode = next(m for m in modes if m["name"] == dst) + engine = build_engine(dst_mode, hf_dir, tag=f"cross_{src}") + engine.load_dcp(root / src / "dcp_step5") + losses = train(engine, range(5, 10), vocab_size, SEQ_LEN) + release(engine) + diff = max_rel_diff(losses, continuous[src]) + self.assertLessEqual(diff, LOSS_RTOL, f"cross load {src} -> {dst}: losses differ by {diff:.2e}") + dist.barrier() + finally: + self._cleanup_shared_tmpdir(root) + + @property + def world_size(self) -> int: + return 8 + + def _check_round_trip( + self, mode: LayoutMode, hf_dir: Path, src_bf16: Path, mode_dir: Path, vocab_size: int + ) -> list[float]: + name = mode["name"] + engine = build_engine(mode, hf_dir, tag="main") + engine.from_hf(hf_path=hf_dir, strict=True) + engine.save_hf(str(mode_dir / "hf_step0")) + dist.barrier() + if dist.get_rank() == 0: + assert_hf_dirs_close( + src_bf16, mode_dir / "hf_step0", rtol=0.0, atol=0.0, label=f"{name} hf_step0 vs source" + ) + + train(engine, range(0, 5), vocab_size, SEQ_LEN) + engine.save_dcp(mode_dir / "dcp_step5") + dist.barrier() + losses_continuous = train(engine, range(5, 10), vocab_size, SEQ_LEN) + engine.save_hf(str(mode_dir / "hf_step10")) + dist.barrier() + release(engine) + + resumed = build_engine(mode, hf_dir, tag="resume") + resumed.load_dcp(mode_dir / "dcp_step5") + losses_resumed = train(resumed, range(5, 10), vocab_size, SEQ_LEN) + resumed.save_hf(str(mode_dir / "hf_step10_resumed")) + dist.barrier() + release(resumed) + + diff = max_rel_diff(losses_resumed, losses_continuous) + self.assertLessEqual(diff, LOSS_RTOL, f"{name}: resumed losses differ from the continuous run by {diff:.2e}") + if dist.get_rank() == 0: + assert_hf_dirs_close( + mode_dir / "hf_step10", + mode_dir / "hf_step10_resumed", + rtol=HF_RESUME_RTOL, + atol=HF_RESUME_ATOL, + label=f"{name} hf_step10 resumed vs continuous", + ) + dist.barrier() + return losses_continuous + + +class TestHfDirComparison: + """CPU checks of the checkpoint comparison the GPU gates rely on.""" + + @staticmethod + def _write(path: Path, tensors: dict[str, torch.Tensor]) -> Path: + path.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(path / "model.safetensors")) + return path + + def test_bit_exact_passes_and_one_ulp_fails(self, tmp_path: Path) -> None: + base = {"w": torch.linspace(-1, 1, 64, dtype=torch.bfloat16)} + lhs = self._write(tmp_path / "lhs", base) + rhs = self._write(tmp_path / "rhs", {"w": base["w"].clone()}) + assert_hf_dirs_close(lhs, rhs, rtol=0.0, atol=0.0, label="same") + + nudged = base["w"].clone() + # One bf16 ulp: bump the bit pattern (a float32 `nextafter` would round back to the same bf16). + nudged[3] = (nudged[3].view(torch.int16) + 1).view(torch.bfloat16) + rhs_ulp = self._write(tmp_path / "rhs_ulp", {"w": nudged}) + with pytest.raises(AssertionError, match="one ulp: w"): + assert_hf_dirs_close(lhs, rhs_ulp, rtol=0.0, atol=0.0, label="one ulp") + # The layout tolerance admits a few ulps. + assert_hf_dirs_close(lhs, rhs_ulp, rtol=HF_LAYOUT_RTOL, atol=HF_LAYOUT_ATOL, label="tolerant") + + def test_key_and_shape_mismatches_fail(self, tmp_path: Path) -> None: + lhs = self._write(tmp_path / "lhs", {"w": torch.ones(4, dtype=torch.bfloat16)}) + with pytest.raises(AssertionError, match="key sets differ"): + assert_hf_dirs_close( + lhs, self._write(tmp_path / "keys", {"v": torch.ones(4, dtype=torch.bfloat16)}), 0, 0, "keys" + ) + with pytest.raises(AssertionError, match="expected \\(4,\\)"): + assert_hf_dirs_close( + lhs, self._write(tmp_path / "shape", {"w": torch.ones(8, dtype=torch.bfloat16)}), 0, 0, "shape" + ) + + def test_max_rel_diff(self) -> None: + assert max_rel_diff([1.0, 2.0], [1.0, 2.0]) == 0.0 + assert max_rel_diff([1.0, 2.2], [1.0, 2.0]) == pytest.approx(0.1) + with pytest.raises(ValueError, match="length mismatch"): + max_rel_diff([1.0], [1.0, 2.0]) diff --git a/tests/model/run_decoupled_ep_fsdp_ckpt.py b/tests/model/run_decoupled_ep_fsdp_ckpt.py index 06f505c66a..50e6ffafe4 100644 --- a/tests/model/run_decoupled_ep_fsdp_ckpt.py +++ b/tests/model/run_decoupled_ep_fsdp_ckpt.py @@ -11,102 +11,35 @@ 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 gc import json import os import shutil -import sys from pathlib import Path from typing import Any import torch import torch.distributed as dist -from safetensors.torch import load_file - - -sys.path.insert(0, str(Path(__file__).resolve().parent)) -from run_decoupled_ep_fsdp_numerics import build_hf_checkpoint, make_batch, parse_mode # noqa: E402 - -from xtuner.v1.config import AdamWConfig, FSDPConfig # noqa: E402 -from xtuner.v1.engine.train_engine import TrainEngine # noqa: E402 -from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig # noqa: E402 - - -def build_engine(mode: dict[str, Any], hf_path: Path, tag: str) -> TrainEngine: - model_cfg = Qwen3MoEConfig.from_hf(hf_path) - model_cfg.ep_size = mode["ep"] - model_cfg.dispatcher = "all2all" - model_cfg.compile_cfg = False - model_cfg.mesh_prefix = f"ckpt_{mode['name']}_{tag}" - fsdp_cfg = FSDPConfig( - ep_size=mode["ep"], - decouple_ep_fsdp=mode["decouple"], - hsdp_sharding_size=mode["hsdp"], - torch_compile=False, - ) - return TrainEngine(model_cfg=model_cfg, optim_cfg=AdamWConfig(lr=1e-4), fsdp_cfg=fsdp_cfg) - - -def train(engine: TrainEngine, steps: range, vocab_size: int, seq_len: int) -> list[float]: - losses = [] - for step in steps: - info = engine.train_step([make_batch(step, dist.get_rank(), vocab_size, seq_len)]) - grad_norm = engine.clip_grad_norm() - engine.step_optimizer(grad_norm) - losses.append(float(info["logs_info"]["reduced_llm_loss"])) - return losses - - -def load_hf_dir(path: Path) -> dict[str, torch.Tensor]: - tensors: dict[str, torch.Tensor] = {} - for file in sorted(path.glob("*.safetensors")): - tensors.update(load_file(str(file))) - return tensors - - -def compare_hf(lhs: Path, rhs: Path) -> dict[str, Any]: - a = load_hf_dir(lhs) - b = load_hf_dir(rhs) - missing = sorted(set(a) - set(b)) - extra = sorted(set(b) - set(a)) - max_abs = 0.0 - max_rel = 0.0 - mismatched = 0 - total = 0 - worst = "" - for key in sorted(set(a) & set(b)): - x = a[key].to(torch.float32) - y = b[key].to(torch.float32) - if x.shape != y.shape: - return {"error": f"shape mismatch for {key}: {tuple(x.shape)} vs {tuple(y.shape)}"} - diff = (x - y).abs() - cur = float(diff.max()) - if cur > max_abs: - max_abs, worst = cur, key - max_rel = max(max_rel, float((diff / x.abs().clamp_min(1e-6)).max())) - mismatched += int((diff > 0).sum()) - total += diff.numel() - return { - "keys": len(set(a) & set(b)), - "missing_in_rhs": missing[:5], - "extra_in_rhs": extra[:5], - "max_abs_diff": max_abs, - "max_rel_diff": max_rel, - "mismatched_elements": mismatched, - "total_elements": total, - "worst_key": worst, - } - - -def release(engine: TrainEngine) -> None: - del engine - gc.collect() - torch.cuda.empty_cache() +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: @@ -134,8 +67,6 @@ def main() -> None: src_bf16 = out / "source_bf16" if rank == 0: src_bf16.mkdir(exist_ok=True) - from safetensors.torch import save_file - 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() @@ -152,7 +83,7 @@ def main() -> None: if rank == 0: print(f"===== {mode}", flush=True) - engine = build_engine(mode, hf_dir, "main") + 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() @@ -168,7 +99,7 @@ def main() -> None: dist.barrier() release(engine) - resumed = build_engine(mode, hf_dir, "resume") + 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")) @@ -178,9 +109,7 @@ def main() -> None: 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( - abs(a - b) / abs(b) for a, b in zip(losses_resumed, losses_cont, strict=True) - ) + 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" @@ -200,14 +129,14 @@ def main() -> None: 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"cross_{src}") + 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(abs(a - b) / abs(b) for a, b in zip(losses, ref, strict=True)), + "max_rel_loss_diff_vs_src_continuous": max_rel_diff(losses, ref), } if rank == 0: diff --git a/tests/model/run_decoupled_ep_fsdp_numerics.py b/tests/model/run_decoupled_ep_fsdp_numerics.py index ed819277c9..d8dc0d7494 100644 --- a/tests/model/run_decoupled_ep_fsdp_numerics.py +++ b/tests/model/run_decoupled_ep_fsdp_numerics.py @@ -1,7 +1,11 @@ """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. +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): @@ -13,203 +17,14 @@ """ import argparse -import gc import json import os -import time from pathlib import Path -from typing import Any import torch import torch.distributed as dist -from torch.distributed.tensor import DTensor -from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM -from xtuner.v1.config import AdamWConfig, FSDPConfig -from xtuner.v1.engine.train_engine import TrainEngine -from xtuner.v1.float8.config import Float8Config, ScalingGranularity -from xtuner.v1.loss.ce_loss import CELossConfig -from xtuner.v1.model.base import ModelItem -from xtuner.v1.model.moe.moe import SequenceContext -from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig -from xtuner.v1.utils.device import get_device -from xtuner.v1.utils.dtensor import cal_total_norm - - -DEVICE = get_device() -GRAD_NORM_STEPS = (0, 1, 25) - - -def parse_mode(spec: str) -> dict[str, Any]: - name, _, kv = spec.partition(":") - mode: dict[str, Any] = {"name": name, "ep": 1, "decouple": False, "hsdp": None} - for item in filter(None, kv.split(",")): - key, value = item.split("=") - if key == "ep": - mode["ep"] = int(value) - elif key == "decouple": - mode["decouple"] = bool(int(value)) - elif key == "hsdp": - mode["hsdp"] = int(value) - else: - raise ValueError(f"unknown mode key {key}") - return mode - - -MODEL_SIZES: dict[str, dict[str, int]] = { - # ~108M params: numerics (fast, 50 steps in < 10 s) - "tiny": { - "vocab_size": 4096, - "hidden_size": 512, - "intermediate_size": 1024, - "moe_intermediate_size": 256, - "num_hidden_layers": 4, - "num_attention_heads": 8, - "num_key_value_heads": 4, - "head_dim": 64, - "num_experts": 64, - "num_experts_per_tok": 4, - }, - # ~3.4B params (3.2B in routed experts): step-time / memory comparison - "medium": { - "vocab_size": 32768, - "hidden_size": 2048, - "intermediate_size": 4096, - "moe_intermediate_size": 1024, - "num_hidden_layers": 8, - "num_attention_heads": 16, - "num_key_value_heads": 4, - "head_dim": 128, - "num_experts": 64, - "num_experts_per_tok": 4, - }, -} - - -def build_hf_checkpoint(path: Path, seed: int, model_size: str) -> None: - config = Qwen3MoeConfig( - **MODEL_SIZES[model_size], - max_position_embeddings=8192, - norm_topk_prob=True, - tie_word_embeddings=False, - bos_token_id=1, - eos_token_id=2, - rope_theta=1000000.0, - use_sliding_window=False, - max_window_layers=MODEL_SIZES[model_size]["num_hidden_layers"], - ) - torch.manual_seed(seed) - model = Qwen3MoeForCausalLM(config) - model.save_pretrained(path, safe_serialization=True) - - -def make_batch(step: int, rank: int, vocab_size: int, seq_len: int) -> ModelItem: - generator = torch.Generator().manual_seed(100_000 + step * 1024 + rank) - tokens = torch.randint(0, vocab_size, (1, seq_len + 1), generator=generator) - input_ids = tokens[:, :-1] - labels = tokens[:, 1:].to(DEVICE) - seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) - loss_cfg = CELossConfig() - loss_ctx = loss_cfg.build(data={"shifted_labels": labels}, sp_mesh=None) - loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] - return ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) - - -def clean_name(name: str) -> str: - return name.replace("_checkpoint_wrapped_module.", "") - - -def param_memory(model: torch.nn.Module) -> dict[str, float]: - expert_bytes = 0 - dense_bytes = 0 - for name, param in model.named_parameters(): - local = param.to_local() if isinstance(param, DTensor) else param - nbytes = local.numel() * local.element_size() - if ".experts" in name: - expert_bytes += nbytes - else: - dense_bytes += nbytes - return {"expert_param_mib": expert_bytes / 2**20, "dense_param_mib": dense_bytes / 2**20} - - -def per_param_grad_norms(model: torch.nn.Module) -> dict[str, float]: - norms: dict[str, float] = {} - for name, param in model.named_parameters(): - if param.grad is None: - continue - if isinstance(param.grad, DTensor): - norm = cal_total_norm([param.grad], norm_type=2.0, foreach=True, dtype=torch.float32) - else: - norm = torch.linalg.vector_norm(param.grad, 2.0, dtype=torch.float32) - norms[clean_name(name)] = float(norm) - return norms - - -def run_mode( - mode: dict[str, Any], hf_path: Path, steps: int, seq_len: int, lr: float, dispatcher: str, fp8: bool -) -> dict[str, Any]: - rank = dist.get_rank() - torch.cuda.reset_peak_memory_stats() - torch.cuda.synchronize() - base_allocated = torch.cuda.memory_allocated() - - model_cfg = Qwen3MoEConfig.from_hf(hf_path) - model_cfg.ep_size = mode["ep"] - model_cfg.dispatcher = dispatcher - model_cfg.compile_cfg = False - model_cfg.mesh_prefix = f"numerics_{mode['name']}" - if fp8: - model_cfg.float8_cfg = Float8Config( - scaling_granularity_gemm=ScalingGranularity.TILEWISE, - scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, - ) - fsdp_kwargs: dict[str, Any] = {"ep_size": mode["ep"], "hsdp_sharding_size": mode["hsdp"], "torch_compile": False} - if mode["decouple"]: - # Only pass the flag when set so the script also runs against checkouts that predate it. - fsdp_kwargs["decouple_ep_fsdp"] = True - fsdp_cfg = FSDPConfig(**fsdp_kwargs) - optim_cfg = AdamWConfig(lr=lr) - - engine = TrainEngine(model_cfg=model_cfg, optim_cfg=optim_cfg, fsdp_cfg=fsdp_cfg) - engine.from_hf(hf_path=hf_path, strict=True) - torch.cuda.synchronize() - result: dict[str, Any] = { - "mode": mode, - "memory": param_memory(engine.model), - "losses": [], - "grad_norms": [], - "param_grad_norms": {}, - "step_time_s": None, - } - result["memory"]["allocated_after_load_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 - - step_times = [] - for step in range(steps): - torch.cuda.synchronize() - t0 = time.perf_counter() - batch = make_batch(step, rank, model_cfg.vocab_size, seq_len) - info = engine.train_step([batch]) - grad_norm = engine.clip_grad_norm() - if step in GRAD_NORM_STEPS: - result["param_grad_norms"][str(step)] = per_param_grad_norms(engine.model) - engine.step_optimizer(grad_norm) - torch.cuda.synchronize() - step_times.append(time.perf_counter() - t0) - result["losses"].append(float(info["logs_info"]["reduced_llm_loss"])) - result["grad_norms"].append(float(grad_norm)) - if step == 0: - result["memory"]["allocated_after_step0_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 - - warm = step_times[5:] if len(step_times) > 10 else step_times - result["step_time_s"] = sum(warm) / len(warm) - result["step_times_s"] = step_times - result["memory"]["peak_allocated_mib"] = (torch.cuda.max_memory_allocated() - base_allocated) / 2**20 - result["memory"]["peak_reserved_mib"] = torch.cuda.max_memory_reserved() / 2**20 - - del engine - gc.collect() - torch.cuda.empty_cache() - return result +from xtuner._testing.decoupled_ep_fsdp import MODEL_SIZES, build_hf_checkpoint, parse_mode, run_mode def main() -> None: diff --git a/xtuner/_testing/decoupled_ep_fsdp.py b/xtuner/_testing/decoupled_ep_fsdp.py new file mode 100644 index 0000000000..fe1611dbd7 --- /dev/null +++ b/xtuner/_testing/decoupled_ep_fsdp.py @@ -0,0 +1,461 @@ +"""Shared helpers for the decoupled EP/FSDP ("dp2ep") numerics and checkpoint checks. + +The GPU pytest gates in ``tests/engine/test_decoupled_ep_fsdp_train_engine.py`` and the manual +experiment scripts under ``tests/model/`` (``run_decoupled_ep_fsdp_numerics.py`` and +``run_decoupled_ep_fsdp_ckpt.py``) build the same tiny random Qwen3-MoE, feed it the same token +stream and compare layouts with the same functions, so a number in a report and a gate threshold +mean the same thing. +""" + +from __future__ import annotations + +import gc +import time +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import Any, TypedDict + +import torch +import torch.distributed as dist +from safetensors.torch import load_file +from torch.distributed.tensor import DTensor + +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.float8.config import Float8Config, ScalingGranularity +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.base import ModelItem +from xtuner.v1.model.moe.moe import SequenceContext +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig +from xtuner.v1.utils.device import get_device +from xtuner.v1.utils.dtensor import cal_total_norm + + +DEVICE = get_device() + +MODEL_SIZES: dict[str, dict[str, int]] = { + # ~108M params: numerics (fast, 50 steps in < 10 s) + "tiny": { + "vocab_size": 4096, + "hidden_size": 512, + "intermediate_size": 1024, + "moe_intermediate_size": 256, + "num_hidden_layers": 4, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "head_dim": 64, + "num_experts": 64, + "num_experts_per_tok": 4, + }, + # ~3.4B params (3.2B in routed experts): step-time / memory comparison + "medium": { + "vocab_size": 32768, + "hidden_size": 2048, + "intermediate_size": 4096, + "moe_intermediate_size": 1024, + "num_hidden_layers": 8, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 128, + "num_experts": 64, + "num_experts_per_tok": 4, + }, +} + + +class LayoutMode(TypedDict): + """One EP/FSDP layout under test. + + ``name`` labels the layout in results, ``ep`` is the expert-parallel size, ``decouple`` selects + ``FSDPConfig.decouple_ep_fsdp`` and ``hsdp`` is ``FSDPConfig.hsdp_sharding_size`` (``None`` for + plain FSDP over the world). + """ + + name: str + ep: int + decouple: bool + hsdp: int | None + + +def parse_mode(spec: str) -> LayoutMode: + """Parse a ``:key=value[,key=value...]`` layout spec. + + Args: + spec (str): Mode spec with keys ``ep`` (int), ``decouple`` (0/1) and ``hsdp`` (int). + + Returns: + LayoutMode: The parsed layout; unspecified keys default to ``ep=1, decouple=0, hsdp=None``. + """ + name, _, kv = spec.partition(":") + mode: LayoutMode = {"name": name, "ep": 1, "decouple": False, "hsdp": None} + for item in filter(None, kv.split(",")): + key, value = item.split("=") + if key == "ep": + mode["ep"] = int(value) + elif key == "decouple": + mode["decouple"] = bool(int(value)) + elif key == "hsdp": + mode["hsdp"] = int(value) + else: + raise ValueError(f"unknown mode key {key}") + return mode + + +def build_hf_checkpoint(path: Path, seed: int, model_size: str) -> None: + """Write a randomly initialised Qwen3-MoE HF checkpoint to ``path``. + + Args: + path (Path): Output directory (created by ``save_pretrained``). + seed (int): Seed for the random initialisation. + model_size (str): A key of :data:`MODEL_SIZES`. + """ + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + + config = Qwen3MoeConfig( + **MODEL_SIZES[model_size], + max_position_embeddings=8192, + norm_topk_prob=True, + tie_word_embeddings=False, + bos_token_id=1, + eos_token_id=2, + rope_theta=1000000.0, + use_sliding_window=False, + max_window_layers=MODEL_SIZES[model_size]["num_hidden_layers"], + ) + torch.manual_seed(seed) + model = Qwen3MoeForCausalLM(config) + model.save_pretrained(path, safe_serialization=True) + + +def make_batch(step: int, rank: int, vocab_size: int, seq_len: int) -> ModelItem: + """Build the deterministic random token batch of ``(step, rank)``. + + Args: + step (int): Training step; together with ``rank`` it seeds the batch. + rank (int): Data-parallel rank. + vocab_size (int): Vocabulary size to sample tokens from. + seq_len (int): Number of input tokens. + + Returns: + ModelItem: Sequence context and CE loss context on :data:`DEVICE`. + """ + generator = torch.Generator().manual_seed(100_000 + step * 1024 + rank) + tokens = torch.randint(0, vocab_size, (1, seq_len + 1), generator=generator) + input_ids = tokens[:, :-1] + labels = tokens[:, 1:].to(DEVICE) + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) + loss_cfg = CELossConfig() + loss_ctx = loss_cfg.build(data={"shifted_labels": labels}, sp_mesh=None) + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + return ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) + + +def build_engine( + mode: LayoutMode, + hf_path: Path, + tag: str, + dispatcher: str = "all2all", + fp8: bool = False, + lr: float = 1e-4, +) -> TrainEngine: + """Build a ``TrainEngine`` for ``mode`` without loading weights. + + Args: + mode (LayoutMode): Layout under test. + hf_path (Path): HF checkpoint the model config is read from. + tag (str): Unique tag for the mesh prefix; every engine alive in a process needs its own. + dispatcher (str): MoE dispatcher (``"all2all"`` or ``"deepep"``). + fp8 (bool): Enable tile-wise float8 for linear and grouped linear layers. + lr (float): AdamW learning rate. + + Returns: + TrainEngine: Sharded engine with freshly built (uninitialised) parameters. + """ + model_cfg = Qwen3MoEConfig.from_hf(hf_path) + model_cfg.ep_size = mode["ep"] + model_cfg.dispatcher = dispatcher + model_cfg.compile_cfg = False + model_cfg.mesh_prefix = f"{tag}_{mode['name']}" + if fp8: + model_cfg.float8_cfg = Float8Config( + scaling_granularity_gemm=ScalingGranularity.TILEWISE, + scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE, + ) + fsdp_cfg = FSDPConfig( + ep_size=mode["ep"], + decouple_ep_fsdp=mode["decouple"], + hsdp_sharding_size=mode["hsdp"], + torch_compile=False, + ) + return TrainEngine(model_cfg=model_cfg, optim_cfg=AdamWConfig(lr=lr), fsdp_cfg=fsdp_cfg) + + +def train(engine: TrainEngine, steps: Iterable[int], vocab_size: int, seq_len: int) -> list[float]: + """Run ``steps`` optimizer steps on the shared token stream. + + Args: + engine (TrainEngine): Engine to train. + steps (Iterable[int]): Step indices; they select the batches via :func:`make_batch`. + vocab_size (int): Vocabulary size of the model. + seq_len (int): Tokens per batch. + + Returns: + list[float]: ``reduced_llm_loss`` of every step. + """ + losses = [] + for step in steps: + info = engine.train_step([make_batch(step, dist.get_rank(), vocab_size, seq_len)]) + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + losses.append(float(info["logs_info"]["reduced_llm_loss"])) + return losses + + +def run_mode( + mode: LayoutMode, + hf_path: Path, + steps: int, + seq_len: int, + lr: float, + dispatcher: str, + fp8: bool, + grad_norm_steps: Sequence[int] = (0, 1, 25), + tag: str = "numerics", +) -> dict[str, Any]: + """Train ``mode`` from the HF checkpoint and record losses, grad norms, memory and step time. + + Args: + mode (LayoutMode): Layout under test. + hf_path (Path): HF checkpoint to load. + steps (int): Number of optimizer steps. + seq_len (int): Tokens per batch. + lr (float): AdamW learning rate. + dispatcher (str): MoE dispatcher. + fp8 (bool): Enable tile-wise float8. + grad_norm_steps (Sequence[int]): Steps whose per-parameter grad norms are recorded. + tag (str): Mesh-prefix tag passed to :func:`build_engine`. + + Returns: + dict[str, Any]: ``losses``, ``grad_norms`` (per step), ``param_grad_norms`` (per recorded + step, keyed by parameter name), ``memory`` (per-rank MiB) and ``step_time_s``. + """ + rank = dist.get_rank() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + base_allocated = torch.cuda.memory_allocated() + + engine = build_engine(mode, hf_path, tag, dispatcher=dispatcher, fp8=fp8, lr=lr) + engine.from_hf(hf_path=hf_path, strict=True) + vocab_size = engine.model.config.vocab_size + torch.cuda.synchronize() + result: dict[str, Any] = { + "mode": mode, + "memory": param_memory(engine.model), + "losses": [], + "grad_norms": [], + "param_grad_norms": {}, + "step_time_s": None, + } + result["memory"]["allocated_after_load_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 + + step_times = [] + for step in range(steps): + torch.cuda.synchronize() + t0 = time.perf_counter() + batch = make_batch(step, rank, vocab_size, seq_len) + info = engine.train_step([batch]) + grad_norm = engine.clip_grad_norm() + if step in grad_norm_steps: + result["param_grad_norms"][str(step)] = per_param_grad_norms(engine.model) + engine.step_optimizer(grad_norm) + torch.cuda.synchronize() + step_times.append(time.perf_counter() - t0) + result["losses"].append(float(info["logs_info"]["reduced_llm_loss"])) + result["grad_norms"].append(float(grad_norm)) + if step == 0: + result["memory"]["allocated_after_step0_mib"] = (torch.cuda.memory_allocated() - base_allocated) / 2**20 + + warm = step_times[5:] if len(step_times) > 10 else step_times + result["step_time_s"] = sum(warm) / len(warm) + result["step_times_s"] = step_times + result["memory"]["peak_allocated_mib"] = (torch.cuda.max_memory_allocated() - base_allocated) / 2**20 + result["memory"]["peak_reserved_mib"] = torch.cuda.max_memory_reserved() / 2**20 + + release(engine) + return result + + +def param_memory(model: torch.nn.Module) -> dict[str, float]: + """Per-rank parameter bytes of the local shards, split into routed experts and everything else. + + Args: + model (torch.nn.Module): Sharded model. + + Returns: + dict[str, float]: ``expert_param_mib`` and ``dense_param_mib``. + """ + expert_bytes = 0 + dense_bytes = 0 + for name, param in model.named_parameters(): + local = param.to_local() if isinstance(param, DTensor) else param + nbytes = local.numel() * local.element_size() + if ".experts" in name: + expert_bytes += nbytes + else: + dense_bytes += nbytes + return {"expert_param_mib": expert_bytes / 2**20, "dense_param_mib": dense_bytes / 2**20} + + +def per_param_grad_norms(model: torch.nn.Module) -> dict[str, float]: + """Global L2 norm of every parameter's gradient, keyed by the unwrapped parameter name. + + Args: + model (torch.nn.Module): Sharded model after ``backward``. + + Returns: + dict[str, float]: Parameter name (without ``_checkpoint_wrapped_module.``) to grad norm. + """ + norms: dict[str, float] = {} + for name, param in model.named_parameters(): + if param.grad is None: + continue + if isinstance(param.grad, DTensor): + norm = cal_total_norm([param.grad], norm_type=2.0, foreach=True, dtype=torch.float32) + else: + norm = torch.linalg.vector_norm(param.grad, 2.0, dtype=torch.float32) + norms[clean_name(name)] = float(norm) + return norms + + +def clean_name(name: str) -> str: + """Strip the activation-checkpoint wrapper prefix from a parameter name. + + Args: + name (str): Name from ``named_parameters``. + + Returns: + str: The name without ``_checkpoint_wrapped_module.``. + """ + return name.replace("_checkpoint_wrapped_module.", "") + + +def rel_diff(a: float, b: float) -> float: + """Relative difference ``|a - b| / |b|`` (``b`` is the reference). + + Args: + a (float): Value under test. + b (float): Reference value. + + Returns: + float: The relative difference; the denominator is clamped to ``1e-12``. + """ + return abs(a - b) / max(abs(b), 1e-12) + + +def max_rel_diff(values: Sequence[float], reference: Sequence[float]) -> float: + """Largest element-wise :func:`rel_diff` of two equally long sequences. + + Args: + values (Sequence[float]): Values under test. + reference (Sequence[float]): Reference values. + + Returns: + float: ``max(rel_diff(v, r))`` over the zipped sequences. + """ + if len(values) != len(reference): + raise ValueError(f"length mismatch: {len(values)} vs {len(reference)}") + return max(rel_diff(a, b) for a, b in zip(values, reference)) + + +def load_hf_dir(path: Path) -> dict[str, torch.Tensor]: + """Load every ``*.safetensors`` file of an HF checkpoint directory. + + Args: + path (Path): Checkpoint directory. + + Returns: + dict[str, torch.Tensor]: All tensors keyed by their checkpoint name. + """ + tensors: dict[str, torch.Tensor] = {} + for file in sorted(path.glob("*.safetensors")): + tensors.update(load_file(str(file))) + return tensors + + +def compare_hf(lhs: Path, rhs: Path) -> dict[str, Any]: + """Compare two HF checkpoint directories tensor by tensor. + + Args: + lhs (Path): Reference checkpoint. + rhs (Path): Checkpoint under test. + + Returns: + dict[str, Any]: Key counts, missing / extra keys, max absolute and relative differences, + the number of mismatched elements and the worst key; ``{"error": ...}`` on a shape mismatch. + """ + a = load_hf_dir(lhs) + b = load_hf_dir(rhs) + missing = sorted(set(a) - set(b)) + extra = sorted(set(b) - set(a)) + max_abs = 0.0 + max_rel = 0.0 + mismatched = 0 + total = 0 + worst = "" + for key in sorted(set(a) & set(b)): + x = a[key].to(torch.float32) + y = b[key].to(torch.float32) + if x.shape != y.shape: + return {"error": f"shape mismatch for {key}: {tuple(x.shape)} vs {tuple(y.shape)}"} + diff = (x - y).abs() + cur = float(diff.max()) + if cur > max_abs: + max_abs, worst = cur, key + max_rel = max(max_rel, float((diff / x.abs().clamp_min(1e-6)).max())) + mismatched += int((diff > 0).sum()) + total += diff.numel() + return { + "keys": len(set(a) & set(b)), + "missing_in_rhs": missing[:5], + "extra_in_rhs": extra[:5], + "max_abs_diff": max_abs, + "max_rel_diff": max_rel, + "mismatched_elements": mismatched, + "total_elements": total, + "worst_key": worst, + } + + +def assert_hf_dirs_close(lhs: Path, rhs: Path, rtol: float, atol: float, label: str) -> None: + """Fail unless two HF checkpoint directories hold the same keys, shapes and close values. + + Args: + lhs (Path): Reference checkpoint. + rhs (Path): Checkpoint under test. + rtol (float): Relative tolerance passed to ``torch.testing.assert_close``. + atol (float): Absolute tolerance passed to ``torch.testing.assert_close``. + label (str): Prefix for the failure message. + """ + a = load_hf_dir(lhs) + b = load_hf_dir(rhs) + if set(a) != set(b): + raise AssertionError( + f"{label}: key sets differ; missing in rhs {sorted(set(a) - set(b))[:5]}, " + f"extra in rhs {sorted(set(b) - set(a))[:5]}" + ) + for key in sorted(a): + x, y = a[key], b[key] + if x.shape != y.shape or x.dtype != y.dtype: + raise AssertionError(f"{label}: {key} has {tuple(y.shape)}/{y.dtype}, expected {tuple(x.shape)}/{x.dtype}") + torch.testing.assert_close(y.float(), x.float(), rtol=rtol, atol=atol, msg=lambda m: f"{label}: {key}: {m}") + + +def release(engine: TrainEngine) -> None: + """Drop an engine and return its device memory to the allocator. + + Args: + engine (TrainEngine): Engine to release; the caller must not use it afterwards. + """ + del engine + gc.collect() + torch.cuda.empty_cache() From 212c05ee7867a001b5962639dae5782d8f68776e Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 17 Sep 2026 10:04:50 +0000 Subject: [PATCH 10/14] [Docs] Record the topology validator and the pytest gates in the decoupled EP/FSDP design note MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - §3.1 shows the `model_validator`; §3.6 and the §5 evidence table link the 8-GPU gates next to the experiment scripts; the two resolved items leave §6.3. - reports/decoupled_ep_fsdp_review_zh.md: status note for N1 and W5. --- docs/design/decouple_ep_fsdp.md | 49 +++++++++++++++++---------------- 1 file changed, 26 insertions(+), 23 deletions(-) diff --git a/docs/design/decouple_ep_fsdp.md b/docs/design/decouple_ep_fsdp.md index 7d1a4476f2..37603d20e7 100644 --- a/docs/design/decouple_ep_fsdp.md +++ b/docs/design/decouple_ep_fsdp.md @@ -18,7 +18,7 @@ flowchart LR |---|---|---| | routed expert | `Shard(0)` on `ep`，再 FSDP over `fsdp` | 正确，且与新路径数学同构 | | dense（attention / router / shared expert / embedding / lm_head） | `Replicate()` on `ep`（`_replicate_other_params`），再 FSDP over `fsdp` | 参数、梯度、优化器状态在 EP 维复制 `ep` 份；每步一次额外的 coalesced all-reduce 补梯度 | -| HSDP | `assert ep_size == 1` | HSDP 与 EP 不能共存 | +| HSDP | 要求 `ep_size == 1`（`FSDPConfig` 校验） | HSDP 与 EP 不能共存 | EP 越大，dense 侧越亏。以 GLM-5.2-30B 单机 8 卡为例，EP=8 时每卡保存一份完整的 dense 参数与 fp32 优化器状态，legacy 路径 EP=8 的 reserved 显存顶到约 132 GiB 的 allocator 上限（allocated 峰值 120 GiB），每步触发 allocator 释放重试，步时退化到 EP=4 的 7 倍（见 §5）。 @@ -50,29 +50,34 @@ expert : Shard(0) over ep + FSDP over efsdp (+ replicate 上 HS ## 3. 实现 -各小节的"位置"行给出文件与行号，行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `5923dc1d`（2026-09-03）。 +各小节的"位置"行给出文件与行号，行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `0a7c6523`（2026-09-17，基于 upstream/main `6c06f96e`）。 ### 3.1 配置与校验 -位置：`xtuner/v1/config/fsdp.py` 48–65，`xtuner/v1/model/moe/moe.py` 1665–1678。 +位置：`xtuner/v1/config/fsdp.py` 44–73，`xtuner/v1/model/moe/moe.py` 1764–1771。 ```python # xtuner/v1/config/fsdp.py decouple_ep_fsdp: bool = False # 默认关闭，旧路径不受影响 -def model_post_init(self, __context): +@model_validator(mode="after") +def _validate_ep_fsdp_topology(self): + # ValueError（pydantic 包装为 ValidationError），不是 assert：python -O 下同样生效 + if self.ep_size < 1 or (self.hsdp_sharding_size is not None and self.hsdp_sharding_size < 1): + raise ValueError(...) # 必须是正整数 if self.hsdp_sharding_size is not None: if self.decouple_ep_fsdp: - assert self.hsdp_sharding_size % self.ep_size == 0 # 必须存在整数 efsdp - else: - assert self.ep_size == 1 # 旧断言只在旧路径保留 + if self.hsdp_sharding_size % self.ep_size != 0: # 必须存在整数 efsdp + raise ValueError(...) + elif self.ep_size != 1: # 旧约束只在旧路径保留 + raise ValueError(...) ``` -`MoE._init_decoupled_device_mesh` 再次校验 `world_size % dp_shard == 0` 与 `dp_shard % ep_size == 0`。ExpertTP（`expert_tp_size > 1`）与解耦同时开启时抛 `NotImplementedError`，ETP 的 `(Shard, InterleavedShard)` 二维 placement 需要单独设计 mesh 维。 +`MoE._init_decoupled_device_mesh` 再次校验 `world_size % dp_shard == 0` 与 `dp_shard % ep_size == 0`，同样抛 `ValueError`。ExpertTP（`expert_tp_size > 1`）与解耦同时开启时抛 `NotImplementedError`，ETP 的 `(Shard, InterleavedShard)` 二维 placement 需要单独设计 mesh 维。 ### 3.2 单一 root mesh -位置：`xtuner/v1/model/moe/moe.py` 1651–1716（`_init_decoupled_device_mesh`），1486–1492（分流）。 +位置：`xtuner/v1/model/moe/moe.py` 1743–1808（`_init_decoupled_device_mesh`），1576–1581（分流）。 ```python # xtuner/v1/model/moe/moe.py::_init_decoupled_device_mesh（节选） @@ -107,7 +112,7 @@ else: ### 3.3 两级 `fully_shard` -位置：`xtuner/v1/model/moe/moe.py` 1194–1225（decoder layer），1250–1260（prefetch），1328–1360（MTP），1631–1649（helper）。 +位置：`xtuner/v1/model/moe/moe.py` 1310–1342（decoder layer），1370–1380（prefetch），1419–1451（MTP），1723–1741（helper）。 ```python # xtuner/v1/model/moe/moe.py::fully_shard（节选） @@ -140,7 +145,7 @@ def _fully_shard_expert_blocks(self, module, mp_policy, reshard_after_forward): ### 3.4 梯度归约 -位置：`xtuner/v1/model/moe/moe.py` 1392–1396（分流），1583–1629（`_scale_and_reduce_grad_decoupled`）。 +位置：`xtuner/v1/model/moe/moe.py` 1482–1486（分流），1673–1722（`_scale_and_reduce_grad_decoupled`）。 不变式：每个参数的最终梯度等于全 data-parallel world 上的均值，与 `ep` / `efsdp` 取值无关。 @@ -178,7 +183,7 @@ for group, grads in grads_by_group.items(): # 每个 group 一次 coalesced al ### 3.5 FP8 -位置：`xtuner/v1/float8/float8_handler.py` 121–161、166–239、316–335；`xtuner/v1/float8/fsdp_utils.py` 127–137；`xtuner/v1/model/moe/moe.py` 1165–1176；`xtuner/v1/model/base.py` 1015。 +位置：`xtuner/v1/float8/float8_handler.py` 121–164、166–238、316–335；`xtuner/v1/float8/fsdp_utils.py` 127–137；`xtuner/v1/model/moe/moe.py` 1281–1292；`xtuner/v1/model/base.py` 1015。 tile-wise FP8 需要知道每个本地权重会被多少个 FSDP rank 切分（决定 padding），以及 scale 的 reduce-max 应跨哪些 rank。解耦后 dense 与 expert 的答案不同： @@ -207,7 +212,7 @@ self.expert_tilewise_reduce_mesh_mapping = self._build_strided_reduce_mesh_mappi ### 3.6 HF 权重与 DCP -位置：无代码改动。依赖 `xtuner/v1/utils/load_spec.py` 的 `LoadSpec.from_tensor`、`RuntimeLayout.from_dtensor`、`plan_hf_save`；验证脚本 `tests/model/run_decoupled_ep_fsdp_ckpt.py`。 +位置：无代码改动。依赖 `xtuner/v1/utils/load_spec.py` 的 `LoadSpec.from_tensor`、`RuntimeLayout.from_dtensor`、`plan_hf_save`；验证用例 `tests/engine/test_decoupled_ep_fsdp_train_engine.py::TestDecoupledEpFsdpCheckpoint`，实验脚本 `tests/model/run_decoupled_ep_fsdp_ckpt.py`。 `LoadSpec.from_tensor` 直接从 DTensor placement 记录 shard 历史，新 placement 无需特殊处理： @@ -224,7 +229,7 @@ dense（解耦 + HSDP） ### 3.7 RL 权重同步 -位置：`xtuner/v1/model/base.py` 1916–1947；`xtuner/v1/rl/weight_update/weight_iterator.py` 140–149、196–208、236–247；测试 `tests/rl/test_weight_iterator.py` 234–418。 +位置：`xtuner/v1/model/base.py` 1916–1947；`xtuner/v1/rl/weight_update/weight_iterator.py` 142–151、198、207–230、240–249；测试 `tests/rl/test_weight_iterator.py` 237–427。 Turbomind 的 layer-wise 更新只 gather FSDP shard、保留 EP-local 的 expert 切片。`BaseModel._fsdp_foreach_allgather` 按每个 `LoadSpec` 选择 gather group： @@ -267,13 +272,13 @@ def _param_owner(model, name): | 项目 | 证据 | 边界 | 链接 | |---|---|---|---| -| mesh / placement / 两级 FSDP（L0） | fake-PG 下 8/16/64 rank 共 29 例；旧布局在开关关闭时逐项固定 | 需要一张空闲 GPU 建 CUDA context，不通信 | [test_decoupled_ep_fsdp_mesh.py](../../tests/model/test_decoupled_ep_fsdp_mesh.py)、[baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md) | -| BF16 训练语义（L1） | tiny Qwen3-MoE 50 步：解耦 EP8 vs 旧 EP8 loss 最大相对差 `1.9e-5`（旧 EP8 vs EP1 为 `2.5e-5`）；3.4B 模型每卡 dense 参数 836 → 105 MiB，峰值 allocated 11861 → 8203 MiB，步时 170 → 158 ms | 单机；reduction order 不同，不是逐 bit 相同 | [L1.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L1.md)、[run_decoupled_ep_fsdp_numerics.py](../../tests/model/run_decoupled_ep_fsdp_numerics.py) | -| `efsdp > 1`、HSDP + EP（L2） | 单机缩比拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)` 均在噪声内；16/64 rank 形状由 L0 固定 | 真正的双节点 collective 未跑 | [L2.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L2.md)、[decisions.md D10](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | -| HF / DCP（L3） | 刚 `from_hf` 后导出 807 个 key 全部 bit-exact；DCP 同布局 resume 与跨布局 reshard 后 loss 连续 | 未覆盖 async save、FP8 checkpoint、GLM/MTP 导出、world size 变化 | [L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[run_decoupled_ep_fsdp_ckpt.py](../../tests/model/run_decoupled_ep_fsdp_ckpt.py) | +| mesh / placement / 两级 FSDP（L0） | fake-PG 下 8/16/64 rank 共 31 例（含配置校验与运行时校验的回归）；旧布局在开关关闭时逐项固定 | 需要一张空闲 GPU 建 CUDA context，不通信 | [test_decoupled_ep_fsdp_mesh.py](../../tests/model/test_decoupled_ep_fsdp_mesh.py)、[baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md) | +| BF16 训练语义（L1） | tiny Qwen3-MoE 50 步：解耦 EP8 vs 旧 EP8 loss 最大相对差 `1.9e-5`（旧 EP8 vs EP1 为 `2.5e-5`）；3.4B 模型每卡 dense 参数 836 → 105 MiB，峰值 allocated 11861 → 8203 MiB，步时 170 → 158 ms | 单机；reduction order 不同，不是逐 bit 相同 | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（8 卡 gate：loss 相对差 ≤ 1e-3、总/逐参数梯度范数 ≤ 1e-2、每卡参数显存 ±25%）、[L1.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L1.md)、[run_decoupled_ep_fsdp_numerics.py](../../tests/model/run_decoupled_ep_fsdp_numerics.py) | +| `efsdp > 1`、HSDP + EP（L2） | 单机缩比拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)` 均在噪声内；16/64 rank 形状由 L0 固定 | 真正的双节点 collective 未跑 | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（同上阈值，拓扑 `(1,2,4)`、`(2,1,4)`、`(2,2,2)`）、[L2.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L2.md)、[decisions.md D10](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | +| HF / DCP（L3） | 刚 `from_hf` 后导出 807 个 key 全部 bit-exact；DCP 同布局 resume 与跨布局 reshard 后 loss 连续 | 未覆盖 async save、FP8 checkpoint、GLM/MTP 导出、world size 变化 | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（8 卡 gate：step-0 导出逐 bit 相同、resume 与跨布局 loss 相对差 ≤ 1e-3、导出差异 ≤ 数个 bf16 ulp）、[L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[run_decoupled_ep_fsdp_ckpt.py](../../tests/model/run_decoupled_ep_fsdp_ckpt.py) | | FP8 | 3.4B 与 GLM-5.2 tile-wise FP8 训练 20 步稳定，差异量级接近同配置重复噪声 | FP8 + HSDP、FP8 + checkpoint 未测 | [L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | | GLM-5.2-30B，8×H200 目标配方 | EP4：1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB；EP8：11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB | EP8 的加速来自离开 allocator 上限（legacy 每步 2–3 次 alloc retry，解耦为 0），不是 collective 本身快 7 倍 | [GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | -| RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | [test_weight_iterator.py](../../tests/rl/test_weight_iterator.py#L234-L418) | +| RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | [test_weight_iterator.py](../../tests/rl/test_weight_iterator.py#L237-L427) | | legacy 兼容 | 开关默认关闭，旧分支执行逻辑未改动，L0 旧布局回归通过 | 未做 base-vs-branch 端到端逐 tensor 对照 | [baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md)、[decisions.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | `reports/` 下的报告不随代码 PR 合入，链接指向 fork 分支 `feat/decouple-ep-fsdp` 上的固定地址；测试与脚本是仓库内相对链接。 @@ -311,9 +316,7 @@ def _param_owner(model, name): ### 6.3 其它 -- 新拓扑约束目前用 `assert`，应改为 Pydantic validator + `ValueError`。 - 两次不同 mesh 的 `fully_shard` 与 `DeviceMesh._flatten(name)` 只在 torch 2.8 / 2.9 验证；`pyproject.toml` 声明 `torch>=2.6`，需要补 CI 或为该开关声明更高的最低版本。 -- L1–L3 是人工实验脚本（写 JSON、打印统计），不是会失败的 pytest gate。 - ExpertTP 与解耦互斥；expert dim-1 分片、`shard_placement_fn` 单次包装留作优化。 ## 7. 开放问题（请架构师决策） @@ -343,10 +346,10 @@ expert_fsdp = False（新增）: Shard(0) on ep + Replicate on (replicate, efsdp ## 8. 合入建议 -按 torchtitan 的演进顺序拆分： +PR 的提交按 torchtitan 的演进顺序排列，维护者若希望拆分，可按同一顺序切分： 1. 基础解耦：配置开关、root mesh、两级 `fully_shard`、梯度归约、L0 fake-PG 测试； 2. 生态：FP8 per-class padding / reduce mesh、`LoadSpec` 驱动的 HF / DCP、RL per-spec gather 与 compose owner 修复； 3. 文档与后续：本文、§6 的限制项。 -`decouple_ep_fsdp` 默认 `False`，两个代码 PR 都不改动旧路径的执行逻辑，出问题时一键回退。 +`decouple_ep_fsdp` 默认 `False`，所有代码提交都不改动旧路径的执行逻辑，出问题时一键回退。 From 90f6e5e8641a650c04932ebb375647bee978d636 Mon Sep 17 00:00:00 2001 From: daixu Date: Wed, 23 Sep 2026 17:04:54 +0000 Subject: [PATCH 11/14] [Fix] Reshard decoupled expert FSDP groups once per layer instead of after every call On the decoupled path each MoEBlock is its own FSDP unit nested inside the layer's reentrant activation checkpoint. With FSDP2's automatic reshard the group was resharded after every call: once per intra-layer micro-batch in forward, again for every call the checkpoint replays in backward and once per micro-batch backward, each followed by a fresh all-gather. With four micro-batches that is 13 expert all-gathers per layer and step where the legacy layout, whose single FSDP unit per layer wraps the checkpoint, needs 2. Expert groups are now sharded with reshard_after_forward=False and reshard_after_backward=False and resharded explicitly: by a forward hook on the outer layer unit (checkpoint replay does not call it) and by a hook on the gradient of the layer inputs, which fires after the whole layer backward. A model forward pre-hook reshards any group left gathered so no step can run on pre-update weights. MTP layers use the same two-level wrap and get the same treatment. The expert groups are all-gathered once in forward and once in backward per layer, and at most one layer's experts are gathered at a time. --- xtuner/v1/model/moe/moe.py | 74 ++++++++++++++++++++++++++++++++------ 1 file changed, 64 insertions(+), 10 deletions(-) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 31ad44d967..b69a160ae6 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -19,6 +19,7 @@ MixedPrecisionPolicy, ) from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor +from torch.utils._pytree import tree_flatten from tqdm import tqdm from typing_extensions import overload, override @@ -1334,11 +1335,6 @@ def fully_shard( self._fully_shard_expert_blocks( layer, mp_policy=mp_policy, - reshard_after_forward=( - False - if layer_idx >= len(self.layers) - 1 and self.mtp_block is None - else self.fsdp_config.reshard_after_forward - ), ) if self._should_recompute( layer_idx=layer_idx, @@ -1362,6 +1358,8 @@ def fully_shard( offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=layer, ) + if decoupled: + self._reshard_expert_blocks_per_layer(layer, reshard_after_forward=reshard_after_forward) for layer_cur, layer_next in zip( list(self.layers.values())[:-1], @@ -1417,9 +1415,7 @@ def fully_shard( reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1 if decoupled: - self._fully_shard_expert_blocks( - mtp_layer, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward - ) + self._fully_shard_expert_blocks(mtp_layer, mp_policy=mp_policy) self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, mp_policy=mp_policy, @@ -1427,6 +1423,8 @@ def fully_shard( offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=mtp_layer, ) + if decoupled: + self._reshard_expert_blocks_per_layer(mtp_layer, reshard_after_forward=reshard_after_forward) if mtp_idx == 0: if decoupled: layer_next.set_modules_to_forward_prefetch( # type: ignore @@ -1457,6 +1455,8 @@ def fully_shard( offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, ) self.set_modules_to_forward_prefetch([self.embed_tokens, self.layers["0"]]) # type: ignore + if decoupled: + self._reshard_expert_blocks_before_forward() for _, module in self.named_modules(): if isinstance(module, nn.Embedding): @@ -1727,18 +1727,72 @@ def _fully_shard_expert_blocks( self, module: nn.Module, mp_policy: MixedPrecisionPolicy, - reshard_after_forward: bool, ) -> None: + # Each routed-expert block becomes its own FSDP unit on `expert_fsdp_mesh`, nested inside the layer + # and therefore inside the layer's reentrant activation checkpoint. FSDP2's automatic reshard would + # then run after every call of the block: once per intra-layer micro-batch in forward, again for every + # call the checkpoint replays in backward, and once per micro-batch backward, each followed by a fresh + # all-gather. Automatic reshard is disabled here; `_reshard_expert_blocks_per_layer` reshards the + # group once after the layer's forward and once after the layer's backward instead. assert self.expert_fsdp_mesh is not None assert self.fsdp_config is not None for expert_block in self._expert_blocks(module): self._fully_shard( mesh=self.expert_fsdp_mesh, mp_policy=mp_policy, - reshard_after_forward=reshard_after_forward, + reshard_after_forward=False, offload_policy=CPUOffloadPolicy() if self.fsdp_config.cpu_offload else None, module=expert_block, ) + expert_block.set_reshard_after_backward(False) # type: ignore[operator] + + def _reshard_expert_blocks_per_layer(self, layer: nn.Module, reshard_after_forward: bool) -> None: + # `layer` is the outer FSDP unit, which wraps the checkpoint boundary: checkpoint replay calls the + # wrapped module, not this one, so the forward hook below runs for the original forward only, and the + # gradient of the layer inputs is produced only after the whole layer backward (replay, every + # micro-batch backward and the experts' reduce-scatter launch). `reshard_after_forward` follows the + # layer's own setting, so the last layer keeps its experts gathered into backward like its dense group. + experts = self._expert_blocks(layer) + + def reshard_experts() -> None: + for expert_block in experts: + expert_block.reshard() # type: ignore[operator] + + def reshard_after_layer_backward(module: nn.Module, args: tuple, kwargs: dict) -> None: + if not torch.is_grad_enabled(): + return None + inputs = [t for t in tree_flatten((args, kwargs))[0] if isinstance(t, torch.Tensor) and t.requires_grad] + resharded = [False] + + def on_input_grad(grad: torch.Tensor) -> None: + if not resharded[0]: + resharded[0] = True + reshard_experts() + + for tensor in inputs: + tensor.register_hook(on_input_grad) + return None + + def reshard_after_layer_forward(module: nn.Module, args: tuple, kwargs: dict, output: object) -> None: + if reshard_after_forward: + reshard_experts() + return None + + layer.register_forward_pre_hook(reshard_after_layer_backward, with_kwargs=True) + layer.register_forward_hook(reshard_after_layer_forward, with_kwargs=True) + + def _reshard_expert_blocks_before_forward(self) -> None: + # Correctness guard: an expert group still gathered when the next forward starts would skip its + # all-gather and compute with the weights from before the optimizer step. A no-op whenever every + # layer's backward hook has fired. + experts = [expert_block for expert_block in self.modules() if isinstance(expert_block, MoEBlock)] + + def reshard_experts(module: nn.Module, args: tuple, kwargs: dict) -> None: + for expert_block in experts: + expert_block.reshard() # type: ignore[operator] + return None + + self.register_forward_pre_hook(reshard_experts, with_kwargs=True) def _init_decoupled_device_mesh(self, fsdp_config: FSDPConfig) -> None: # Decoupled ("dp2ep") layout: EP is a sub-dimension of the FSDP shard dimension. From 1edc86d9aab500e9603185c01b2efff313271e58 Mon Sep 17 00:00:00 2001 From: daixu Date: Wed, 23 Sep 2026 17:05:30 +0000 Subject: [PATCH 12/14] [Fix] Point decoupled expert groups' backward prefetch at the layer below FSDP2 picks each unit's default backward prefetch target from the order in which units finished forward. The decoupled expert units sit inside the reentrant activation checkpoint, so checkpoint replay re-runs their forward hooks during backward and appends to that order. Their default target then becomes the experts of the layer above, whose backward is already done: one all-gather per layer that is never used and whose output buffer stays allocated until the end of backward (0.3 GB per layer on Qwen3-30B-A3B at EP4, ~14 GB in total, which set the peak memory of the decoupled layout). Nothing prefetches the dense unit of the layer below either, so its all-gather is issued on demand and queues behind this layer's dense reduce-scatter on the same process group. Every expert group now prefetches the layer below (for decoder layers and MTP layers alike); an explicit list replaces the default target. The first layer of each stack keeps the default. --- xtuner/v1/model/moe/moe.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index b69a160ae6..2da5783dfc 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1376,6 +1376,7 @@ def fully_shard( if decoupled: last_layer = list(self.layers.values())[-1] last_layer.set_modules_to_forward_prefetch(self._expert_blocks(last_layer)) # type: ignore + self._set_expert_backward_prefetch(list(self.layers.values())) self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, @@ -1447,6 +1448,7 @@ def fully_shard( if decoupled: last_mtp_layer = list(self.mtp_block.layers)[-1] last_mtp_layer.set_modules_to_forward_prefetch(self._expert_blocks(last_mtp_layer)) # type: ignore + self._set_expert_backward_prefetch(list(self.mtp_block.layers)) self._fully_shard( mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, @@ -1746,6 +1748,18 @@ def _fully_shard_expert_blocks( ) expert_block.set_reshard_after_backward(False) # type: ignore[operator] + def _set_expert_backward_prefetch(self, layers: list[nn.Module]) -> None: + # FSDP2 picks each unit's default backward prefetch target from the order in which units finished + # forward. Checkpoint replay re-runs the expert units' forward hooks during backward and appends to that + # order, so an expert group's default target becomes the experts of the layer above, whose backward is + # already done: an all-gather that is never used and stays allocated until the end of backward. Nothing + # prefetches the dense unit of the layer below either, so its all-gather is issued on demand behind this + # layer's dense reduce-scatter on the same process group. Point every expert group at the layer below; + # an explicit list replaces the default target. + for layer_prev, layer_cur in zip(layers[:-1], layers[1:]): + for expert_block in self._expert_blocks(layer_cur): + expert_block.set_modules_to_backward_prefetch([layer_prev]) # type: ignore[operator] + def _reshard_expert_blocks_per_layer(self, layer: nn.Module, reshard_after_forward: bool) -> None: # `layer` is the outer FSDP unit, which wraps the checkpoint boundary: checkpoint replay calls the # wrapped module, not this one, so the forward hook below runs for the original forward only, and the From d86987d3c2590cfe11e432f7ebfbe85b48e5928f Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 24 Sep 2026 00:00:35 +0000 Subject: [PATCH 13/14] [Test] Count decoupled expert all-gathers under intra-layer micro-batches and recompute The existing gates compare losses, gradient norms and per-rank parameter memory, none of which changes when an FSDP group is gathered more often than needed, and they run without intra-layer micro-batches. The new 8-GPU gate trains the tiny Qwen3-MoE at EP4 (efsdp == 2) with two intra-layer micro-batches and the default full recompute, counts per FSDP unit the all-gathers issued and the ones whose result was used, and asserts at most two used all-gathers per expert group and step (forward + backward) and at most one unused all-gather in total. Before the fix it fails with six used all-gathers per expert group. `count_fsdp_all_gathers` wraps `FSDPParamGroup.unshard` / `wait_for_unshard` to collect the counts; `build_engine` / `run_mode` take the number of intra-layer micro-batches. --- .../test_decoupled_ep_fsdp_train_engine.py | 35 +++++++++- xtuner/_testing/decoupled_ep_fsdp.py | 67 +++++++++++++++++-- 2 files changed, 96 insertions(+), 6 deletions(-) diff --git a/tests/engine/test_decoupled_ep_fsdp_train_engine.py b/tests/engine/test_decoupled_ep_fsdp_train_engine.py index f3cbe177c3..fb4e15fc0c 100644 --- a/tests/engine/test_decoupled_ep_fsdp_train_engine.py +++ b/tests/engine/test_decoupled_ep_fsdp_train_engine.py @@ -8,6 +8,8 @@ that AdamW would hide from the loss curve; after the first update bf16 weight differences flip individual top-k routing decisions and router gradients legitimately differ by percents) and per-rank parameter memory, for ``efsdp == 1``, ``efsdp > 1`` and HSDP + EP. + With intra-layer micro-batches and full recompute it also counts the expert groups' FSDP + all-gathers per step, which the numeric checks cannot see. * ``TestDecoupledEpFsdpCheckpoint`` (L3): bit-exact HF export right after ``from_hf``, DCP resume, HF export after resume and cross-layout DCP resharding. @@ -120,11 +122,40 @@ def test_efsdp_and_hsdp_layouts_match_legacy_ep4(self) -> None: self._assert_param_memory(results["H41"], results["B4"], dense_ratio=(1 / 4) / (1 / 2), expert_ratio=2.0) self._assert_param_memory(results["H22"], results["B4"], dense_ratio=(1 / 4) / (1 / 2), expert_ratio=2.0) + def test_expert_groups_gather_once_per_layer_with_micro_batches_and_recompute(self) -> None: + # The decoupled expert groups are FSDP units nested inside each layer's reentrant activation checkpoint. + # With two intra-layer micro-batches, full recompute (the default `recompute_ratio`) and `efsdp == 2`, + # FSDP2's per-call reshard re-gathered them for every forward call, replayed call and backward call + # (6 all-gathers per layer and step), and checkpoint replay pointed their default backward prefetch at + # the already-finished experts of the layer above (one unused all-gather per layer). Expected now: one + # all-gather in forward and one in backward per layer, and no unused one except from the first layer, + # which keeps FSDP2's default backward prefetch target. + # + # Only communication is asserted here; loss and gradient parity with the legacy layout is covered by the + # tests above. With intra-layer micro-batches the all2all dispatcher's asynchronous path occasionally + # produces non-finite gradients on both layouts, independent of this change, which would make a numeric + # comparison flaky. The counts do not depend on the values. + results = self._run_layouts( + ("C4:ep=4,decouple=1",), steps=3, intra_layer_micro_batch=2, count_all_gathers_at=2 + ) + experts = {fqn: counts for fqn, counts in results["C4"]["all_gathers"].items() if fqn.endswith("experts")} + self.assertEqual(len(experts), 4, f"expected one expert group per layer, got {sorted(experts)}") + for fqn, (issued, used) in experts.items(): + self.assertLessEqual(used, 2, f"{fqn}: {used} all-gathers used in one step, expected forward + backward") + unused = sum(issued - used for issued, used in experts.values()) + self.assertLessEqual(unused, 1, f"{unused} expert all-gathers were issued but never used: {experts}") + @property def world_size(self) -> int: return 8 - def _run_layouts(self, specs: tuple[str, ...], steps: int) -> dict[str, dict[str, Any]]: + def _run_layouts( + self, + specs: tuple[str, ...], + steps: int, + intra_layer_micro_batch: int = 1, + count_all_gathers_at: int | None = None, + ) -> dict[str, dict[str, Any]]: self.create_pg("cuda") root = self._make_shared_tmpdir() try: @@ -142,6 +173,8 @@ def _run_layouts(self, specs: tuple[str, ...], steps: int) -> dict[str, dict[str fp8=False, grad_norm_steps=(0,), tag=self._testMethodName, + intra_layer_micro_batch=intra_layer_micro_batch, + count_all_gathers_at=count_all_gathers_at, ) dist.barrier() return results diff --git a/xtuner/_testing/decoupled_ep_fsdp.py b/xtuner/_testing/decoupled_ep_fsdp.py index fe1611dbd7..f6e12f92d9 100644 --- a/xtuner/_testing/decoupled_ep_fsdp.py +++ b/xtuner/_testing/decoupled_ep_fsdp.py @@ -11,13 +11,17 @@ import gc import time -from collections.abc import Iterable, Sequence +from collections import defaultdict +from collections.abc import Iterable, Iterator, Sequence +from contextlib import contextmanager, nullcontext from pathlib import Path from typing import Any, TypedDict +from unittest import mock import torch import torch.distributed as dist from safetensors.torch import load_file +from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup from torch.distributed.tensor import DTensor from xtuner.v1.config import AdamWConfig, FSDPConfig @@ -157,6 +161,7 @@ def build_engine( dispatcher: str = "all2all", fp8: bool = False, lr: float = 1e-4, + intra_layer_micro_batch: int = 1, ) -> TrainEngine: """Build a ``TrainEngine`` for ``mode`` without loading weights. @@ -167,6 +172,7 @@ def build_engine( dispatcher (str): MoE dispatcher (``"all2all"`` or ``"deepep"``). fp8 (bool): Enable tile-wise float8 for linear and grouped linear layers. lr (float): AdamW learning rate. + intra_layer_micro_batch (int): Number of micro-batches interleaved inside every decoder layer. Returns: TrainEngine: Sharded engine with freshly built (uninitialised) parameters. @@ -187,7 +193,45 @@ def build_engine( hsdp_sharding_size=mode["hsdp"], torch_compile=False, ) - return TrainEngine(model_cfg=model_cfg, optim_cfg=AdamWConfig(lr=lr), fsdp_cfg=fsdp_cfg) + return TrainEngine( + model_cfg=model_cfg, + optim_cfg=AdamWConfig(lr=lr), + fsdp_cfg=fsdp_cfg, + intra_layer_micro_batch=intra_layer_micro_batch, + ) + + +@contextmanager +def count_fsdp_all_gathers() -> Iterator[dict[str, list[int]]]: + """Count, per FSDP unit, the parameter all-gathers issued and the ones whose result was used. + + An all-gather that is issued but never waited on was prefetched for nothing: FSDP2 only waits on it in + ``finalize_backward`` at the end of backward, and its output buffer stays allocated until then. Units whose + shard mesh has a single rank do not communicate and are not counted. + + Yields: + dict[str, list[int]]: ``{module_fqn: [issued, used]}``, filled while the context is active. + """ + counts: dict[str, list[int]] = defaultdict(lambda: [0, 0]) + unshard = FSDPParamGroup.unshard + wait_for_unshard = FSDPParamGroup.wait_for_unshard + + def counting_unshard(self: FSDPParamGroup, async_op: bool = False) -> None: + pending = self._all_gather_result is not None + unshard(self, async_op) + if not pending and self._all_gather_result is not None and self._all_gather_process_group.size() > 1: + counts[self._module_fqn][0] += 1 + + def counting_wait_for_unshard(self: FSDPParamGroup) -> None: + if self._all_gather_result is not None and self._all_gather_process_group.size() > 1: + counts[self._module_fqn][1] += 1 + wait_for_unshard(self) + + with ( + mock.patch.object(FSDPParamGroup, "unshard", counting_unshard), + mock.patch.object(FSDPParamGroup, "wait_for_unshard", counting_wait_for_unshard), + ): + yield counts def train(engine: TrainEngine, steps: Iterable[int], vocab_size: int, seq_len: int) -> list[float]: @@ -221,6 +265,8 @@ def run_mode( fp8: bool, grad_norm_steps: Sequence[int] = (0, 1, 25), tag: str = "numerics", + intra_layer_micro_batch: int = 1, + count_all_gathers_at: int | None = None, ) -> dict[str, Any]: """Train ``mode`` from the HF checkpoint and record losses, grad norms, memory and step time. @@ -234,6 +280,9 @@ def run_mode( fp8 (bool): Enable tile-wise float8. grad_norm_steps (Sequence[int]): Steps whose per-parameter grad norms are recorded. tag (str): Mesh-prefix tag passed to :func:`build_engine`. + intra_layer_micro_batch (int): Micro-batches per step, interleaved inside every decoder layer. + count_all_gathers_at (int | None): Step whose FSDP all-gathers are counted with + :func:`count_fsdp_all_gathers`; the counts are returned under ``all_gathers``. Returns: dict[str, Any]: ``losses``, ``grad_norms`` (per step), ``param_grad_norms`` (per recorded @@ -244,7 +293,9 @@ def run_mode( torch.cuda.synchronize() base_allocated = torch.cuda.memory_allocated() - engine = build_engine(mode, hf_path, tag, dispatcher=dispatcher, fp8=fp8, lr=lr) + engine = build_engine( + mode, hf_path, tag, dispatcher=dispatcher, fp8=fp8, lr=lr, intra_layer_micro_batch=intra_layer_micro_batch + ) engine.from_hf(hf_path=hf_path, strict=True) vocab_size = engine.model.config.vocab_size torch.cuda.synchronize() @@ -262,8 +313,14 @@ def run_mode( for step in range(steps): torch.cuda.synchronize() t0 = time.perf_counter() - batch = make_batch(step, rank, vocab_size, seq_len) - info = engine.train_step([batch]) + batches = [ + make_batch(step * intra_layer_micro_batch + i, rank, vocab_size, seq_len) + for i in range(intra_layer_micro_batch) + ] + with count_fsdp_all_gathers() if step == count_all_gathers_at else nullcontext() as all_gathers: + info = engine.train_step(batches) + if all_gathers is not None: + result["all_gathers"] = {fqn: list(counts) for fqn, counts in all_gathers.items()} grad_norm = engine.clip_grad_norm() if step in grad_norm_steps: result["param_grad_norms"][str(step)] = per_param_grad_norms(engine.model) From 7c53b7da727d36ea143a7c1c14de0293b719c1a6 Mon Sep 17 00:00:00 2001 From: daixu Date: Thu, 24 Sep 2026 01:13:52 +0000 Subject: [PATCH 14/14] [Docs] Describe the decoupled expert groups' reshard timing and backward prefetch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add §3.3.1 to the design note: why the expert FSDP groups, nested inside the layer's reentrant activation checkpoint, need per-layer reshard and an explicit backward prefetch target, and record the measurements in §5. Refresh the moe.py line references. --- docs/design/decouple_ep_fsdp.md | 46 ++++++++++++++++++++++++++------- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/docs/design/decouple_ep_fsdp.md b/docs/design/decouple_ep_fsdp.md index 37603d20e7..aceb0fc70c 100644 --- a/docs/design/decouple_ep_fsdp.md +++ b/docs/design/decouple_ep_fsdp.md @@ -50,11 +50,11 @@ expert : Shard(0) over ep + FSDP over efsdp (+ replicate 上 HS ## 3. 实现 -各小节的"位置"行给出文件与行号，行号对应分支 `feat/decouple-ep-fsdp` 的代码提交 `0a7c6523`（2026-09-17，基于 upstream/main `6c06f96e`）。 +各小节的"位置"行给出文件与行号，行号对应加入专家组 reshard 修复（§3.3.1）之后的代码（2026-09-24；分支 `feat/decouple-ep-fsdp` 与 PR 分支上这些文件内容相同）。 ### 3.1 配置与校验 -位置：`xtuner/v1/config/fsdp.py` 44–73，`xtuner/v1/model/moe/moe.py` 1764–1771。 +位置：`xtuner/v1/config/fsdp.py` 44–73，`xtuner/v1/model/moe/moe.py` 1832–1839。 ```python # xtuner/v1/config/fsdp.py @@ -77,7 +77,7 @@ def _validate_ep_fsdp_topology(self): ### 3.2 单一 root mesh -位置：`xtuner/v1/model/moe/moe.py` 1743–1808（`_init_decoupled_device_mesh`），1576–1581（分流）。 +位置：`xtuner/v1/model/moe/moe.py` 1811–1876（`_init_decoupled_device_mesh`），1578–1583（分流）。 ```python # xtuner/v1/model/moe/moe.py::_init_decoupled_device_mesh（节选） @@ -112,7 +112,7 @@ else: ### 3.3 两级 `fully_shard` -位置：`xtuner/v1/model/moe/moe.py` 1310–1342（decoder layer），1370–1380（prefetch），1419–1451（MTP），1723–1741（helper）。 +位置：`xtuner/v1/model/moe/moe.py` 1311–1362（decoder layer），1364–1379（prefetch），1405–1451（MTP），1459–1461（模型 forward pre-hook），1724–1808（helper）。 ```python # xtuner/v1/model/moe/moe.py::fully_shard（节选） @@ -122,30 +122,57 @@ if not decoupled and (self.ep_mesh.size() > 1 or tp_enabled): for layer_idx, layer in self.layers.items(): if decoupled: # 1) 先包 MoEBlock -> expert_fsdp_mesh - self._fully_shard_expert_blocks(layer, mp_policy=mp_policy, reshard_after_forward=...) + self._fully_shard_expert_blocks(layer, mp_policy=mp_policy) ... self._fully_shard( # 2) 再包整层 -> dense 的 fsdp_mesh / hsdp_mesh mesh=self.fsdp_mesh if self.hsdp_mesh is None else self.hsdp_mesh, module=layer, ... ) + if decoupled: # 3) 专家组按层 reshard（§3.3.1） + self._reshard_expert_blocks_per_layer(layer, reshard_after_forward=reshard_after_forward) for layer_cur, layer_next in zip(layers[:-1], layers[1:]): if decoupled: # expert 的 all-gather 与下一层 dense 一起预取 layer_cur.set_modules_to_forward_prefetch([*self._expert_blocks(layer_cur), layer_next]) +if decoupled: # 专家组的反向预取指向下一层（§3.3.1） + self._set_expert_backward_prefetch(layers) @staticmethod def _expert_blocks(module): return [m for m in module.modules() if isinstance(m, MoEBlock)] -def _fully_shard_expert_blocks(self, module, mp_policy, reshard_after_forward): +def _fully_shard_expert_blocks(self, module, mp_policy): for expert_block in self._expert_blocks(module): - self._fully_shard(mesh=self.expert_fsdp_mesh, module=expert_block, mp_policy=mp_policy, ...) + self._fully_shard(mesh=self.expert_fsdp_mesh, module=expert_block, reshard_after_forward=False, ...) + expert_block.set_reshard_after_backward(False) # reshard 时机见 §3.3.1 ``` 解耦路径跳过 `_replicate_other_params`：dense 参数进入 `fully_shard` 时是普通 tensor，由 FSDP 在完整 `dp_shard` 上分片。每个 decoder layer 先把所有 `MoEBlock`（只含 routed expert 的 grouped linear 与激活）在 `expert_fsdp_mesh` 上 `fully_shard`，再在 dense mesh 上包装整层；FSDP2 会保留内层的 expert wrapper，把 expert 参数从外层 param group 中排除。MTP layer 使用同样的两级包装。forward prefetch 列表同时包含下一层与其 expert block，让 expert 的 all-gather 与 attention 重叠，而不是串行等在 MoE block 前。 +#### 3.3.1 专家组的 reshard 时机与反向预取 + +专家组的 FSDP 单元嵌在整层的 reentrant activation checkpoint **里面**（外层 dense 单元则包在它外面）。重算会把这一层的前向 +再跑一遍，途中重新触发专家组的 FSDP hook，这带来两个问题： + +1. **重复 all-gather**：按 FSDP2 默认的 `reshard_after_forward=True` / `reshard_after_backward=True`，专家组每次调用后都会释放， + 前向每个层内微批、重算每次调用、反向每个微批都要重新 all-gather。mb=4、全重算时每层每步 13 次，旧布局的整层单元只要 2 次。 +2. **反向预取目标错位**：重算在反向期间往 FSDP2 的前向完成顺序里追加记录，专家组的默认预取目标变成上一层（反向已经做完的） + 专家，发出一次永远不会被用到的 all-gather，其输出缓冲要到反向结束才释放；同时没有单元预取下一层的 dense 单元， + 它的 all-gather 只能排在本层 dense 的 reduce-scatter 后面。 + +处理方式，两者都不新增配置，解耦路径下默认生效： + +- 专家组关闭 FSDP2 的自动 reshard，由 `_reshard_expert_blocks_per_layer` 在外层单元上挂两个 hook：前向结束后 reshard 一次 + （forward post-hook，重算不会调用外层单元）；该层输入梯度产生时——即重算、所有微批反向与 reduce-scatter 都已发出之后—— + 再 reshard 一次。最后一层（无 MTP 时）沿用外层的 `reshard_after_forward=False`。模型的 forward pre-hook 在每步前向前 + reshard 仍处于拼好状态的专家组，保证不会用到优化器更新前的权重。 +- `_set_expert_backward_prefetch` 让第 i 层（i ≥ 1）的专家组显式预取第 i−1 层；显式列表替换默认目标。 + +decoder layer 与 MTP layer 都按此处理。每层专家组每步 all-gather 两次，同一时刻最多一层专家处于拼好状态。 +机制、trace 计数与测量见 [decoupled_ep_fsdp_expert_reshard_zh.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decoupled_ep_fsdp_expert_reshard_zh.md)。 + ### 3.4 梯度归约 -位置：`xtuner/v1/model/moe/moe.py` 1482–1486（分流），1673–1722（`_scale_and_reduce_grad_decoupled`）。 +位置：`xtuner/v1/model/moe/moe.py` 1484–1488（分流），1675–1722（`_scale_and_reduce_grad_decoupled`）。 不变式：每个参数的最终梯度等于全 data-parallel world 上的均值，与 `ep` / `efsdp` 取值无关。 @@ -183,7 +210,7 @@ for group, grads in grads_by_group.items(): # 每个 group 一次 coalesced al ### 3.5 FP8 -位置：`xtuner/v1/float8/float8_handler.py` 121–164、166–238、316–335；`xtuner/v1/float8/fsdp_utils.py` 127–137；`xtuner/v1/model/moe/moe.py` 1281–1292；`xtuner/v1/model/base.py` 1015。 +位置：`xtuner/v1/float8/float8_handler.py` 121–164、166–238、316–335；`xtuner/v1/float8/fsdp_utils.py` 127–137；`xtuner/v1/model/moe/moe.py` 1282–1293；`xtuner/v1/model/base.py` 1015。 tile-wise FP8 需要知道每个本地权重会被多少个 FSDP rank 切分（决定 padding），以及 scale 的 reduce-max 应跨哪些 rank。解耦后 dense 与 expert 的答案不同： @@ -279,6 +306,7 @@ def _param_owner(model, name): | FP8 | 3.4B 与 GLM-5.2 tile-wise FP8 训练 20 步稳定，差异量级接近同配置重复噪声 | FP8 + HSDP、FP8 + checkpoint 未测 | [L3.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/L3.md)、[GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | | GLM-5.2-30B，8×H200 目标配方 | EP4：1.688 s / 99.5 GiB → 1.658 s / 83.5 GiB；EP8：11.8 s / 120.0 GiB → 1.64 s / 76.3 GiB | EP8 的加速来自离开 allocator 上限（legacy 每步 2–3 次 alloc retry，解耦为 0），不是 collective 本身快 7 倍 | [GLM52.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/GLM52.md) | | RL 权重同步 | plain MoE 的 per-spec gather 逻辑成立；compose 的 owner 错误已修复并有单测 | 未做真实 Turbomind 端到端对拍 | [test_weight_iterator.py](../../tests/rl/test_weight_iterator.py#L237-L427) | +| 专家组 reshard 时机与反向预取（§3.3.1） | Qwen3-30B-A3B EP4、mb4、64K token/卡/步（8×H200 同批 3 次交错，DeepEP v1）：专家组 all-gather 每步 614 → 96 次、发出后从未被使用 46 → 1；解耦对旧布局由 −5.3%、峰值 +3.2 GB 变为 +0.6%、−8.3 GB。EP8（efsdp=1）−1.6% → +1.0%。GLM-5.2-30B（含 MTP 层）EP4 mb2：对修复前 +4.4%，峰值 87.0 → 83.2 GB，MTP 层专家组每步同样 ≤ 2 次。两机 16 卡（专家组跨机）：对修复前 EP8（efsdp=2）+8.0%、EP4（efsdp=4）+11.2%、GLM +4.7%，与旧布局持平，显存少 8.6 / 3.8 GB | 单机与两机；更多节点、efsdp > 4 未测；主干第 0 层保留默认反向预取，每步 1 次未被使用的 all-gather | [test_decoupled_ep_fsdp_train_engine.py](../../tests/engine/test_decoupled_ep_fsdp_train_engine.py)（8 卡 gate：EP4、2 个层内微批、全重算下每个专家组每步被使用的 all-gather ≤ 2、发出后未被使用的合计 ≤ 1；只数通信，不比数值）、[decoupled_ep_fsdp_expert_reshard_zh.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decoupled_ep_fsdp_expert_reshard_zh.md) | | legacy 兼容 | 开关默认关闭，旧分支执行逻辑未改动，L0 旧布局回归通过 | 未做 base-vs-branch 端到端逐 tensor 对照 | [baseline.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/baseline.md)、[decisions.md](https://github.com/silencelamb/xtuner/blob/feat/decouple-ep-fsdp/reports/decisions.md) | `reports/` 下的报告不随代码 PR 合入，链接指向 fork 分支 `feat/decouple-ep-fsdp` 上的固定地址；测试与脚本是仓库内相对链接。