[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training

tests/acceptance/run_qwen35_moonep_acceptance.sh

@@ -0,0 +1,90 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -lt 3 || $# -gt 4 ]]; then
+    echo "Usage: $0 <deepep|moonep> <mtp:0|1> <pack-length> [acceptance-root]" >&2
+    exit 2
+fi
+
+backend=$1
+mtp=$2
+pack_length=$3
+acceptance_root=${4:-work_dirs/moonep_qwen35_acceptance}
+# Resolve the checkout that owns this script so commit-level comparisons can
+# run in isolated worktrees without importing the developer's dirty checkout.
+repo_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)
+gpu_lock=/mnt/shared-storage-user/zhaopenghao/github/xtuner/zdev/gpu_lock.sh
+
+if [[ $backend != "deepep" && $backend != "moonep" ]]; then
+    echo "backend must be deepep or moonep" >&2
+    exit 2
+fi
+if [[ $mtp != "0" && $mtp != "1" ]]; then
+    echo "mtp must be 0 or 1" >&2
+    exit 2
+fi
+if ! [[ $pack_length =~ ^[1-9][0-9]*$ ]]; then
+    echo "pack-length must be a positive integer" >&2
+    exit 2
+fi
+
+# Re-enter the exact same command while holding the repository-wide 8-GPU
+# lock. The marker prevents recursively acquiring the non-reentrant lock.
+if [[ ${MOONEP_ACCEPTANCE_LOCK_HELD:-0} != "1" ]]; then
+    exec "$gpu_lock" env MOONEP_ACCEPTANCE_LOCK_HELD=1 "$0" "$@"
+fi
+
+source /mnt/shared-storage-user/zhaopenghao/miniconda3/etc/profile.d/conda.sh
+conda activate pt212_cu132
+cd "$repo_root"
+
+run_dir="$acceptance_root/${backend}_mtp${mtp}_pack${pack_length}"
+if [[ ${MOONEP_ACCEPTANCE_MICRO_BATCH:-1} != "1" ]]; then
+    run_dir="${run_dir}_micro${MOONEP_ACCEPTANCE_MICRO_BATCH}"
+fi
+if [[ -e $run_dir ]]; then
+    echo "refusing to mix acceptance attempts in existing directory: $run_dir" >&2
+    exit 2
+fi
+mkdir -p "$run_dir"
+
+export PYTHONPATH="$repo_root"
+export MOONEP_ACCEPTANCE_BACKEND=$backend
+export MOONEP_ACCEPTANCE_MTP=$mtp
+export MOONEP_ACCEPTANCE_PACK_LENGTH=$pack_length
+export MOONEP_ACCEPTANCE_WORK_DIR=$run_dir
+export MOONEP_ACCEPTANCE_MODEL_PATH=${MOONEP_ACCEPTANCE_MODEL_PATH:-/mnt/shared-storage-user/llmrazor-share/model/Qwen3.5-35B-A3B}
+export MOONEP_ACCEPTANCE_DATA_PATH=${MOONEP_ACCEPTANCE_DATA_PATH:-/mnt/shared-storage-user/llmrazor-share/data/alpaca}
+export MODEL_COMPILE=1
+export XTUNER_USE_FA3="${XTUNER_USE_FA3:-0}"
+# FA3's Hopper backward has no deterministic path for head_dim 256 (Qwen3.5),
+# so the FA3 comparison runs non-deterministically; keep the default otherwise.
+export XTUNER_DETERMINISTIC="${XTUNER_DETERMINISTIC:-true}"
+export XTUNER_ACTIVATION_OFFLOAD=0
+export XTUNER_COMPILE_NO_INPLACE_BUFFERS=1
+export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export PROFILE_RANKS=${PROFILE_RANKS:-0}
+export TRITON_CACHE_DIR="$run_dir/triton_cache"
+export TORCHINDUCTOR_CACHE_DIR="$run_dir/torchinductor_cache"
+unset XTUNER_USE_CUTLASS_GROUP_GEMM
+unset GROUPED_GEMM_USE_CUTLASS
+
+config=tests/acceptance/sft_qwen35_moonep_acceptance.py
+python -m xtuner._testing.moonep_acceptance capture \
+    --config "$config" \
+    --output "$run_dir/acceptance_manifest.json"
+
+# VMM/NCCL allocations are not all visible to PyTorch's allocator counters.
+# Sample all eight devices while the same GPU lock covers the training job.
+nvidia-smi --query-gpu=timestamp,index,memory.used --format=csv,noheader,nounits \
+    --loop-ms=500 > "$run_dir/device_memory.csv" &
+memory_monitor_pid=$!
+trap 'kill "$memory_monitor_pid" 2>/dev/null || true; wait "$memory_monitor_pid" 2>/dev/null || true' EXIT
+
+torchrun \
+    --nproc-per-node 8 \
+    --master-port "${MOONEP_ACCEPTANCE_MASTER_PORT:-29618}" \
+    xtuner/v1/train/cli/sft.py \
+    --config "$config" \
+    2>&1 | tee "$run_dir/stdout.log"

