diff --git a/tests/acceptance/run_qwen35_moonep_acceptance.sh b/tests/acceptance/run_qwen35_moonep_acceptance.sh new file mode 100755 index 0000000000..6ad56df3ca --- /dev/null +++ b/tests/acceptance/run_qwen35_moonep_acceptance.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -lt 3 || $# -gt 4 ]]; then + echo "Usage: $0 [acceptance-root]" >&2 + exit 2 +fi + +backend=$1 +mtp=$2 +pack_length=$3 +acceptance_root=${4:-work_dirs/moonep_qwen35_acceptance} +repo_root=/mnt/shared-storage-user/zhaopenghao/github/xtuner_moonep +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 [[ -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_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" + +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" diff --git a/tests/acceptance/sft_qwen35_moonep_acceptance.py b/tests/acceptance/sft_qwen35_moonep_acceptance.py new file mode 100644 index 0000000000..ad7b547bab --- /dev/null +++ b/tests/acceptance/sft_qwen35_moonep_acceptance.py @@ -0,0 +1,108 @@ +"""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")) +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 +# The installed FlashAttention package metadata has no importable extension in +# pt212_cu132. Flex attention keeps both dispatcher runs on the same real-model +# workload instead of depending on that broken optional binary. +text_cfg.attention.attn_impl = "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, + intra_layer_micro_batch=1, + 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, +) diff --git a/tests/engine/test_moe_train_engine_deepep_expert_tp.py b/tests/engine/test_moe_train_engine_deepep_expert_tp.py index ffe195c9f1..38d64ad9a7 100644 --- a/tests/engine/test_moe_train_engine_deepep_expert_tp.py +++ b/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() diff --git a/tests/engine/test_moe_train_engine_tpep.py b/tests/engine/test_moe_train_engine_tpep.py index 4580ab680c..e3b5aff730 100644 --- a/tests/engine/test_moe_train_engine_tpep.py +++ b/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) diff --git a/tests/engine/test_moonep_acceptance.py b/tests/engine/test_moonep_acceptance.py new file mode 100644 index 0000000000..f5025b0a3e --- /dev/null +++ b/tests/engine/test_moonep_acceptance.py @@ -0,0 +1,119 @@ +import json +import runpy + +import pytest + +from xtuner._testing.moonep_acceptance import AcceptanceRun, compare_runs + + +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]) +def test_qwen35_acceptance_config_locks_the_formal_workload(monkeypatch, tmp_path, backend, mtp) -> None: + monkeypatch.setenv("MOONEP_ACCEPTANCE_BACKEND", backend) + monkeypatch.setenv("MOONEP_ACCEPTANCE_MTP", str(int(mtp))) + 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 + assert trainer.intra_layer_micro_batch == 1 + 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) + + 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) diff --git a/tests/engine/test_moonep_forward.py b/tests/engine/test_moonep_forward.py new file mode 100644 index 0000000000..a4fff5bbec --- /dev/null +++ b/tests/engine/test_moonep_forward.py @@ -0,0 +1,878 @@ +import unittest + +import torch +import torch.distributed as dist +from torch.distributed.tensor import DTensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.loss import CELossConfig +from xtuner.v1.model.base import ModelItem +from xtuner.v1.model.moe.glm52 import Glm52MoEConfig +from xtuner.v1.model.moe.moe import MoEConfig +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig +from xtuner.v1.module.attention import DSAMLAConfig, MHAConfig +from xtuner.v1.module.mtp import MTPConfig +from xtuner.v1.module.router import GreedyRouterConfig, NoAuxRouterConfig +from xtuner.v1.utils.test_utils import init_data_mesh + + +def _tiny_config( + family: str, + dispatcher: str, + *, + compile: bool, + router_compute_dtype: str = "float32", + staging_reference: bool | None = None, + mtp_config: MTPConfig | None = None, + n_shared_experts: int = 1, + with_shared_expert_gate: bool = False, +) -> MoEConfig: + common = dict( + vocab_size=256, + max_position_embeddings=64, + pad_token_id=0, + eos_token_id=1, + num_hidden_layers=3, + first_k_dense_replace=1, + # With EP4/E8 each home chunk must satisfy CUDA's 2 MiB VMM + # granularity for both fused projections. + hidden_size=512, + intermediate_size=1024, + rms_norm_eps=1e-6, + hidden_act="silu", + n_routed_experts=8, + n_shared_experts=n_shared_experts, + with_shared_expert_gate=with_shared_expert_gate, + num_experts_per_tok=2, + moe_intermediate_size=1024, + ep_size=4, + dispatcher=dispatcher, + router_compute_dtype=router_compute_dtype, + moonep_staging_reference=False if staging_reference is None else staging_reference, + balancing_loss_cfg=None, + mtp_config=mtp_config, + compile_cfg=( + {"xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEBlock.forward": {"fullgraph": True}} + if compile + else False + ), + ) + if family == "qwen": + return Qwen3MoEConfig( + **common, + bos_token_id=2, + attention=MHAConfig( + num_attention_heads=8, + num_key_value_heads=8, + head_dim=64, + qk_norm=True, + attn_impl="flex_attention", + ), + router=GreedyRouterConfig( + scoring_func="softmax", + norm_topk_prob=True, + router_scaling_factor=1.0, + ), + ) + if family == "glm52": + return Glm52MoEConfig( + **common, + hf_eos_token_id=[1], + attention=DSAMLAConfig( + num_attention_heads=2, + head_dim=4, + kv_lora_rank=4, + q_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=4, + index_topk=4, + index_head_dim=4, + index_n_heads=2, + indexer_types=["full", "shared", "shared"], + sparse_mla_backend="torch", + ), + hf_head_dim=4, + qk_head_dim=8, + router=NoAuxRouterConfig( + n_group=1, + topk_group=1, + scoring_func="sigmoid", + norm_topk_prob=True, + router_scaling_factor=2.5, + ), + mlp_layer_types=["dense", "sparse", "sparse"], + num_nextn_predict_layers=None, + ) + raise AssertionError(f"unknown tiny model family: {family}") + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestMoonEPStagingForward(DeterministicDDPTestCase): + @staticmethod + def _training_item() -> ModelItem: + input_ids = torch.arange(2, 18, device="cuda").view(1, -1) + labels = (input_ids + 1) % 256 + loss_cfg = CELossConfig() + loss_ctx = loss_cfg.build(data={"shifted_labels": labels}) + assert loss_ctx is not None + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + return ModelItem( + seq_ctx=SequenceContext.from_input_ids((input_ids,), device="cuda"), + loss_ctx={"lm": loss_ctx}, + ) + + @staticmethod + def _model_training_item( + engine: TrainEngine, + *, + offset: int = 0, + sequence_length: int = 16, + sp_mesh=None, + ) -> ModelItem: + input_ids = (torch.arange(2, 2 + sequence_length, device="cuda") + offset).view(1, -1) % 256 + full_seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cuda") + loss_ctx = engine.model.build_loss_ctx_batch( + [{"seq_ctx": full_seq_ctx, "shifted_labels": (input_ids + 1) % 256}], + sp_mesh=sp_mesh, + )[0] + seq_ctx = full_seq_ctx if sp_mesh is None else full_seq_ctx.split(sp_mesh) + return ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx) + + def _forward(self, family: str, dispatcher: str) -> torch.Tensor: + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config(family, dispatcher, compile=True), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + intra_layer_micro_batch=1, + ) + engine.init_model_weights() + if dispatcher == "moonep": + assert engine.model.config.intra_layer_micro_batch == 1 + input_ids = torch.arange(2, 18, device="cuda").view(1, -1) + + try: + engine.model.eval() + with torch.no_grad(): + first = engine.model( + seq_ctx=SequenceContext.from_input_ids((input_ids,), device="cuda"), + loss_ctx=None, + ).logits + + assert first is not None + assert torch.isfinite(first).all() + repeats = 3 if dispatcher == "moonep" else int(dispatcher != "deepep") + for _ in range(repeats): + with torch.no_grad(): + repeated = engine.model( + seq_ctx=SequenceContext.from_input_ids((input_ids,), device="cuda"), + loss_ctx=None, + ).logits + torch.testing.assert_close(first, repeated, rtol=0, atol=0) + return first.clone() + finally: + # Resource teardown may unmap VMM landings, so the test must first + # complete queued output copies. This is lifecycle-only, not a hot-path sync. + torch.cuda.synchronize() + if dispatcher == "moonep": + engine.model.close_ep_runtime() + del engine + # DeepEP owns a process-scoped C++ Buffer. Forcing cyclic GC here + # can destruct it on only a subset of ranks; leave that resource + # to the distributed process teardown. + torch.cuda.empty_cache() + dist.barrier() + + def _assert_matches_reference(self, family: str, reference: str) -> None: + self.create_pg("cuda") + expected = self._forward(family, reference) + moonep = self._forward(family, "moonep") + torch.testing.assert_close(moonep, expected, rtol=1e-2, atol=1e-2) + + @staticmethod + def _selected_training_tensors(engine: TrainEngine, *, gradients: bool) -> dict[str, torch.Tensor]: + selected = {} + for name, parameter in engine.model.named_parameters(): + if not any( + marker in name for marker in (".experts.", ".shared_experts.", ".shared_expert_gate.", ".gate.") + ): + continue + value = parameter.grad if gradients else parameter + assert value is not None + if isinstance(value, DTensor): + value = value.to_local() + selected[name] = value.detach().clone() + return selected + + def _train_two_steps( + self, + dispatcher: str, + *, + staging_reference: bool | None = None, + ) -> tuple[list[float], list[torch.Tensor], list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]: + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config( + "qwen", + dispatcher, + compile=True, + staging_reference=staging_reference, + ), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + ) + engine.init_model_weights() + losses = [] + grad_norms = [] + gradients = [] + try: + for _ in range(2): + step = engine.train_step([self._training_item()]) + losses.append(step["total_loss"]) + grad_norms.append(engine.clip_grad_norm(do_clip=False).detach().clone()) + gradients.append(self._selected_training_tensors(engine, gradients=True)) + engine.step_optimizer(grad_norms[-1]) + parameters = self._selected_training_tensors(engine, gradients=False) + return losses, grad_norms, gradients, parameters + finally: + torch.cuda.synchronize() + if dispatcher == "moonep": + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + @staticmethod + def _assert_training_runs_close( + actual: tuple[list[float], list[torch.Tensor], list[dict[str, torch.Tensor]], dict[str, torch.Tensor]], + expected: tuple[list[float], list[torch.Tensor], list[dict[str, torch.Tensor]], dict[str, torch.Tensor]], + ) -> None: + actual_losses, actual_norms, actual_gradients, actual_parameters = actual + expected_losses, expected_norms, expected_gradients, expected_parameters = expected + torch.testing.assert_close( + torch.tensor(actual_losses, device="cuda"), + torch.tensor(expected_losses, device="cuda"), + rtol=1e-2, + atol=1e-3, + ) + for actual_norm, expected_norm in zip(actual_norms, expected_norms): + torch.testing.assert_close(actual_norm, expected_norm, rtol=1e-2, atol=1e-3) + for actual_step, expected_step in zip(actual_gradients, expected_gradients): + assert actual_step.keys() == expected_step.keys() + for name in actual_step: + torch.testing.assert_close(actual_step[name], expected_step[name], rtol=1e-2, atol=1e-3) + assert actual_parameters.keys() == expected_parameters.keys() + for name in actual_parameters: + torch.testing.assert_close(actual_parameters[name], expected_parameters[name], rtol=1e-2, atol=1e-3) + + def test_qwen_fixed_length_fused_expert_forward_matches_deepep(self) -> None: + self._assert_matches_reference("qwen", "deepep") + + def test_glm52_fixed_length_fused_expert_forward_matches_all2all(self) -> None: + self._assert_matches_reference("glm52", "all2all") + + def test_qwen_backward_updates_routed_expert_fsdp_shards(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config("qwen", "moonep", compile=True, router_compute_dtype="native"), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + ) + engine.init_model_weights() + routed_parameter = next( + parameter for name, parameter in engine.model.named_parameters() if ".experts." in name + ) + routed_parameter.grad = torch.full_like(routed_parameter, 15) + engine.model.scale_and_reduce_grad() + torch.testing.assert_close( + routed_parameter.grad.to_local(), + torch.full_like(routed_parameter.grad.to_local(), 15 / 4), + rtol=0, + atol=0, + ) + engine.optimizer.zero_grad() + before = { + name: parameter.to_local().detach().clone() + for name, parameter in engine.model.named_parameters() + if ".experts." in name + } + + try: + step = engine.train_step([self._training_item()]) + assert torch.isfinite(torch.tensor(step["total_loss"], device="cuda")) + routed = {name: parameter for name, parameter in engine.model.named_parameters() if ".experts." in name} + assert routed + assert all(parameter.grad is not None for parameter in routed.values()) + assert all(torch.isfinite(parameter.grad.to_local()).all() for parameter in routed.values()) + + grad_norm = engine.clip_grad_norm(do_clip=False) + engine.step_optimizer(grad_norm) + assert any(not torch.equal(before[name], parameter.to_local()) for name, parameter in routed.items()) + finally: + torch.cuda.synchronize() + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def test_qwen_two_step_training_matches_deepep(self) -> None: + self.create_pg("cuda") + expected = self._train_two_steps("deepep") + actual = self._train_two_steps("moonep") + repeated = self._train_two_steps("moonep") + self._assert_training_runs_close(actual, expected) + self._assert_training_runs_close(repeated, actual) + + def test_qwen_direct_landing_matches_staging_training(self) -> None: + self.create_pg("cuda") + staging = self._train_two_steps("moonep", staging_reference=True) + direct = self._train_two_steps("moonep", staging_reference=False) + self._assert_training_runs_close(direct, staging) + + def test_qwen_direct_hot_path_has_no_full_weight_copy_or_host_sync(self) -> None: + self.create_pg("cuda") + + def profile_mode( + staging_reference: bool, + *, + mtp_micro2_sp4: bool = False, + ) -> tuple[int, int, list[str], int]: + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config( + "qwen", + "moonep", + compile=True, + staging_reference=staging_reference, + mtp_config=(MTPConfig(num_layers=2, share_weights=True) if mtp_micro2_sp4 else None), + ), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig( + ep_size=4, + recompute_ratio=1.0 if mtp_micro2_sp4 else 0.0, + torch_compile=True, + mtp_checkpoint_use_reentrant=True, + ), + intra_layer_micro_batch=2 if mtp_micro2_sp4 else 1, + ) + engine.init_model_weights() + if mtp_micro2_sp4: + sp_mesh = init_data_mesh("cuda", sp_size=4)["sp"] + train_items = [ + self._model_training_item( + engine, + offset=micro_batch_idx * 32, + sequence_length=32, + sp_mesh=sp_mesh, + ) + for micro_batch_idx in range(2) + ] + else: + train_items = [self._training_item()] + try: + torch.cuda.reset_peak_memory_stats() + # Compile/autotune before profiling so their setup-only CUDA + # synchronization cannot be confused with the steady hot path. + engine.train_step(train_items) + engine.optimizer.zero_grad() + + with torch.profiler.profile( + activities=(torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA), + record_shapes=True, + ) as profiler: + engine.train_step(train_items) + + full_home_shapes = {(2, 2048, 512), (2, 512, 1024)} + full_local_dw_shapes = {(4, 2048, 512), (4, 512, 1024)} + full_weight_copies = 0 + full_dw_materializations = 0 + host_syncs: list[str] = [] + for event in profiler.events(): + tensor_shapes = { + tuple(shape) + for shape in event.input_shapes + if isinstance(shape, list) and all(isinstance(dim, int) for dim in shape) + } + if event.name == "aten::copy_" and tensor_shapes & full_home_shapes: + full_weight_copies += 1 + if event.name in {"aten::clone", "aten::copy_", "aten::zeros_like"} and ( + tensor_shapes & full_local_dw_shapes + ): + full_dw_materializations += 1 + + parent = event.cpu_parent + inside_gate = False + while parent is not None: + if parent.name.startswith("MoonEP::"): + inside_gate = True + break + parent = parent.cpu_parent + if inside_gate and event.name in { + "cudaDeviceSynchronize", + "cudaEventSynchronize", + "cudaStreamSynchronize", + }: + ancestry = [] + parent = event.cpu_parent + while parent is not None: + ancestry.append(parent.name) + parent = parent.cpu_parent + host_syncs.append(f"{event.name} <- {' <- '.join(ancestry)}") + return ( + full_weight_copies, + full_dw_materializations, + host_syncs, + torch.cuda.max_memory_allocated(), + ) + finally: + torch.cuda.synchronize() + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + staging_copies, _, _, _ = profile_mode(True) + direct_copies, direct_dw_materializations, direct_host_syncs, _ = profile_mode(False) + assert staging_copies > 0 # Calibrates the shape-based copy detector. + assert direct_copies == 0 + assert direct_dw_materializations == 0 + assert direct_host_syncs == [], direct_host_syncs + combo_copies, combo_dw_materializations, combo_host_syncs, combo_peak_bytes = profile_mode( + False, + mtp_micro2_sp4=True, + ) + assert combo_copies == 0 + assert combo_dw_materializations == 0 + assert combo_host_syncs == [], combo_host_syncs + # The fixed tiny fallback measured 0.185 GiB/rank on H200; leave ample + # headroom while still catching an accidental full-model materialization. + assert combo_peak_bytes < 2**30 + + def _train_mtp_micro2( + self, + dispatcher: str, + *, + share_weights: bool, + ) -> tuple[list[tuple[float, float]], list[torch.Tensor], list[dict[str, torch.Tensor]]]: + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config( + "qwen", + dispatcher, + compile=True, + mtp_config=MTPConfig(num_layers=2, share_weights=share_weights), + ), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig( + ep_size=4, + recompute_ratio=1.0, + torch_compile=True, + mtp_checkpoint_use_reentrant=True, + ), + intra_layer_micro_batch=2, + ) + engine.init_model_weights() + losses = [] + grad_norms = [] + gradients = [] + try: + if dispatcher == "moonep": + # 两次 forward-only 调用必须各自完成并释放 main/MTP plan；随后 + # 同一个 runtime 直接进入正常 reentrant training/replay。 + engine.model.eval() + forward_only_losses = [] + with torch.no_grad(): + for _ in range(2): + item = self._model_training_item(engine) + output = engine.model(seq_ctx=item["seq_ctx"], loss_ctx=item["loss_ctx"]) + assert output.loss is not None and output.mtp_loss is not None + forward_only_losses.append(torch.stack((output.loss, output.mtp_loss))) + torch.testing.assert_close(forward_only_losses[1], forward_only_losses[0], rtol=1e-2, atol=3e-3) + engine.model.train() + for step_idx in range(3): + step = engine.train_step( + [ + self._model_training_item(engine, offset=step_idx * 32), + self._model_training_item(engine, offset=step_idx * 32 + 16), + ] + ) + losses.append((step["total_loss"], step["logs_info"]["reduced_mtp_loss"])) + grad_norms.append(engine.clip_grad_norm(do_clip=False).detach().clone()) + if step_idx == 0: + gradients.append(self._selected_training_tensors(engine, gradients=True)) + engine.step_optimizer(grad_norms[-1]) + return losses, grad_norms, gradients + finally: + torch.cuda.synchronize() + if dispatcher == "moonep": + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def _assert_mtp_micro2_matches_deepep(self, *, share_weights: bool) -> None: + self.create_pg("cuda") + expected = self._train_mtp_micro2("deepep", share_weights=share_weights) + actual = self._train_mtp_micro2("moonep", share_weights=share_weights) + try: + expected_losses, expected_norms, expected_gradients = expected + actual_losses, actual_norms, actual_gradients = actual + torch.testing.assert_close( + torch.tensor(actual_losses, device="cuda"), + torch.tensor(expected_losses, device="cuda"), + rtol=1e-2, + atol=1e-3, + ) + for actual_norm, expected_norm in zip(actual_norms, expected_norms): + torch.testing.assert_close(actual_norm, expected_norm, rtol=1e-2, atol=1e-3) + # DeepEP/MoonEP 的 BF16 前向舍入会通过后续 router 放大，不适合在 + # micro2 多层图上逐元素判梯度；slot 累加由 identical-items 测试精确覆盖。 + for actual_step, expected_step in zip(actual_gradients, expected_gradients): + assert actual_step.keys() == expected_step.keys() + for gradients in (actual_step, expected_step): + assert all(torch.isfinite(tensor).all() for tensor in gradients.values()) + assert any(torch.count_nonzero(tensor) > 0 for tensor in gradients.values()) + finally: + # All ranks must leave numerical assertions together before the + # distributed test harness destroys the world process group. + torch.cuda.synchronize() + dist.barrier() + + def test_qwen_unshared_mtp_reentrant_micro2_matches_deepep(self) -> None: + self._assert_mtp_micro2_matches_deepep(share_weights=False) + + def test_qwen_shared_mtp_reentrant_micro2_matches_deepep(self) -> None: + self._assert_mtp_micro2_matches_deepep(share_weights=True) + + def test_qwen_mtp_reentrant_micro2_completes_forward_and_backward(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config( + "qwen", + "moonep", + compile=True, + mtp_config=MTPConfig(num_layers=2, share_weights=True), + ), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig( + ep_size=4, + recompute_ratio=1.0, + torch_compile=True, + mtp_checkpoint_use_reentrant=True, + ), + intra_layer_micro_batch=2, + ) + engine.init_model_weights() + try: + result = engine.train_step( + [ + self._model_training_item(engine, offset=0), + self._model_training_item(engine, offset=16), + ] + ) + gradients = self._selected_training_tensors(engine, gradients=True) + assert torch.isfinite(torch.tensor(result["total_loss"], device="cuda")) + assert gradients and all(torch.isfinite(gradient).all() for gradient in gradients.values()) + assert any(torch.count_nonzero(gradient) > 0 for gradient in gradients.values()) + finally: + torch.cuda.synchronize() + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def test_qwen_two_live_forwards_support_reverse_backward(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config("qwen", "moonep", compile=True), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + intra_layer_micro_batch=2, + ) + engine.init_model_weights() + try: + items = [self._model_training_item(engine, offset=offset) for offset in (0, 16)] + outputs = [engine.model(seq_ctx=item["seq_ctx"], loss_ctx=item["loss_ctx"]) for item in items] + for output in reversed(outputs): + assert output.loss is not None + output.loss.backward() + + gradients = self._selected_training_tensors(engine, gradients=True) + assert gradients and all(torch.isfinite(gradient).all() for gradient in gradients.values()) + assert any(torch.count_nonzero(gradient) > 0 for gradient in gradients.values()) + finally: + torch.cuda.synchronize() + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def test_qwen_requires_the_configured_domino_width(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config("qwen", "moonep", compile=True), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + intra_layer_micro_batch=2, + ) + engine.init_model_weights() + items = [self._model_training_item(engine, offset=idx * 16) for idx in range(3)] + try: + for actual_width in (1, 3): + with ( + self.subTest(actual_width=actual_width), + torch.no_grad(), + self.assertRaisesRegex(ValueError, f"width {actual_width} does not match configured width 2"), + ): + engine.model( + seq_ctx=[item["seq_ctx"] for item in items[:actual_width]], + loss_ctx=[item["loss_ctx"] for item in items[:actual_width]], + ) + finally: + torch.cuda.synchronize() + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def _train_microbatches_without_mtp( + self, + dispatcher: str, + *, + recompute_ratio: float, + offsets: tuple[int, ...] = (0, 16), + n_shared_experts: int = 1, + with_shared_expert_gate: bool = False, + routed_only: bool = True, + ) -> tuple[float, torch.Tensor, dict[str, torch.Tensor]]: + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config( + "qwen", + dispatcher, + compile=True, + n_shared_experts=n_shared_experts, + with_shared_expert_gate=with_shared_expert_gate, + ), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig( + ep_size=4, + recompute_ratio=recompute_ratio, + torch_compile=True, + mtp_checkpoint_use_reentrant=True, + ), + intra_layer_micro_batch=len(offsets), + ) + engine.init_model_weights() + try: + step = engine.train_step([self._model_training_item(engine, offset=offset) for offset in offsets]) + grad_norm = engine.clip_grad_norm(do_clip=False).detach().clone() + gradients = { + name: tensor + for name, tensor in self._selected_training_tensors(engine, gradients=True).items() + if not routed_only or ".experts." in name + } + return step["total_loss"], grad_norm, gradients + finally: + torch.cuda.synchronize() + if dispatcher == "moonep": + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def test_qwen_reentrant_micro2_routed_gradients_match_deepep(self) -> None: + self.create_pg("cuda") + expected_loss, expected_norm, expected_gradients = self._train_microbatches_without_mtp( + "deepep", recompute_ratio=1.0 + ) + actual_loss, actual_norm, actual_gradients = self._train_microbatches_without_mtp( + "moonep", recompute_ratio=1.0 + ) + torch.testing.assert_close( + torch.tensor(actual_loss, device="cuda"), + torch.tensor(expected_loss, device="cuda"), + rtol=1e-2, + atol=1e-3, + ) + torch.testing.assert_close(actual_norm, expected_norm, rtol=1e-2, atol=1e-3) + assert actual_gradients.keys() == expected_gradients.keys() + for gradients in (actual_gradients, expected_gradients): + assert all(torch.isfinite(tensor).all() for tensor in gradients.values()) + assert any(torch.count_nonzero(tensor) > 0 for tensor in gradients.values()) + + def test_qwen_micro2_identical_items_accumulate_like_micro1(self) -> None: + self.create_pg("cuda") + expected_loss, expected_norm, expected_gradients = self._train_microbatches_without_mtp( + "moonep", recompute_ratio=0.0, offsets=(0,) + ) + actual_loss, actual_norm, actual_gradients = self._train_microbatches_without_mtp( + "moonep", recompute_ratio=0.0, offsets=(0, 0) + ) + torch.testing.assert_close( + torch.tensor(actual_loss, device="cuda"), + torch.tensor(2 * expected_loss, device="cuda"), + rtol=1e-2, + atol=1e-3, + ) + torch.testing.assert_close(actual_norm, 2 * expected_norm, rtol=1e-2, atol=1e-3) + assert actual_gradients.keys() == expected_gradients.keys() + for name in actual_gradients: + max_error = (actual_gradients[name].float() - 2 * expected_gradients[name].float()).abs().max() + dist.all_reduce(max_error, op=dist.ReduceOp.MAX) + assert max_error <= 1e-3, f"{name}: max_abs={max_error.item()}" + + def test_qwen_micro4_trains_with_the_configured_width(self) -> None: + self.create_pg("cuda") + loss, grad_norm, gradients = self._train_microbatches_without_mtp( + "moonep", + recompute_ratio=0.0, + offsets=(0, 16, 32, 48), + ) + assert torch.isfinite(torch.tensor(loss, device="cuda")) + assert torch.isfinite(grad_norm) + assert gradients and all(torch.isfinite(gradient).all() for gradient in gradients.values()) + + def test_qwen_micro8_trains_with_the_configured_width(self) -> None: + self.create_pg("cuda") + loss, grad_norm, gradients = self._train_microbatches_without_mtp( + "moonep", + recompute_ratio=0.0, + offsets=tuple(range(0, 128, 16)), + ) + assert torch.isfinite(torch.tensor(loss, device="cuda")) + assert torch.isfinite(grad_norm) + assert gradients and all(torch.isfinite(gradient).all() for gradient in gradients.values()) + + def test_qwen_shared_expert_variants_train(self) -> None: + self.create_pg("cuda") + _, no_shared_norm, no_shared_gradients = self._train_microbatches_without_mtp( + "moonep", + recompute_ratio=0.0, + offsets=(0,), + n_shared_experts=0, + routed_only=False, + ) + _, gated_norm, gated_gradients = self._train_microbatches_without_mtp( + "moonep", + recompute_ratio=0.0, + offsets=(0,), + n_shared_experts=1, + with_shared_expert_gate=True, + routed_only=False, + ) + assert torch.isfinite(no_shared_norm) and torch.isfinite(gated_norm) + assert not any("shared_expert" in name for name in no_shared_gradients) + assert any("shared_experts" in name for name in gated_gradients) + assert any("shared_expert_gate" in name for name in gated_gradients) + assert all(torch.isfinite(gradient).all() for gradient in gated_gradients.values()) + + def test_qwen_shared_expert_gradient_uses_fp32_ep_mean(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config( + "qwen", + "moonep", + compile=True, + n_shared_experts=1, + with_shared_expert_gate=True, + ), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + ) + engine.init_model_weights() + shared_gate = next( + parameter for name, parameter in engine.model.named_parameters() if ".shared_expert_gate." in name + ) + assert isinstance(shared_gate, DTensor) + shared_gate.grad = torch.full_like(shared_gate, dist.get_rank() + 1) + + try: + engine.model.scale_and_reduce_grad() + local_grad = shared_gate.grad.to_local() + # Model mesh is [FSDP2, EP4]. Shared parameters are sharded on the + # first dimension and replicated on each contiguous EP4 row. + expected_mean = (dist.get_rank() // 4) * 4 + 2.5 + assert local_grad.dtype is torch.float32 + torch.testing.assert_close( + local_grad, + torch.full_like(local_grad, expected_mean), + rtol=0, + atol=0, + ) + finally: + torch.cuda.synchronize() + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def _train_sp_once(self, dispatcher: str, *, sp_size: int) -> tuple[float, torch.Tensor]: + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config("qwen", dispatcher, compile=True), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0, torch_compile=True), + ) + engine.init_model_weights() + sp_mesh = init_data_mesh("cuda", sp_size=sp_size)["sp"] + item = self._model_training_item( + engine, + sequence_length=32, + sp_mesh=sp_mesh, + ) + + try: + step = engine.train_step([item]) + grad_norm = engine.clip_grad_norm(do_clip=False).detach().clone() + routed_gradients = { + name: gradient + for name, gradient in self._selected_training_tensors(engine, gradients=True).items() + if ".experts." in name + } + assert routed_gradients + assert all(torch.isfinite(gradient).all() for gradient in routed_gradients.values()) + assert any(torch.count_nonzero(gradient) > 0 for gradient in routed_gradients.values()) + engine.step_optimizer(grad_norm) + return step["total_loss"], grad_norm + finally: + torch.cuda.synchronize() + if dispatcher == "moonep": + engine.model.close_ep_runtime() + del engine + torch.cuda.empty_cache() + dist.barrier() + + def _assert_sp_matches_deepep(self, sp_size: int) -> None: + self.create_pg("cuda") + expected_loss, expected_norm = self._train_sp_once("deepep", sp_size=sp_size) + actual_loss, actual_norm = self._train_sp_once("moonep", sp_size=sp_size) + torch.testing.assert_close( + torch.tensor(actual_loss, device="cuda"), + torch.tensor(expected_loss, device="cuda"), + rtol=1e-2, + atol=1e-3, + ) + torch.testing.assert_close(actual_norm, expected_norm, rtol=1e-2, atol=1e-3) + + def test_qwen_sp2_ep4_matches_deepep(self) -> None: + self._assert_sp_matches_deepep(2) + + def test_qwen_sp4_ep4_matches_deepep(self) -> None: + self._assert_sp_matches_deepep(4) + + def test_qwen_sp8_ep4_matches_deepep(self) -> None: + self._assert_sp_matches_deepep(8) + + @property + def world_size(self) -> int: + return 8 diff --git a/tests/engine/test_moonep_persistence.py b/tests/engine/test_moonep_persistence.py new file mode 100644 index 0000000000..5add5d2e60 --- /dev/null +++ b/tests/engine/test_moonep_persistence.py @@ -0,0 +1,398 @@ +import gc +import os +import shutil +import tempfile +import unittest +import warnings +from pathlib import Path +from unittest.mock import patch + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.tensor import DTensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.config import AdamWConfig, FSDPConfig +from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.engine.train_engine import TrainEngine +from xtuner.v1.model.base import ModelItem +from xtuner.v1.model.moe.qwen3 import Qwen3MoEConfig +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.router import GreedyRouterConfig + + +def _tiny_moonep_config( + *, + return_router_results: bool = False, + router_async_offload: bool = False, +) -> Qwen3MoEConfig: + return Qwen3MoEConfig( + vocab_size=256, + max_position_embeddings=64, + pad_token_id=0, + eos_token_id=1, + bos_token_id=2, + num_hidden_layers=3, + first_k_dense_replace=1, + hidden_size=512, + intermediate_size=1024, + rms_norm_eps=1e-6, + hidden_act="silu", + n_routed_experts=8, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=1024, + ep_size=4, + dispatcher="moonep", + router_compute_dtype="float32", + moonep_staging_reference=False, + balancing_loss_cfg=None, + return_router_results=return_router_results, + router_async_offload=router_async_offload, + compile_cfg=False, + attention=MHAConfig( + num_attention_heads=8, + num_key_value_heads=8, + head_dim=64, + qk_norm=True, + attn_impl="flex_attention", + ), + router=GreedyRouterConfig( + scoring_func="softmax", + norm_topk_prob=True, + router_scaling_factor=1.0, + ), + ) + + +def _training_item(engine: TrainEngine, *, offset: int = 0) -> ModelItem: + input_ids = (torch.arange(2, 18, device="cuda") + offset).view(1, -1) % 256 + seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cuda") + loss_ctx = engine.model.build_loss_ctx_batch( + [{"seq_ctx": seq_ctx, "shifted_labels": (input_ids + 1) % 256}], + sp_mesh=None, + )[0] + return ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx) + + +def _local_model_state(engine: TrainEngine) -> dict[str, torch.Tensor]: + state = {} + for name, value in engine.model.state_dict().items(): + if isinstance(value, DTensor): + value = value.to_local() + state[name] = value.detach().clone() + return state + + +def _optimizer_tensor_state(engine: TrainEngine) -> dict[str, torch.Tensor]: + tensors = {} + parameter_names = {id(parameter): name for name, parameter in engine.model.named_parameters()} + state_dict = engine.optimizer.state_dict() + for saved_group, live_group in zip(state_dict["param_groups"], engine.optimizer.param_groups, strict=True): + for parameter_id, parameter in zip(saved_group["params"], live_group["params"], strict=True): + parameter_name = parameter_names[id(parameter)] + for name, value in state_dict["state"][parameter_id].items(): + if isinstance(value, torch.Tensor): + if isinstance(value, DTensor): + value = value.to_local() + tensors[f"{parameter_name}.{name}"] = value.detach().clone() + return tensors + + +def _optimizer_step(engine: TrainEngine, *, offset: int) -> tuple[float, torch.Tensor]: + step = engine.train_step([_training_item(engine, offset=offset)]) + grad_norm = engine.clip_grad_norm(do_clip=False).detach().clone() + engine.step_optimizer(grad_norm) + return step["total_loss"], grad_norm + + +@torch.no_grad() +def _probe_logits(engine: TrainEngine, *, offset: int = 48) -> torch.Tensor: + item = _training_item(engine, offset=offset) + engine.model.eval() + output = engine.model(seq_ctx=item["seq_ctx"], loss_ctx=None) + assert output.logits is not None + return output.logits.detach().clone() + + +def _shared_temporary_directory() -> Path: + directory = tempfile.mkdtemp() if dist.get_rank() == 0 else None + shared = [directory] + dist.broadcast_object_list(shared, src=0) + assert shared[0] is not None + return Path(shared[0]) + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestMoonEPPersistence(DeterministicDDPTestCase): + @staticmethod + def _build_engine(optim_cfg=None, *, model_cfg=None) -> TrainEngine: + return TrainEngine( + model_cfg=model_cfg or _tiny_moonep_config(), + optim_cfg=optim_cfg or AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=0.0), + ) + + def test_engine_close_is_idempotent_and_rejects_further_forward(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = self._build_engine() + engine.init_model_weights() + item = _training_item(engine) + + try: + engine.model.eval() + with torch.no_grad(): + output = engine.model(seq_ctx=item["seq_ctx"], loss_ctx=None) + assert output.logits is not None + + engine.close() + engine.close() + + with self.assertRaisesRegex(RuntimeError, "closed"): + engine.train_step([item]) + with self.assertRaisesRegex(RuntimeError, "closed"): + engine.model(seq_ctx=item["seq_ctx"], loss_ctx=None) + finally: + # 红测阶段 close 尚不存在，仍需显式释放 collective 资源。 + if not getattr(engine, "_closed", False): + torch.cuda.synchronize() + engine.model.close_ep_runtime() + dist.barrier() + + def test_rank_divergent_destructor_only_warns(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = self._build_engine() + engine.init_model_weights() + + # Rank 0 drops the engine first. If __del__ enters a Buffer barrier, + # CUDA synchronize, or VMM teardown, this world barrier cannot finish. + if dist.get_rank() == 0: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", ResourceWarning) + del engine + gc.collect() + assert any("TrainEngine.close" in str(item.message) for item in caught) + dist.barrier() + + if dist.get_rank() != 0: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", ResourceWarning) + del engine + gc.collect() + assert any("TrainEngine.close" in str(item.message) for item in caught) + dist.barrier() + + def test_sync_dcp_cold_resume_matches_uninterrupted_step(self) -> None: + self.create_pg("cuda") + checkpoint_root = _shared_temporary_directory() + weights_dir = checkpoint_root / "weights" + torch.manual_seed(20260805) + reference = self._build_engine() + reference.init_model_weights() + + try: + _optimizer_step(reference, offset=0) + _optimizer_step(reference, offset=16) + checkpoint_model = _local_model_state(reference) + checkpoint_optimizer = _optimizer_tensor_state(reference) + reference.save_dcp(weights_dir) + dist.barrier() + + metadata_keys = dcp.FileSystemReader(weights_dir).read_metadata().state_dict_metadata.keys() + assert any(key.startswith("model.") for key in metadata_keys) + assert any(key.startswith("optimizer.") for key in metadata_keys) + transient_markers = ("moonep", "workspace", "landing", "invocation", "gradient_slot", "event") + assert not any(marker in key.lower() for key in metadata_keys for marker in transient_markers) + + expected_loss, expected_norm = _optimizer_step(reference, offset=32) + expected_model = _local_model_state(reference) + expected_optimizer = _optimizer_tensor_state(reference) + updated_names = { + name for name, value in expected_model.items() if not torch.equal(value, checkpoint_model[name]) + } + assert updated_names + reference.close() + + # Load occurs before this fresh runtime's first forward/AllGather. + torch.manual_seed(17) + resumed = self._build_engine() + resumed.init_model_weights() + resumed.load_dcp(weights_dir) + + actual_checkpoint_model = _local_model_state(resumed) + actual_checkpoint_optimizer = _optimizer_tensor_state(resumed) + assert actual_checkpoint_model.keys() == checkpoint_model.keys() + assert actual_checkpoint_optimizer.keys() == checkpoint_optimizer.keys() + for name, expected in checkpoint_model.items(): + torch.testing.assert_close(actual_checkpoint_model[name], expected, rtol=0, atol=0) + for name, expected in checkpoint_optimizer.items(): + torch.testing.assert_close(actual_checkpoint_optimizer[name], expected, rtol=0, atol=0) + + actual_loss, actual_norm = _optimizer_step(resumed, offset=32) + actual_model = _local_model_state(resumed) + actual_optimizer = _optimizer_tensor_state(resumed) + torch.testing.assert_close( + torch.tensor(actual_loss, device="cuda"), + torch.tensor(expected_loss, device="cuda"), + rtol=1e-5, + atol=1e-6, + ) + torch.testing.assert_close(actual_norm, expected_norm, rtol=1e-5, atol=1e-6) + for name in updated_names: + torch.testing.assert_close(actual_model[name], expected_model[name], rtol=0, atol=0) + assert actual_optimizer.keys() == expected_optimizer.keys() + for name, expected in expected_optimizer.items(): + torch.testing.assert_close(actual_optimizer[name], expected, rtol=0, atol=0) + resumed.close() + finally: + if not reference._closed: + reference.close() + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(checkpoint_root) + + def test_sync_hf_export_loads_into_fresh_moonep_runtime(self) -> None: + self.create_pg("cuda") + export_root = _shared_temporary_directory() + hf_dir = export_root / "hf" + torch.manual_seed(20260805) + reference = self._build_engine() + reference.init_model_weights() + restored = None + + try: + _optimizer_step(reference, offset=0) + _optimizer_step(reference, offset=16) + expected_state = _local_model_state(reference) + expected_logits = _probe_logits(reference) + reference.save_hf(str(hf_dir)) + reference.close() + + restored = self._build_engine() + restored.from_hf(hf_dir, strict=True) + actual_state = _local_model_state(restored) + assert actual_state.keys() == expected_state.keys() + for name, expected in expected_state.items(): + # HF is deliberately a BF16 interchange format while FSDP + # optimizer shards remain FP32 in the live engine. + torch.testing.assert_close(actual_state[name].bfloat16(), expected.bfloat16(), rtol=0, atol=0) + torch.testing.assert_close(_probe_logits(restored), expected_logits, rtol=0, atol=0) + restored.close() + finally: + if not reference._closed: + reference.close() + if restored is not None and not restored._closed: + restored.close() + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(export_root) + + def test_activation_offload_preserves_moonep_training(self) -> None: + self.create_pg("cuda") + results = [] + for enabled in (False, True): + torch.manual_seed(20260805) + engine = self._build_engine() + engine.init_model_weights() + try: + with patch.dict(os.environ, {"XTUNER_ACTIVATION_OFFLOAD": str(int(enabled))}): + loss, grad_norm = _optimizer_step(engine, offset=0) + results.append((loss, grad_norm, _local_model_state(engine))) + finally: + engine.close() + + expected_loss, expected_norm, expected_state = results[0] + actual_loss, actual_norm, actual_state = results[1] + torch.testing.assert_close( + torch.tensor(actual_loss, device="cuda"), + torch.tensor(expected_loss, device="cuda"), + rtol=0, + atol=0, + ) + torch.testing.assert_close(actual_norm, expected_norm, rtol=0, atol=0) + for name, expected in expected_state.items(): + torch.testing.assert_close(actual_state[name], expected, rtol=0, atol=0) + dist.barrier() + + def test_router_async_offload_only_changes_detached_logging_outputs(self) -> None: + self.create_pg("cuda") + results = [] + for enabled in (False, True): + torch.manual_seed(20260805) + config = _tiny_moonep_config(return_router_results=True, router_async_offload=enabled) + engine = self._build_engine(model_cfg=config) + engine.init_model_weights() + try: + item = _training_item(engine) + with torch.no_grad(): + output = engine.model(seq_ctx=item["seq_ctx"], loss_ctx=None) + assert output.router_logits + assert output.router_weights + logging_tensors = [*output.router_logits.values(), *output.router_weights.values()] + expected_device = "cpu" if enabled else "cuda" + assert all(tensor.device.type == expected_device for tensor in logging_tensors) + assert all(not tensor.requires_grad for tensor in logging_tensors) + + loss, grad_norm = _optimizer_step(engine, offset=16) + results.append((loss, grad_norm, _local_model_state(engine))) + finally: + engine.close() + + expected_loss, expected_norm, expected_state = results[0] + actual_loss, actual_norm, actual_state = results[1] + torch.testing.assert_close( + torch.tensor(actual_loss, device="cuda"), + torch.tensor(expected_loss, device="cuda"), + rtol=0, + atol=0, + ) + torch.testing.assert_close(actual_norm, expected_norm, rtol=0, atol=0) + for name, expected in expected_state.items(): + torch.testing.assert_close(actual_state[name], expected, rtol=0, atol=0) + dist.barrier() + + def test_async_hf_export_is_immutable_and_close_waits_for_writer(self) -> None: + self.create_pg("cuda") + export_root = _shared_temporary_directory() + hf_dir = export_root / "hf" + torch.manual_seed(20260805) + reference = self._build_engine() + reference.init_model_weights() + restored = None + + try: + _optimizer_step(reference, offset=0) + _optimizer_step(reference, offset=16) + expected_state = _local_model_state(reference) + expected_logits = _probe_logits(reference) + + save_future = reference.async_save_hf(str(hf_dir)) + _optimizer_step(reference, offset=32) + reference.close() + assert save_future.done() + assert hf_dir.is_dir() + + restored = self._build_engine() + restored.from_hf(hf_dir, strict=True) + actual_state = _local_model_state(restored) + for name, expected in expected_state.items(): + torch.testing.assert_close(actual_state[name].bfloat16(), expected.bfloat16(), rtol=0, atol=0) + torch.testing.assert_close(_probe_logits(restored), expected_logits, rtol=0, atol=0) + restored.close() + finally: + if not reference._closed: + if "save_future" in locals(): + save_future.result() + reference.close() + if restored is not None and not restored._closed: + restored.close() + dist.barrier() + if dist.get_rank() == 0: + shutil.rmtree(export_root) + + @property + def world_size(self) -> int: + return 8 diff --git a/tests/module/dispatcher/test_agrs_all2all.py b/tests/module/dispatcher/test_agrs_all2all.py index 1d0d68f530..499630e8dd 100644 --- a/tests/module/dispatcher/test_agrs_all2all.py +++ b/tests/module/dispatcher/test_agrs_all2all.py @@ -1,16 +1,15 @@ -import unittest +import os + +import parametrize import torch from torch.testing._internal.common_distributed import DistributedTestBase -from xtuner.v1.module.dispatcher.base import NaiveDispatcher, DispacherInterface -from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher -import parametrize + from xtuner.v1.module.dispatcher.agrs import MoEAGRSDispatcher +from xtuner.v1.module.dispatcher.base import DispacherInterface +from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher from xtuner.v1.module.router.greedy import GreedyGroupedRouter -import os - - EP_SIZE = 8 @@ -26,16 +25,10 @@ def test_dispatch_and_combine(self, dtype, device): num_experts = 128 all2all_dispatcher = TorchAll2AllDispatcher( - n_routed_experts=num_experts, - training_dtype="bf16", - process_group=torch.distributed.group.WORLD + n_routed_experts=num_experts, process_group=torch.distributed.group.WORLD ) - agrs_dispatcher = MoEAGRSDispatcher( - n_routed_experts=num_experts, - training_dtype="bf16", - process_group=torch.distributed.group.WORLD - ) + agrs_dispatcher = MoEAGRSDispatcher(n_routed_experts=num_experts, process_group=torch.distributed.group.WORLD) seq_len = 32 hidden_size = 128 @@ -57,29 +50,32 @@ def test_dispatch_and_combine(self, dtype, device): dispatcher=all2all_dispatcher, hidden_states=hidden_states, topk_ids=router_out["topk_ids"], - topk_weights=router_out["topk_weights"] + topk_weights=router_out["topk_weights"], ) agrs_results = self._dispatcher_call( dispatcher=agrs_dispatcher, hidden_states=hidden_states, topk_ids=router_out["topk_ids"], - topk_weights=router_out["topk_weights"] + topk_weights=router_out["topk_weights"], ) - self.assertTrue(torch.allclose(all2all_results["hidden_states"], agrs_results["hidden_states"], atol=1e-2, rtol=1e-2)) + self.assertTrue( + torch.allclose(all2all_results["hidden_states"], agrs_results["hidden_states"], atol=1e-2, rtol=1e-2) + ) def _dispatcher_call( - self, - dispatcher: DispacherInterface, - hidden_states: torch.Tensor, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor + self, + dispatcher: DispacherInterface, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, ): pre_dispatched = dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=128), ) dispatched = dispatcher.dispatch( pre_dispatched=pre_dispatched, diff --git a/tests/module/dispatcher/test_deepep.py b/tests/module/dispatcher/test_deepep.py index 4be7cec866..149ee56b26 100644 --- a/tests/module/dispatcher/test_deepep.py +++ b/tests/module/dispatcher/test_deepep.py @@ -1,17 +1,13 @@ -from unittest.mock import Mock +import os from typing import cast +import parametrize import torch import torch.distributed as dist from torch.testing._internal.common_distributed import DistributedTestBase +from xtuner.v1.module.dispatcher.base import GenericDispatcher, NaiveDispatcher from xtuner.v1.module.dispatcher.deepep import DeepEPDispatcher -from xtuner.v1.model.base import TransformerConfig -from xtuner.v1.module.dispatcher.base import NaiveDispatcher, GenericDispatcher -import parametrize - - -import os def mock_experts(hidden_states: torch.Tensor, tokens_per_exprts: torch.Tensor): @@ -24,7 +20,7 @@ class TestMoETorchAll2AllDispatcher(DistributedTestBase): [ (torch.bfloat16, "cuda", False), (torch.bfloat16, "cuda", True), - ] + ], ) def test_dispatch_and_combine(self, dtype, device, async_op): self.create_pg(device) @@ -32,13 +28,10 @@ def test_dispatch_and_combine(self, dtype, device, async_op): noep_dispatcher = NaiveDispatcher( n_routed_experts=num_experts, - training_dtype="bf16", ) all2all_dispatcher = DeepEPDispatcher( - n_routed_experts=num_experts, - training_dtype="bf16", - process_group=cast(dist.ProcessGroup, dist.group.WORLD) + n_routed_experts=num_experts, process_group=cast(dist.ProcessGroup, dist.group.WORLD) ) seq_len = 32 @@ -49,10 +42,7 @@ def test_dispatch_and_combine(self, dtype, device, async_op): topk_weights = torch.ones(seq_len, topk_experts).to(device).to(torch.float32) noep_results = self._dispatcher_call( - dispatcher=noep_dispatcher, - hidden_states=hidden_states, - topk_ids=topk_idx, - topk_weights=topk_weights + dispatcher=noep_dispatcher, hidden_states=hidden_states, topk_ids=topk_idx, topk_weights=topk_weights ) all2all_results = self._dispatcher_call( dispatcher=all2all_dispatcher, @@ -70,12 +60,13 @@ def _dispatcher_call( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, - async_op: bool=False + async_op: bool = False, ): pre_dispatched = dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=16), async_op=async_op, ) dispatched = dispatcher.dispatch( diff --git a/tests/module/dispatcher/test_deepep_expert_tp.py b/tests/module/dispatcher/test_deepep_expert_tp.py index 0fa3728ebf..ef6ce9eba7 100644 --- a/tests/module/dispatcher/test_deepep_expert_tp.py +++ b/tests/module/dispatcher/test_deepep_expert_tp.py @@ -70,6 +70,7 @@ def test_sync_virtual_expert_path_preserves_output_and_gradients(self) -> None: hidden_states=hidden_leaf, topk_ids=local_topk_ids, topk_weights=topk_weights_leaf, + tokens_per_expert=torch.bincount(local_topk_ids.flatten(), minlength=4), ) expected_virtual_ids = torch.tensor( [0, 2, 1, 3, 4, 6, 5, 7], @@ -265,6 +266,7 @@ def _run_public_api( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=4), async_op=async_op, ) dispatched = dispatcher.dispatch( diff --git a/tests/module/dispatcher/test_fsdp_vmm_landing.py b/tests/module/dispatcher/test_fsdp_vmm_landing.py new file mode 100644 index 0000000000..dc53746493 --- /dev/null +++ b/tests/module/dispatcher/test_fsdp_vmm_landing.py @@ -0,0 +1,81 @@ +import unittest + +import torch +from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard +from torch.distributed.tensor import DTensor + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.module.dispatcher.fsdp_vmm_landing import ( + accumulate_fsdp_unsharded_expert_gradients, + fsdp_current_unsharded_expert_parameters, + install_fsdp_vmm_landing, + uninstall_fsdp_vmm_landing, +) + + +class _ExpertLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fused_w1w3 = nn.Linear(8, 16, bias=False) + self.fused_w2 = nn.Linear(8, 8, bias=False) + + +@unittest.skipUnless(torch.cuda.device_count() >= 2, "requires 2 CUDA devices") +class TestFSDPVMMDirectLanding(DeterministicDDPTestCase): + def test_installation_is_atomic_and_preserves_the_unsharded_contract(self) -> None: + self.create_pg("cuda") + root = nn.ModuleList([_ExpertLayer(), _ExpertLayer()]).cuda() + mesh = init_device_mesh("cuda", (2,)) + policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16) + for layer in root: + fully_shard(layer.fused_w1w3, mesh=mesh, mp_policy=policy, reshard_after_forward=False) + fully_shard(layer.fused_w2, mesh=mesh, mp_policy=policy, reshard_after_forward=False) + + targets = [] + for layer_idx, layer in enumerate(root): + projections = (layer.fused_w1w3, layer.fused_w2) + landings = tuple( + torch.empty(projection.weight.shape, dtype=torch.bfloat16, device="cuda") for projection in projections + ) + targets.append((f"layers.{layer_idx}.experts", projections, landings)) + + # A failure in the last target must leave the earlier, valid targets + # untouched. A corrected retry through the public API proves that no + # binding attributes or method overrides leaked from validation. + invalid_targets = [*targets] + last_fqn, last_projections, last_landings = invalid_targets[-1] + invalid_targets[-1] = ( + last_fqn, + last_projections, + (last_landings[0], torch.empty(last_landings[1].shape, dtype=torch.float32, device="cuda")), + ) + with self.assertRaisesRegex(RuntimeError, "metadata mismatch"): + install_fsdp_vmm_landing(fsdp_root=root, targets=invalid_targets) + + fsdp_params = install_fsdp_vmm_landing(fsdp_root=root, targets=targets) + assert len(fsdp_params) == 4 + for layer in root: + layer.fused_w1w3.unshard() + layer.fused_w2.unshard() + parameters = fsdp_current_unsharded_expert_parameters((layer.fused_w1w3, layer.fused_w2)) + gradients = tuple( + torch.ones_like(parameter.to_local() if isinstance(parameter, DTensor) else parameter) + for parameter in parameters + ) + accumulate_fsdp_unsharded_expert_gradients(parameters, gradients) + for parameter in parameters: + gradient = parameter.grad + assert gradient is not None + gradient = gradient.to_local() if isinstance(gradient, DTensor) else gradient + assert gradient.dtype is torch.bfloat16 + torch.testing.assert_close(gradient, torch.ones_like(gradient), rtol=0, atol=0) + + uninstall_fsdp_vmm_landing(fsdp_params) + with self.assertRaisesRegex(RuntimeError, "is not installed"): + fsdp_current_unsharded_expert_parameters((root[0].fused_w1w3, root[0].fused_w2)) + + @property + def world_size(self) -> int: + return 2 diff --git a/tests/module/dispatcher/test_moonep_contract.py b/tests/module/dispatcher/test_moonep_contract.py new file mode 100644 index 0000000000..3fd75db1a0 --- /dev/null +++ b/tests/module/dispatcher/test_moonep_contract.py @@ -0,0 +1,303 @@ +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from xtuner.v1.data_proto import SequenceContext +from xtuner.v1.float8.config import Float8Config, ScalingGranularity +from xtuner.v1.model.moe.moe import MoEConfig +from xtuner.v1.module.attention import MHAConfig +from xtuner.v1.module.dispatcher import NaiveDispatcher +from xtuner.v1.module.router import GreedyRouter, GreedyRouterConfig + + +def _moe_config(**overrides) -> MoEConfig: + values = dict( + vocab_size=128, + max_position_embeddings=32, + pad_token_id=0, + eos_token_id=1, + num_hidden_layers=2, + hidden_size=128, + intermediate_size=128, + rms_norm_eps=1e-6, + hidden_act="silu", + attention=MHAConfig( + num_attention_heads=2, + num_key_value_heads=2, + head_dim=64, + ), + n_routed_experts=8, + n_shared_experts=0, + num_experts_per_tok=2, + moe_intermediate_size=128, + router=GreedyRouterConfig( + scoring_func="softmax", + norm_topk_prob=True, + router_scaling_factor=1.0, + ), + compile_cfg=False, + ) + values.update(overrides) + return MoEConfig(**values) + + +def test_moonep_is_a_standard_model_config_choice() -> None: + config = _moe_config(dispatcher="moonep", moonep_staging_reference=True) + + assert config.dispatcher == "moonep" + assert config.moonep_staging_reference is True + assert config.moonep_num_sms == 64 + assert config.intra_layer_micro_batch == 1 + + +def test_moonep_rejects_fp8_at_model_construction() -> None: + config = _moe_config( + dispatcher="moonep", + float8_cfg=Float8Config(scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE), + ) + + with pytest.raises(ValueError, match="requires BF16 expert compute; FP8 is not supported"): + config.build() + + +def test_non_moonep_build_does_not_import_optional_backend() -> None: + # A fresh interpreter with an explicit missing-module sentinel models an + # XTuner installation that does not have MoonEP installed. + source = """ +import sys +sys.modules[\"moonep\"] = None +from xtuner.v1.module.dispatcher import NaiveDispatcher, build_dispatcher +dispatcher = build_dispatcher(None, n_routed_experts=4) +assert isinstance(dispatcher, NaiveDispatcher) +""" + subprocess.run([sys.executable, "-c", source], check=True, capture_output=True, text=True) + + +def test_selecting_moonep_reports_the_missing_optional_backend() -> None: + source = """ +import sys +sys.modules["moonep"] = None +from types import SimpleNamespace +from xtuner.v1.module.dispatcher.moonep import MoonEPModelRuntime +try: + MoonEPModelRuntime( + ep_group=SimpleNamespace(size=lambda: 4), + hidden_size=128, + intermediate_size=128, + num_experts=8, + top_k=2, + intra_layer_micro_batch=1, + staging_reference=False, + ) +except RuntimeError as exc: + assert "requires the MoonEP-mod integration package" in str(exc) +else: + raise AssertionError("selecting MoonEP unexpectedly succeeded") +""" + subprocess.run([sys.executable, "-c", source], check=True, capture_output=True, text=True) + + +def test_missing_grouped_gemm_does_not_disable_triton_backend() -> None: + source = """ +import sys +sys.modules["grouped_gemm"] = None +sys.modules["grouped_gemm_backend"] = None +from xtuner.v1.ops.moe.cuda import cutlass_group_gemm, triton_group_gemm +assert cutlass_group_gemm is None +assert callable(triton_group_gemm) +""" + subprocess.run([sys.executable, "-c", source], check=True, capture_output=True, text=True) + + +def test_fully_shard_private_api_is_isolated_to_the_landing_module() -> None: + dispatcher_dir = Path(__file__).parents[3] / "xtuner" / "v1" / "module" / "dispatcher" + users = [ + path.name for path in dispatcher_dir.glob("*.py") if "torch.distributed.fsdp._fully_shard" in path.read_text() + ] + + assert users == ["fsdp_vmm_landing.py"] + + +def test_existing_dispatcher_keeps_public_preprocess_behavior() -> None: + dispatcher = NaiveDispatcher(n_routed_experts=4) + hidden_states = torch.randn(3, 8) + topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]]) + topk_weights = torch.full((3, 2), 0.5) + + result = dispatcher.dispatch_preprocess( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=4), + ) + + assert result["hidden_states"] is hidden_states + assert result["topk_ids"] is topk_ids + + +def test_router_owns_logical_tokens_per_expert() -> None: + router = GreedyRouter( + n_routed_experts=4, + num_experts_per_tok=2, + norm_topk_prob=True, + ) + result = router(torch.tensor([[4.0, 3.0, 2.0, 1.0], [1.0, 2.0, 3.0, 4.0]])) + + assert set(result) == {"logits", "router_weights", "topk_weights", "topk_ids", "tokens_per_expert"} + torch.testing.assert_close( + result["tokens_per_expert"], + torch.bincount(result["topk_ids"].flatten(), minlength=4), + ) + + +@pytest.mark.parametrize("width", [1, 2, 4]) +def test_moe_list_forward_rejects_a_different_width(width: int) -> None: + model = _moe_config(intra_layer_micro_batch=width).build() + input_ids = torch.tensor([[2, 3, 4]]) + contexts = [ + SequenceContext.from_input_ids((input_ids.clone(),), device="cpu") + for _ in range(width + 1) + ] + + with pytest.raises(ValueError, match=f"width {width + 1} does not match configured width {width}"): + model( + seq_ctx=contexts, + loss_ctx=[{} for _ in contexts], + ) + + +def test_runtime_meta_build_does_not_require_or_allocate_a_backend_workspace(monkeypatch) -> None: + from xtuner.v1.module.dispatcher import moonep as moonep_integration + from xtuner.v1.module.dispatcher.moonep import MoonEPModelRuntime + from xtuner.v1.module.grouped_linear import moe_group_linear + from xtuner.v1.ops.moe.cuda.group_gemm import triton_group_gemm + + # Workspace policy belongs to XTuner and allocation happens only after + # FSDP installation, so the optional backend needs no workspace interface. + backend = SimpleNamespace( + __file__="/tmp/MoonEP-mod/moonep/__init__.py", + XTUNER_INTEGRATION_API_VERSION=3, + Buffer=object, + ) + monkeypatch.setattr(moonep_integration, "_moonep_backend", backend) + monkeypatch.setattr(moonep_integration, "_MOONEP_IMPORT_ERROR", None) + monkeypatch.setattr(moe_group_linear, "group_gemm", triton_group_gemm) + ep_group = SimpleNamespace(size=lambda: 4) + + runtime = MoonEPModelRuntime( + ep_group=ep_group, + hidden_size=128, + intermediate_size=128, + num_experts=8, + top_k=2, + intra_layer_micro_batch=2, + staging_reference=False, + ) + + assert not hasattr(backend, "ExpertVMMWorkspace") + assert isinstance(runtime, MoonEPModelRuntime) + + +def test_runtime_allows_triton_grouped_gemm(monkeypatch) -> None: + from xtuner.v1.module.dispatcher import moonep as moonep_integration + from xtuner.v1.module.dispatcher.moonep import MoonEPModelRuntime + from xtuner.v1.module.grouped_linear import moe_group_linear + from xtuner.v1.ops.moe.cuda.group_gemm import triton_group_gemm + + monkeypatch.setattr( + moonep_integration, + "_moonep_backend", + SimpleNamespace( + __file__="/tmp/MoonEP-mod/moonep/__init__.py", XTUNER_INTEGRATION_API_VERSION=3, Buffer=object + ), + ) + monkeypatch.setattr(moonep_integration, "_MOONEP_IMPORT_ERROR", None) + monkeypatch.setattr(moe_group_linear, "group_gemm", triton_group_gemm) + + runtime = MoonEPModelRuntime( + ep_group=SimpleNamespace(size=lambda: 4), + hidden_size=128, + intermediate_size=128, + num_experts=8, + top_k=2, + intra_layer_micro_batch=1, + staging_reference=False, + ) + + assert isinstance(runtime, MoonEPModelRuntime) + + +@pytest.mark.parametrize( + ("environment_value", "effective_cutlass", "valid"), + [(None, True, False), ("1", False, False), ("1", True, True)], +) +def test_runtime_requires_grouped_gemm_cutlass_backend( + monkeypatch, environment_value: str | None, effective_cutlass: bool, valid: bool +) -> None: + pytest.importorskip("grouped_gemm") + from grouped_gemm import backend as grouped_gemm_backend + + from xtuner.v1.module.dispatcher import moonep as moonep_integration + from xtuner.v1.module.dispatcher.moonep import MoonEPModelRuntime + from xtuner.v1.module.grouped_linear import moe_group_linear + from xtuner.v1.ops.moe.cuda import cutlass_group_gemm + + assert cutlass_group_gemm is not None + monkeypatch.setattr( + moonep_integration, + "_moonep_backend", + SimpleNamespace( + __file__="/tmp/MoonEP-mod/moonep/__init__.py", + XTUNER_INTEGRATION_API_VERSION=3, + Buffer=object, + ), + ) + monkeypatch.setattr(moonep_integration, "_MOONEP_IMPORT_ERROR", None) + monkeypatch.setattr(moe_group_linear, "group_gemm", cutlass_group_gemm) + monkeypatch.setattr(grouped_gemm_backend, "use_cutlass", effective_cutlass) + if environment_value is None: + monkeypatch.delenv("GROUPED_GEMM_USE_CUTLASS", raising=False) + else: + monkeypatch.setenv("GROUPED_GEMM_USE_CUTLASS", environment_value) + + kwargs = dict( + ep_group=SimpleNamespace(size=lambda: 4), + hidden_size=128, + intermediate_size=128, + num_experts=8, + top_k=2, + intra_layer_micro_batch=1, + staging_reference=False, + ) + if valid: + assert isinstance(MoonEPModelRuntime(**kwargs), MoonEPModelRuntime) + else: + with pytest.raises(RuntimeError, match="grouped_gemm requires GROUPED_GEMM_USE_CUTLASS=1"): + MoonEPModelRuntime(**kwargs) + + +def test_runtime_reports_optional_backend_source_on_capability_mismatch(monkeypatch) -> None: + from xtuner.v1.module.dispatcher import moonep as moonep_integration + from xtuner.v1.module.dispatcher.moonep import MoonEPModelRuntime + + backend = SimpleNamespace( + __file__="/wrong/worktree/moonep/__init__.py", + XTUNER_INTEGRATION_API_VERSION=0, + ) + monkeypatch.setattr(moonep_integration, "_moonep_backend", backend) + monkeypatch.setattr(moonep_integration, "_MOONEP_IMPORT_ERROR", None) + + with pytest.raises(RuntimeError, match="/wrong/worktree/moonep/__init__.py"): + MoonEPModelRuntime( + ep_group=SimpleNamespace(size=lambda: 4), + hidden_size=128, + intermediate_size=128, + num_experts=8, + top_k=2, + intra_layer_micro_batch=1, + staging_reference=False, + ) diff --git a/tests/module/dispatcher/test_moonep_dispatcher.py b/tests/module/dispatcher/test_moonep_dispatcher.py new file mode 100644 index 0000000000..6e714d6326 --- /dev/null +++ b/tests/module/dispatcher/test_moonep_dispatcher.py @@ -0,0 +1,413 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from xtuner.v1.module.dispatcher import build_dispatcher +from xtuner.v1.module.dispatcher.moonep import ( + MoonEPDispatcher, + MoonEPModelRuntime, + _MoonEPGradReduceJoin, + _MoonEPGradReduceStart, +) + + +class _Event: + def wait(self) -> None: + return None + + +class _Stream: + def record_event(self): + return _Event() + + def wait_event(self, event) -> None: + del event + + def synchronize(self) -> None: + return None + + +class _Buffer: + def __init__( + self, + *, + S, + H, + K, + E, + num_ep_ranks, + group, + explicitly_destroy, + num_sms, + ): + self.S = S + self.K = K + self.E = E + self.B = E // num_ep_ranks + self.num_sms = num_sms + self.destroyed = False + self.prefetch_calls = 0 + + def dispatch( + self, + hidden_states, + route_weights_sk=None, + topk_experts_sk=None, + tokens_per_expert=None, + plan=None, + async_finish=False, + zero_copy=False, + ): + if plan is None: + plan = object() + cu_seqlens = torch.full((self.E + self.B,), hidden_states.shape[0], dtype=torch.int32) + else: + cu_seqlens = None + result = (hidden_states.clone(), route_weights_sk[:, 0].contiguous(), cu_seqlens, plan) + return (*result, _Event()) if async_finish else result + + def prefetch_weight(self, **kwargs): + assert kwargs["async_finish"] is False + self.prefetch_calls += 1 + return None + + def combine( + self, + *, + plan, + hidden_nvsh, + route_weights_nvs=None, + hidden_scales_nvs=None, + async_finish=False, + zero_copy=False, + ): + output = hidden_nvsh + if hidden_scales_nvs is not None: + output = output * hidden_scales_nvs[:, None].to(output.dtype) + result = (output, None, _Event() if async_finish else None) + return result + + def destroy(self) -> None: + self.destroyed = True + + +class _Workspace: + allocated = [] + + @classmethod + def allocate( + cls, + *, + projection_shapes, + num_experts, + ep_group, + gradient_slots, + **kwargs, + ): + instance = cls() + b = num_experts // ep_group.size() + instance._landings = tuple( + tuple(torch.zeros(b, *shape, dtype=torch.bfloat16) for shape in projection_shapes) for _ in range(2) + ) + instance._slots = tuple( + tuple(torch.zeros(2 * b, *shape, dtype=torch.bfloat16) for shape in projection_shapes) + for _ in range(gradient_slots) + ) + instance.destroyed = False + cls.allocated.append(instance) + return instance + + def landing(self, generation): + return self._landings[generation] + + def prefetch_weights(self, *, buffer, plan, generation, grad_slot): + landings = self.landing(generation) + local_weights = tuple(torch.cat((weight, torch.zeros_like(weight))) for weight in landings) + buffer.prefetch_weight(plan=plan, projections=landings, async_finish=False) + return local_weights, self._slots[grad_slot] + + def local_token_counts(self, cu_seqlens): + return torch.tensor([cu_seqlens[-1], 0, 0, 0], dtype=torch.int32) + + def destroy(self) -> None: + self.destroyed = True + + +class _Experts(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fused_w1w3 = nn.Linear(128, 2 * 2 * 128, bias=False, dtype=torch.bfloat16) + self.fused_w2 = nn.Linear(128, 2 * 128, bias=False, dtype=torch.bfloat16) + + +@pytest.fixture +def backend(monkeypatch): + from xtuner.v1.module.dispatcher import moonep as moonep_integration + from xtuner.v1.module.grouped_linear import moe_group_linear + from xtuner.v1.ops.moe.cuda.group_gemm import triton_group_gemm + + _Workspace.allocated.clear() + module = SimpleNamespace( + __file__="/tmp/MoonEP-mod/moonep/__init__.py", + XTUNER_INTEGRATION_API_VERSION=3, + Buffer=_Buffer, + ) + monkeypatch.setattr(moonep_integration, "_moonep_backend", module) + monkeypatch.setattr(moonep_integration, "_MOONEP_IMPORT_ERROR", None) + stream = _Stream() + monkeypatch.setattr(moonep_integration.torch.cuda, "Stream", lambda **kwargs: stream) + monkeypatch.setattr(moonep_integration.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(moonep_integration.torch.cuda, "current_stream", lambda: stream) + monkeypatch.setattr( + MoonEPModelRuntime, + "_enqueue", + lambda self, operation, inputs=(): (operation(), _Event()), + ) + monkeypatch.setattr( + "xtuner.v1.module.dispatcher.moonep._ExpertVMMWorkspace", + _Workspace, + ) + monkeypatch.setattr(moe_group_linear, "group_gemm", triton_group_gemm) + return module + + +def test_staging_dispatcher_runs_the_public_forward_path(backend) -> None: + ep_group = SimpleNamespace(size=lambda: 2) + runtime = MoonEPModelRuntime( + ep_group=ep_group, + hidden_size=128, + intermediate_size=128, + num_experts=4, + top_k=2, + intra_layer_micro_batch=1, + staging_reference=True, + ) + experts = _Experts() + dispatcher = build_dispatcher( + dispatcher="moonep", + n_routed_experts=4, + ep_group=ep_group, + moonep_runtime=runtime, + layer_fqn="layers.0.experts", + projections=(experts.fused_w1w3, experts.fused_w2), + ) + runtime.validate_before_fsdp( + SimpleNamespace( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + requires_grad=True, + cpu_offload=False, + reshard_after_forward=True, + ) + ) + runtime.install_after_fsdp(fsdp_root=experts) + + hidden_states = torch.randn(3, 128, dtype=torch.bfloat16, requires_grad=True) + topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int64) + source_counts = torch.tensor([1, 2, 2, 1], dtype=torch.int64) + route_weights = torch.full((3, 2), 0.5, dtype=torch.float32) + + with torch.no_grad(): + layer_input, layer_state = dispatcher.prepare_layer_input(hidden_states) + assert layer_input.grad_fn is None + pre = dispatcher.dispatch_preprocess( + hidden_states=layer_input, + topk_ids=topk_ids, + topk_weights=route_weights, + tokens_per_expert=source_counts, + layer_state=layer_state, + ) + dispatched = dispatcher.dispatch(pre_dispatched=pre, topk_weights=route_weights) + assert runtime._buffer.prefetch_calls == 1 + post = dispatcher.dispatch_postprocess(pre_dispatched=pre, dispatched=dispatched) + pre_combined = dispatcher.combine_preprocess( + hidden_states=post["hidden_states"], + pre_dispatched=pre, + dispatched=dispatched, + post_dispatched=post, + ) + combined = dispatcher.combine( + pre_dispatched=pre, + dispatched=dispatched, + post_dispatched=post, + pre_combined=pre_combined, + ) + result = dispatcher.combine_postprocess( + pre_dispatched=pre, + dispatched=dispatched, + post_dispatched=post, + pre_combined=pre_combined, + combined=combined, + ) + + assert isinstance(dispatcher, MoonEPDispatcher) + assert pre["topk_ids"].dtype == torch.int32 + assert pre["tokens_per_expert"].dtype == torch.int32 + assert torch.equal(pre["tokens_per_expert"], source_counts.to(torch.int32)) + assert post["tokens_per_expert"].shape == (4,) + assert post["expert_weight_layout"].trainable_weights is not None + assert all(isinstance(weight, nn.Parameter) for weight in post["expert_weight_layout"].trainable_weights) + assert post["expert_weight_layout"].trainable_weights[0].shape == (4, 256, 128) + assert post["expert_weight_layout"].trainable_wgrad_outs is not None + assert torch.equal(result["hidden_states"], hidden_states * 0.5) + assert not result["hidden_states"].requires_grad + assert runtime._buffer.num_sms == 64 + + invocation_ref = weakref.ref(layer_state) + del layer_state, pre, dispatched, post, pre_combined, combined + assert invocation_ref() is None + + with pytest.raises(RuntimeError, match="requires layer_state from prepare_layer_input"): + dispatcher.dispatch_preprocess( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=route_weights, + tokens_per_expert=source_counts, + ) + + with pytest.raises(RuntimeError, match="fixed S changed"): + invalid_hidden, invalid_state = dispatcher.prepare_layer_input(torch.randn(4, 128, dtype=torch.bfloat16)) + dispatcher.dispatch_preprocess( + hidden_states=invalid_hidden, + topk_ids=torch.zeros(4, 2, dtype=torch.int64), + topk_weights=torch.full((4, 2), 0.5), + tokens_per_expert=torch.tensor([8, 0, 0, 0]), + layer_state=invalid_state, + ) + + +def test_direct_install_failure_is_explicit_and_never_falls_back_to_staging(backend) -> None: + ep_group = SimpleNamespace(size=lambda: 2) + runtime = MoonEPModelRuntime( + ep_group=ep_group, + hidden_size=128, + intermediate_size=128, + num_experts=4, + top_k=2, + intra_layer_micro_batch=1, + staging_reference=False, + ) + experts = _Experts() + build_dispatcher( + dispatcher="moonep", + n_routed_experts=4, + ep_group=ep_group, + moonep_runtime=runtime, + layer_fqn="layers.0.experts", + projections=(experts.fused_w1w3, experts.fused_w2), + ) + + runtime.validate_before_fsdp( + SimpleNamespace( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + requires_grad=True, + cpu_offload=False, + reshard_after_forward=True, + ) + ) + with pytest.raises(RuntimeError, match="could not find FSDPParam"): + runtime.install_after_fsdp(fsdp_root=experts) + + assert _Workspace.allocated[-1].destroyed + + +def test_gradient_reduce_start_uses_the_workspace_targets_once() -> None: + from xtuner.v1.module.dispatcher import moonep as moonep_integration + + class _CompletionEvent: + def __init__(self) -> None: + self.waits = 0 + + def wait(self) -> None: + self.waits += 1 + + local_grads = tuple(torch.arange(24, dtype=torch.bfloat16).view(4, 2, 3) + projection for projection in range(2)) + fallback = tuple(gradient.clone() for gradient in local_grads) + calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + event = _CompletionEvent() + + class _Workspace: + def complete_gradients(self, *, local_grads, **kwargs): + del kwargs + calls.append(local_grads) + return local_grads[0][:2], local_grads[1][:2] + + runtime = SimpleNamespace(_workspace=_Workspace(), _buffer=object()) + + def enqueue(operation, inputs=()): + del inputs + return operation(), event + + runtime._enqueue = enqueue + invocation = moonep_integration._MoonEPLayerInvocation( + runtime=runtime, + layer_fqn="layers.0.experts", + projections=(nn.Linear(3, 3), nn.Linear(3, 3)), + generation=0, + grad_slot=0, + ) + invocation._plan = object() + home_parameters = ( + nn.Parameter(torch.zeros_like(local_grads[0][:2])), + nn.Parameter(torch.zeros_like(local_grads[1][:2])), + ) + invocation._home_parameters = home_parameters + invocation._fallback_gradient_targets = fallback + + invocation._start_gradient_completion() + + assert len(calls) == 1 + assert all(actual is expected for actual, expected in zip(calls[0], fallback, strict=True)) + assert event.waits == 0 + + invocation._finish_gradient_completion() + assert event.waits == 1 + for parameter, expected in zip(home_parameters, fallback, strict=True): + torch.testing.assert_close(parameter.grad, expected[:2]) + assert invocation._fallback_gradient_targets is None + + +def test_gradient_reduce_start_and_join_preserve_device_order() -> None: + events: list[str] = [] + targets = (torch.zeros(4), torch.zeros(4)) + + class _Invocation: + def _start_gradient_completion(self) -> None: + assert all( + torch.equal(target, torch.full_like(target, projection + 1)) + for projection, target in enumerate(targets) + ) + events.append("start") + + def _finish_gradient_completion(self) -> None: + events.append("finish") + + class _WriteWGrad(torch.autograd.Function): + @staticmethod + def forward(ctx, value, target, projection): + ctx.target = target + ctx.projection = projection + return value + + @staticmethod + def backward(ctx, grad): + ctx.target.fill_(ctx.projection + 1) + events.append(f"projection-{ctx.projection}") + return grad, None, None + + invocation = _Invocation() + source = torch.ones(4, requires_grad=True) + + joined = _MoonEPGradReduceJoin.apply(source, invocation) + started = _MoonEPGradReduceStart.apply(joined, invocation) + projection_0 = _WriteWGrad.apply(started, targets[0], 0) + projection_1 = _WriteWGrad.apply(projection_0, targets[1], 1) + projection_1.sum().backward() + + assert joined.data_ptr() == source.data_ptr() + assert started.data_ptr() == source.data_ptr() + assert events == ["projection-1", "projection-0", "start", "finish"] diff --git a/tests/module/dispatcher/test_moonep_workspace.py b/tests/module/dispatcher/test_moonep_workspace.py new file mode 100644 index 0000000000..6a6217049e --- /dev/null +++ b/tests/module/dispatcher/test_moonep_workspace.py @@ -0,0 +1,116 @@ +import unittest + +import torch +import torch.distributed as dist + +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.module.dispatcher.moonep_workspace import _ExpertVMMWorkspace + + +@unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices") +class TestMoonEPOneSegmentWorkspace(DeterministicDDPTestCase): + def test_ep2_ep4_ep8_share_the_one_segment_contract(self) -> None: + """Exercise the real VMM/transport path for every supported EP size.""" + self.create_pg("cuda") + global_rank = dist.get_rank() + world_size = dist.get_world_size() + + for ep_size in (2, 4, 8): + rank_lists = [list(range(start, start + ep_size)) for start in range(0, world_size, ep_size)] + groups = [dist.new_group(ranks=ranks) for ranks in rank_lists] + ep_group = groups[global_rank // ep_size] + ep_rank = dist.get_rank(ep_group) + experts_per_rank = 2 + num_experts = experts_per_rank * ep_size + device = torch.device("cuda", global_rank) + + workspace = _ExpertVMMWorkspace.allocate( + projection_shapes=((512, 1024), (1024, 512)), + num_experts=num_experts, + ep_group=ep_group, + gradient_slots=2, + ) + from moonep import Buffer + + buffer = Buffer( + S=64, + H=128, + K=1, + E=num_experts, + num_ep_ranks=ep_size, + B=experts_per_rank, + num_sms=8, + token_padding=16, + group=ep_group, + explicitly_destroy=True, + ) + try: + for projection, landing in enumerate(workspace.landing(0)): + for local_expert in range(experts_per_rank): + expert = ep_rank * experts_per_rank + local_expert + landing[local_expert].fill_(100 * projection + expert + 1) + + # A globally hot expert forces duplicate weight placement. + topk_ids = torch.zeros((64, 1), dtype=torch.int32, device=device) + tokens_per_expert = torch.bincount(topk_ids.flatten(), minlength=num_experts).to(torch.int32) + hidden = torch.randn(64, 128, dtype=torch.bfloat16, device=device) + _, _, cu_seqlens, plan = buffer.dispatch( + hidden, + topk_experts_sk=topk_ids, + tokens_per_expert=tokens_per_expert, + ) + + local_weights, gradients_0 = workspace.prefetch_weights( + buffer=buffer, + plan=plan, + generation=0, + grad_slot=0, + ) + _, gradients_1 = workspace.prefetch_weights( + buffer=buffer, + plan=plan, + generation=0, + grad_slot=1, + ) + + # One grouped GEMM receives exactly one contiguous [B+B] + # segment. Its home prefix aliases the current FSDP landing. + assert workspace.local_token_counts(cu_seqlens).shape == (2 * experts_per_rank,) + for projection, weight in enumerate(local_weights): + assert weight.is_contiguous() + assert weight.shape[0] == 2 * experts_per_rank + assert torch.equal( + weight[:experts_per_rank], + workspace.landing(0)[projection], + ) + + # Gradient slots are independent. Duplicate BF16 partials are + # returned to their home chunk and cleared without repacking. + for gradient in gradients_0: + gradient.zero_() + gradient[experts_per_rank:].fill_(1) + for gradient in gradients_1: + gradient.fill_(7) + home_grads = workspace.complete_gradients( + buffer=buffer, + plan=plan, + local_grads=gradients_0, + grad_slot=0, + ) + assert all(torch.count_nonzero(gradient[experts_per_rank:]) == 0 for gradient in gradients_0) + assert all(torch.all(gradient == 7) for gradient in gradients_1) + + copied = torch.count_nonzero(plan.experts_to_copy >= 0) + dist.all_reduce(copied, group=ep_group) + assert copied > 0 + hot_gradient = home_grads[0][0].float().sum() if ep_rank == 0 else torch.zeros((), device=device) + dist.all_reduce(hot_gradient, group=ep_group) + assert hot_gradient > 0 + finally: + buffer.destroy() + workspace.destroy() + dist.barrier() + + @property + def world_size(self) -> int: + return 8 diff --git a/tests/module/dispatcher/test_noep.py b/tests/module/dispatcher/test_noep.py index 7790c96733..7655921522 100644 --- a/tests/module/dispatcher/test_noep.py +++ b/tests/module/dispatcher/test_noep.py @@ -49,6 +49,7 @@ def test_dispatch_and_combine(self, dtype, device): hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=4), ) dispatched = self.dispatcher.dispatch( pre_dispatched=pre_dispatched, diff --git a/tests/module/dispatcher/test_noep_expert_tp.py b/tests/module/dispatcher/test_noep_expert_tp.py index 415ca11965..cf07292e3b 100644 --- a/tests/module/dispatcher/test_noep_expert_tp.py +++ b/tests/module/dispatcher/test_noep_expert_tp.py @@ -38,6 +38,7 @@ def _run_dispatcher( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=4), async_op=async_op, ) dispatched = dispatcher.dispatch( @@ -167,6 +168,7 @@ def test_async_path_exposes_events_at_stage_boundaries(self) -> None: hidden_states=hidden, topk_ids=local_topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(local_topk_ids.flatten(), minlength=4), async_op=True, ) _assert_cuda_event(pre_dispatched["forward_finished_event"]) diff --git a/tests/module/dispatcher/test_torch_all2all.py b/tests/module/dispatcher/test_torch_all2all.py index 802542c450..3f553f552c 100644 --- a/tests/module/dispatcher/test_torch_all2all.py +++ b/tests/module/dispatcher/test_torch_all2all.py @@ -1,12 +1,12 @@ +import os import unittest + +import parametrize import torch from torch.testing._internal.common_distributed import DistributedTestBase -from xtuner.v1.module.dispatcher.base import NaiveDispatcher, DispacherInterface -from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher -import parametrize - -import os +from xtuner.v1.module.dispatcher.base import DispacherInterface, NaiveDispatcher +from xtuner.v1.module.dispatcher.torch_all2all import TorchAll2AllDispatcher EP_SIZE = 8 @@ -24,13 +24,10 @@ def test_dispatch_and_combine(self, dtype, device): num_experts = 16 noep_dispatcher = NaiveDispatcher( n_routed_experts=num_experts, - training_dtype="bf16", ) all2all_dispatcher = TorchAll2AllDispatcher( - n_routed_experts=num_experts, - training_dtype="bf16", - process_group=torch.distributed.group.WORLD + n_routed_experts=num_experts, process_group=torch.distributed.group.WORLD ) seq_len = 32 @@ -41,31 +38,28 @@ def test_dispatch_and_combine(self, dtype, device): topk_weights = torch.ones(seq_len, topk_experts).to(device).to(torch.float32) noep_results = self._dispatcher_call( - dispatcher=noep_dispatcher, - hidden_states=hidden_states, - topk_ids=topk_idx, - topk_weights=topk_weights + dispatcher=noep_dispatcher, hidden_states=hidden_states, topk_ids=topk_idx, topk_weights=topk_weights ) all2all_results = self._dispatcher_call( - dispatcher=all2all_dispatcher, - hidden_states=hidden_states, - topk_ids=topk_idx, - topk_weights=topk_weights + dispatcher=all2all_dispatcher, hidden_states=hidden_states, topk_ids=topk_idx, topk_weights=topk_weights ) - self.assertTrue(torch.allclose(noep_results["hidden_states"], all2all_results["hidden_states"], atol=1e-6, rtol=1e-4)) + self.assertTrue( + torch.allclose(noep_results["hidden_states"], all2all_results["hidden_states"], atol=1e-6, rtol=1e-4) + ) def _dispatcher_call( - self, - dispatcher: DispacherInterface, - hidden_states: torch.Tensor, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor + self, + dispatcher: DispacherInterface, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, ): pre_dispatched = dispatcher.dispatch_preprocess( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=16), ) dispatched = dispatcher.dispatch( pre_dispatched=pre_dispatched, diff --git a/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py b/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py index db5528635f..495485bd24 100644 --- a/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py +++ b/tests/module/dispatcher/test_torch_all2all_shared_expert_tp.py @@ -56,6 +56,7 @@ def _run_dispatcher( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, + tokens_per_expert=torch.bincount(topk_ids.flatten(), minlength=4), async_op=async_op, ) dispatched = dispatcher.dispatch( diff --git a/tests/module/test_grouped_linear.py b/tests/module/test_grouped_linear.py index 1da49fecd0..668ae29995 100644 --- a/tests/module/test_grouped_linear.py +++ b/tests/module/test_grouped_linear.py @@ -5,6 +5,7 @@ """ import pytest +import torch from xtuner.v1.float8.config import Float8Config, ScalingGranularity from xtuner.v1.float8.float8_gmm_tile_wise import ADAPTIVEGEMM_INSTALLED, TileWiseFloat8GroupedLinear @@ -42,3 +43,50 @@ def test_grouped_gemm_switch_selects_implementation( ) assert type(layer) is expected_type + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_grouped_linear_returns_natural_gradient_for_call_local_weight() -> None: + layer = GroupedLinear(in_features=128, out_features=128, num_routed_experts=2).cuda().bfloat16() + original_parameter = layer.weight + override = torch.randn(2, 128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + hidden_states = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + counts = torch.tensor([2, 2], device="cuda", dtype=torch.int32) + override_ref = override.detach().clone().requires_grad_() + hidden_states_ref = hidden_states.detach().clone().requires_grad_() + + output = layer(hidden_states, counts, trainable_weight=override) + expected = torch.cat( + ( + hidden_states_ref[:2] @ override_ref[0].T, + hidden_states_ref[2:] @ override_ref[1].T, + ) + ) + grad_output = torch.randn_like(output) + output.backward(grad_output) + expected.backward(grad_output) + + assert layer.weight is original_parameter + assert original_parameter.grad is None + torch.testing.assert_close(output, expected, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(hidden_states.grad, hidden_states_ref.grad, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(override.grad, override_ref.grad, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available() or not ADAPTIVEGEMM_INSTALLED, reason="requires CUDA AdaptiveGEMM") +def test_non_moonep_fp8_accepts_the_uniform_grouped_linear_interface() -> None: + layer = build_grouped_linear( + in_features=128, + out_features=128, + num_routed_experts=2, + float8_cfg=Float8Config(scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE), + ).cuda().bfloat16() + hidden_states = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) + counts = torch.tensor([2, 2], device="cuda", dtype=torch.int64) + + output = layer(hidden_states, counts, trainable_wgrad_out=None) + + assert output.shape == (4, 128) + assert torch.isfinite(output).all() + with pytest.raises(NotImplementedError, match="preallocated trainable WGrad"): + layer(hidden_states, counts, trainable_wgrad_out=torch.empty_like(layer.weight)) diff --git a/tests/module/test_router_counts.py b/tests/module/test_router_counts.py new file mode 100644 index 0000000000..e6fc034363 --- /dev/null +++ b/tests/module/test_router_counts.py @@ -0,0 +1,109 @@ +import pytest +import torch + +from xtuner.v1.module.router.greedy import GreedyGroupedRouter, GreedyRouter +from xtuner.v1.module.router.noaux_router import NoAuxGroupedRouter, NoAuxRouter + + +ROUTER_KINDS = ("greedy", "greedy_grouped", "noaux", "noaux_grouped") + + +def _build_router(kind: str, device: torch.device): + if kind == "greedy": + router = GreedyRouter( + n_routed_experts=16, + num_experts_per_tok=2, + norm_topk_prob=True, + ) + elif kind == "greedy_grouped": + router = GreedyGroupedRouter( + n_routed_experts=16, + num_experts_per_tok=8, + router_n_groups=4, + norm_topk_prob=True, + ) + elif kind == "noaux": + router = NoAuxRouter( + n_routed_experts=16, + num_experts_per_tok=2, + router_scaling_factor=1.0, + scoring_func="sigmoid", + n_group=1, + topk_group=1, + ) + elif kind == "noaux_grouped": + router = NoAuxGroupedRouter( + n_routed_experts=16, + num_experts_per_tok=8, + router_scaling_factor=1.0, + router_n_groups=4, + scoring_func="sigmoid", + n_group=1, + topk_group=1, + ) + else: + raise AssertionError(f"unknown router kind: {kind}") + + router = router.to(device) + if isinstance(router, NoAuxRouter): + router.e_score_correction_bias.zero_() + return router + + +@pytest.mark.parametrize("kind", ROUTER_KINDS) +def test_router_returns_fixed_device_counts(kind: str) -> None: + device = torch.device("cpu") + router = _build_router(kind, device) + logits = torch.arange(48, dtype=torch.float32, device=device).view(3, 16) + top_k = router.top_k + routed_experts = torch.arange(3 * top_k, device=device).view(3, top_k).remainder(16) + + result = router(logits, rollout_routed_experts=routed_experts) + expected = torch.bincount(routed_experts.flatten(), minlength=16) + + assert result["tokens_per_expert"].shape == (16,) + assert result["tokens_per_expert"].dtype == torch.int64 + assert result["tokens_per_expert"].device == logits.device + assert torch.equal(result["tokens_per_expert"], expected) + assert result["tokens_per_expert"].sum() == routed_experts.numel() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA profiler regression requires a GPU") +@pytest.mark.parametrize("kind", ROUTER_KINDS) +def test_router_counting_has_no_host_synchronization(kind: str) -> None: + device = torch.device("cuda") + router = _build_router(kind, device) + logits = torch.randn(512, 16, dtype=torch.float32, device=device) + routed_experts = torch.arange( + 512 * router.top_k, + dtype=torch.int64, + device=device, + ).view(512, router.top_k).remainder(16) + + # Warm allocations before profiling so the marker contains only the real + # router path, including its device-side count production. + router(logits, rollout_routed_experts=routed_experts) + torch.cuda.synchronize() + marker = f"{kind}_router_forward" + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as profile, torch.profiler.record_function(marker): + result = router(logits, rollout_routed_experts=routed_experts) + torch.cuda.synchronize() + + blocking_events = [] + for event in profile.events(): + parent = event.cpu_parent + while parent is not None and parent.name != marker: + parent = parent.cpu_parent + if parent is not None and ( + "synchronize" in event.name.lower() + or event.name in {"aten::item", "aten::_local_scalar_dense"} + ): + blocking_events.append(event.name) + + assert not blocking_events + assert result["tokens_per_expert"].sum() == routed_experts.numel() diff --git a/tests/ops/test_grouped_gemm_cutlass.py b/tests/ops/test_grouped_gemm_cutlass.py new file mode 100644 index 0000000000..0c4f3c3ede --- /dev/null +++ b/tests/ops/test_grouped_gemm_cutlass.py @@ -0,0 +1,68 @@ +import pytest +import torch + + +pytest.importorskip("grouped_gemm_backend") + +from grouped_gemm import backend + +from xtuner.v1.ops.moe.cuda.group_gemm_cutlass import cutlass_group_gemm + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("compile", [False, True]) +@pytest.mark.parametrize("use_cutlass", [False, True]) +def test_grouped_gemm_wrapper_supports_natural_gradients( + compile: bool, + use_cutlass: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(backend, "use_cutlass", use_cutlass) + grouped_gemm = torch.compile(cutlass_group_gemm, fullgraph=True) if compile else cutlass_group_gemm + + for sizes in ([2, 0, 5], [1, 3, 3]): + torch.manual_seed(0) + counts = torch.tensor(sizes, device="cuda", dtype=torch.int32) + hidden_states = torch.randn(7, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(3, 256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + hidden_states_ref = hidden_states.detach().clone().requires_grad_() + weight_ref = weight.detach().clone().requires_grad_() + + output = grouped_gemm(hidden_states, weight, counts) + expected_groups = [] + offset = 0 + for expert, size in enumerate(sizes): + expected_groups.append(hidden_states_ref[offset : offset + size] @ weight_ref[expert].T) + offset += size + expected = torch.cat(expected_groups) + grad_output = torch.randn_like(output) + output.backward(grad_output) + expected.backward(grad_output) + + torch.testing.assert_close(output, expected) + torch.testing.assert_close(hidden_states.grad, hidden_states_ref.grad) + torch.testing.assert_close(weight.grad, weight_ref.grad) + if sizes[1] == 0: + torch.testing.assert_close(weight.grad[1], torch.zeros_like(weight.grad[1]), rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("compile", [False, True]) +def test_cublas_path_accepts_preallocated_gradient( + compile: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(backend, "use_cutlass", False) + hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(1, 256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + counts = torch.tensor([2], device="cuda", dtype=torch.int32) + grad_weight_out = torch.empty_like(weight) + grouped_gemm = torch.compile(cutlass_group_gemm, fullgraph=True) if compile else cutlass_group_gemm + + output = grouped_gemm(hidden_states, weight, counts, grad_weight_out=grad_weight_out) + grad_output = torch.randn_like(output) + output.backward(grad_output) + + expected = grad_output.T @ hidden_states.detach() + torch.testing.assert_close(grad_weight_out[0], expected) + torch.testing.assert_close(weight.grad, grad_weight_out) diff --git a/tests/ops/test_grouped_gemm_out.py b/tests/ops/test_grouped_gemm_out.py new file mode 100644 index 0000000000..0adc628d72 --- /dev/null +++ b/tests/ops/test_grouped_gemm_out.py @@ -0,0 +1,158 @@ +import pytest +import torch +from torch import nn + +from xtuner.v1.ops.moe.cuda import cutlass_group_gemm +from xtuner.v1.ops.moe.cuda.group_gemm import triton_group_gemm +from xtuner.v1.ops.moe.cuda.route_weight import route_weight_rows_backward + + +@pytest.fixture(params=[triton_group_gemm, cutlass_group_gemm], ids=["triton", "cutlass"]) +def grouped_gemm(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch): + implementation = request.param + if implementation is None: + pytest.skip("requires grouped_gemm") + if implementation is cutlass_group_gemm: + from grouped_gemm import backend + + monkeypatch.setattr(backend, "use_cutlass", True) + return implementation + + +@pytest.mark.parametrize("compile", [False, True]) +def test_grouped_gemm_backward_returns_natural_bf16_weight_gradient(compile: bool) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + torch.manual_seed(17) + counts = torch.tensor([2, 0, 3, 1], device="cuda", dtype=torch.int32) + x = torch.randn(6, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(4, 256, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + grad_output = torch.randn(6, 256, device="cuda", dtype=torch.bfloat16) + + x_ref = x.detach().clone().requires_grad_() + weight_ref = weight.detach().clone().requires_grad_() + expected = torch.cat( + ( + x_ref[:2] @ weight_ref[0].T, + x_ref[2:5] @ weight_ref[2].T, + x_ref[5:] @ weight_ref[3].T, + ) + ) + expected.backward(grad_output) + + grouped_gemm = triton_group_gemm + if compile: + grouped_gemm = torch.compile(grouped_gemm, fullgraph=True) + actual = grouped_gemm(x, weight, counts) + actual.backward(grad_output) + + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(x.grad, x_ref.grad, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(weight.grad, weight_ref.grad, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(weight.grad[1], torch.zeros_like(weight.grad[1]), rtol=0, atol=0) + + +@pytest.mark.parametrize("compile", [False, True]) +def test_parameter_owns_preallocated_grouped_gemm_gradient_without_copy(grouped_gemm, compile: bool) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + counts = torch.tensor([2, 0, 3, 1], device="cuda", dtype=torch.int32) + hidden_states = torch.randn(6, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = nn.Parameter(torch.randn(4, 256, 128, device="cuda", dtype=torch.bfloat16)) + target_storage = torch.empty_like(weight) + target = target_storage.new_empty(0).set_( + target_storage.untyped_storage(), + target_storage.storage_offset(), + target_storage.shape, + target_storage.stride(), + ) + seen_grad_pointers: list[int] = [] + weight.register_post_accumulate_grad_hook(lambda parameter: seen_grad_pointers.append(parameter.grad.data_ptr())) + + if compile: + grouped_gemm = torch.compile(grouped_gemm, fullgraph=True) + output = grouped_gemm(hidden_states, weight, counts, grad_weight_out=target) + del target + output.sum().backward() + + assert weight.grad is not None + assert weight.grad.data_ptr() == target_storage.data_ptr() + assert seen_grad_pointers == [target_storage.data_ptr()] + torch.testing.assert_close(weight.grad[1], torch.zeros_like(weight.grad[1]), rtol=0, atol=0) + + +@pytest.mark.parametrize("compile", [False, True]) +def test_retained_gradient_target_uses_parameter_copy_path(grouped_gemm, compile: bool) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + counts = torch.tensor([2, 1], device="cuda", dtype=torch.int32) + hidden_states = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = nn.Parameter(torch.randn(2, 256, 128, device="cuda", dtype=torch.bfloat16)) + target = torch.empty_like(weight) + + if compile: + grouped_gemm = torch.compile(grouped_gemm, fullgraph=True) + output = grouped_gemm(hidden_states, weight, counts, grad_weight_out=target) + output.sum().backward() + + assert weight.grad is not None + assert weight.grad.data_ptr() != target.data_ptr() + torch.testing.assert_close(weight.grad, target, rtol=0, atol=0) + + +@pytest.mark.parametrize("compile", [False, True]) +def test_preallocated_grouped_gemm_gradient_covers_zero_token_batch(grouped_gemm, compile: bool) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + counts = torch.zeros(4, device="cuda", dtype=torch.int32) + hidden_states = torch.empty(0, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = nn.Parameter(torch.randn(4, 256, 128, device="cuda", dtype=torch.bfloat16)) + target_storage = torch.full_like(weight, 7) + target = target_storage.new_empty(0).set_( + target_storage.untyped_storage(), + target_storage.storage_offset(), + target_storage.shape, + target_storage.stride(), + ) + + if compile: + grouped_gemm = torch.compile(grouped_gemm, fullgraph=True) + output = grouped_gemm(hidden_states, weight, counts, grad_weight_out=target) + del target + output.sum().backward() + + assert output.shape == (0, 256) + assert weight.grad is not None + assert weight.grad.data_ptr() == target_storage.data_ptr() + torch.testing.assert_close(weight.grad, torch.zeros_like(weight.grad), rtol=0, atol=0) + + +def test_fused_route_weight_backward_returns_bf16_rows_and_fp32_weights() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + torch.manual_seed(23) + grad_weighted = torch.randn(7, 512, device="cuda", dtype=torch.bfloat16) + expert_output = torch.randn_like(grad_weighted) + route_weights = torch.randn(7, device="cuda", dtype=torch.float32) + + grad_expert, grad_route = route_weight_rows_backward( + grad_weighted, + expert_output, + route_weights, + ) + # grouped-gemm's BF16 unpermute backward first rounds the FP32 router + # weight to BF16, and its route-gradient dot rounds every product to BF16 + # before the FP32 reduction. MoonEP must preserve that public numerical + # contract when it fuses the same operation into combine backward. + expected_expert = grad_weighted * route_weights.bfloat16()[:, None] + expected_route = (grad_weighted * expert_output).float().sum(dim=-1) + + assert grad_expert.dtype is torch.bfloat16 + assert grad_route.dtype is torch.float32 + torch.testing.assert_close(grad_expert, expected_expert, rtol=0, atol=0) + torch.testing.assert_close(grad_route, expected_route, rtol=1e-5, atol=1e-4) diff --git a/tests/train/test_trainer.py b/tests/train/test_trainer.py index aa2b559807..7441aa95da 100644 --- a/tests/train/test_trainer.py +++ b/tests/train/test_trainer.py @@ -49,6 +49,7 @@ def __init__(self): self.grad_norm_calls = 0 self.optimizer_step_calls = 0 self.optimizer_device_calls = [] + self.close_calls = 0 self.model = model = nn.Linear(10, 10) self.optimizer = torch.optim.Adam(model.parameters(), lr=0.001) @@ -108,6 +109,9 @@ def async_save_dcp(self, weights_dir: Path) -> Future: def destroy_async_checkpoint_pg(self) -> None: pass + def close(self) -> None: + self.close_calls += 1 + def prepare(fn): def wrapper(self, *args, **kwargs): @@ -196,6 +200,7 @@ def test_save_hf_interval(self): expected_dirs = {"hf-9", "hf-10", "hf-latest"} actual_dirs = {d.name for d in hf_dirs} self.assertEqual(actual_dirs, expected_dirs) + self.assertEqual(trainer._engine.close_calls, 1) # Verify the files were actually created and contain expected content for hf_dir in hf_dirs: diff --git a/tests/utils/test_interleaved_shard.py b/tests/utils/test_interleaved_shard.py index c85f54ab56..fcbc1c61b5 100644 --- a/tests/utils/test_interleaved_shard.py +++ b/tests/utils/test_interleaved_shard.py @@ -18,6 +18,7 @@ from torch.distributed.tensor import DTensor, Shard, distribute_tensor from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.utils.dtensor import cal_total_norm from xtuner.v1.utils.interleaved_shard import ( InterleavedShard, RuntimeLayout, @@ -192,3 +193,44 @@ def test_reconstruct_and_load(self) -> None: @property def world_size(self) -> int: return 8 + + +class TestNestedShardGradNorm(DeterministicDDPTestCase): + def test_cal_total_norm_for_fsdp2_ep4(self) -> None: + """FSDP's prepended shard and EP's shard must both contribute.""" + self.create_pg("cuda") + mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("fsdp", "ep")) + + global_weight = torch.zeros( + GLOBAL_ROWS, + IN_FEATURES, + device="cuda", + dtype=torch.bfloat16, + ) + tensor = distribute_tensor(global_weight, mesh["ep"], (Shard(0),)) + model = _ToyGroupedLinear(tensor).cuda() + fully_shard( + model, + mesh=mesh["fsdp"], + mp_policy=MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ), + reshard_after_forward=True, + ) + + inputs = torch.ones(6, IN_FEATURES, device="cuda", dtype=torch.bfloat16) + model(inputs).sum().backward() + + assert isinstance(model.weight.grad, DTensor) + total_norm = cal_total_norm([model.weight.grad], foreach=True) + expected = torch.tensor( + 6 * (GLOBAL_ROWS * IN_FEATURES) ** 0.5, + device="cuda", + dtype=torch.float32, + ) + torch.testing.assert_close(total_norm, expected) + + @property + def world_size(self) -> int: + return 8 diff --git a/xtuner/_testing/moonep_acceptance.py b/xtuner/_testing/moonep_acceptance.py new file mode 100644 index 0000000000..72d65bcc5e --- /dev/null +++ b/xtuner/_testing/moonep_acceptance.py @@ -0,0 +1,249 @@ +"""Parse and judge matched MoonEP/DeepEP Qwen3.5 acceptance runs.""" + +from __future__ import annotations + +import argparse +import importlib +import json +import math +import os +import runpy +import statistics +import subprocess +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +_LOSS_PREFIX = "loss/" + + +@dataclass(frozen=True) +class AcceptanceRun: + backend: str + mtp: bool + pack_length: int + records: tuple[dict[str, Any], ...] + + @classmethod + def from_tracker( + cls, + tracker: str | Path, + *, + backend: str, + mtp: bool, + pack_length: int, + ) -> "AcceptanceRun": + tracker = Path(tracker) + records = tuple(json.loads(line) for line in tracker.read_text(encoding="utf-8").splitlines() if line) + steps = [record.get("step") for record in records] + if steps != list(range(1, 21)): + raise ValueError(f"{tracker} must contain exactly steps 1..20, got {steps}") + + required = {"runtime_info/text_tokens", "runtime_info/tgs", "loss/reduced_llm_loss", "grad_norm"} + if mtp: + required.add("loss/reduced_mtp_loss") + for record in records: + missing = required - record.keys() + if missing: + raise ValueError(f"step {record['step']} is missing metrics: {sorted(missing)}") + return cls(backend=backend, mtp=mtp, pack_length=pack_length, records=records) + + @classmethod + def from_work_dir(cls, work_dir: str | Path) -> "AcceptanceRun": + work_dir = Path(work_dir) + manifest = json.loads((work_dir / "acceptance_manifest.json").read_text(encoding="utf-8")) + trackers = list(work_dir.glob("**/exp_tracking/rank0/tracker.jsonl")) + if len(trackers) != 1: + raise ValueError(f"expected one rank0 tracker below {work_dir}, got {trackers}") + return cls.from_tracker( + trackers[0], + backend=manifest["backend"], + mtp=manifest["mtp"], + pack_length=manifest["pack_length"], + ) + + @property + def steps(self) -> list[int]: + return [int(record["step"]) for record in self.records] + + @property + def tokens(self) -> list[int]: + return [int(record["runtime_info/text_tokens"]) for record in self.records] + + @property + def throughput(self) -> list[float]: + return [float(record["runtime_info/tgs"]) for record in self.records] + + def curves(self) -> dict[str, list[float]]: + names = { + key.removeprefix(_LOSS_PREFIX) + for record in self.records + for key in record + if key.startswith(f"{_LOSS_PREFIX}reduced_") and key.endswith("loss") + } + curves = {name: [float(record[f"{_LOSS_PREFIX}{name}"]) for record in self.records] for name in sorted(names)} + curves["total_loss"] = [ + sum( + float(value) + for key, value in record.items() + if key.startswith(f"{_LOSS_PREFIX}reduced_") and key.endswith("loss") + ) + for record in self.records + ] + curves["grad_norm"] = [float(record["grad_norm"]) for record in self.records] + return curves + + +@dataclass(frozen=True) +class CurveComparison: + cosine_similarity: float + mean_relative_difference: float + finite: bool + passed: bool + + +@dataclass(frozen=True) +class PairComparison: + throughput_steps: list[int] + deepep_throughput: list[float] + moonep_throughput: list[float] + deepep_median: float + moonep_median: float + throughput_ratio: float + curves: dict[str, CurveComparison] + passed: bool + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _compare_curve( + reference: list[float], + actual: list[float], + *, + minimum_cosine: float, + maximum_relative_difference: float, +) -> CurveComparison: + finite = all(math.isfinite(value) for value in [*reference, *actual]) + dot = sum(expected * observed for expected, observed in zip(reference, actual, strict=True)) + reference_norm = math.sqrt(sum(value * value for value in reference)) + actual_norm = math.sqrt(sum(value * value for value in actual)) + if reference_norm == 0 or actual_norm == 0: + cosine = 1.0 if reference == actual else 0.0 + else: + cosine = dot / (reference_norm * actual_norm) + relative = statistics.fmean( + abs(observed - expected) / max(abs(expected), 1e-12) + for expected, observed in zip(reference, actual, strict=True) + ) + return CurveComparison( + cosine_similarity=cosine, + mean_relative_difference=relative, + finite=finite, + passed=finite and cosine >= minimum_cosine and relative < maximum_relative_difference, + ) + + +def compare_runs(deepep: AcceptanceRun, moonep: AcceptanceRun) -> PairComparison: + if deepep.backend != "deepep" or moonep.backend != "moonep": + raise ValueError(f"expected deepep/moonep pair, got {deepep.backend}/{moonep.backend}") + for field in ("mtp", "pack_length"): + if getattr(deepep, field) != getattr(moonep, field): + raise ValueError(f"workload mismatch for {field}: {getattr(deepep, field)} != {getattr(moonep, field)}") + if deepep.tokens != moonep.tokens: + raise ValueError("workload mismatch for per-step text tokens") + + throughput_slice = slice(5, 20) + deepep_throughput = deepep.throughput[throughput_slice] + moonep_throughput = moonep.throughput[throughput_slice] + deepep_median = statistics.median(deepep_throughput) + moonep_median = statistics.median(moonep_throughput) + throughput_ratio = moonep_median / deepep_median + + deepep_curves = deepep.curves() + moonep_curves = moonep.curves() + if deepep_curves.keys() != moonep_curves.keys(): + raise ValueError(f"metric mismatch: deepep={sorted(deepep_curves)}, moonep={sorted(moonep_curves)}") + curves = { + name: _compare_curve( + deepep_curves[name], + moonep_curves[name], + minimum_cosine=0.98 if name == "grad_norm" else 0.99, + maximum_relative_difference=0.05 if name == "grad_norm" else 0.03, + ) + for name in deepep_curves + } + return PairComparison( + throughput_steps=list(range(6, 21)), + deepep_throughput=deepep_throughput, + moonep_throughput=moonep_throughput, + deepep_median=deepep_median, + moonep_median=moonep_median, + throughput_ratio=throughput_ratio, + curves=curves, + passed=throughput_ratio >= 0.95 and all(curve.passed for curve in curves.values()), + ) + + +def _git_commit(directory: Path) -> str: + return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=directory, text=True).strip() + + +def capture_manifest(config_path: Path, output: Path) -> None: + trainer = runpy.run_path(str(config_path))["trainer"] + moonep = importlib.import_module("moonep") + torch = importlib.import_module("torch") + moonep_source = Path(moonep.__file__).resolve() + repo_root = Path(__file__).resolve().parents[2] + payload = { + "backend": os.environ["MOONEP_ACCEPTANCE_BACKEND"], + "mtp": bool(int(os.environ["MOONEP_ACCEPTANCE_MTP"])), + "pack_length": int(os.environ["MOONEP_ACCEPTANCE_PACK_LENGTH"]), + "xtuner_commit": _git_commit(repo_root), + "moonep_commit": _git_commit(moonep_source.parents[1]), + "moonep_module": str(moonep_source), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "gpu_names": [torch.cuda.get_device_name(index) for index in range(torch.cuda.device_count())], + "configuration": json.loads(trainer.model_dump_json(serialize_as_any=True)), + "environment": { + name: os.environ.get(name) + for name in ( + "CUDA_VISIBLE_DEVICES", + "MODEL_COMPILE", + "XTUNER_DETERMINISTIC", + "XTUNER_ACTIVATION_OFFLOAD", + "XTUNER_USE_CUTLASS_GROUP_GEMM", + "XTUNER_COMPILE_NO_INPLACE_BUFFERS", + ) + }, + } + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + + +def main() -> int: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command", required=True) + capture = subparsers.add_parser("capture") + capture.add_argument("--config", type=Path, required=True) + capture.add_argument("--output", type=Path, required=True) + compare = subparsers.add_parser("compare") + compare.add_argument("--deepep", type=Path, required=True) + compare.add_argument("--moonep", type=Path, required=True) + compare.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + if args.command == "capture": + capture_manifest(args.config, args.output) + return 0 + + result = compare_runs(AcceptanceRun.from_work_dir(args.deepep), AcceptanceRun.from_work_dir(args.moonep)) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + return 0 if result.passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index fffec2a442..43ba08c79e 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -7,6 +7,7 @@ import threading import time import traceback +import warnings from concurrent.futures import Future, ThreadPoolExecutor, wait from pathlib import Path from typing import Any, Dict, List, cast @@ -150,16 +151,25 @@ def __init__( fsdp_cfg: FSDPConfig, intra_layer_micro_batch: int = 1, ) -> None: + if intra_layer_micro_batch < 1: + raise ValueError("intra_layer_micro_batch must be positive") + self.intra_layer_micro_batch = intra_layer_micro_batch + execution_cfg = getattr(model_cfg, "text_config", model_cfg) + if hasattr(execution_cfg, "intra_layer_micro_batch"): + # This is the number of forwards issued consecutively inside one + # layer call; it is independent of EP size and MTP depth. + setattr(execution_cfg, "intra_layer_micro_batch", intra_layer_micro_batch) self.model_cfg = model_cfg self.optim_cfg = optim_cfg self.fsdp_cfg = fsdp_cfg self.model = self.build_model() self.optimizer = self.build_optimizer(optim_cfg) - self.intra_layer_micro_batch = intra_layer_micro_batch self._count = 0 self.has_freeze_params = self.__has_freeze_params() self._async_checkpoint_pg: dist.ProcessGroup | None = None self._async_state_dict_cache: dict[str, Any] | None = None + self._pending_async_saves: list[Future[Any]] = [] + self._closed = False def __has_freeze_params(self) -> bool: has_freeze_params = False @@ -189,6 +199,7 @@ def data_replicate_size(self) -> int: @torch.no_grad() def forward_only(self, seq_ctx: SequenceContext, loss_ctx: LogProbContext): + self._ensure_open() output = self.model(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx}) # type: ignore[call-overload] return output @@ -202,6 +213,7 @@ def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo: Args: data_batches (List[Dict]): The input data batches for the training step. """ + self._ensure_open() self._maybe_precompute_float8_dynamic_scale_for_fsdp() intra_layer_micro_batch = self.intra_layer_micro_batch @@ -347,11 +359,14 @@ def async_save_hf( hf_dir: str, save_dtype: torch.dtype = torch.bfloat16, ) -> Future[Path]: + self._ensure_open() with profile_time_and_memory(f"[Async saving HF to {hf_dir} launch cost]"): - return self.model.async_save_hf( + future = self.model.async_save_hf( hf_dir=hf_dir, save_dtype=save_dtype, ) + self._pending_async_saves.append(future) + return future def _get_dcp_state_dict( self, @@ -404,6 +419,7 @@ def async_save_dcp( weights_dir: Path, save_optimizer: bool = True, ) -> Future: + self._ensure_open() async_checkpoint_pg = self._get_async_checkpoint_pg() # Match async HF export semantics: write the DCP payload into a @@ -474,6 +490,7 @@ def commit_async_save() -> None: commit_executor = ThreadPoolExecutor(max_workers=1) commit_future = commit_executor.submit(commit_async_save) commit_future.add_done_callback(lambda _: commit_executor.shutdown(wait=False)) + self._pending_async_saves.append(commit_future) return commit_future def _build_async_storage_writer(self, weights_dir: Path, *, save_optimizer: bool) -> XtunerCacheWriter: @@ -500,15 +517,38 @@ def destroy_async_checkpoint_pg(self) -> None: dist.destroy_process_group(self._async_checkpoint_pg) self._async_checkpoint_pg = None + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("TrainEngine is closed") + + def close(self) -> None: + """Collectively release model execution resources before PG teardown. + + All ranks in the training group must call this method at the same quiescent boundary. Exception paths + intentionally leave cleanup to process exit, because a rank-divergent destructor cannot run collectives safely. + """ + if self._closed: + return + # Both persistence paths snapshot before returning, but their writers + # still own CPU storage and auxiliary process groups until completion. + # Propagate background failures before releasing either resource. + for future in self._pending_async_saves: + future.result() + self._pending_async_saves.clear() + self.model.close_ep_runtime() + self.model.destroy_async_hf_resources() + self.destroy_async_checkpoint_pg() + self._closed = True + def __del__(self) -> None: - try: - self.model.destroy_async_hf_resources() - except Exception: - pass - try: - self.destroy_async_checkpoint_pg() - except Exception: - pass + if not getattr(self, "_closed", True): + # A rank-divergent finalizer must never enter process-group or + # CUDA/VMM teardown. Normal clients own the coordinated close. + warnings.warn( + "TrainEngine.close() was not called; distributed resources are left for process exit", + ResourceWarning, + stacklevel=2, + ) def load_dcp( self, diff --git a/xtuner/v1/float8/float8_gmm_tile_wise.py b/xtuner/v1/float8/float8_gmm_tile_wise.py index d9704d5180..8d80bd4331 100644 --- a/xtuner/v1/float8/float8_gmm_tile_wise.py +++ b/xtuner/v1/float8/float8_gmm_tile_wise.py @@ -346,7 +346,23 @@ def _check_shape(self, weight): f"but got {weight.shape}." ) - def forward(self, input: torch.Tensor, tokens_per_expert, decoding: bool = False) -> torch.Tensor: + def forward( + self, + input: torch.Tensor, + tokens_per_expert: torch.Tensor, + *, + trainable_weight: torch.Tensor | None = None, + trainable_wgrad_out: torch.Tensor | None = None, + external_weight: torch.Tensor | None = None, + external_wgrad_out: torch.Tensor | None = None, + ) -> torch.Tensor: + if trainable_weight is not None: + raise NotImplementedError("FP8 grouped linear does not support a trainable weight override yet") + if trainable_wgrad_out is not None: + raise NotImplementedError("FP8 grouped linear does not support a preallocated trainable WGrad") + if external_weight is not None or external_wgrad_out is not None: + raise NotImplementedError("FP8 two-segment grouped linear is not implemented") + weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight self._check_shape(weight) diff --git a/xtuner/v1/model/base.py b/xtuner/v1/model/base.py index 50e4ee59dc..1ac4c1afda 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -610,6 +610,10 @@ def cal_grad_norm(self, grads: list[DTensor], dtype=torch.float32): return cal_grad_norm(grads, dtype=dtype) + def close_ep_runtime(self) -> None: + """Release optional dynamic-EP resources at a coordinated boundary.""" + return + def to_hf_key_list(self, key: str) -> list[str]: raise NotImplementedError() diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index b0111afa73..29754790f1 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -148,7 +148,18 @@ class MoEConfig(TransformerConfig): moe_intermediate_size: Annotated[int, Parameter(group="moe")] ep_size: Annotated[int, Parameter(group="moe")] = 1 expert_tp_size: Annotated[int, Parameter(group="moe")] = 1 - dispatcher: Annotated[Literal["deepep", "all2all", "agrs"] | None, Parameter(group="moe")] = None + dispatcher: Annotated[Literal["deepep", "all2all", "agrs", "moonep"] | None, Parameter(group="moe")] = None + # Staging keeps the native FSDP unsharded tensor and copies its BF16 home + # experts into MoonEP VMM after AllGather. It is an explicit bring-up path; + # production direct landing is installed by the later FSDP adapter. + moonep_staging_reference: bool = False + # MoonEP reserves this many SMs for its communication kernels. The + # H200 acceptance workload is measurably faster at 64 than its upstream + # default of 32; keep it model-scoped so other deployments can tune it. + moonep_num_sms: int = 64 + # TrainEngine resolves this scalar before model build. MoonEP uses it to + # size per-invocation resources without depending on TrainerConfig. + intra_layer_micro_batch: int = 1 router: GreedyRouterConfig | NoAuxRouterConfig balancing_loss_cfg: BalancingLossConfig | None = BalancingLossConfig() z_loss_cfg: ZLossConfig | None = None @@ -194,6 +205,8 @@ class MoE(BaseModel): def __init__(self, config: MoEConfig): # Concrete MoE configs override build(), so validate dispatcher support # at the shared model-construction boundary. + if config.dispatcher == "moonep" and config.float8_cfg is not None and config.float8_cfg.enable_float8: + raise ValueError("MoonEP currently requires BF16 expert compute; FP8 is not supported") if config.dispatcher == "agrs": if config.expert_tp_size > 1: raise NotImplementedError("AGRS with ExpertTP is not supported") @@ -239,6 +252,25 @@ def __init__(self, config: MoEConfig): self.expert_tp_mesh = None self.ep_tp_mesh = None + self._moonep_runtime = None + if config.dispatcher == "moonep": + if self.ep_mesh is None: + raise ValueError("MoonEP requires expert parallelism") + if config.moe_bias: + raise ValueError("MoonEP does not support routed-expert linear bias") + from xtuner.v1.module.dispatcher.moonep import MoonEPModelRuntime + + self._moonep_runtime = MoonEPModelRuntime( + ep_group=self.ep_mesh.get_group(), + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + intra_layer_micro_batch=config.intra_layer_micro_batch, + staging_reference=config.moonep_staging_reference, + num_sms=config.moonep_num_sms, + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, type=config.rms_norm_type) self.lm_head = LMHead(config.hidden_size, config.vocab_size, bias=False) @@ -483,6 +515,13 @@ def forward( assert isinstance(loss_ctx, list) and len(loss_ctx) == len(seq_ctx), ( "seq_ctx_list and loss_ctx_list must be lists of the same length" ) + # The configured width is the exact number of forwards co-scheduled + # in one layer call, not EP world size or the number of MTP layers. + if len(seq_ctx) != self.config.intra_layer_micro_batch: + raise ValueError( + f"intra-layer micro-batch width {len(seq_ctx)} does not match " + f"configured width {self.config.intra_layer_micro_batch}" + ) if loss_ctx is None: raise NotImplementedError("loss_ctx must be provided for intra-layer bsz > 1") @@ -607,8 +646,9 @@ def _micro_batch_forward( d2h_stream=self.offload_stream, block_idx=layer_idx - self.config.first_k_dense_replace, group="text", - custom_check_fn=lambda x: x.data_ptr() - in [hidden_states.data_ptr() for hidden_states in hidden_states_list], + custom_check_fn=lambda x: ( + x.data_ptr() in [hidden_states.data_ptr() for hidden_states in hidden_states_list] + ), prefetch=True, reserve_pin_memory=True, ): @@ -1040,6 +1080,8 @@ def build_layers(self, config: MoEConfig) -> nn.ModuleDict: ep_mesh=self.ep_mesh, expert_tp_mesh=self.expert_tp_mesh, ep_tp_mesh=self.ep_tp_mesh, + moonep_runtime=self._moonep_runtime, + layer_fqn=f"layers.{layer_idx}.experts", ) if self.config.freeze_routers: layers[str(layer_idx)].gate.requires_grad_(False) @@ -1107,6 +1149,8 @@ def build_mtp_block(self, config: MoEConfig) -> MTPBlock: ep_mesh=self.ep_mesh, expert_tp_mesh=self.expert_tp_mesh, ep_tp_mesh=self.ep_tp_mesh, + moonep_runtime=self._moonep_runtime, + layer_fqn=f"mtp_block.layers.{i}.decoder_layer.experts", ) # Wrap decoder layer in MTPLayer @@ -1147,6 +1191,8 @@ def fully_shard( ) -> Self: if fsdp_config.hsdp_sharding_size is not None and self.config.expert_tp_size > 1: raise NotImplementedError("HSDP with ExpertTP is not supported") + if self._moonep_runtime is not None: + self._moonep_runtime.validate_before_fsdp(fsdp_config) self.fsdp_config = fsdp_config assert self.fsdp_config.ep_size == self.config.ep_size @@ -1319,6 +1365,8 @@ def fully_shard( self._init_load_spec() self._to_empty_meta() + if self._moonep_runtime is not None: + self._moonep_runtime.install_after_fsdp(fsdp_root=self) return self @property @@ -1329,6 +1377,11 @@ def default_compile_cfg(self) -> dict[str, TorchCompileOption]: else: return MOE_NON_EP_COMPILE_CFG + def close_ep_runtime(self) -> None: + """Release optional dynamic-EP resources before PG teardown.""" + if self._moonep_runtime is not None: + self._moonep_runtime.close() + @property def need_update_bias(self) -> bool: router_config = self.config.router diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py index 862d09c030..8188fff2a7 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -28,6 +28,7 @@ from xtuner.v1.module.dispatcher import ( CombineResult, DispatchResult, + ExpertWeightLayout, PostDispatchResult, PreCombineResult, PreDispatchResult, @@ -193,10 +194,34 @@ def __init__( ) self.moe_act = moe_act_fn_cfg.build() - def forward(self, x, tokens_per_expert, decoding): - gate_up_out = self.fused_w1w3(x, tokens_per_expert, decoding) + def forward( + self, + x: torch.Tensor, + tokens_per_expert: torch.Tensor, + *, + weight_layout: ExpertWeightLayout, + ) -> torch.Tensor: + trainable = weight_layout.trainable_weights or (None, None) + trainable_wgrad_outs = weight_layout.trainable_wgrad_outs or (None, None) + external = weight_layout.external_weights or (None, None) + external_wgrad_outs = weight_layout.external_wgrad_outs or (None, None) + gate_up_out = self.fused_w1w3( + x, + tokens_per_expert, + trainable_weight=trainable[0], + trainable_wgrad_out=trainable_wgrad_outs[0], + external_weight=external[0], + external_wgrad_out=external_wgrad_outs[0], + ) out = self.moe_act(gate_up_out, split_dim=-1) - res = self.fused_w2(out, tokens_per_expert, decoding) + res = self.fused_w2( + out, + tokens_per_expert, + trainable_weight=trainable[1], + trainable_wgrad_out=trainable_wgrad_outs[1], + external_weight=external[1], + external_wgrad_out=external_wgrad_outs[1], + ) return res @@ -229,10 +254,12 @@ def __init__( moe_act_fn_cfg: MoEActFnConfig, float8_cfg: Float8Config | None = None, layer_idx: int = 0, - dispatcher: Literal["deepep", "all2all", "agrs"] | None, + dispatcher: Literal["deepep", "all2all", "agrs", "moonep"] | None, ep_mesh: DeviceMesh | None = None, expert_tp_mesh: DeviceMesh | None = None, ep_tp_mesh: DeviceMesh | None = None, + moonep_runtime=None, + layer_fqn: str | None = None, ): super().__init__() self.ep_mesh = ep_mesh @@ -297,14 +324,18 @@ def __init__( process_group = ep_mesh.get_group() if ep_mesh is not None else None tp_group = expert_tp_mesh.get_group() if expert_tp_mesh is not None else None ep_tp_group = ep_tp_mesh._flatten().get_group() if ep_tp_mesh is not None else None + # EP membership is static for the layer. Keep the decision outside the + # compiled forward; Naive EP=1 execution retains its synchronous API. + self._async_combine = process_group is not None and process_group.size() > 1 self.dispatcher = build_dispatcher( dispatcher=dispatcher, n_routed_experts=n_routed_experts, ep_group=process_group, tp_group=tp_group, ep_tp_group=ep_tp_group, - training_dtype="fp8" if float8_cfg is not None else "bf16", - generate_dtype=generate_config.dtype if generate_config is not None else "bf16", + moonep_runtime=moonep_runtime, + layer_fqn=layer_fqn, + projections=(self.experts.fused_w1w3, self.experts.fused_w2), ) def forward( @@ -395,6 +426,9 @@ def _forward( seq_ctx: SequenceContext, position_embeddings: tuple[torch.Tensor, torch.Tensor], ) -> tuple[HiddenStates, RouterLogits, RouterWeights, RouterTopKIds]: + # MoonEP uses this identity seam to place Join before attention and + # carries its opaque invocation token into dispatch phase 1. + hidden_states, layer_state = self.dispatcher.prepare_layer_input(hidden_states) residual, hidden_states, router_results = self._pre_moe_forward( hidden_states=hidden_states, seq_ctx=seq_ctx, @@ -412,6 +446,8 @@ def _forward( hidden_states=hidden_states.view(-1, hidden_states.shape[-1]), topk_ids=router_results["topk_ids"], topk_weights=router_results["topk_weights"], + tokens_per_expert=router_results["tokens_per_expert"], + layer_state=layer_state, ) dispatched = self.dispatcher.dispatch( pre_dispatched=pre_dispatched, @@ -432,7 +468,7 @@ def _forward( experts_out = self.experts( post_dispatched["hidden_states"], post_dispatched["tokens_per_expert"], - decoding=False, + weight_layout=post_dispatched["expert_weight_layout"], ) # ProberList.before_combine( # self.layer_idx, @@ -445,6 +481,7 @@ def _forward( pre_dispatched=pre_dispatched, dispatched=dispatched, post_dispatched=post_dispatched, + async_op=self._async_combine, decoding=False, ) @@ -453,14 +490,24 @@ def _forward( dispatched=dispatched, post_dispatched=post_dispatched, pre_combined=pre_combined, + async_op=self._async_combine, decoding=False, ) + + # EP combine 已经在通信流上异步启动；共享专家在默认流计算，只在 + # routed + shared 相加前建立设备侧依赖，不引入 host sync。 + if self.n_shared_experts > 0: + shared_experts_out = self._shared_experts_forward(hidden_states=hidden_states) + else: + shared_experts_out = None + post_combined = self.dispatcher.combine_postprocess( pre_dispatched=pre_dispatched, dispatched=dispatched, post_dispatched=post_dispatched, pre_combined=pre_combined, combined=combined, + async_op=self._async_combine, ) combined_hidden_states = post_combined["hidden_states"] combined_hidden_states = combined_hidden_states.view(*origin_shape) @@ -470,11 +517,6 @@ def _forward( # ProberList.after_combine(self.layer_idx, combined_hidden_states) - if self.n_shared_experts > 0: - shared_experts_out = self._shared_experts_forward(hidden_states=hidden_states) - else: - shared_experts_out = None - hidden_states = self._post_moe_forward( combined_hidden_states=combined_hidden_states, residual=residual, @@ -515,6 +557,7 @@ def _micro_batch_forward( seq_ctx_list, position_embeddings_list, ): + hidden_states, layer_state = self.dispatcher.prepare_layer_input(hidden_states) residual, hidden_states, router_results = self._pre_moe_forward( hidden_states=hidden_states, seq_ctx=seq_ctx, @@ -527,6 +570,8 @@ def _micro_batch_forward( hidden_states=hidden_states, topk_ids=router_results["topk_ids"], topk_weights=router_results["topk_weights"], + tokens_per_expert=router_results["tokens_per_expert"], + layer_state=layer_state, async_op=True, ) pre_dispatched_list.append(pre_dispatched) @@ -557,7 +602,7 @@ def _micro_batch_forward( experts_out = self.experts( post_dispatched["hidden_states"], post_dispatched["tokens_per_expert"], - decoding=False, + weight_layout=post_dispatched["expert_weight_layout"], ) pre_combined = self.dispatcher.combine_preprocess( diff --git a/xtuner/v1/module/dispatcher/__init__.py b/xtuner/v1/module/dispatcher/__init__.py index d5d3c96860..c64255a9d6 100644 --- a/xtuner/v1/module/dispatcher/__init__.py +++ b/xtuner/v1/module/dispatcher/__init__.py @@ -5,6 +5,7 @@ XTUNER_DISPATCHER_DEBUG = os.getenv("XTUNER_DISPATCHER_DEBUG", "0") == "1" import torch.distributed as dist +from torch import nn from xtuner.v1.utils import get_logger, log_rank0 @@ -13,6 +14,7 @@ CombineResult, DispacherInterface, DispatchResult, + ExpertWeightLayout, NaiveDispatcher, PostCombineResult, PostDispatchResult, @@ -28,14 +30,28 @@ # TODO: (yehaochen) This interface declaration does not follow the Liskov Substitution Principle. # Maybe we should find a better way to handle the dispatchers. def build_dispatcher( - dispatcher: Literal["deepep", "all2all", "agrs"] | None, + dispatcher: Literal["deepep", "all2all", "agrs", "moonep"] | None, n_routed_experts: int, ep_group: dist.ProcessGroup | None = None, tp_group: dist.ProcessGroup | None = None, ep_tp_group: dist.ProcessGroup | None = None, - training_dtype: Literal["bf16", "fp8"] = "bf16", - generate_dtype: Literal["bf16", "fp8"] = "bf16", + *, + moonep_runtime=None, + layer_fqn: str | None = None, + projections: tuple[nn.Module, nn.Module] | None = None, ) -> DispacherInterface: + if dispatcher == "moonep": + if ep_group is None or ep_group.size() not in (2, 4, 8): + raise ValueError("MoonEP requires ep_size in {2, 4, 8}") + if n_routed_experts % ep_group.size(): + raise ValueError("MoonEP requires n_routed_experts divisible by ep_size") + if moonep_runtime is None or layer_fqn is None or projections is None: + raise ValueError("MoonEP runtime, layer_fqn, and expert projections are required") + return moonep_runtime.build_dispatcher( + layer_fqn=layer_fqn, + projections=projections, + ) # type: ignore[return-value] + if ep_group is None or ep_group.size() == 1: if dispatcher is not None: log_rank0.warning(f"{dispatcher} will not be used because the ep group is None.") @@ -43,8 +59,6 @@ def build_dispatcher( n_routed_experts=n_routed_experts, process_group=ep_group, tp_group=tp_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) # type: ignore[return-value] if dispatcher is None: @@ -72,8 +86,6 @@ def build_dispatcher( n_routed_experts=n_routed_experts, process_group=process_group, tp_size=tp_size, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) # type: ignore elif dispatcher == "all2all": assert ep_group is not None, "TorchAll2AllDispatcher requires a non-null ep_group." @@ -81,19 +93,17 @@ def build_dispatcher( n_routed_experts=n_routed_experts, process_group=ep_group, tp_group=tp_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) # type: ignore[return-value] elif dispatcher == "agrs": assert ep_group is not None, "MoEAGRSDispatcher requires a non-null process group." return MoEAGRSDispatcher( n_routed_experts=n_routed_experts, process_group=ep_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) # type: ignore[return-value] else: - raise ValueError(f"Unknown dispatcher name: {dispatcher}, name must be one of 'deepep' or 'all2all'.") + raise ValueError( + f"Unknown dispatcher name: {dispatcher}, name must be one of 'deepep', 'all2all', 'agrs', or 'moonep'." + ) __all__ = [ @@ -104,6 +114,7 @@ def build_dispatcher( "build_dispatcher", "PreDispatchResult", "DispatchResult", + "ExpertWeightLayout", "PostDispatchResult", "PreCombineResult", "CombineResult", diff --git a/xtuner/v1/module/dispatcher/agrs.py b/xtuner/v1/module/dispatcher/agrs.py index 255070a6b9..3a42194a23 100644 --- a/xtuner/v1/module/dispatcher/agrs.py +++ b/xtuner/v1/module/dispatcher/agrs.py @@ -1,4 +1,4 @@ -from typing import Literal, TypeAlias, cast +from typing import TypeAlias, cast import torch import torch.distributed as dist @@ -19,6 +19,7 @@ from .base import ( CombineResult, DispatchResult, + ExpertWeightLayout, GenericDispatcher, PostCombineResult, PostDispatchResult, @@ -235,14 +236,10 @@ def __init__( *, n_routed_experts: int, process_group: torch.distributed.ProcessGroup, - training_dtype: Literal["fp8", "bf16"] = "bf16", - generate_dtype: Literal["fp8", "bf16"] = "bf16", ): super().__init__( n_routed_experts=n_routed_experts, process_group=process_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) assert self._process_group is not None, ( "Process group must be provided for `DeepEPDispatcher`. " @@ -259,8 +256,11 @@ def dispatch_preprocess( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, # noqa: ARG002 — kept for interface compatibility; not used here + tokens_per_expert: torch.Tensor, + layer_state: object | None = None, async_op: bool = False, ) -> MoEAGRSPreDispatchResult: + del tokens_per_expert, layer_state if async_op: forward_finished_event = cast(torch.cuda.Event, torch.cuda.Event()) forward_finished_event.record() @@ -396,6 +396,7 @@ def dispatch_postprocess( hidden_states=permuted_hidden_states, row_ids_map=row_ids_map, tokens_per_expert=tokens_per_expert, + expert_weight_layout=ExpertWeightLayout(), ) @override diff --git a/xtuner/v1/module/dispatcher/base.py b/xtuner/v1/module/dispatcher/base.py index 81bc94d919..f0113c44cf 100644 --- a/xtuner/v1/module/dispatcher/base.py +++ b/xtuner/v1/module/dispatcher/base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import ( Generic, - Literal, + NamedTuple, TypeAlias, TypeVar, ) @@ -15,6 +15,16 @@ HiddenStates: TypeAlias = torch.Tensor +ProjectionPair: TypeAlias = tuple[torch.Tensor, torch.Tensor] + + +class ExpertWeightLayout(NamedTuple): + """Call-local expert weight ownership at the dispatcher/MLP seam.""" + + trainable_weights: ProjectionPair | None = None + trainable_wgrad_outs: ProjectionPair | None = None + external_weights: ProjectionPair | None = None + external_wgrad_outs: ProjectionPair | None = None def _get_backward_pre_hook(backward_previous_event: torch.cuda.Event): @@ -59,6 +69,7 @@ class PostDispatchResult(TypedDict): # TODO: hidden_states: torch.Tensor tokens_per_expert: torch.Tensor + expert_weight_layout: ExpertWeightLayout class PreCombineResult(TypedDict): @@ -102,13 +113,21 @@ def __init__( *, n_routed_experts: int, process_group: torch.distributed.ProcessGroup | None = None, - training_dtype: Literal["fp8", "bf16"] = "bf16", - generate_dtype: Literal["fp8", "bf16"] = "bf16", ): self._process_group = process_group self._n_routed_experts = n_routed_experts - self._training_dtype = training_dtype - self._generate_dtype = generate_dtype + + def prepare_layer_input( + self, + layer_input: torch.Tensor, + ) -> tuple[torch.Tensor, object | None]: + """Return the layer input and optional backend ordering state. + + Most dispatchers have no work to schedule before attention, so they + keep the identity behavior. A backend that needs an autograd ordering + seam may return an opaque token for ``dispatch_preprocess``. + """ + return layer_input, None @abstractmethod def dispatch( @@ -136,6 +155,10 @@ def dispatch_preprocess( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, + # Source-logical counts are owned by Router. Post-dispatch counts have + # a different meaning: local physical groups consumed by expert GMM. + tokens_per_expert: torch.Tensor, + layer_state: object | None = None, async_op: bool = False, ) -> PreDispatch: ... @@ -237,14 +260,10 @@ def __init__( n_routed_experts: int, process_group: torch.distributed.ProcessGroup | None = None, tp_group: torch.distributed.ProcessGroup | None = None, - training_dtype: Literal["fp8", "bf16"] = "bf16", - generate_dtype: Literal["fp8", "bf16"] = "bf16", ): super().__init__( n_routed_experts=n_routed_experts, process_group=process_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) if self._process_group is not None: assert self._process_group.size() == 1, "Naive dispatcher is only for ep=1." @@ -259,8 +278,11 @@ def dispatch_preprocess( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, + tokens_per_expert: torch.Tensor, + layer_state: object | None = None, async_op: bool = False, ) -> NaivePreDispatchResult: + del tokens_per_expert, layer_state if async_op: if self._expert_tp is None: raise NotImplementedError("Naive dispatcher async_op=True requires ExpertTP.") @@ -409,6 +431,7 @@ def dispatch_postprocess( hidden_states=hidden_states, row_ids_map=row_id_maps, tokens_per_expert=tokens_per_expert, + expert_weight_layout=ExpertWeightLayout(), ) @override diff --git a/xtuner/v1/module/dispatcher/deepep.py b/xtuner/v1/module/dispatcher/deepep.py index de264bf4ab..2cc39aa003 100644 --- a/xtuner/v1/module/dispatcher/deepep.py +++ b/xtuner/v1/module/dispatcher/deepep.py @@ -1,4 +1,4 @@ -from typing import Literal, TypeAlias, cast +from typing import TypeAlias, cast import torch import torch.distributed as dist @@ -20,6 +20,7 @@ from .base import ( CombineResult, DispatchResult, + ExpertWeightLayout, GenericDispatcher, PostCombineResult, PostDispatchResult, @@ -264,8 +265,6 @@ def __init__( n_routed_experts: int, process_group: torch.distributed.ProcessGroup, tp_size: int = 1, - training_dtype: Literal["fp8", "bf16"] = "bf16", - generate_dtype: Literal["fp8", "bf16"] = "bf16", ): """DeepEP-backed MoE dispatcher. @@ -294,8 +293,6 @@ def __init__( super().__init__( n_routed_experts=n_routed_experts, process_group=process_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) assert self._process_group is not None, ( "Process group must be provided for `DeepEPDispatcher`. " @@ -324,8 +321,11 @@ def dispatch_preprocess( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, + tokens_per_expert: torch.Tensor, + layer_state: object | None = None, async_op: bool = False, ) -> DeepEPPreDispatchResult: + del tokens_per_expert, layer_state if async_op: backward_previous_event = EventOverlap(None) if hidden_states.grad_fn is not None: @@ -501,6 +501,7 @@ def dispatch_postprocess( hidden_states=permuted_hidden_states, row_ids_map=row_ids_map, tokens_per_expert=tokens_per_expert, + expert_weight_layout=ExpertWeightLayout(), ) @override diff --git a/xtuner/v1/module/dispatcher/fsdp_vmm_landing.py b/xtuner/v1/module/dispatcher/fsdp_vmm_landing.py new file mode 100644 index 0000000000..87ccb3c747 --- /dev/null +++ b/xtuner/v1/module/dispatcher/fsdp_vmm_landing.py @@ -0,0 +1,239 @@ +"""Version-pinned FSDP2 landing adapter for MoonEP expert weights. + +This is the only XTuner module allowed to know about ``_fully_shard`` +internals. The rest of the MoonEP integration only sees installation, +current-view, and uninstallation functions from this module. +""" + +from __future__ import annotations + +import types +from collections.abc import Sequence +from typing import Any, cast + +import torch +from torch import nn +from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState +from torch.distributed.fsdp._fully_shard._fsdp_state import _get_module_fsdp_state +from torch.distributed.tensor import DTensor + + +_TARGET_TORCH_VERSION = "2.12.1+cu132" +_BINDING_ATTR = "_xtuner_moonep_landing" +_OWNER_ATTR = "_xtuner_moonep_fsdp_owner" +_PROJECTION_ATTR = "_xtuner_moonep_projection" +_FSDP_PARAM_ATTR = "_xtuner_moonep_fsdp_param" + + +def _init_direct_all_gather_outputs( + fsdp_param: FSDPParam, + all_gather_input_numels: list[int], + all_gather_input_dtypes: list[torch.dtype], + world_size: int, + device: torch.device, + force_recreate: bool = False, +) -> None: + """Point FSDP's final per-parameter unpack at the fixed VMM landing.""" + del force_recreate + landing = getattr(fsdp_param, _BINDING_ATTR) + if ( + len(all_gather_input_numels) != 1 + or len(all_gather_input_dtypes) != 1 + or all_gather_input_numels[0] * world_size != landing.numel() + or all_gather_input_dtypes[0] is not landing.dtype + or device != landing.device + ): + raise RuntimeError("MoonEP direct landing no longer matches FSDP AllGather metadata") + fsdp_param.all_gather_outputs = [landing.view(-1)] + + +def _keep_direct_all_gather_storage(fsdp_param: FSDPParam) -> None: + """FSDP must not resize/free runtime-owned, non-resizable VMM storage.""" + del fsdp_param + + +def _resolve_and_validate_targets( + fsdp_root: nn.Module, + targets: Sequence[tuple[str, tuple[nn.Module, nn.Module], tuple[torch.Tensor, torch.Tensor]]], +) -> list[tuple[FSDPParam, nn.Module, nn.Module, torch.Tensor]]: + """Resolve every target and validate the landing ABI before mutation.""" + by_identity: dict[tuple[int, str], tuple[FSDPParam, nn.Module]] = {} + for fsdp_owner in fsdp_root.modules(): + state = _get_module_fsdp_state(fsdp_owner) + if state is None: + continue + # The 2.12 runtime has the plural list; its bundled type stub still + # exposes only the deprecated singular compatibility property. + for param_group in cast(Any, state)._fsdp_param_groups: + for fsdp_param in param_group.fsdp_params: + key = (id(fsdp_param._module_info.module), fsdp_param._module_info.param_name) + if key in by_identity: + raise RuntimeError("MoonEP found duplicate FSDP parameter identity") + by_identity[key] = fsdp_param, fsdp_owner + + selected: list[tuple[FSDPParam, nn.Module, nn.Module, torch.Tensor]] = [] + for layer_fqn, projections, landings in targets: + for projection_name, projection, landing in zip( + ("fused_w1w3", "fused_w2"), projections, landings, strict=True + ): + if hasattr(projection, _FSDP_PARAM_ATTR): + raise RuntimeError(f"MoonEP direct landing is already installed for {layer_fqn}.{projection_name}") + match = by_identity.get((id(projection), "weight")) + if match is None: + raise RuntimeError(f"MoonEP could not find FSDPParam for {layer_fqn}.{projection_name}.weight") + fsdp_param, fsdp_owner = match + expected_dtype = fsdp_param.mp_policy.param_dtype or fsdp_param.sharded_param.dtype + shard_world_size = fsdp_param.mesh_info.shard_mesh_size + unpadded_numel = fsdp_param._orig_size.numel() + gathered_numel = fsdp_param.padded_sharded_param_size.numel() * shard_world_size + if fsdp_param.fsdp_placement.dim != 0 or not landing.is_contiguous(): + raise RuntimeError(f"MoonEP requires contiguous dim-0 FSDP layout for {layer_fqn}.{projection_name}") + if unpadded_numel != gathered_numel: + raise RuntimeError( + f"MoonEP direct landing does not support FSDP padding for {layer_fqn}.{projection_name}" + ) + if landing.numel() != unpadded_numel or landing.dtype is not expected_dtype: + raise RuntimeError(f"MoonEP VMM landing metadata mismatch for {layer_fqn}.{projection_name}") + if landing.device != fsdp_param.device or shard_world_size <= 1: + raise RuntimeError( + f"MoonEP direct landing requires multi-rank CUDA FSDP for {layer_fqn}.{projection_name}" + ) + if hasattr(fsdp_param._sharded_local_tensor, "fsdp_post_all_gather"): + raise RuntimeError(f"MoonEP direct landing does not support post-AllGather extensions: {layer_fqn}") + if fsdp_param.sharded_state is not ShardedState.SHARDED or fsdp_param.all_gather_outputs: + raise RuntimeError(f"MoonEP direct landing must be installed before first AllGather: {layer_fqn}") + if any( + name in fsdp_param.__dict__ + for name in ("init_all_gather_outputs", "alloc_all_gather_outputs", "free_unsharded_param") + ): + raise RuntimeError(f"MoonEP refuses an already customized FSDPParam: {layer_fqn}") + selected.append((fsdp_param, fsdp_owner, projection, landing)) + + if len({id(item[0]) for item in selected}) != len(selected): + raise RuntimeError("MoonEP routed expert targets must map to distinct FSDPParams") + return selected + + +def install_fsdp_vmm_landing( + *, + fsdp_root: nn.Module, + targets: Sequence[tuple[str, tuple[nn.Module, nn.Module], tuple[torch.Tensor, torch.Tensor]]], +) -> tuple[FSDPParam, ...]: + """Bind routed expert FSDPParams to their two-generation VMM landings. + + Each target is ``(layer_fqn, projections, landings)``. + Matching uses the original module and parameter-name identities recorded + by FSDP, never an FQN guess. + """ + if torch.__version__ != _TARGET_TORCH_VERSION: + raise RuntimeError( + f"MoonEP direct FSDP landing requires torch {_TARGET_TORCH_VERSION}, got {torch.__version__}" + ) + + selected = _resolve_and_validate_targets(fsdp_root, targets) + + for fsdp_param, fsdp_owner, projection, landing in selected: + setattr(projection, _FSDP_PARAM_ATTR, fsdp_param) + setattr(fsdp_param, _BINDING_ATTR, landing) + setattr(fsdp_param, _OWNER_ATTR, fsdp_owner) + setattr(fsdp_param, _PROJECTION_ATTR, projection) + fsdp_param.init_all_gather_outputs = types.MethodType( # type: ignore[method-assign] + _init_direct_all_gather_outputs, fsdp_param + ) + fsdp_param.alloc_all_gather_outputs = types.MethodType( # type: ignore[method-assign] + _keep_direct_all_gather_storage, fsdp_param + ) + fsdp_param.free_unsharded_param = types.MethodType( # type: ignore[method-assign] + _keep_direct_all_gather_storage, fsdp_param + ) + return tuple(item[0] for item in selected) + + +def fsdp_current_unsharded_expert_parameters( + projections: tuple[nn.Module, nn.Module], +) -> tuple[nn.Parameter, nn.Parameter]: + """Return the current FSDP leaf Parameters without starting an + AllGather.""" + current_parameters: list[nn.Parameter] = [] + for projection in projections: + fsdp_param = getattr(projection, _FSDP_PARAM_ATTR, None) + if fsdp_param is None: + raise RuntimeError("MoonEP direct FSDP landing is not installed for this expert projection") + if fsdp_param.sharded_state is not ShardedState.UNSHARDED: + raise RuntimeError("MoonEP expert weight was read outside its FSDP unsharded window") + registered = getattr(fsdp_param._module_info.module, fsdp_param._module_info.param_name) + if registered is not fsdp_param.unsharded_param: + raise RuntimeError("MoonEP observed an unexpected FSDP Parameter switch") + if not isinstance(registered, nn.Parameter): + raise RuntimeError("MoonEP expected FSDP to expose an unsharded Parameter") + local = registered.to_local() if isinstance(registered, DTensor) else registered + landing = getattr(fsdp_param, _BINDING_ATTR) + if local.data_ptr() != landing.data_ptr() or local.numel() != landing.numel(): + raise RuntimeError("MoonEP unsharded FSDP view no longer aliases its VMM landing") + current_parameters.append(registered) + return current_parameters[0], current_parameters[1] + + +def accumulate_fsdp_unsharded_expert_gradients( + parameters: tuple[nn.Parameter, nn.Parameter], + local_gradients: tuple[torch.Tensor, torch.Tensor], +) -> None: + """Hand completed home gradients to the Parameters consumed by FSDP.""" + with torch.no_grad(): + for parameter, local_gradient in zip(parameters, local_gradients, strict=True): + local_parameter = parameter.to_local() if isinstance(parameter, DTensor) else parameter + local_gradient = local_gradient.reshape(local_parameter.shape) + if isinstance(parameter, DTensor): + gradient: torch.Tensor = DTensor.from_local( + local_gradient, + parameter.device_mesh, + parameter.placements, + run_check=False, + shape=parameter.shape, + stride=parameter.stride(), + ) + else: + gradient = local_gradient + + # The first producer transfers storage ownership without a copy; + # Domino's later producers use the same accumulation semantics as + # native AccumulateGrad. + if parameter.grad is None: + parameter.grad = gradient + else: + parameter.grad.add_(gradient) + + +def uninstall_fsdp_vmm_landing(fsdp_params: tuple[FSDPParam, ...]) -> None: + """Restore native instance methods and release every FSDP VMM reference.""" + owners = {getattr(fsdp_param, _OWNER_ATTR) for fsdp_param in fsdp_params} + for owner in owners: + state = _get_module_fsdp_state(owner) + if state is None or state._training_state.name != "IDLE": + raise RuntimeError("MoonEP direct landing may only be removed at an idle FSDP boundary") + # Public FSDPModule operation: swap modules back to their sharded + # Parameters before removing the VMM-backed unsharded Parameter. + owner.reshard() + + projections = {getattr(fsdp_param, _PROJECTION_ATTR) for fsdp_param in fsdp_params} + for fsdp_param in fsdp_params: + if fsdp_param.sharded_state is not ShardedState.SHARDED: + raise RuntimeError("MoonEP failed to reshard a bound FSDP parameter") + fsdp_param.all_gather_outputs.clear() + fsdp_param._unsharded_inner_tensors.clear() + if hasattr(fsdp_param, "_unsharded_param"): + del fsdp_param._unsharded_param + for method_name in ("init_all_gather_outputs", "alloc_all_gather_outputs", "free_unsharded_param"): + delattr(fsdp_param, method_name) + for attr_name in (_BINDING_ATTR, _OWNER_ATTR, _PROJECTION_ATTR): + delattr(fsdp_param, attr_name) + for projection in projections: + delattr(projection, _FSDP_PARAM_ATTR) + + +__all__ = [ + "accumulate_fsdp_unsharded_expert_gradients", + "fsdp_current_unsharded_expert_parameters", + "install_fsdp_vmm_landing", + "uninstall_fsdp_vmm_landing", +] diff --git a/xtuner/v1/module/dispatcher/moonep.py b/xtuner/v1/module/dispatcher/moonep.py new file mode 100644 index 0000000000..d82d5a1fe3 --- /dev/null +++ b/xtuner/v1/module/dispatcher/moonep.py @@ -0,0 +1,959 @@ +"""MoonEP's model-scoped XTuner integration. + +The backend import remains lazy so unrelated dispatchers do not require +MoonEP. ``MoonEPModelRuntime`` owns model resources, ``MoonEPDispatcher`` owns +one routed layer's static policy, and ``_MoonEPLayerInvocation`` owns one +dispatch/combine transaction. The private VMM workspace remains the deep +module for physical expert layout. +""" + +from __future__ import annotations + +import os +from typing import Any, cast + +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.tensor import DTensor +from typing_extensions import TypedDict, override + +from xtuner.v1.ops.moe.cuda.route_weight import route_weight_rows_backward +from xtuner.v1.utils import log_rank0 + +from .base import ExpertWeightLayout, GenericDispatcher, PostDispatchResult, ProjectionPair +from .fsdp_vmm_landing import ( + accumulate_fsdp_unsharded_expert_gradients, + fsdp_current_unsharded_expert_parameters, + install_fsdp_vmm_landing, + uninstall_fsdp_vmm_landing, +) +from .moonep_workspace import _ExpertVMMWorkspace + + +_INTEGRATION_API_VERSION = 3 +_MOONEP_IMPORT_ERROR: ImportError | None + +try: + import moonep as _moonep_backend +except ImportError as exc: + _moonep_backend = None + _MOONEP_IMPORT_ERROR = exc +else: + _MOONEP_IMPORT_ERROR = None + + +def require_moonep_backend() -> Any: + """Validate the optional MoonEP-mod package when MoonEP is selected.""" + if _moonep_backend is None: + raise RuntimeError("dispatcher='moonep' requires the MoonEP-mod integration package") from _MOONEP_IMPORT_ERROR + + source = getattr(_moonep_backend, "__file__", "") + if getattr(_moonep_backend, "XTUNER_INTEGRATION_API_VERSION", None) != _INTEGRATION_API_VERSION: + raise RuntimeError( + f"incompatible MoonEP integration API; expected {_INTEGRATION_API_VERSION}; loaded module: {source}" + ) + return _moonep_backend + + +class MoonEPModelRuntime: + """Own the resources shared by all routed layers in one model/EP group.""" + + def __init__( + self, + *, + ep_group: dist.ProcessGroup, + hidden_size: int, + intermediate_size: int, + num_experts: int, + top_k: int, + intra_layer_micro_batch: int, + staging_reference: bool, + num_sms: int = 64, + ) -> None: + require_moonep_backend() + if intra_layer_micro_batch < 1: + raise ValueError("intra_layer_micro_batch must be positive") + + # MoonEP keeps token counts device-resident. Triton already satisfies + # that contract; grouped_gemm does so only with its CUTLASS backend. + # Validate this once here instead of branching in every GMM call. + from xtuner.v1.module.grouped_linear import moe_group_linear + from xtuner.v1.ops.moe.cuda import cutlass_group_gemm + + if cutlass_group_gemm is not None and moe_group_linear.group_gemm is cutlass_group_gemm: + from grouped_gemm import backend as grouped_gemm_backend + + if os.environ.get("GROUPED_GEMM_USE_CUTLASS") != "1" or not grouped_gemm_backend.use_cutlass: + raise RuntimeError( + "MoonEP with grouped_gemm requires GROUPED_GEMM_USE_CUTLASS=1 before importing grouped_gemm" + ) + + self._ep_group = ep_group + self._hidden_size = hidden_size + self._intermediate_size = intermediate_size + self._num_experts = num_experts + self._top_k = top_k + self._intra_layer_micro_batch = intra_layer_micro_batch + self._staging_reference = staging_reference + self._num_sms = num_sms + + # Model-owned execution resources. Buffer is created lazily once S is + # known; workspace/stream and FSDP bindings are installed as one unit. + self._buffer: Any | None = None + self._workspace: _ExpertVMMWorkspace | None = None + self._comm_stream: torch.cuda.Stream | None = None + self._fsdp_params: tuple[Any, ...] = () + + # Physical routed layers in FSDP execution order. Each entry is + # (layer FQN, (fused_w1w3, fused_w2), landing generation). + self._layers: list[tuple[str, tuple[nn.Module, nn.Module], int]] = [] + # MoonEP allocates fixed-size communication buffers on the first call; + # later calls must retain the same source-token count S per EP rank. + self._fixed_tokens_per_rank: int | None = None + self._closed = False + + def build_dispatcher( + self, + *, + layer_fqn: str, + projections: tuple[nn.Module, nn.Module], + ) -> MoonEPDispatcher: + """Register one physical routed layer in FSDP execution order.""" + if any(registered_fqn == layer_fqn for registered_fqn, _, _ in self._layers): + raise ValueError(f"duplicate MoonEP routed layer: {layer_fqn}") + generation = len(self._layers) % 2 + self._layers.append((layer_fqn, projections, generation)) + return MoonEPDispatcher( + runtime=self, + layer_fqn=layer_fqn, + projections=projections, + generation=generation, + ) + + def validate_before_fsdp(self, fsdp_config: Any) -> None: + """Validate the build-time FSDP policy without retaining its config.""" + if fsdp_config.param_dtype is not torch.bfloat16 or fsdp_config.reduce_dtype is not torch.bfloat16: + raise ValueError("MoonEP requires BF16 FSDP param and reduce dtypes") + if fsdp_config.cpu_offload: + raise ValueError("MoonEP VMM weights cannot use FSDP CPU offload") + if not fsdp_config.requires_grad: + raise ValueError("MoonEP v1 requires trainable FSDP parameters") + if not fsdp_config.reshard_after_forward: + raise ValueError("MoonEP requires reshard_after_forward=True") + + def install_after_fsdp(self, *, fsdp_root: nn.Module) -> None: + """Allocate execution resources after native FSDP has been + installed.""" + if self._workspace is not None: + raise RuntimeError("MoonEP FSDP resources are already installed") + if not self._layers: + raise TypeError("MoonEP requires at least one physical routed-expert layer") + if self._staging_reference: + log_rank0.warning( + "moonep_staging_reference=True copies complete BF16 home expert " + "weights after every FSDP AllGather; it is a numerical reference, " + "not the production performance path." + ) + workspace = _ExpertVMMWorkspace.allocate( + projection_shapes=( + (2 * self._intermediate_size, self._hidden_size), + (self._hidden_size, self._intermediate_size), + ), + num_experts=self._num_experts, + ep_group=self._ep_group, + gradient_slots=self._intra_layer_micro_batch, + ) + # Keep MoonEP collectives in FSDP's device-side launch order. A + # separate high-priority stream forms an orthogonal progress wave with + # NCCL and stalls at MoonEP's rank barriers under a full model. + comm_stream = torch.cuda.current_stream() + if not self._staging_reference: + try: + self._fsdp_params = install_fsdp_vmm_landing( + fsdp_root=fsdp_root, + targets=tuple( + ( + layer_fqn, + projections, + workspace.landing(generation), + ) + for layer_fqn, projections, generation in self._layers + ), + ) + except Exception: + workspace.destroy() + raise + self._workspace = workspace + self._comm_stream = comm_stream + + def _validate_tokens_per_rank(self, tokens_per_rank: int) -> None: + if self._closed: + raise RuntimeError("MoonEP runtime was closed") + if self._fixed_tokens_per_rank is None: + self._fixed_tokens_per_rank = tokens_per_rank + elif tokens_per_rank != self._fixed_tokens_per_rank: + raise RuntimeError(f"MoonEP fixed S changed: {self._fixed_tokens_per_rank} -> {tokens_per_rank}") + + def _buffer_for(self, tokens_per_rank: int) -> Any: + self._validate_tokens_per_rank(tokens_per_rank) + if self._workspace is None: + raise RuntimeError("MoonEP FSDP resources must be installed before forward") + if self._buffer is None: + assert _moonep_backend is not None + self._buffer = _moonep_backend.Buffer( + S=tokens_per_rank, + H=self._hidden_size, + K=self._top_k, + E=self._num_experts, + num_ep_ranks=self._ep_group.size(), + group=self._ep_group, + explicitly_destroy=True, + num_sms=self._num_sms, + ) + return self._buffer + + def _enqueue(self, operation, *, inputs: tuple[torch.Tensor | None, ...] = ()): + """Run one MoonEP transaction on XTuner's stream and return its + event.""" + stream = self._comm_stream + if stream is None: + raise RuntimeError("MoonEP FSDP resources must be installed before execution") + caller_stream = torch.cuda.current_stream() + stream.wait_event(caller_stream.record_event()) + for tensor in inputs: + if tensor is not None: + tensor.record_stream(stream) + with torch.cuda.stream(stream): + result = operation() + done = stream.record_event() + return result, done + + def close(self) -> None: + """Release Buffer before VMM workspace at a coordinated boundary.""" + if self._closed: + return + if self._comm_stream is not None: + self._comm_stream.synchronize() + self._comm_stream = None + if self._buffer is not None: + self._buffer.destroy() + self._buffer = None + if self._fsdp_params: + uninstall_fsdp_vmm_landing(self._fsdp_params) + self._fsdp_params = () + if self._workspace is not None: + self._workspace.destroy() + self._workspace = None + self._layers.clear() + self._closed = True + + +# Dispatcher shape legend: +# S: source tokens on this EP rank, K: routed experts per token, +# NvS: MoonEP's padded VM-group rows, E: global experts, B=E/R: home +# experts per EP rank, H: hidden size. + + +class MoonEPPreDispatchResult(TypedDict): + """Stage 1: device-resident router-space inputs normalized for MoonEP.""" + + hidden_states: torch.Tensor # [S, H], BF16 source-token order. + topk_ids: torch.Tensor # [S, K], contiguous int32 global expert IDs. + tokens_per_expert: torch.Tensor # [E], contiguous int32 source histogram. + # The invocation is opaque control state for the remaining five stages; + # it never crosses into the compiled tensor-only expert block. + _moonep_invocation: _MoonEPLayerInvocation + + +class MoonEPDispatchResult(TypedDict): + """Stage 2: global dispatch outputs plus its eager control-plane state.""" + + hidden_states: torch.Tensor # [NvS, H], BF16 physical VM-group order. + topk_weights: torch.Tensor # [NvS], FP32 weights in the same row order. + # [E+B], int32 padded group ends; stays on device and is non-differentiable. + cu_seqlens: torch.Tensor + # Opaque per-call state shared by later stages; never enters compiled expert compute. + _moonep_invocation: _MoonEPLayerInvocation + + +class MoonEPPostDispatchResult(PostDispatchResult): + """Stage 3: tensor-only local ``[2B]`` expert-compute bundle. + + ``hidden_states`` is ``[NvS, H]``; ``tokens_per_expert`` is device int32 + ``[2B]`` for home then duplicate groups; ``expert_weight_layout`` holds + projection-paired ``[2B, O_p, I_p]`` weights and direct BF16 WGrad targets. + """ + + +class MoonEPPreCombineResult(TypedDict): + """Stage 4: expert outputs before route scaling and reverse transport.""" + + hidden_states: torch.Tensor # [NvS, H], physical VM-group order. + + +class MoonEPCombineResult(TypedDict): + """Stage 5: fused route-scaled output restored to source-token order.""" + + hidden_states: torch.Tensor # [S, H]. + + +class MoonEPPostCombineResult(TypedDict): + """Stage 6: final tensor bundle returned through the generic interface.""" + + hidden_states: torch.Tensor # [S, H]. + + +class _MoonEPLayerInvocation: + """Own one routed layer's complete forward/backward transaction. + + The invocation borrows model resources and layer projections, but it never + owns or calls back into ``MoonEPDispatcher``. All behavior that mutates + call-local plan, event, weight, and gradient state stays here. + """ + + def __init__( + self, + *, + runtime: MoonEPModelRuntime, + layer_fqn: str, + projections: tuple[nn.Module, nn.Module], + generation: int, + grad_slot: int, + ) -> None: + self._runtime = runtime + self._layer_fqn = layer_fqn + self._projections = projections + self._generation = generation + self._grad_slot = grad_slot + + # One MoonEP communication plan and its device-side dependency chain. + # Events are recorded once their named producer has been enqueued. + self._plan: Any | None = None + self._dispatch_done: Any | None = None + self._weights_ready: Any | None = None + self._combine_done: Any | None = None + + # Borrowed VMM aliases for this call. Weights and direct WGrad targets + # are projection pairs, each tensor laid out as [2B, O_p, I_p]. + self._local_weights: ProjectionPair | None = None + self._gradient_targets: ProjectionPair | None = None + self._fallback_gradient_targets: ProjectionPair | None = None + + # Current FSDP unsharded home Parameters [B, O_p, I_p] receive the + # returned BF16 home gradients after both local projections complete. + self._home_parameters: tuple[nn.Parameter, nn.Parameter] | None = None + # Completed home views and the event covering the pair reduction. + self._gradient_completion: tuple[ProjectionPair, torch.cuda.Event] | None = None + + def begin_dispatch( + self, + *, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + tokens_per_expert: torch.Tensor, + topk_weights: torch.Tensor, + async_op: bool, + ) -> MoonEPDispatchResult: + """Create the activation autograd edge and start weight prefetch.""" + hidden_nvsh, weights_nvs, cu_seqlens = _DispatchAutograd.apply( + hidden_states, + topk_ids, + tokens_per_expert, + topk_weights, + self, + async_op, + ) + return MoonEPDispatchResult( + hidden_states=hidden_nvsh, + topk_weights=weights_nvs, + cu_seqlens=cu_seqlens, + _moonep_invocation=self, + ) + + def begin_combine( + self, + *, + expert_output: torch.Tensor, + route_weights: torch.Tensor, + async_op: bool, + ) -> torch.Tensor: + """Create the fused route-scaled combine autograd edge.""" + return _CombineAutograd.apply( + expert_output, + route_weights, + self, + async_op, + ) + + def finish_combine(self, combined: torch.Tensor, *, async_op: bool) -> torch.Tensor: + """Establish the final device dependency.""" + if async_op: + assert self._combine_done is not None + self._combine_done.wait() + return combined + + def prepare_experts(self, dispatched: MoonEPDispatchResult) -> MoonEPPostDispatchResult: + """Wait at the first weight consumer and expose tensor-only layout.""" + workspace = self._runtime._workspace + assert workspace is not None + assert self._weights_ready is not None + assert self._local_weights is not None and self._gradient_targets is not None + + with torch.profiler.record_function("MoonEP::prepare_experts"): + # This inserts a device dependency; it never waits on the host. + self._weights_ready.wait() + local_counts = workspace.local_token_counts(dispatched["cu_seqlens"]) + covered_rows = local_counts.sum() + row_is_covered = ( + torch.arange( + dispatched["hidden_states"].shape[0], + device=dispatched["hidden_states"].device, + ) + < covered_rows + ) + hidden_states = dispatched["hidden_states"] * row_is_covered.unsqueeze(-1) + local_counts = torch.cat( + ( + local_counts[:-1], + local_counts[-1:] + dispatched["hidden_states"].shape[0] - covered_rows, + ) + ) + + local_weights = self._local_weights + gradient_targets = self._gradient_targets + self._local_weights = None + self._gradient_targets = None + # Leaf Parameters make grouped GEMM compute dW. Its backward writes the + # pair directly into the VMM targets; Start consumes the replay aliases, + # so no AccumulateGrad hook or Parameter.grad handoff is required. + differentiable_weights: ProjectionPair = ( + nn.Parameter(local_weights[0]), + nn.Parameter(local_weights[1]), + ) + return MoonEPPostDispatchResult( + hidden_states=hidden_states, + tokens_per_expert=local_counts, + expert_weight_layout=ExpertWeightLayout( + trainable_weights=differentiable_weights, + trainable_wgrad_outs=gradient_targets, + ), + ) + + def _current_home_parameters(self) -> tuple[nn.Parameter, nn.Parameter]: + """Return current FSDP leaves, staging only in reference mode.""" + if not self._runtime._staging_reference: + return fsdp_current_unsharded_expert_parameters(self._projections) + + workspace = self._runtime._workspace + assert workspace is not None + parameters: list[nn.Parameter] = [] + for linear, landing in zip( + self._projections, + workspace.landing(self._generation), + strict=True, + ): + weight = cast(torch.Tensor, linear.weight) + if not isinstance(weight, nn.Parameter): + raise RuntimeError(f"{self._layer_fqn} staging expected an unsharded expert Parameter") + source = weight.to_local() if isinstance(weight, DTensor) else weight + if source.dtype is not torch.bfloat16 or source.numel() != landing.numel(): + raise RuntimeError(f"{self._layer_fqn} staging expected an unsharded BF16 expert weight") + with torch.no_grad(): + landing.copy_(source.view_as(landing)) + parameters.append(weight) + return parameters[0], parameters[1] + + def _dispatch_forward( + self, + source_hidden: torch.Tensor, + topk_ids: torch.Tensor, + tokens_per_expert: torch.Tensor, + source_route_weights: torch.Tensor, + *, + async_op: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + runtime = self._runtime + buffer = runtime._buffer_for(source_hidden.shape[0]) + workspace = runtime._workspace + assert workspace is not None + + def dispatch_and_prefetch(): + # The staging copy precedes dispatch's device barrier. The fresh + # plan then starts both projection prefetches in this transaction. + self._home_parameters = self._current_home_parameters() + hidden_nvsh, route_weights_nvs, cu_seqlens, plan = buffer.dispatch( + source_hidden, + route_weights_sk=source_route_weights, + topk_experts_sk=topk_ids, + tokens_per_expert=tokens_per_expert, + async_finish=False, + zero_copy=False, + ) + assert route_weights_nvs is not None and cu_seqlens is not None + self._plan = plan + self._dispatch_done = torch.cuda.current_stream().record_event() + self._local_weights, self._gradient_targets = workspace.prefetch_weights( + buffer=buffer, + plan=plan, + generation=self._generation, + grad_slot=self._grad_slot, + ) + return hidden_nvsh, route_weights_nvs, cu_seqlens + + with torch.profiler.record_function("MoonEP::dispatch_forward"): + result, self._weights_ready = runtime._enqueue( + dispatch_and_prefetch, + inputs=(source_hidden, source_route_weights, topk_ids, tokens_per_expert), + ) + assert self._dispatch_done is not None + if not async_op: + self._dispatch_done.wait() + return result + + def _dispatch_backward( + self, + grad_hidden_nvsh: torch.Tensor, + grad_route_weights_nvs: torch.Tensor, + ) -> ProjectionPair: + runtime = self._runtime + buffer = runtime._buffer + assert self._plan is not None and buffer is not None + grad_hidden_nvsh = grad_hidden_nvsh.contiguous() + grad_route_weights_nvs = grad_route_weights_nvs.contiguous() + + def combine_gradients(): + grad_hidden, grad_route_weights, no_event = buffer.combine( + plan=self._plan, + hidden_nvsh=grad_hidden_nvsh, + route_weights_nvs=grad_route_weights_nvs, + async_finish=False, + zero_copy=False, + ) + assert grad_route_weights is not None and no_event is None + return grad_hidden, grad_route_weights + + with torch.profiler.record_function("MoonEP::dispatch_backward"): + result, done = runtime._enqueue( + combine_gradients, + inputs=(grad_hidden_nvsh, grad_route_weights_nvs), + ) + done.wait() + return result + + def _combine_forward( + self, + expert_output: torch.Tensor, + route_weights: torch.Tensor, + *, + async_op: bool, + ) -> torch.Tensor: + runtime = self._runtime + buffer = runtime._buffer + assert self._plan is not None and buffer is not None + + def combine_output(): + output, gathered_weights, no_event = buffer.combine( + plan=self._plan, + hidden_nvsh=expert_output, + hidden_scales_nvs=route_weights, + route_weights_nvs=None, + async_finish=False, + zero_copy=False, + ) + assert gathered_weights is None and no_event is None + return output + + with torch.profiler.record_function("MoonEP::combine_forward"): + output, self._combine_done = runtime._enqueue( + combine_output, + inputs=(expert_output, route_weights), + ) + if not async_op: + self._combine_done.wait() + return output + + def _combine_backward(self, grad_output: torch.Tensor) -> tuple[torch.Tensor, Any]: + runtime = self._runtime + buffer = runtime._buffer + workspace = runtime._workspace + assert self._plan is not None and buffer is not None and workspace is not None + grad_output = grad_output.contiguous() + + def dispatch_gradient_and_prefetch(): + # FSDP pre-backward has restored this generation. Stage it before + # dispatch's barrier, then replay remote weights on the same stream. + replay_home_parameters = self._current_home_parameters() + if self._home_parameters is None: + raise RuntimeError("MoonEP backward has no forward home Parameters") + if any( + replay is not forward + for replay, forward in zip(replay_home_parameters, self._home_parameters, strict=True) + ): + raise RuntimeError("MoonEP backward observed a different FSDP unsharded Parameter") + grad_weighted, no_weights, no_cu, reused_plan = buffer.dispatch( + grad_output, + plan=self._plan, + async_finish=False, + zero_copy=False, + ) + assert no_weights is None and no_cu is None and reused_plan is self._plan + gradient_dispatch_done = torch.cuda.current_stream().record_event() + _, self._fallback_gradient_targets = workspace.prefetch_weights( + buffer=buffer, + plan=self._plan, + generation=self._generation, + grad_slot=self._grad_slot, + ) + return grad_weighted, gradient_dispatch_done + + with torch.profiler.record_function("MoonEP::combine_backward"): + (grad_weighted, gradient_dispatch_done), replay_done = runtime._enqueue( + dispatch_gradient_and_prefetch, + inputs=(grad_output,), + ) + # Route-scale backward overlaps weight replay but cannot read the + # dispatched gradient before this device event. + gradient_dispatch_done.wait() + return grad_weighted, replay_done + + def _start_gradient_completion(self) -> None: + """Enqueue the one pair reduction after both GMM backwards.""" + if self._gradient_completion is not None: + raise RuntimeError("MoonEP gradient completion was started twice") + if self._plan is None or self._fallback_gradient_targets is None: + raise RuntimeError("MoonEP gradient completion has no plan/targets") + # Both grouped-GEMM backwards write these VMM slots before propagating + # activation gradients to Start. The fresh aliases share that storage + # without retaining the forward TensorImpl/version counters. + reduction_pair = self._fallback_gradient_targets + runtime = self._runtime + workspace = runtime._workspace + buffer = runtime._buffer + if workspace is None or buffer is None: + raise RuntimeError("MoonEP gradient completion has no runtime workspace") + + # Only the pair reduction is enqueued here. The completion event is + # consumed later by ``_finish_gradient_completion`` on the caller's + # stream, so upstream backward can overlap the communication. + with torch.profiler.record_function("MoonEP::gradient_handoff"): + home_grads, done = runtime._enqueue( + lambda: workspace.complete_gradients( + buffer=buffer, + plan=self._plan, + local_grads=reduction_pair, + grad_slot=self._grad_slot, + ), + inputs=reduction_pair, + ) + self._gradient_completion = (home_grads, done) + + def _finish_gradient_completion(self) -> None: + """Wait on the device event and hand BF16 home grads to FSDP.""" + completion = self._gradient_completion + if completion is None: + raise RuntimeError("MoonEP gradient completion was not started") + home_grads, done = completion + done.wait() + + if self._home_parameters is None: + raise RuntimeError("MoonEP gradient completion has no home Parameters") + accumulate_fsdp_unsharded_expert_gradients(self._home_parameters, home_grads) + + self._fallback_gradient_targets = None + self._home_parameters = None + self._gradient_completion = None + + +class _MoonEPGradReduceStart(torch.autograd.Function): + """Identity edge whose backward starts duplicate-gradient transport.""" + + @staticmethod + def forward( + ctx: Any, + hidden_states: torch.Tensor, + invocation: _MoonEPLayerInvocation, + ) -> torch.Tensor: + ctx.invocation = invocation + return hidden_states + + @staticmethod + def backward(ctx: Any, grad_hidden: torch.Tensor) -> tuple[torch.Tensor, None]: + # Only enqueue work here. The current stream does not wait, allowing + # attention/router backward to overlap the MoonEP communication. + cast(_MoonEPLayerInvocation, ctx.invocation)._start_gradient_completion() + return grad_hidden, None + + +class _MoonEPGradReduceJoin(torch.autograd.Function): + """Identity edge whose backward waits and hands home grads to FSDP.""" + + @staticmethod + def forward( + ctx: Any, + layer_input: torch.Tensor, + invocation: _MoonEPLayerInvocation, + ) -> torch.Tensor: + ctx.invocation = invocation + return layer_input + + @staticmethod + def backward(ctx: Any, grad_input: torch.Tensor) -> tuple[torch.Tensor, None]: + # Event.wait inserts a dependency into the current CUDA stream; it does + # not block the Python host or poll event readiness. + cast(_MoonEPLayerInvocation, ctx.invocation)._finish_gradient_completion() + return grad_input, None + + +class _DispatchAutograd(torch.autograd.Function): + """Bridge the dispatch/combine pair into PyTorch autograd.""" + + @staticmethod + def forward( + ctx: Any, + source_hidden: torch.Tensor, + topk_ids: torch.Tensor, + tokens_per_expert: torch.Tensor, + source_route_weights: torch.Tensor, + invocation: _MoonEPLayerInvocation, + async_op: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ctx.invocation = invocation + # TODO: 所有外部调用的函数， 去掉 _ 前缀 + hidden_nvsh, route_weights_nvs, cu_seqlens = invocation._dispatch_forward( + source_hidden, + topk_ids, + tokens_per_expert, + source_route_weights, + async_op=async_op, + ) + ctx.mark_non_differentiable(cu_seqlens) + return hidden_nvsh, route_weights_nvs, cu_seqlens + + @staticmethod + def backward( + ctx: Any, + grad_hidden_nvsh: torch.Tensor, + grad_route_weights_nvs: torch.Tensor, + grad_cu_seqlens: None, + ) -> tuple[torch.Tensor, None, None, torch.Tensor, None, None]: + del grad_cu_seqlens + grad_hidden, grad_route_weights = cast(_MoonEPLayerInvocation, ctx.invocation)._dispatch_backward( + grad_hidden_nvsh, + grad_route_weights_nvs, + ) + return grad_hidden, None, None, grad_route_weights, None, None + + +class _CombineAutograd(torch.autograd.Function): + """Bridge fused combine and saved-plan dispatch into autograd.""" + + @staticmethod + def forward( + ctx: Any, + expert_output: torch.Tensor, + route_weights: torch.Tensor, + invocation: _MoonEPLayerInvocation, + async_op: bool, + ) -> torch.Tensor: + ctx.invocation = invocation + ctx.save_for_backward(expert_output, route_weights) + return invocation._combine_forward(expert_output, route_weights, async_op=async_op) + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None, None]: + grad_weighted, replay_done = ctx.invocation._combine_backward(grad_output) + expert_output, route_weights = ctx.saved_tensors + grad_expert, grad_route_weights = route_weight_rows_backward( + grad_weighted, + expert_output, + route_weights, + ) + # The next autograd node immediately reads duplicated weights. + replay_done.wait() + return grad_expert, grad_route_weights, None, None + + +class MoonEPDispatcher( + GenericDispatcher[ + MoonEPPreDispatchResult, + MoonEPDispatchResult, + MoonEPPostDispatchResult, + MoonEPPreCombineResult, + MoonEPCombineResult, + MoonEPPostCombineResult, + ] +): + """Adapt one routed layer to XTuner's six-stage dispatcher interface. + + This class owns only layer-static policy. Every dispatch creates a fresh + ``_MoonEPLayerInvocation`` for plan, event, weight, and gradient state. + """ + + def __init__( + self, + *, + runtime: MoonEPModelRuntime, + layer_fqn: str, + projections: tuple[nn.Module, nn.Module], + generation: int, + ) -> None: + super().__init__( + n_routed_experts=runtime._num_experts, + process_group=runtime._ep_group, + ) + self._runtime = runtime + self._layer_fqn = layer_fqn + self._projections = projections + self._generation = generation + self._next_gradient_slot = 0 + + def _new_invocation(self) -> _MoonEPLayerInvocation: + """Allocate the next call-local slot and invocation token.""" + grad_slot = self._next_gradient_slot + self._next_gradient_slot = (grad_slot + 1) % self._runtime._intra_layer_micro_batch + return _MoonEPLayerInvocation( + runtime=self._runtime, + layer_fqn=self._layer_fqn, + projections=self._projections, + generation=self._generation, + grad_slot=grad_slot, + ) + + @override + def prepare_layer_input( + self, + layer_input: torch.Tensor, + ) -> tuple[torch.Tensor, object | None]: + """Create an invocation before attention/router backward is built.""" + invocation = self._new_invocation() + # Join is an identity in forward and gives the invocation a stable + # autograd edge immediately before attention/router computation. Under + # no-grad, apply remains an identity without recording a backward node. + return _MoonEPGradReduceJoin.apply(layer_input, invocation), invocation + + @override + def dispatch_preprocess( + self, + *, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + tokens_per_expert: torch.Tensor, + layer_state: object | None = None, + async_op: bool = False, + ) -> MoonEPPreDispatchResult: + del topk_weights, async_op + if layer_state is None: + raise RuntimeError("MoonEP dispatch_preprocess requires layer_state from prepare_layer_input") + if not isinstance(layer_state, _MoonEPLayerInvocation): + raise TypeError("MoonEP layer_state must be a _MoonEPLayerInvocation") + self._runtime._validate_tokens_per_rank(hidden_states.shape[0]) + invocation = layer_state + # As with Join, outer no-grad suppresses the backward node naturally. + hidden_states = _MoonEPGradReduceStart.apply(hidden_states, invocation) + return MoonEPPreDispatchResult( + hidden_states=hidden_states, + topk_ids=topk_ids.to(dtype=torch.int32).contiguous(), + tokens_per_expert=tokens_per_expert.to(dtype=torch.int32).contiguous(), + _moonep_invocation=invocation, + ) + + @override + def dispatch( + self, + *, + pre_dispatched: MoonEPPreDispatchResult, + topk_weights: torch.Tensor, + async_op: bool = False, + decoding: bool = False, + ) -> MoonEPDispatchResult: + if decoding: + raise NotImplementedError("MoonEP fixed-S training dispatch does not implement decoding") + invocation = pre_dispatched["_moonep_invocation"] + return invocation.begin_dispatch( + hidden_states=pre_dispatched["hidden_states"], + topk_ids=pre_dispatched["topk_ids"], + tokens_per_expert=pre_dispatched["tokens_per_expert"], + topk_weights=topk_weights.to(dtype=torch.float32).contiguous(), + async_op=async_op, + ) + + @override + def dispatch_postprocess( + self, + *, + pre_dispatched: MoonEPPreDispatchResult, + dispatched: MoonEPDispatchResult, + async_op: bool = False, + ) -> MoonEPPostDispatchResult: + del pre_dispatched, async_op + return dispatched["_moonep_invocation"].prepare_experts(dispatched) + + @override + def combine_preprocess( + self, + *, + hidden_states: torch.Tensor, + pre_dispatched: MoonEPPreDispatchResult, + dispatched: MoonEPDispatchResult, + post_dispatched: MoonEPPostDispatchResult, + async_op: bool = False, + decoding: bool = False, + ) -> MoonEPPreCombineResult: + del pre_dispatched, dispatched, post_dispatched, async_op, decoding + return MoonEPPreCombineResult(hidden_states=hidden_states) + + @override + def combine( + self, + *, + pre_dispatched: MoonEPPreDispatchResult, + dispatched: MoonEPDispatchResult, + post_dispatched: MoonEPPostDispatchResult, + pre_combined: MoonEPPreCombineResult, + async_op: bool = False, + decoding: bool = False, + ) -> MoonEPCombineResult: + del pre_dispatched, post_dispatched, decoding + invocation = dispatched["_moonep_invocation"] + return MoonEPCombineResult( + hidden_states=invocation.begin_combine( + expert_output=pre_combined["hidden_states"], + route_weights=dispatched["topk_weights"], + async_op=async_op, + ) + ) + + @override + def combine_postprocess( + self, + *, + pre_dispatched: MoonEPPreDispatchResult, + dispatched: MoonEPDispatchResult, + post_dispatched: MoonEPPostDispatchResult, + pre_combined: MoonEPPreCombineResult, + combined: MoonEPCombineResult, + async_op: bool = False, + ) -> MoonEPPostCombineResult: + del pre_dispatched, post_dispatched, pre_combined + invocation = dispatched["_moonep_invocation"] + return MoonEPPostCombineResult( + hidden_states=invocation.finish_combine( + combined["hidden_states"], + async_op=async_op, + ) + ) + + +__all__ = [ + "MoonEPDispatcher", + "MoonEPModelRuntime", + "MoonEPPreDispatchResult", + "MoonEPDispatchResult", + "MoonEPPostDispatchResult", + "MoonEPPreCombineResult", + "MoonEPCombineResult", + "MoonEPPostCombineResult", + "require_moonep_backend", +] diff --git a/xtuner/v1/module/dispatcher/moonep_workspace.py b/xtuner/v1/module/dispatcher/moonep_workspace.py new file mode 100644 index 0000000000..b8ea38a8d4 --- /dev/null +++ b/xtuner/v1/module/dispatcher/moonep_workspace.py @@ -0,0 +1,522 @@ +"""XTuner-owned VMM layout for MoonEP expert weights and gradients. + +MoonEP owns the transport and low-level VMM primitives. XTuner owns this +layout because it is coupled to XTuner's FSDP lifecycle and grouped-GEMM +contract: communication addresses ``E + B`` expert chunks while the single +grouped GEMM consumes one contiguous ``2B`` alias (home followed by duplicate). +""" + +from __future__ import annotations + +import os +import socket +import warnings +from collections.abc import Sequence +from contextlib import ExitStack +from typing import TypeAlias, cast + +import torch +import torch.distributed as dist +from typing_extensions import TypedDict + + +# Shape legend used by every workspace structure below: +# E: global experts, R: EP ranks, B=E/R: home experts per rank, +# P=2: fused projections, G=2: FSDP home generations, N: gradient slots. +# Process-lifetime quarantine for an abandoned workspace's complete tensor +# reference graph. Rank-divergent ``__del__`` must not unmap VMM storage that +# a surviving peer may still use; explicit ``destroy()`` never appends here. +_UNDISPOSED_WORKSPACE_TENSORS: list[object] = [] + +# (physical tensor [B, O_p, I_p], still-open local export FD). The matching +# CUDA allocation handle is deliberately owned separately by ``ExitStack``. +_VMMAllocation: TypeAlias = tuple[torch.Tensor, int] +# Rank-ordered imported FDs: [P][G][R] for home weights or [P][N][R] for +# duplicate gradients. Every descriptor closes after all views are mapped. +_FDGraph: TypeAlias = tuple[tuple[tuple[int, ...], ...], ...] + + +class _WorkspaceAllocations(TypedDict): + """Temporary ownership graph for physical chunks and local export FDs. + + It exists only inside ``allocate()`` while its ``ExitStack`` is open. + Collections are projection-first: ``P=2`` projections, ``G=2`` home + generations, ``N`` gradient slots, and every chunk is ``[B, O_p, I_p]``. + """ + + # [P], each (B, O_p, I_p): allocation and mapping granularity. + chunk_shapes: tuple[tuple[int, ...], ...] + # [P][G]: local home chunks/FDS that become FSDP AllGather landings. + home_weights: tuple[tuple[_VMMAllocation, ...], ...] + # [P]: local duplicate-weight destination/FDS shared by both generations. + duplicate_weights: tuple[_VMMAllocation, ...] + # [P][N]: local-home WGrad chunks/FDS, one per invocation slot. + home_gradients: tuple[tuple[_VMMAllocation, ...], ...] + # [P][N]: duplicate WGrad chunks/FDS published to every EP owner. + duplicate_gradients: tuple[tuple[_VMMAllocation, ...], ...] + # Flat strong references to all physical tensors; excludes mapped views. + keepalives: tuple[torch.Tensor, ...] + + +class _WorkspaceLayout(TypedDict): + """Completed runtime view graph, transposed consumer-first. + + ``P=2`` projections, ``G=2`` home generations, ``N`` gradient slots, + ``R`` EP ranks, ``B=E/R``, and projection ``p`` has shape ``[O_p, I_p]``. + This structure crosses the allocation commit point and initializes the + long-lived ``_ExpertVMMWorkspace``; it contains no open descriptors. + """ + + # [G][P], each [B, O_p, I_p]: FSDP AllGather output targets. + landings: tuple[tuple[torch.Tensor, ...], ...] + # [G][P], each [E+B, O_p, I_p]: MoonEP prefetch addresses all home experts plus local duplicates. + global_weights: tuple[tuple[torch.Tensor, ...], ...] + # [G][P], each [2B, O_p, I_p]: zero-copy [home, duplicate] weights consumed by grouped GEMM. + local_weights: tuple[tuple[torch.Tensor, ...], ...] + # [N][P], each [2B, O_p, I_p]: direct grouped-GEMM WGrad targets for one invocation slot. + local_grad_outputs: tuple[tuple[torch.Tensor, ...], ...] + # [N][P], each [R, B, O_p, I_p]: every rank's duplicate WGrad, mapped for return to local home. + distributed_duplicate_grads: tuple[tuple[torch.Tensor, ...], ...] + # Physical chunks that own storage backing every non-owning mapped view. + keepalives: tuple[torch.Tensor, ...] + + +class _ExpertVMMWorkspace: + """Own one model/EP group's completed ``_WorkspaceLayout``. + + ``_WorkspaceLayout`` is the single source of truth for every runtime + view's shape, indexing, and storage ownership. This object adds the EP + metadata and explicit distributed lifecycle around that layout. + """ + + def __init__( + self, + *, + layout: _WorkspaceLayout, + ep_group: dist.ProcessGroup, + ep_rank: int, + num_experts: int, + experts_per_rank: int, + ) -> None: + # Publish the completed layout under lifecycle-managed private names; + # their detailed contracts remain centralized on _WorkspaceLayout. + self._landings = layout["landings"] + self._global_weights = layout["global_weights"] + self._local_weights = layout["local_weights"] + self._local_grad_outputs = layout["local_grad_outputs"] + self._distributed_duplicate_grads = layout["distributed_duplicate_grads"] + self._keepalives = layout["keepalives"] + self._ep_group = ep_group + self._ep_rank = ep_rank + self._num_experts = num_experts + self._experts_per_rank = experts_per_rank + self._gradient_slots = len(layout["local_grad_outputs"]) + self._destroyed = False + + @classmethod + def allocate( + cls, + *, + projection_shapes: Sequence[tuple[int, int]], + num_experts: int, + ep_group: dist.ProcessGroup, + gradient_slots: int, + ) -> _ExpertVMMWorkspace: + """Validate, allocate, and publish one complete VMM workspace.""" + ep_size, ep_rank = cls._validate_and_resolve_topology( + projection_shapes=projection_shapes, + num_experts=num_experts, + ep_group=ep_group, + ) + experts_per_rank = num_experts // ep_size + + # Keep descriptors and allocation handles alive across all three + # setup phases. They are released together after every VMM view has + # imported them, including when a later phase fails. + with ExitStack() as resources: + allocations = cls._allocate_physical_chunks( + projection_shapes=projection_shapes, + experts_per_rank=experts_per_rank, + gradient_slots=gradient_slots, + resources=resources, + ) + home_weight_graph, duplicate_gradient_graph = cls._build_ipc_fd_graph( + allocations=allocations, + ep_group=ep_group, + ep_size=ep_size, + ep_rank=ep_rank, + resources=resources, + ) + layout = cls._map_workspace_views( + allocations=allocations, + home_weight_graph=home_weight_graph, + duplicate_gradient_graph=duplicate_gradient_graph, + ep_size=ep_size, + ep_rank=ep_rank, + ) + + return cls( + layout=layout, + ep_group=ep_group, + ep_rank=ep_rank, + num_experts=num_experts, + experts_per_rank=experts_per_rank, + ) + + @staticmethod + def _validate_and_resolve_topology( + *, + projection_shapes: Sequence[tuple[int, int]], + num_experts: int, + ep_group: dist.ProcessGroup, + ) -> tuple[int, int]: + """Resolve the EP coordinates after group-wide topology checks.""" + if not dist.is_initialized(): + raise RuntimeError("MoonEP workspace requires an initialized process group") + if ep_group is None: + raise ValueError("ep_group must be provided explicitly") + + # These are implementation preconditions, not a second configuration + # validation layer. Dispatcher construction already validates EP size, + # dtype, top-k, and model metadata before this allocation boundary. + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size: + raise ValueError(f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size})") + if len(projection_shapes) != 2: + raise ValueError("MoonEP requires fused w1/w3 and w2 projections") + + # Host coordination is restricted to this one-time initialization. + # The forward/backward hot path uses only VMM aliases and CUDA events. + # ``members[R]`` is ordered by EP rank and stores (hostname, device). + members: list[tuple[str, int] | None] = [None] * ep_size + dist.all_gather_object( + members, + (socket.gethostname(), torch.cuda.current_device()), + group=ep_group, + ) + hosts = {member[0] for member in members if member is not None} + if len(hosts) != 1: + raise ValueError("MoonEP requires a node-local ep_group") + devices = [member[1] for member in members if member is not None] + if len(set(devices)) != ep_size: + raise ValueError("each EP rank must use a distinct CUDA device") + local_device = torch.cuda.current_device() + if any(peer != local_device and not torch.cuda.can_device_access_peer(local_device, peer) for peer in devices): + raise ValueError("all EP devices must be CUDA peer-accessible") + return ep_size, ep_rank + + @staticmethod + def _allocate_physical_chunks( + *, + projection_shapes: Sequence[tuple[int, int]], + experts_per_rank: int, + gradient_slots: int, + resources: ExitStack, + ) -> _WorkspaceAllocations: + """Allocate the physical chunks before any cross-rank mapping.""" + # Import the optional backend only at resource installation time. A + # normal XTuner import or meta model build remains MoonEP-independent. + from moonep._C import get_vmm_granularity, nvl_dist_alloc, nvl_release_mem_handle + + dtype = torch.bfloat16 + chunk_shapes = tuple((experts_per_rank, *projection_shape) for projection_shape in projection_shapes) + granularity = get_vmm_granularity() + element_size = torch.empty((), dtype=dtype).element_size() + for chunk_shape in chunk_shapes: + chunk_bytes = element_size + for dim in chunk_shape: + chunk_bytes *= dim + if chunk_bytes % granularity: + raise ValueError( + f"expert home chunk requires {granularity}-byte VMM alignment, " + f"got {chunk_bytes} bytes for shape {chunk_shape}" + ) + + # Flat ownership list for every physical tensor. Mapped views do not + # themselves keep the underlying CUDA physical allocations alive. + keepalives: list[torch.Tensor] = [] + + def allocate_chunk(chunk_shape: tuple[int, ...]) -> _VMMAllocation: + tensor, fd, handle = nvl_dist_alloc(shape=list(chunk_shape), dtype=dtype) + # ExitStack runs callbacks in reverse: close the exported FD before + # releasing its allocation handle, matching MoonEP's lifecycle. + resources.callback(nvl_release_mem_handle, handle) + resources.callback(os.close, fd) + keepalives.append(tensor) + return tensor, fd + + # Build projection-first ``[P][G/N]`` collections because each + # projection has its own (O_p, I_p) chunk shape. + home_weights: list[tuple[_VMMAllocation, ...]] = [] + duplicate_weights: list[_VMMAllocation] = [] + home_gradients: list[tuple[_VMMAllocation, ...]] = [] + duplicate_gradients: list[tuple[_VMMAllocation, ...]] = [] + for chunk_shape in chunk_shapes: + duplicate_weights.append(allocate_chunk(chunk_shape)) + home_weights.append(tuple(allocate_chunk(chunk_shape) for _ in range(2))) + home_gradients.append(tuple(allocate_chunk(chunk_shape) for _ in range(gradient_slots))) + duplicate_gradients.append(tuple(allocate_chunk(chunk_shape) for _ in range(gradient_slots))) + + return _WorkspaceAllocations( + chunk_shapes=chunk_shapes, + home_weights=tuple(home_weights), + duplicate_weights=tuple(duplicate_weights), + home_gradients=tuple(home_gradients), + duplicate_gradients=tuple(duplicate_gradients), + keepalives=tuple(keepalives), + ) + + @staticmethod + def _build_ipc_fd_graph( + *, + allocations: _WorkspaceAllocations, + ep_group: dist.ProcessGroup, + ep_size: int, + ep_rank: int, + resources: ExitStack, + ) -> tuple[_FDGraph, _FDGraph]: + """Exchange the FDs needed by the global weight and gradient views.""" + from moonep.buffer import _exchange_ipc_fds + + sender_ranks = list(range(ep_size)) + + def exchange(local_fd: int) -> tuple[int, ...]: + exchanged = _exchange_ipc_fds( + local_fd, + sender_ranks, + ep_rank, + ep_size, + ep_group, + ) + ordered_fds = tuple(exchanged[rank] for rank in sender_ranks) + for fd in ordered_fds: + resources.callback(os.close, fd) + return ordered_fds + + # Both graphs preserve projection and generation/slot ordering; each + # innermost tuple is ordered by EP rank for direct ``nvl_dist_map`` use. + home_weight_graph = tuple( + tuple(exchange(home_fd) for _, home_fd in projection) for projection in allocations["home_weights"] + ) + duplicate_gradient_graph = tuple( + tuple(exchange(duplicate_fd) for _, duplicate_fd in projection) + for projection in allocations["duplicate_gradients"] + ) + return home_weight_graph, duplicate_gradient_graph + + @staticmethod + def _map_workspace_views( + *, + allocations: _WorkspaceAllocations, + home_weight_graph: _FDGraph, + duplicate_gradient_graph: _FDGraph, + ep_size: int, + ep_rank: int, + ) -> _WorkspaceLayout: + """Map the physical/IPC graph into the views consumed at runtime.""" + from moonep._C import nvl_dist_map + + dtype = torch.bfloat16 + # Mapping is easiest projection-first (shape varies with p). The final + # return transposes these builders to the runtime's [G/N][P] indexing. + projection_landings: list[tuple[torch.Tensor, ...]] = [] + projection_globals: list[tuple[torch.Tensor, ...]] = [] + projection_locals: list[tuple[torch.Tensor, ...]] = [] + projection_grad_locals: list[tuple[torch.Tensor, ...]] = [] + projection_distributed_grads: list[tuple[torch.Tensor, ...]] = [] + + for projection, chunk_shape in enumerate(allocations["chunk_shapes"]): + duplicate_weight_fd = allocations["duplicate_weights"][projection][1] + projection_landings.append(tuple(tensor for tensor, _ in allocations["home_weights"][projection])) + + global_generations: list[torch.Tensor] = [] + local_generations: list[torch.Tensor] = [] + for all_home_fds in home_weight_graph[projection]: + # Communication addresses all E home chunks followed by this + # rank's B duplicate chunks. Grouped GEMM instead receives the + # zero-copy local [home B, duplicate B] alias. + global_generations.append( + nvl_dist_map( + chunk_shape=list(chunk_shape), + dtype=dtype, + fds=[*all_home_fds, duplicate_weight_fd], + local_rank=ep_rank, + world_size=ep_size + 1, + ) + ) + local_generations.append( + nvl_dist_map( + chunk_shape=list(chunk_shape), + dtype=dtype, + fds=[all_home_fds[ep_rank], duplicate_weight_fd], + local_rank=0, + world_size=2, + ) + ) + projection_globals.append(tuple(global_generations)) + projection_locals.append(tuple(local_generations)) + + grad_locals: list[torch.Tensor] = [] + distributed_grads: list[torch.Tensor] = [] + for slot, all_duplicate_fds in enumerate(duplicate_gradient_graph[projection]): + home_fd = allocations["home_gradients"][projection][slot][1] + duplicate_fd = allocations["duplicate_gradients"][projection][slot][1] + # Grouped GEMM writes one contiguous [home, duplicate] segment. + # Owners also map every rank's duplicate chunk for in-place + # BF16 gradient return. + grad_locals.append( + nvl_dist_map( + chunk_shape=list(chunk_shape), + dtype=dtype, + fds=[home_fd, duplicate_fd], + local_rank=0, + world_size=2, + ) + ) + distributed_grads.append( + nvl_dist_map( + chunk_shape=list(chunk_shape), + dtype=dtype, + fds=list(all_duplicate_fds), + local_rank=ep_rank, + world_size=ep_size, + ).view(ep_size, *chunk_shape) + ) + projection_grad_locals.append(tuple(grad_locals)) + projection_distributed_grads.append(tuple(distributed_grads)) + + return _WorkspaceLayout( + landings=tuple(zip(*projection_landings, strict=True)), + global_weights=tuple(zip(*projection_globals, strict=True)), + local_weights=tuple(zip(*projection_locals, strict=True)), + local_grad_outputs=tuple(zip(*projection_grad_locals, strict=True)), + distributed_duplicate_grads=tuple(zip(*projection_distributed_grads, strict=True)), + keepalives=allocations["keepalives"], + ) + + @property + def destroyed(self) -> bool: + return self._destroyed + + def landing(self, generation: int) -> tuple[torch.Tensor, torch.Tensor]: + """Return projection-paired FSDP targets, each ``[B, O_p, I_p]``.""" + if self._destroyed: + raise RuntimeError("MoonEP workspace has been destroyed") + if generation not in (0, 1): + raise ValueError(f"generation must be 0 or 1, got {generation}") + return cast(tuple[torch.Tensor, torch.Tensor], self._landings[generation]) + + def local_token_counts(self, cu_seqlens: torch.Tensor) -> torch.Tensor: + """Convert ``[E+B]`` endpoints to ``[2B]`` home/duplicate counts.""" + if self._destroyed: + raise RuntimeError("MoonEP workspace has been destroyed") + if cu_seqlens.dtype != torch.int32 or cu_seqlens.numel() != (self._num_experts + self._experts_per_rank): + raise ValueError("cu_seqlens must be int32 with E+B cumulative endpoints") + + # Static device slices select home B and duplicate B. No route value + # reaches the host and activation rows are not repacked. + counts = torch.diff(cu_seqlens, prepend=torch.zeros_like(cu_seqlens[:1])) + home_start = self._ep_rank * self._experts_per_rank + return torch.cat( + ( + counts[home_start : home_start + self._experts_per_rank], + counts[self._num_experts : self._num_experts + self._experts_per_rank], + ) + ) + + def prefetch_weights(self, *, buffer, plan, generation: int, grad_slot: int): + """Prefetch ``[E+B]`` weights and return local ``[2B]`` weight/WGrad + pairs.""" + if self._destroyed: + raise RuntimeError("MoonEP workspace has been destroyed") + if generation not in (0, 1): + raise ValueError(f"generation must be 0 or 1, got {generation}") + if not 0 <= grad_slot < self._gradient_slots: + raise ValueError(f"gradient slot out of range: {grad_slot}") + buffer.prefetch_weight( + plan=plan, + projections=self._global_weights[generation], + async_finish=False, + ) + # A slot is reused sequentially across physical layers. Each + # invocation receives a fresh TensorImpl/version counter over the same + # VMM storage, avoiding both payload allocation and AOT version clashes. + grad_outputs = tuple( + target.new_empty(0).set_( + target.untyped_storage(), + target.storage_offset(), + target.shape, + target.stride(), + ) + for target in self._local_grad_outputs[grad_slot] + ) + return self._local_weights[generation], grad_outputs + + def complete_gradients(self, *, buffer, plan, local_grads, grad_slot: int): + """Reduce ``[2B]`` WGrads and return their local-home ``[B]`` + prefixes.""" + if self._destroyed: + raise RuntimeError("MoonEP workspace has been destroyed") + if not 0 <= grad_slot < self._gradient_slots: + raise ValueError(f"gradient slot out of range: {grad_slot}") + local_grads = tuple(local_grads) + targets = self._local_grad_outputs[grad_slot] + if len(local_grads) != 2: + raise ValueError("local_grads must contain the two fused projections") + for actual, target in zip(local_grads, targets, strict=True): + if ( + actual.dtype != torch.bfloat16 + or actual.shape != target.shape + or actual.data_ptr() != target.data_ptr() + ): + raise ValueError("local_grads must be the selected workspace gradient slot") + + buffer.reduce_grad_bf16( + plan=plan, + local_grads=local_grads, + distributed_duplicate_grads=self._distributed_duplicate_grads[grad_slot], + async_finish=False, + ) + b = self._experts_per_rank + return local_grads[0][:b], local_grads[1][:b] + + def destroy(self) -> None: + """Release mappings at an explicit, rank-coordinated boundary.""" + if self._destroyed: + return + torch.cuda.synchronize() + dist.barrier(group=self._ep_group) + self._landings = () + self._global_weights = () + self._local_weights = () + self._local_grad_outputs = () + self._distributed_duplicate_grads = () + self._keepalives = () + self._destroyed = True + + def __del__(self) -> None: + if getattr(self, "_destroyed", True): + return + warnings.warn( + "MoonEP workspace was not destroyed explicitly; resources may leak.", + ResourceWarning, + ) + # Keep mappings alive instead of tearing CUDA/VMM state down after + # distributed ranks may have diverged during interpreter shutdown. + _UNDISPOSED_WORKSPACE_TENSORS.append( + ( + self._landings, + self._global_weights, + self._local_weights, + self._local_grad_outputs, + self._distributed_duplicate_grads, + self._keepalives, + ) + ) + self._landings = () + self._global_weights = () + self._local_weights = () + self._local_grad_outputs = () + self._distributed_duplicate_grads = () + self._keepalives = () diff --git a/xtuner/v1/module/dispatcher/torch_all2all.py b/xtuner/v1/module/dispatcher/torch_all2all.py index 6edc6002be..d2f286c5d3 100644 --- a/xtuner/v1/module/dispatcher/torch_all2all.py +++ b/xtuner/v1/module/dispatcher/torch_all2all.py @@ -1,4 +1,4 @@ -from typing import Literal, TypeAlias, cast +from typing import TypeAlias, cast import torch import torch.distributed as dist @@ -13,6 +13,7 @@ from .base import ( CombineResult, DispatchResult, + ExpertWeightLayout, GenericDispatcher, PostCombineResult, PostDispatchResult, @@ -296,14 +297,10 @@ def __init__( n_routed_experts: int, process_group: torch.distributed.ProcessGroup, tp_group: torch.distributed.ProcessGroup | None = None, - training_dtype: Literal["fp8", "bf16"] = "bf16", - generate_dtype: Literal["fp8", "bf16"] = "bf16", ): super().__init__( n_routed_experts=n_routed_experts, process_group=process_group, - training_dtype=training_dtype, - generate_dtype=generate_dtype, ) assert self._process_group is not None, ( "Process group must be provided for `TorchAll2AllDispatcher`. " @@ -332,8 +329,11 @@ def dispatch_preprocess( hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, # noqa: ARG002 — kept for interface compatibility; not used here + tokens_per_expert: torch.Tensor, + layer_state: object | None = None, async_op: bool = False, ) -> TorchAll2AllPreDispatchResult: + del tokens_per_expert, layer_state permuted_hidden_states, row_ids_map = permute(hidden_states, topk_ids.to(torch.int32)) if async_op: @@ -513,6 +513,7 @@ def dispatch_postprocess( hidden_states=global_input_tokens, row_ids_map=row_ids_map, tokens_per_expert=tokens_per_expert, + expert_weight_layout=ExpertWeightLayout(), ) @override diff --git a/xtuner/v1/module/grouped_linear/moe_group_linear.py b/xtuner/v1/module/grouped_linear/moe_group_linear.py index e00a129e34..24c82632f9 100644 --- a/xtuner/v1/module/grouped_linear/moe_group_linear.py +++ b/xtuner/v1/module/grouped_linear/moe_group_linear.py @@ -159,10 +159,29 @@ def __init__( else: self.bias = nn.Parameter(bias) - def forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False): - weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight - weight = weight.view(-1, self.local_out_features, self.local_in_features) - out = group_gemm(x, weight, tokens_per_expert) + def forward( + self, + x: torch.Tensor, + tokens_per_expert: torch.Tensor, + *, + trainable_weight: torch.Tensor | None = None, + trainable_wgrad_out: torch.Tensor | None = None, + external_weight: torch.Tensor | None = None, + external_wgrad_out: torch.Tensor | None = None, + ): + if external_weight is not None or external_wgrad_out is not None: + raise NotImplementedError("two-segment grouped linear is not implemented") + + # A dynamic EP backend may supply a differentiable call-local alias. + # The selected one-segment op still returns its dW through autograd. + if trainable_weight is None: + weight = self.weight.to_local() if isinstance(self.weight, DTensor) else self.weight + weight = weight.view(-1, self.local_out_features, self.local_in_features) + else: + weight = trainable_weight + if trainable_wgrad_out is not None and not isinstance(weight, nn.Parameter): + raise TypeError("a preallocated trainable WGrad requires a leaf Parameter weight") + out = group_gemm(x, weight, tokens_per_expert, grad_weight_out=trainable_wgrad_out) if self.moe_bias: bias = self.bias.to_local() if isinstance(self.bias, DTensor) else self.bias diff --git a/xtuner/v1/module/router/greedy.py b/xtuner/v1/module/router/greedy.py index b1a34b1de0..dda6d2d2c7 100644 --- a/xtuner/v1/module/router/greedy.py +++ b/xtuner/v1/module/router/greedy.py @@ -87,14 +87,22 @@ def forward(self, logits: torch.Tensor, rollout_routed_experts: torch.Tensor | N # moe forward # (e, ) - tokens_per_expert = torch.histc(topk_ids, bins=self.n_routed_experts, min=0, max=self.n_routed_experts) + # bincount determines its output size through a scalar readback. The + # explicit histogram range keeps routing counts entirely on device. + histogram_ids = topk_ids if topk_ids.device.type == "cuda" else topk_ids.float() + tokens_per_expert = torch.histc( + histogram_ids, + bins=self.n_routed_experts, + min=0, + max=self.n_routed_experts, + ).to(torch.int64) return { "logits": logits, "router_weights": routing_weights, "topk_weights": topk_weights, "topk_ids": topk_ids, - "topkens_per_expert": tokens_per_expert, + "tokens_per_expert": tokens_per_expert, } @@ -163,12 +171,18 @@ def forward(self, logits: torch.Tensor, rollout_routed_experts: torch.Tensor | N # moe forward # (e, ) - tokens_per_expert = torch.histc(topk_ids, bins=self.n_routed_experts, min=0, max=self.n_routed_experts) + histogram_ids = topk_ids if topk_ids.device.type == "cuda" else topk_ids.float() + tokens_per_expert = torch.histc( + histogram_ids, + bins=self.n_routed_experts, + min=0, + max=self.n_routed_experts, + ).to(torch.int64) return { "logits": logits, "router_weights": routing_weights, "topk_weights": topk_weights, "topk_ids": topk_ids, - "topkens_per_expert": tokens_per_expert, + "tokens_per_expert": tokens_per_expert, } diff --git a/xtuner/v1/module/router/noaux_router.py b/xtuner/v1/module/router/noaux_router.py index 1fbe864f6c..1847eb79f2 100644 --- a/xtuner/v1/module/router/noaux_router.py +++ b/xtuner/v1/module/router/noaux_router.py @@ -133,20 +133,22 @@ def forward(self, logits, rollout_routed_experts: torch.Tensor | None = None) -> topk_weight = topk_weight / denominator topk_weight = topk_weight * self.router_scaling_factor # must multiply the scaling factor - # TODO: (yehaochen) `Dispatcher` calculate the distribution duplicatedly + # An explicit histogram range avoids bincount's output-size readback + # while preserving the integer count contract consumed by dispatchers. + histogram_ids = topk_ids if topk_ids.device.type == "cuda" else topk_ids.float() tokens_per_expert = torch.histc( - topk_ids.float(), + histogram_ids, bins=self.n_routed_experts, min=0, max=self.n_routed_experts, - ) # .view(self.ep_mesh.size(), -1) + ).to(torch.int64) return { "logits": logits, "router_weights": scores_for_choice, "topk_weights": topk_weight, "topk_ids": topk_ids, - "topkens_per_expert": tokens_per_expert, + "tokens_per_expert": tokens_per_expert, } @@ -225,17 +227,18 @@ def forward(self, logits, rollout_routed_experts: torch.Tensor | None = None) -> topk_weight = topk_weight / denominator topk_weight = topk_weight * self.router_scaling_factor # must multiply the scaling factor + histogram_ids = topk_ids if topk_ids.device.type == "cuda" else topk_ids.float() tokens_per_expert = torch.histc( - topk_ids.float(), + histogram_ids, bins=self.n_routed_experts, min=0, max=self.n_routed_experts, - ) # .view(self.ep_mesh.size(), -1) + ).to(torch.int64) return { "logits": logits, "router_weights": scores_for_choice, "topk_weights": topk_weight, "topk_ids": topk_ids, - "topkens_per_expert": tokens_per_expert, + "tokens_per_expert": tokens_per_expert, } diff --git a/xtuner/v1/module/router/protocol.py b/xtuner/v1/module/router/protocol.py index b5cb1f293c..40815924c5 100644 --- a/xtuner/v1/module/router/protocol.py +++ b/xtuner/v1/module/router/protocol.py @@ -9,7 +9,7 @@ class RouterResults(TypedDict): router_weights: torch.Tensor topk_weights: torch.Tensor topk_ids: torch.Tensor - topkens_per_expert: torch.Tensor + tokens_per_expert: torch.Tensor class RouterProtocol(Protocol): diff --git a/xtuner/v1/ops/moe/__init__.py b/xtuner/v1/ops/moe/__init__.py index 15640fb692..94fc32b64c 100644 --- a/xtuner/v1/ops/moe/__init__.py +++ b/xtuner/v1/ops/moe/__init__.py @@ -1,5 +1,6 @@ import os import traceback +from typing import cast import torch from mmengine import digit_version @@ -19,7 +20,7 @@ def get_group_gemm() -> GroupGemmProtocol: device = get_device() if device == "cpu": - return cpu_group_gemm + return cast(GroupGemmProtocol, cpu_group_gemm) elif device == "cuda": if os.environ.get("XTUNER_USE_CUTLASS_GROUP_GEMM", "0") == "1": from .cuda import cutlass_group_gemm as cuda_group_gemm @@ -28,12 +29,12 @@ def get_group_gemm() -> GroupGemmProtocol: else: from .cuda import triton_group_gemm as cuda_group_gemm - return cuda_group_gemm + return cast(GroupGemmProtocol, cuda_group_gemm) elif device == "npu": from .npu import npu_group_gemm - return npu_group_gemm + return cast(GroupGemmProtocol, npu_group_gemm) else: raise NotImplementedError diff --git a/xtuner/v1/ops/moe/cuda/group_gemm.py b/xtuner/v1/ops/moe/cuda/group_gemm.py index bcd5313904..0e34e12611 100644 --- a/xtuner/v1/ops/moe/cuda/group_gemm.py +++ b/xtuner/v1/ops/moe/cuda/group_gemm.py @@ -2,25 +2,34 @@ import torch -from .triton_kernels import k_grouped_gemm, m_grouped_gemm +from .triton_kernels import k_grouped_gemm, k_grouped_gemm_out, m_grouped_gemm class GroupedGemm(torch.autograd.Function): @staticmethod - def forward(ctx, x, w, tokens_per_expert): - out = m_grouped_gemm(x, w, tokens_per_expert, trans_b=True) - ctx.save_for_backward(x, w, tokens_per_expert) - return out + def forward(ctx, x, w, tokens_per_expert, grad_weight_out=None): + ctx.save_for_backward(x, w, tokens_per_expert, grad_weight_out) + if x.shape[0] == 0: + return x.new_empty((0, w.shape[1])) + return m_grouped_gemm(x, w, tokens_per_expert, trans_b=True) @staticmethod def backward(ctx, grad_output): - x, w, tokens_per_expert = ctx.saved_tensors - dx = m_grouped_gemm(grad_output, w, tokens_per_expert, trans_b=False) - dw = k_grouped_gemm(grad_output, x, tokens_per_expert) - return dx, dw, None - - -def triton_group_gemm(x, w, tokens_per_expert): + grad_output = grad_output.contiguous() + x, w, tokens_per_expert, grad_weight_out = ctx.saved_tensors + if x.shape[0] == 0: + dx = torch.empty_like(x) + else: + dx = m_grouped_gemm(grad_output, w, tokens_per_expert, trans_b=False) + if grad_weight_out is None: + dw = k_grouped_gemm(grad_output, x, tokens_per_expert) + else: + k_grouped_gemm_out(grad_output, x, tokens_per_expert, grad_weight_out) + dw = grad_weight_out + return dx, dw, None, None + + +def triton_group_gemm(x, w, tokens_per_expert, *, grad_weight_out=None): """Grouped matrix multiplication (GMM) for expert models. Args: @@ -31,7 +40,4 @@ def triton_group_gemm(x, w, tokens_per_expert): Returns: Tensor: Output tensor of shape (batch_size, seq_len, dout). """ - if x.shape[0] == 0: - # put x and w to the pytorch graph - return torch.matmul(x, w[0].T) - return GroupedGemm.apply(x, w, tokens_per_expert) + return GroupedGemm.apply(x, w, tokens_per_expert, grad_weight_out) diff --git a/xtuner/v1/ops/moe/cuda/group_gemm_cutlass.py b/xtuner/v1/ops/moe/cuda/group_gemm_cutlass.py index 2aff977328..ed4d11dd5a 100644 --- a/xtuner/v1/ops/moe/cuda/group_gemm_cutlass.py +++ b/xtuner/v1/ops/moe/cuda/group_gemm_cutlass.py @@ -4,15 +4,10 @@ Support torch compile.""" import torch +from grouped_gemm import backend from torch import Tensor -try: - from grouped_gemm import backend -except ImportError: - backend = None - - @torch.library.custom_op("moe::gmm", mutates_args=()) def moe_grouped_gemm( a: Tensor, @@ -20,8 +15,26 @@ def moe_grouped_gemm( batch_sizes: Tensor, trans_a: bool = False, trans_b: bool = True, + grad_weight_out: Tensor | None = None, ) -> Tensor: - return backend.gmm(a, b, batch_sizes, trans_a=trans_a, trans_b=trans_b) + del grad_weight_out + output_shape: tuple[int, ...] + if trans_a: + output_shape = (batch_sizes.shape[0], a.shape[-1], b.shape[-1]) + else: + output_shape = (a.shape[0], b.shape[1] if trans_b else b.shape[2]) + output = torch.empty(output_shape, device=a.device, dtype=a.dtype) + + backend.gmm( + a, + b, + batch_sizes, + trans_a=trans_a, + trans_b=trans_b, + c=output, + num_sm=-1, + ) + return output @moe_grouped_gemm.register_fake @@ -31,7 +44,9 @@ def _( batch_sizes: Tensor, trans_a: bool = False, trans_b: bool = True, + grad_weight_out: Tensor | None = None, ) -> Tensor: + del grad_weight_out if trans_a: return torch.empty( (batch_sizes.shape[0], a.shape[-1], b.shape[-1]), @@ -47,14 +62,15 @@ def _( def setup_context(ctx, inputs, output) -> None: a, b, batch_sizes = inputs[:3] - trans_b = inputs[-1] - ctx.save_for_backward(a, b, batch_sizes) + grad_weight_out = inputs[-1] + trans_b = inputs[-2] + ctx.save_for_backward(a, b, batch_sizes, grad_weight_out) ctx.trans_b = trans_b -def backward(ctx, grad) -> tuple[Tensor | None, Tensor | None, None, None, None]: +def backward(ctx, grad) -> tuple[Tensor | None, Tensor | None, None, None, None, None]: grad = grad.contiguous() - a, b, batch_sizes = ctx.saved_tensors + a, b, batch_sizes, grad_weight_out = ctx.saved_tensors trans_b = ctx.trans_b agrad = None @@ -64,14 +80,50 @@ def backward(ctx, grad) -> tuple[Tensor | None, Tensor | None, None, None, None] bgrad = None if ctx.needs_input_grad[1]: lhs, rhs = (grad, a) if trans_b else (a, grad) - bgrad = moe_grouped_gemm(lhs, rhs, batch_sizes, trans_a=True, trans_b=False) - return agrad, bgrad, None, None, None + if grad_weight_out is None: + bgrad = moe_grouped_gemm(lhs, rhs, batch_sizes, trans_a=True, trans_b=False) + else: + moe_grouped_gemm_out(lhs, rhs, batch_sizes, grad_weight_out, trans_a=True, trans_b=False) + bgrad = grad_weight_out + return agrad, bgrad, None, None, None, None moe_grouped_gemm.register_autograd(backward, setup_context=setup_context) -def cutlass_group_gemm(x, w, tokens_per_expert): +@torch.library.custom_op("moe::gmm_out", mutates_args={"out"}) +def moe_grouped_gemm_out( + a: Tensor, + b: Tensor, + batch_sizes: Tensor, + out: Tensor, + trans_a: bool = False, + trans_b: bool = True, +) -> None: + backend.gmm( + a, + b, + batch_sizes, + trans_a=trans_a, + trans_b=trans_b, + c=out, + num_sm=-1, + ) + + +@moe_grouped_gemm_out.register_fake +def _( + a: Tensor, + b: Tensor, + batch_sizes: Tensor, + out: Tensor, + trans_a: bool = False, + trans_b: bool = True, +) -> None: + return None + + +def cutlass_group_gemm(x, w, tokens_per_expert, *, grad_weight_out=None): """Grouped matrix multiplication (GMM) for expert models. Args: @@ -82,7 +134,5 @@ def cutlass_group_gemm(x, w, tokens_per_expert): Returns: Tensor: Output tensor of shape (batch_size, seq_len, dout). """ - if x.shape[0] == 0: - # put x and w to the pytorch graph - return torch.matmul(x, w[0].T) - return moe_grouped_gemm(x, w, tokens_per_expert.cpu(), trans_b=True) + device_counts = tokens_per_expert.to(device=x.device, dtype=torch.int64) + return moe_grouped_gemm(x, w, device_counts, trans_b=True, grad_weight_out=grad_weight_out) diff --git a/xtuner/v1/ops/moe/cuda/route_weight.py b/xtuner/v1/ops/moe/cuda/route_weight.py new file mode 100644 index 0000000000..edf9ae2c95 --- /dev/null +++ b/xtuner/v1/ops/moe/cuda/route_weight.py @@ -0,0 +1,70 @@ +import torch +import triton +import triton.language as tl +from torch import Tensor + + +@triton.jit +def _route_weight_rows_backward_kernel( + grad_weighted, + expert_output, + route_weights, + grad_expert, + grad_route, + hidden_size: tl.constexpr, + block_size: tl.constexpr, +): + row = tl.program_id(0) + # Match grouped-gemm BF16 unpermute backward: the FP32 router weight is + # rounded before multiplication, and each route-gradient product is + # rounded before its FP32 reduction. + route_weight = tl.load(route_weights + row).to(tl.bfloat16).to(tl.float32) + route_grad = 0.0 + for start in tl.static_range(0, hidden_size, block_size): + offsets = start + tl.arange(0, block_size) + mask = offsets < hidden_size + grad = tl.load(grad_weighted + row * hidden_size + offsets, mask=mask, other=0.0).to(tl.float32) + output = tl.load(expert_output + row * hidden_size + offsets, mask=mask, other=0.0).to(tl.float32) + tl.store( + grad_expert + row * hidden_size + offsets, + (grad * route_weight).to(tl.bfloat16), + mask=mask, + ) + route_grad += tl.sum((grad * output).to(tl.bfloat16).to(tl.float32), axis=0) + tl.store(grad_route + row, route_grad) + + +@torch.library.custom_op("moe::route_weight_rows_backward", mutates_args=()) +def route_weight_rows_backward( + grad_weighted: Tensor, + expert_output: Tensor, + route_weights: Tensor, +) -> tuple[Tensor, Tensor]: + """Differentiate fused BF16 row scaling without a full FP32 activation.""" + assert grad_weighted.dtype is torch.bfloat16 and grad_weighted.is_contiguous() + assert expert_output.dtype is torch.bfloat16 and expert_output.is_contiguous() + assert route_weights.dtype is torch.float32 and route_weights.is_contiguous() + assert grad_weighted.shape == expert_output.shape + assert route_weights.shape == grad_weighted.shape[:1] + + grad_expert = torch.empty_like(grad_weighted) + grad_route = torch.empty_like(route_weights) + _route_weight_rows_backward_kernel[(grad_weighted.shape[0],)]( + grad_weighted, + expert_output, + route_weights, + grad_expert, + grad_route, + hidden_size=grad_weighted.shape[1], + block_size=256, + num_warps=4, + ) + return grad_expert, grad_route + + +@route_weight_rows_backward.register_fake +def _(grad_weighted: Tensor, expert_output: Tensor, route_weights: Tensor) -> tuple[Tensor, Tensor]: + return torch.empty_like(grad_weighted), torch.empty_like(route_weights) + + +__all__ = ["route_weight_rows_backward"] diff --git a/xtuner/v1/ops/moe/cuda/triton_kernels/__init__.py b/xtuner/v1/ops/moe/cuda/triton_kernels/__init__.py index 2214f0eddd..5f26b09e94 100644 --- a/xtuner/v1/ops/moe/cuda/triton_kernels/__init__.py +++ b/xtuner/v1/ops/moe/cuda/triton_kernels/__init__.py @@ -11,20 +11,22 @@ import triton if triton.__version__ >= "3.4.0": - from .k_grouped_gemm_TMA_triton3_4 import k_grouped_gemm + from .k_grouped_gemm_TMA_triton3_4 import k_grouped_gemm, k_grouped_gemm_out from .m_grouped_gemm_TMA_triton3_4 import m_grouped_gemm elif triton.__version__ >= "3.2.0": - from .k_grouped_gemm_TMA import k_grouped_gemm + from .k_grouped_gemm_TMA import k_grouped_gemm, k_grouped_gemm_out from .m_grouped_gemm_TMA import m_grouped_gemm else: env_not_available_func = get_env_not_available_func(["torch.accelerator", "triton"]) k_grouped_gemm = env_not_available_func + k_grouped_gemm_out = env_not_available_func m_grouped_gemm = env_not_available_func else: env_not_available_func = get_env_not_available_func(["torch.accelerator", "triton"]) k_grouped_gemm = env_not_available_func + k_grouped_gemm_out = env_not_available_func m_grouped_gemm = env_not_available_func -__all__ = ["k_grouped_gemm", "m_grouped_gemm"] +__all__ = ["k_grouped_gemm", "k_grouped_gemm_out", "m_grouped_gemm"] diff --git a/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA.py b/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA.py index a6662fab55..bebacd7f48 100644 --- a/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA.py +++ b/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA.py @@ -127,21 +127,24 @@ def k_grouped_gemm_kernel( tl._experimental_descriptor_store(c_desc_ptr, c, [off_row, off_col]) -@torch.library.custom_op("moe::k_grouped_gemm", mutates_args=()) -def k_grouped_gemm(A: Tensor, B: Tensor, size_per_group: torch.Tensor) -> Tensor: +def _launch_k_grouped_gemm(A: Tensor, B: Tensor, size_per_group: torch.Tensor, C: Tensor) -> None: assert A.dim() == 2 assert B.dim() == 2 K, M = A.shape K_, N = B.shape - assert A.stride(-1) == 1, "Please make sure A is K-major" - assert B.stride(-1) == 1, "Please make sure B is K-major" assert K == K_, "Please make sure that A and B have the same seqlen" # assert K * A.element_size() % 128 == 0, "A and B should be 128-byte aligned" num_groups = size_per_group.shape[0] - C = A.new_empty(num_groups, M, N) + assert C.shape == (num_groups, M, N) + assert C.dtype == A.dtype and C.device == A.device and C.is_contiguous() + if K == 0: + C.zero_() + return + assert A.stride(-1) == 1, "Please make sure A is K-major" + assert B.stride(-1) == 1, "Please make sure B is K-major" group_end = size_per_group.cumsum(0) - size_per_group + size_per_group group_start = size_per_group.cumsum(0) - size_per_group @@ -217,6 +220,12 @@ def grid(META): dtype_b, dtype_c, ) + + +@torch.library.custom_op("moe::k_grouped_gemm", mutates_args=()) +def k_grouped_gemm(A: Tensor, B: Tensor, size_per_group: torch.Tensor) -> Tensor: + C = A.new_empty(size_per_group.shape[0], A.shape[1], B.shape[1]) + _launch_k_grouped_gemm(A, B, size_per_group, C) return C @@ -229,6 +238,17 @@ def _(A: Tensor, B: Tensor, size_per_group: torch.Tensor) -> Tensor: return C +@torch.library.custom_op("moe::k_grouped_gemm_out", mutates_args={"out"}) +def k_grouped_gemm_out(A: Tensor, B: Tensor, size_per_group: torch.Tensor, out: Tensor) -> None: + """Write grouped WGrad directly into caller-owned storage.""" + _launch_k_grouped_gemm(A, B, size_per_group, out) + + +@k_grouped_gemm_out.register_fake +def _(A: Tensor, B: Tensor, size_per_group: torch.Tensor, out: Tensor) -> None: + return None + + if __name__ == "__main__": from torch.profiler import ProfilerActivity, profile, record_function from utils import generate_random_list, row_max_normalization diff --git a/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA_triton3_4.py b/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA_triton3_4.py index 9cff25e030..58d8986b36 100644 --- a/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA_triton3_4.py +++ b/xtuner/v1/ops/moe/cuda/triton_kernels/k_grouped_gemm_TMA_triton3_4.py @@ -138,21 +138,24 @@ def k_grouped_gemm_kernel( c_desc.store([offs_cm, offs_cn], c) -@torch.library.custom_op("moe::k_grouped_gemm", mutates_args=()) -def k_grouped_gemm(A: Tensor, B: Tensor, size_per_group: torch.Tensor) -> Tensor: +def _launch_k_grouped_gemm(A: Tensor, B: Tensor, size_per_group: torch.Tensor, C: Tensor) -> None: assert A.dim() == 2 assert B.dim() == 2 K, M = A.shape K_, N = B.shape - assert A.stride(-1) == 1, "Please make sure A is K-major" - assert B.stride(-1) == 1, "Please make sure B is K-major" assert K == K_, "Please make sure that A and B have the same seqlen" # assert K * A.element_size() % 128 == 0, "A and B should be 128-byte aligned" num_groups = size_per_group.shape[0] - C = A.new_empty(num_groups, M, N) + assert C.shape == (num_groups, M, N) + assert C.dtype == A.dtype and C.device == A.device and C.is_contiguous() + if K == 0: + C.zero_() + return + assert A.stride(-1) == 1, "Please make sure A is K-major" + assert B.stride(-1) == 1, "Please make sure B is K-major" group_end = size_per_group.cumsum(0) - size_per_group + size_per_group group_start = size_per_group.cumsum(0) - size_per_group @@ -193,6 +196,12 @@ def alloc_fn(size: int, alignment: int, stream: Optional[int]): dtype_b, dtype_c, ) + + +@torch.library.custom_op("moe::k_grouped_gemm", mutates_args=()) +def k_grouped_gemm(A: Tensor, B: Tensor, size_per_group: torch.Tensor) -> Tensor: + C = A.new_empty(size_per_group.shape[0], A.shape[1], B.shape[1]) + _launch_k_grouped_gemm(A, B, size_per_group, C) return C @@ -205,6 +214,17 @@ def _(A: Tensor, B: Tensor, size_per_group: torch.Tensor) -> Tensor: return C +@torch.library.custom_op("moe::k_grouped_gemm_out", mutates_args={"out"}) +def k_grouped_gemm_out(A: Tensor, B: Tensor, size_per_group: torch.Tensor, out: Tensor) -> None: + """Write grouped WGrad directly into caller-owned storage.""" + _launch_k_grouped_gemm(A, B, size_per_group, out) + + +@k_grouped_gemm_out.register_fake +def _(A: Tensor, B: Tensor, size_per_group: torch.Tensor, out: Tensor) -> None: + return None + + if __name__ == "__main__": from torch.profiler import ProfilerActivity, profile, record_function from utils import generate_random_list, row_max_normalization diff --git a/xtuner/v1/ops/moe/protocol.py b/xtuner/v1/ops/moe/protocol.py index f51fba0d81..446e3a7c26 100644 --- a/xtuner/v1/ops/moe/protocol.py +++ b/xtuner/v1/ops/moe/protocol.py @@ -9,6 +9,8 @@ def __call__( x: torch.Tensor, weights: torch.Tensor, split_sizes: torch.Tensor, + *, + grad_weight_out: torch.Tensor | None = None, ) -> torch.Tensor: ... @@ -33,7 +35,10 @@ def cpu_group_gemm( x: torch.Tensor, weights: torch.Tensor, split_sizes: torch.Tensor, + *, + grad_weight_out: torch.Tensor | None = None, ) -> torch.Tensor: + del grad_weight_out raise NotImplementedError("CPU GroupGemm is not implemented yet.") diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index 5beadb6d9d..9d0a0a8552 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -924,11 +924,9 @@ def fit(self): if self._async_hf_export: self._wait_for_pending_async_hf() - self._engine.model.destroy_async_hf_resources() if self._async_checkpoint: self._wait_for_pending_checkpoint() - self._engine.destroy_async_checkpoint_pg() # TODO: Should use flush rather than close if self._async_hf_export or self._async_checkpoint: @@ -937,7 +935,10 @@ def fit(self): if self._metrics_recorder: self._metrics_recorder.close() log_rank0.info(f"Training finished in {time.time() - train_begin:.2f} seconds") + # MoonEP destroy contains EP-group coordination. Enter it only after + # every rank has finished training and all async saves are quiescent. dist.barrier() + self._engine.close() def _prepare_model_input(self, data_batch) -> list[ModelItem]: seq_ctx_list: list[SequenceContext] = [] diff --git a/xtuner/v1/utils/dtensor.py b/xtuner/v1/utils/dtensor.py index f035b5fac1..b9d7462dbb 100644 --- a/xtuner/v1/utils/dtensor.py +++ b/xtuner/v1/utils/dtensor.py @@ -10,6 +10,8 @@ _has_foreach_support, ) +from .interleaved_shard import RuntimeLayout + def group_tensors_by_device_mesh_and_placements( tensors: list[DTensor], @@ -78,9 +80,9 @@ def cal_total_norm( if norm_type == 2: local_norm_squared = local_norm**2 for i, placement in enumerate(placements): - if isinstance(placement, Shard): - # FSDP's strided bookkeeping placement is a Shard subclass, so - # RuntimeLayout owns the only concrete private-type dependency. + if RuntimeLayout.is_sharded_placement(placement): + # FSDP's strided bookkeeping placement changes hierarchy across + # PyTorch versions; RuntimeLayout owns that private-type detail. dist.all_reduce(local_norm_squared, group=device_mesh.get_group(i)) elif isinstance(placement, Replicate): pass diff --git a/xtuner/v1/utils/interleaved_shard.py b/xtuner/v1/utils/interleaved_shard.py index 137bf75c88..78adfc24ec 100644 --- a/xtuner/v1/utils/interleaved_shard.py +++ b/xtuner/v1/utils/interleaved_shard.py @@ -29,12 +29,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import NamedTuple +from typing import NamedTuple, TypeGuard import torch import torch.distributed as dist from torch.distributed.tensor import DTensor, Shard -from torch.distributed.tensor.placement_types import _StridedShard +from torch.distributed.tensor.placement_types import Placement, _StridedShard __all__ = [ @@ -117,6 +117,11 @@ class RuntimeLayout: global_shape: tuple[int, ...] ordered_shards: tuple[RuntimeShard, ...] + @staticmethod + def is_sharded_placement(placement: Placement) -> TypeGuard[Shard | _StridedShard]: + """Hide PyTorch's version-dependent strided-shard hierarchy.""" + return isinstance(placement, (Shard, _StridedShard)) + @classmethod def from_dtensor(cls, tensor: DTensor) -> RuntimeLayout: mesh = tensor.device_mesh @@ -129,7 +134,7 @@ def from_dtensor(cls, tensor: DTensor) -> RuntimeLayout: chain_supported = True for mesh_dim in reversed(range(len(placements))): placement = placements[mesh_dim] - if not isinstance(placement, (Shard, _StridedShard)): + if not cls.is_sharded_placement(placement): continue order = tensor_dim_to_order.setdefault(placement.dim, []) split_factor = placement.split_factor if isinstance(placement, _StridedShard) else 1 @@ -153,7 +158,7 @@ def from_dtensor(cls, tensor: DTensor) -> RuntimeLayout: # placements first, then FSDP's bookkeeping placement. fsdp_prepended: list[tuple[int, int, int]] = [] for mesh_dim, placement in enumerate(placements): - if not isinstance(placement, (Shard, _StridedShard)): + if not cls.is_sharded_placement(placement): continue is_fsdp_prepended = _is_fsdp_prepended_strided(placement, mesh_dim) item = ( diff --git a/xtuner_moonep_acceptance.md b/xtuner_moonep_acceptance.md new file mode 100644 index 0000000000..d542af7963 --- /dev/null +++ b/xtuner_moonep_acceptance.md @@ -0,0 +1,157 @@ +# XTuner MoonEP 09-02 最终验收报告 + +## 结论 + +MoonEP 第一版接入通过最终验收。真实 Qwen3.5-35B-A3B 在 8×H200、BF16、FSDP2×EP4、TP1、`torch.compile` 下完成 MTP0 与 MTP1 各一组 DeepEP/MoonEP 20-step 对照训练。MoonEP 稳态吞吐分别为 DeepEP 的 `108.49%` 和 `112.09%`；全部 loss 与 grad-norm 曲线满足冻结门槛。 + +```mermaid +flowchart LR + A["同 checkpoint、seed、数据与配置"] --> D["DeepEP
MTP0 / MTP1 各 20 steps"] + A --> M["MoonEP
MTP0 / MTP1 各 20 steps"] + D --> P["steps 6-20
吞吐中位数"] + M --> P + D --> N["全部 20 steps
loss / grad norm"] + M --> N + P --> G["吞吐比 >= 95%"] + N --> G2["loss: cosine >= 0.99 且 rel < 3%
grad: cosine >= 0.98 且 rel < 5%"] + G --> R["PASS"] + G2 --> R +``` + +## 版本与配置 + +正式训练 manifest 记录 XTuner `6d5f4b82aa3a9337537ebd7941942c22312afd8a`；整理提交历史后的等价产品代码位于 `3f5414f8`，两者在 `xtuner/v1` 下没有差异。MoonEP-mod 为 `d4494473fb0932dcc35f3a44a0e3b31827f5e282`。 + +| 项目 | 验收值 | +| ----------------------- | --------------------------------------------------------------------------------- | +| Python / PyTorch / CUDA | `pt212_cu132` / `2.12.1+cu132` / `13.2` | +| 硬件 | 单节点 `8 × NVIDIA H200` | +| 模型 / 数据 | Qwen3.5-35B-A3B / Alpaca | +| 训练 | 20 steps、seed 0、global batch 8、pack length 65536 | +| 并行 | FSDP2、EP4、TP1、SP1、micro1、reshard after forward | +| dtype | BF16 param、forward/backward、EP 通信与 FSDP reduce-scatter；FP32 optimizer shard | +| MoE | 256 routed experts、1 shared expert、top-k 8、router FP32 | +| MoonEP | Direct VMM landing、staging reference 关闭、`moonep_num_sms=64` | +| MTP0 / MTP1 | `None` / `MTPConfig(num_layers=1, share_weights=False)` | +| grouped linear | Triton；未启用 `grouped_gemm` CUTLASS 路径 | +| optimizer | AdamW，lr `6e-5`，wd `0.01`，betas `(0.9, 0.95)` | + +完整配置、环境和 GPU 列表保存在 `work_dirs/moonep_0902_final_pass/*/acceptance_manifest.json`。四个 run 只切换 dispatcher 或 MTP 配置；checkpoint、seed、数据顺序、有效 token 数、packing、compile、optimizer 与 grouped-linear backend 保持一致。 + +## 20-step 结果 + +### 吞吐 + +吞吐只统计 compile/warm-up 后的 steps 6～20，单位为 text tokens/s。 + +| Gate | DeepEP median | MoonEP median | MoonEP / DeepEP | 判定 | +| ---- | ------------: | ------------: | --------------: | ---- | +| MTP0 | 5447.278 | 5909.833 | 108.491% | PASS | +| MTP1 | 3035.483 | 3402.410 | 112.088% | PASS | + +### 数值曲线 + +| Gate / curve | Cosine | Mean relative difference | 判定 | +| ------------------- | ----------: | -----------------------: | ---- | +| MTP0 balancing loss | 0.999998824 | 0.114532% | PASS | +| MTP0 LM loss | 0.999996055 | 0.269665% | PASS | +| MTP0 total loss | 0.999996339 | 0.258404% | PASS | +| MTP0 grad norm | 0.999883706 | 2.536973% | PASS | +| MTP1 balancing loss | 0.999999748 | 0.071091% | PASS | +| MTP1 LM loss | 0.999999130 | 0.121277% | PASS | +| MTP1 MTP loss | 0.999999091 | 0.105620% | PASS | +| MTP1 total loss | 0.999999269 | 0.110629% | PASS | +| MTP1 grad norm | 0.999980890 | 1.107751% | PASS | + +机器判定结果： + +- `work_dirs/moonep_0902_final_pass/mtp0_comparison.json` +- `work_dirs/moonep_0902_final_pass/mtp1_comparison.json` + +两份结果均为 `passed=true`。 + +## 关键语义与验证 + +### BF16 duplicate-gradient return + +正式路径由 MoonEP `Buffer.reduce_grad_bf16(...)` 归还 duplicate expert 梯度。两份 projection 共用 publication/reuse barrier；CuTe DSL kernel 从 BF16 VMM slot 读取，在 FP32 寄存器中执行 `SUM`，最后一次舍入写回 BF16 home gradient。原有 FP32 `Buffer.reduce_grad(...)` 保持不变。 + +```mermaid +flowchart LR + W["Grouped GEMM
BF16 local dW [2B]"] --> V["BF16 duplicate slots
[R, B]"] + V --> C["CuTe DSL reduce
FP32 register SUM"] + C --> H["BF16 home full grad"] + H --> F["FSDP BF16 reduce-scatter"] + F --> O["FP32 shard grad + weight
optimizer update"] +``` + +EP2/4/8 public API 测试覆盖空 slot、重复 expert、连续 slot 复用和精确参与集合；最终真实训练使用 EP4。FSDP2×EP4 tiny 测试验证 duplicate contribution 在 reduce-scatter 前落入完整 BF16 gradient。 + +### 生命周期与并发 + +Runtime 只拥有模型级 Buffer/workspace，Dispatcher 拥有层级行为，InvocationState 只保存单次 forward/backward 的 plan、event 与 gradient slot。权重预取在 `dispatch` 阶段发起，`dispatch_postprocess` 只等待 device event 并消费 alias。 + +真实路径已覆盖： + +- 同一 layer 多个 forward 后按原序和逆序 backward； +- Domino micro-batch=2； +- MTP reentrant original forward 与 replay； +- Direct VMM landing 的 validate-before-mutate、安装与卸载； +- 普通 AdamW 同步 DCP、同步/异步 HF export 与 router offload。 + +### 无 host sync + +profiler 对真实 router 与 MoonEP direct hot path 的门禁结果为： + +- router count 路径无 `cudaStreamSynchronize`； +- dispatch、weight prefetch、combine、duplicate-grad return 区间无 host/device/stream synchronize； +- 完整 home weight copy 为 `0`； +- 完整 local `[2B]` dW temporary/copy 为 `0`。 + +FSDP/NCCL 与 MoonEP 统一在 caller current stream 上形成一致的 device-side 顺序；没有用 host synchronize 修补跨 rank 依赖。 + +### 回归结果 + +| Suite | 结果 | +| --------------------------------------------------------- | ------------------------------------ | +| MoonEP-mod EP2/4/8、BF16 reduce、peer-skew 与连续调用 | PASS | +| XTuner acceptance/config contracts | `23 passed` | +| persistence/router/landing/workspace/GroupedGEMM 组合回归 | `38 passed` | +| pre-commit（ruff、format、mypy 等） | PASS；mypy 检查 367 个 source files | +| Qwen3.5 MTP0/MTP1 DeepEP/MoonEP | `4 × 20 steps`，两个 compare 均 PASS | + +Muon cold DCP resume、SwapAdamW cold DCP resume 与 process-mode async DCP 对应 `todo_bug_fix/0003～0005`。它们是已有通用基线问题，不影响本报告的普通 AdamW、`debug_skip_save=True` 正式训练，相关实现与回归测试留给独立 PR。 + +## 复现命令 + +```bash +tests/acceptance/run_qwen35_moonep_acceptance.sh deepep 0 65536 work_dirs/moonep_0902_final_pass +tests/acceptance/run_qwen35_moonep_acceptance.sh moonep 0 65536 work_dirs/moonep_0902_final_pass +tests/acceptance/run_qwen35_moonep_acceptance.sh deepep 1 65536 work_dirs/moonep_0902_final_pass +tests/acceptance/run_qwen35_moonep_acceptance.sh moonep 1 65536 work_dirs/moonep_0902_final_pass + +PYTHONPATH=. python -m xtuner._testing.moonep_acceptance compare \ + --deepep work_dirs/moonep_0902_final_pass/deepep_mtp0_pack65536 \ + --moonep work_dirs/moonep_0902_final_pass/moonep_mtp0_pack65536 \ + --output work_dirs/moonep_0902_final_pass/mtp0_comparison.json + +PYTHONPATH=. python -m xtuner._testing.moonep_acceptance compare \ + --deepep work_dirs/moonep_0902_final_pass/deepep_mtp1_pack65536 \ + --moonep work_dirs/moonep_0902_final_pass/moonep_mtp1_pack65536 \ + --output work_dirs/moonep_0902_final_pass/mtp1_comparison.json +``` + +所有 GPU 测试均先通过 `/mnt/shared-storage-user/zhaopenghao/github/xtuner/zdev/gpu_lock.sh` 获取 8 卡文件锁。 + +## 首版边界与待办 + +本报告只声明 BF16、TP1、单节点 EP2/4/8（正式性能测试 EP4）、FSDP2、MTP 与已覆盖的 Domino micro-batch 能力。以下内容继续作为后续待办： + +- 跨节点 MoonEP； +- XTuner 自管 expert 参数及其 DCP save/load； +- Muon、SwapAdamW 与 process-mode async DCP 的通用基线修复； +- EP+TP、FP8、pipeline parallel、FSDP `no_sync` 与 decoding。 + +## 最终结论 + +MoonEP dispatcher 已达到首版 Definition of Done：FSDP2×EP4 的 BF16 权重与梯度生命周期正确，MTP0/MTP1 真实 20-step 数值和性能门槛全部通过，关键热路径保持无 host sync；未实现能力已明确留在后续待办中。