tests/acceptance/sft_qwen35_moonep_acceptance.py

@@ -0,0 +1,114 @@
+"""Matched Qwen3.5 configuration for the MoonEP/DeepEP 20-step gate.
+
+The dispatcher, MTP switch, fixed pack length and output directory are the
+only run-varying inputs.  All other workload choices are deliberately shared.
+"""
+
+import os
+
+import torch
+
+from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig
+from xtuner.v1.datasets import FTDPTokenizeFnConfig
+from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
+from xtuner.v1.loss.ce_loss import CELossConfig
+from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config
+from xtuner.v1.model.moe.qwen3_5_text import MOE_EP_COMPILE_CFG
+from xtuner.v1.module.mtp import MTPConfig
+from xtuner.v1.train import TrainerConfig
+
+
+def _required_env(name: str) -> str:
+    value = os.environ.get(name)
+    if not value:
+        raise ValueError(f"{name} must be set")
+    return value
+
+
+backend = _required_env("MOONEP_ACCEPTANCE_BACKEND")
+if backend not in {"deepep", "moonep"}:
+    raise ValueError(f"MOONEP_ACCEPTANCE_BACKEND must be deepep or moonep, got {backend!r}")
+
+mtp_enabled = bool(int(_required_env("MOONEP_ACCEPTANCE_MTP")))
+pack_length = int(_required_env("MOONEP_ACCEPTANCE_PACK_LENGTH"))
+micro_batch = int(os.environ.get("MOONEP_ACCEPTANCE_MICRO_BATCH", "1"))
+if micro_batch not in (1, 2):
+    raise ValueError("MOONEP_ACCEPTANCE_MICRO_BATCH must be 1 or 2")
+if pack_length <= 0:
+    raise ValueError("MOONEP_ACCEPTANCE_PACK_LENGTH must be positive")
+if os.environ.get("XTUNER_ACTIVATION_OFFLOAD", "0") != "0":
+    raise ValueError("formal MoonEP acceptance runs require XTUNER_ACTIVATION_OFFLOAD=0")
+if os.environ.get("XTUNER_USE_CUTLASS_GROUP_GEMM", "0") == "1":
+    raise ValueError("formal MoonEP acceptance runs require the Triton grouped-GEMM backend")
+
+model_cfg = Qwen3_5_VLMoE35BA3Config(only_llm_forward=True)
+text_cfg = model_cfg.text_config
+text_cfg.ep_size = 4
+text_cfg.dispatcher = backend
+text_cfg.moonep_staging_reference = False
+text_cfg.router_compute_dtype = "float32"
+text_cfg.router_async_offload = False
+# FA2's package metadata has no importable extension in pt212_cu132, so the
+# default keeps both dispatcher runs on flex attention. Set XTUNER_USE_FA3=1 to
+# compare against FlashAttention 3, whose varlen path skips BlockMask
+# construction entirely (no per-packed-document Dynamo/SymPy mask subgraph).
+text_cfg.attention.attn_impl = (
+    "flash_attention" if os.environ.get("XTUNER_USE_FA3", "0") == "1" else "flex_attention"
+)
+# FlexAttention intentionally compiles behind a graph break so its BlockMask
+# tensors become fixed-layout inputs to the kernel graph. Keep all default
+# Qwen3.5 compile targets, but let MHA form that one required boundary.
+text_cfg.compile_cfg = MOE_EP_COMPILE_CFG | {
+    "xtuner.v1.module.attention.mha.MultiHeadAttention.forward": {"fullgraph": False}
+}
+text_cfg.mtp_config = MTPConfig(num_layers=1) if mtp_enabled else None
+
+dataset_cfg = [
+    {
+        "dataset": DatasetConfig(
+            name="alpaca",
+            anno_path=_required_env("MOONEP_ACCEPTANCE_DATA_PATH"),
+            sample_ratio=1.0,
+        ),
+        "tokenize_fn": FTDPTokenizeFnConfig(max_length=262144),
+    }
+]
+dataloader_cfg = DataloaderConfig(
+    dataset_config_list=dataset_cfg,
+    pack_to_max_length=True,
+    pack_max_length=pack_length,
+    pack_level="hard",
+)
+
+profile_step_env = os.environ.get("MOONEP_ACCEPTANCE_PROFILE_STEP")
+profile_step = int(profile_step_env) if profile_step_env else None
+
+trainer = TrainerConfig(
+    load_from=_required_env("MOONEP_ACCEPTANCE_MODEL_PATH"),
+    tokenizer_path=_required_env("MOONEP_ACCEPTANCE_MODEL_PATH"),
+    model_cfg=model_cfg,
+    optim_cfg=AdamWConfig(lr=6e-5, foreach=False),
+    lr_cfg=LRConfig(lr_type="cosine", lr_min=1e-6),
+    loss_cfg=CELossConfig(mode="chunk", chunk_size=1024),
+    fsdp_cfg=FSDPConfig(
+        ep_size=4,
+        param_dtype=torch.bfloat16,
+        reduce_dtype=torch.bfloat16,
+        torch_compile=True,
+        cpu_offload=False,
+    ),
+    dataloader_cfg=dataloader_cfg,
+    global_batch_size=8 * micro_batch,
+    intra_layer_micro_batch=micro_batch,
+    sp_size=1,
+    total_step=20,
+    work_dir=_required_env("MOONEP_ACCEPTANCE_WORK_DIR"),
+    seed=0,
+    strict_load=False,
+    auto_resume=False,
+    debug_skip_save=True,
+    exp_tracker="jsonl",
+    profile_step=profile_step,
+    profile_time=profile_step is not None,
+    profile_memory=False,
+)

tests/engine/test_moe_train_engine_deepep_expert_tp.py

@@ -291,6 +291,7 @@ def test_deepep_expert_tp_domino_micro_batch_matches_sync_baseline(self) -> None
             expert_tp_size=expert_tp_size,
             intra_layer_micro_batch=2,
         )
+        assert engine_domino.model.config.intra_layer_micro_batch == 2
         engine_domino.init_model_weights()
         _copy_matching_engine_weights(engine_ref, engine_domino)
         dist.barrier()

tests/engine/test_moe_train_engine_tpep.py

@@ -685,6 +685,7 @@ def test_expert_tp_only_domino_micro_batch_matches_sync_baseline(self, device: s
             expert_tp_size=expert_tp_size,
             intra_layer_micro_batch=2,
         )
+        assert engine_domino.model.config.intra_layer_micro_batch == 2
         engine_domino.init_model_weights()
         _copy_matching_engine_weights(engine_ref, engine_domino)
         collective_stages = _record_expert_tp_collective_stages(engine_domino)

tests/engine/test_moonep_acceptance.py

@@ -0,0 +1,151 @@
+import json
+import runpy
+
+import pytest
+
+from xtuner._testing.moonep_acceptance import AcceptanceRun, compare_runs, summarize_memory
+
+
+def _write_tracker(path, *, tgs_scale: float, mtp: bool) -> None:
+    path.parent.mkdir(parents=True)
+    with path.open("w", encoding="utf-8") as output:
+        for step in range(1, 21):
+            record = {
+                "step": step,
+                "runtime_info/text_tokens": 65536,
+                "runtime_info/tgs": (1000 + step) * tgs_scale,
+                "loss/reduced_llm_loss": 2.0 - step / 100,
+                "loss/reduced_balancing_loss": 0.01 + step / 10000,
+                "loss/local_loss": 2.01 - step / 100 + step / 10000,
+                "grad_norm": 0.5 + step / 1000,
+            }
+            if mtp:
+                record["loss/reduced_mtp_loss"] = 0.2 - step / 1000
+                record["loss/local_loss"] += record["loss/reduced_mtp_loss"]
+            output.write(json.dumps(record) + "\n")
+
+
+@pytest.mark.parametrize("backend", ["deepep", "moonep"])
+@pytest.mark.parametrize("mtp", [False, True])
+@pytest.mark.parametrize("micro_batch", [1, 2])
+def test_qwen35_acceptance_config_locks_the_formal_workload(monkeypatch, tmp_path, backend, mtp, micro_batch) -> None:
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_BACKEND", backend)
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_MTP", str(int(mtp)))
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_MICRO_BATCH", str(micro_batch))
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_PACK_LENGTH", "65536")
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_WORK_DIR", str(tmp_path / "run"))
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_MODEL_PATH", "/model")
+    monkeypatch.setenv("MOONEP_ACCEPTANCE_DATA_PATH", "/data")
+
+    trainer = runpy.run_path("tests/acceptance/sft_qwen35_moonep_acceptance.py")["trainer"]
+    model = trainer.model_cfg
+
+    assert trainer.total_step == 20
+    assert trainer.global_batch_size == 8 * micro_batch
+    assert trainer.intra_layer_micro_batch == micro_batch
+    assert trainer.sp_size == 1
+    assert trainer.debug_skip_save is True
+    assert trainer.dataloader_cfg.pack_to_max_length is True
+    assert trainer.dataloader_cfg.pack_max_length == 65536
+    assert trainer.fsdp_cfg.ep_size == 4
+    assert trainer.fsdp_cfg.param_dtype.__str__() == "torch.bfloat16"
+    assert trainer.fsdp_cfg.reduce_dtype.__str__() == "torch.bfloat16"
+    assert trainer.fsdp_cfg.torch_compile is True
+    assert trainer.fsdp_cfg.cpu_offload is False
+    assert model.only_llm_forward is True
+    assert model.text_config.ep_size == 4
+    assert model.text_config.dispatcher == backend
+    assert model.text_config.moonep_staging_reference is False
+    assert model.text_config.moonep_num_sms == 64
+    assert model.text_config.router_async_offload is False
+    assert model.text_config.router_compute_dtype == "float32"
+    assert (
+        model.text_config.compile_cfg["xtuner.v1.module.attention.mha.MultiHeadAttention.forward"]["fullgraph"]
+        is False
+    )
+    assert (model.text_config.mtp_config is not None) is mtp
+    if mtp:
+        assert model.text_config.mtp_config.num_layers == 1
+        assert model.text_config.mtp_config.share_weights is False
+
+
+def test_acceptance_report_compares_all_steps_and_warm_throughput(tmp_path) -> None:
+    deepep_tracker = tmp_path / "deepep" / "tracker.jsonl"
+    moonep_tracker = tmp_path / "moonep" / "tracker.jsonl"
+    _write_tracker(deepep_tracker, tgs_scale=1.0, mtp=True)
+    _write_tracker(moonep_tracker, tgs_scale=0.96, mtp=True)
+    records = [json.loads(line) for line in moonep_tracker.read_text().splitlines()]
+    for record in records:
+        for name in tuple(record):
+            if name.startswith("loss/reduced_") and name.endswith("loss"):
+                record[name] *= 1.02
+        record["grad_norm"] *= 1.04
+    moonep_tracker.write_text("".join(f"{json.dumps(record)}\n" for record in records))
+
+    deepep = AcceptanceRun.from_tracker(deepep_tracker, backend="deepep", mtp=True, pack_length=65536)
+    moonep = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=True, pack_length=65536)
+    result = compare_runs(deepep, moonep)
+
+    assert result.passed
+    assert result.throughput_ratio == pytest.approx(0.96)
+    assert result.throughput_steps == list(range(6, 21))
+    assert set(result.curves) == {
+        "reduced_llm_loss",
+        "reduced_mtp_loss",
+        "reduced_balancing_loss",
+        "total_loss",
+        "grad_norm",
+    }
+    assert all(curve.cosine_similarity >= 0.99 for curve in result.curves.values())
+    assert all(curve.mean_relative_difference < 0.03 for name, curve in result.curves.items() if name != "grad_norm")
+    assert result.curves["grad_norm"].mean_relative_difference == pytest.approx(0.04)
+
+
+def test_acceptance_report_rejects_incomplete_or_mismatched_runs(tmp_path) -> None:
+    deepep_tracker = tmp_path / "deepep" / "tracker.jsonl"
+    moonep_tracker = tmp_path / "moonep" / "tracker.jsonl"
+    _write_tracker(deepep_tracker, tgs_scale=1.0, mtp=False)
+    _write_tracker(moonep_tracker, tgs_scale=0.94, mtp=False)
+
+    deepep = AcceptanceRun.from_tracker(deepep_tracker, backend="deepep", mtp=False, pack_length=65536)
+    slow = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=65536)
+    assert not compare_runs(deepep, slow).passed
+
+    mismatched = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=32768)
+    with pytest.raises(ValueError, match="pack_length"):
+        compare_runs(deepep, mismatched)
+
+    micro2 = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=65536, micro_batch=2)
+    with pytest.raises(ValueError, match="micro_batch"):
+        compare_runs(deepep, micro2)
+
+    lines = moonep_tracker.read_text().splitlines()
+    moonep_tracker.write_text("\n".join(lines[:-1]) + "\n")
+    with pytest.raises(ValueError, match="exactly steps 1..20"):
+        AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=65536)
+
+
+def test_memory_report_includes_all_ranks_and_external_device_allocations(tmp_path) -> None:
+    for rank in range(2):
+        tracker = tmp_path / "logs" / "exp_tracking" / f"rank{rank}" / "tracker.jsonl"
+        tracker.parent.mkdir(parents=True)
+        tracker.write_text(
+            "\n".join(
+                json.dumps(
+                    {
+                        "step": step,
+                        "memory/max_memory_GB": (10 if step == 1 else 6) + rank,
+                        "memory/reserved_memory_GB": (12 if step == 1 else 8) + rank,
+                    }
+                )
+                for step in range(1, 21)
+            )
+        )
+    (tmp_path / "device_memory.csv").write_text(
+        "2026/09/07 00:00:00.000, 0, 15360\n2026/09/07 00:00:00.000, 1, 16384\n2026/09/07 00:00:00.500, 0, 14336\n"
+    )
+    result = summarize_memory(tmp_path)
+    assert result["recorded_ranks"] == 2
+    assert result["all_steps"] == {"allocated_gib": 11, "reserved_gib": 13}
+    assert result["steps_6_20"] == {"allocated_gib": 7, "reserved_gib": 9}
+    assert result["sampled_device_peak_gib"] == {"0": 15, "1": 16}