From 1c6746fca88de7bd959f26b5784a11d746cd7ce2 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 9 Sep 2026 17:24:00 +0000 Subject: [PATCH 1/3] [Fix] Recognize PyTorch 2.12 FSDP strided shards in distributed grad norm On PyTorch 2.12 the FSDP2 + EP nested-shard bookkeeping placement is _StridedShard, which is no longer a Shard subclass, so cal_total_norm rejected it and grad-norm computation raised on FSDP2 + EP4. Route the placement check through RuntimeLayout.is_sharded_placement so both Shard and _StridedShard contribute their all-reduce. --- tests/utils/test_interleaved_shard.py | 42 +++++++++++++++++++++++++++ xtuner/v1/utils/dtensor.py | 8 +++-- xtuner/v1/utils/interleaved_shard.py | 13 ++++++--- 3 files changed, 56 insertions(+), 7 deletions(-) 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/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 = ( From d28a96906190d5a4d55d6f0f1175d9e89b842fe9 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 9 Sep 2026 17:24:00 +0000 Subject: [PATCH 2/3] [Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training Add an optional dispatcher="moonep" path for node-local BF16 MoE training. Native FSDP2 stays the sole owner of expert parameters, optimizer state, and checkpoint identity; MoonEP owns only the communication/VMM execution workspace. - Model-scoped MoonEP runtime adapted to the six-stage dispatcher API, with a versioned backend capability check and lazy import. - FSDP all-gathered BF16 expert weights land directly into MoonEP VMM aliases; the private VMM workspace owns the physical expert layout. - Dispatch activations and weights, reuse the existing grouped GEMM path (extended for dynamic expert gradients), and return duplicated expert gradients in BF16 before FSDP reduce-scatter, joined once per layer. - Support MTP (shared/unshared), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile. - Unify the EP dispatcher weight/routing contracts; keep routing counts on device via torch.histc. - Integrate DCP/HF persistence, activation/router offload, optimizer swap, Muon, and explicit runtime teardown. --- .../test_moe_train_engine_deepep_expert_tp.py | 1 + tests/engine/test_moe_train_engine_tpep.py | 1 + tests/engine/test_moonep_forward.py | 993 ++++++++++++++++ tests/engine/test_moonep_persistence.py | 398 +++++++ tests/model/test_fsdp_model.py | 17 + tests/module/dispatcher/test_agrs_all2all.py | 42 +- tests/module/dispatcher/test_deepep.py | 25 +- .../dispatcher/test_deepep_expert_tp.py | 2 + .../dispatcher/test_fsdp_vmm_landing.py | 81 ++ .../module/dispatcher/test_moonep_contract.py | 295 +++++ .../dispatcher/test_moonep_dispatcher.py | 436 +++++++ .../dispatcher/test_moonep_workspace.py | 119 ++ tests/module/dispatcher/test_noep.py | 1 + .../module/dispatcher/test_noep_expert_tp.py | 2 + tests/module/dispatcher/test_torch_all2all.py | 40 +- .../test_torch_all2all_shared_expert_tp.py | 1 + tests/module/test_grouped_linear.py | 48 + tests/module/test_router_counts.py | 109 ++ tests/ops/test_grouped_gemm_cutlass.py | 68 ++ tests/ops/test_grouped_gemm_out.py | 158 +++ tests/train/test_trainer.py | 5 + xtuner/v1/engine/train_engine.py | 60 +- xtuner/v1/float8/float8_gmm_tile_wise.py | 11 +- xtuner/v1/model/base.py | 10 + xtuner/v1/model/moe/moe.py | 64 +- .../module/decoder_layer/moe_decoder_layer.py | 61 +- xtuner/v1/module/dispatcher/__init__.py | 100 +- xtuner/v1/module/dispatcher/agrs.py | 11 +- xtuner/v1/module/dispatcher/base.py | 43 +- xtuner/v1/module/dispatcher/deepep.py | 11 +- .../v1/module/dispatcher/fsdp_vmm_landing.py | 239 ++++ xtuner/v1/module/dispatcher/moonep.py | 1045 +++++++++++++++++ .../v1/module/dispatcher/moonep_capability.py | 106 ++ .../v1/module/dispatcher/moonep_workspace.py | 539 +++++++++ xtuner/v1/module/dispatcher/torch_all2all.py | 11 +- .../module/grouped_linear/moe_group_linear.py | 17 +- xtuner/v1/module/router/greedy.py | 22 +- xtuner/v1/module/router/noaux_router.py | 17 +- xtuner/v1/module/router/protocol.py | 2 +- xtuner/v1/ops/moe/__init__.py | 7 +- xtuner/v1/ops/moe/cuda/group_gemm.py | 38 +- xtuner/v1/ops/moe/cuda/group_gemm_cutlass.py | 86 +- xtuner/v1/ops/moe/cuda/route_weight.py | 70 ++ .../ops/moe/cuda/triton_kernels/__init__.py | 8 +- .../cuda/triton_kernels/k_grouped_gemm_TMA.py | 30 +- .../k_grouped_gemm_TMA_triton3_4.py | 30 +- xtuner/v1/ops/moe/protocol.py | 5 + xtuner/v1/train/trainer.py | 5 +- xtuner/v1/utils/fsdp.py | 11 + 49 files changed, 5308 insertions(+), 193 deletions(-) create mode 100644 tests/engine/test_moonep_forward.py create mode 100644 tests/engine/test_moonep_persistence.py create mode 100644 tests/module/dispatcher/test_fsdp_vmm_landing.py create mode 100644 tests/module/dispatcher/test_moonep_contract.py create mode 100644 tests/module/dispatcher/test_moonep_dispatcher.py create mode 100644 tests/module/dispatcher/test_moonep_workspace.py create mode 100644 tests/module/test_router_counts.py create mode 100644 tests/ops/test_grouped_gemm_cutlass.py create mode 100644 tests/ops/test_grouped_gemm_out.py create mode 100644 xtuner/v1/module/dispatcher/fsdp_vmm_landing.py create mode 100644 xtuner/v1/module/dispatcher/moonep.py create mode 100644 xtuner/v1/module/dispatcher/moonep_capability.py create mode 100644 xtuner/v1/module/dispatcher/moonep_workspace.py create mode 100644 xtuner/v1/ops/moe/cuda/route_weight.py 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_forward.py b/tests/engine/test_moonep_forward.py new file mode 100644 index 0000000000..5106319d4d --- /dev/null +++ b/tests/engine/test_moonep_forward.py @@ -0,0 +1,993 @@ +import unittest + +import torch +import torch.distributed as dist +from torch.distributed.fsdp import FSDPModule +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): + def test_shared_home_micro1_micro2_micro4_accumulation_and_updates(self) -> None: + """Identical items isolate accumulation from routing/rounding changes. + + Compare eight synchronized backwards with four/two layer joins over + the same effective batch, through full recompute and three AdamW steps. + Different-input/backend parity remains covered by the training tests. + """ + self.create_pg("cuda") + reference = None + for width in (1, 2, 4): + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config("qwen", "moonep", compile=True), + optim_cfg=AdamWConfig(lr=6e-5, foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=1.0, torch_compile=True), + intra_layer_micro_batch=width, + ) + engine.init_model_weights() + observed = [] + for step in range(3): + engine.optimizer.zero_grad() + ids = (torch.arange(2, 18, device="cuda") + 16 * step).view(1, -1) % 256 + seq_ctxs = [SequenceContext.from_input_ids((ids,), device="cuda") for _ in range(8)] + # Normalize the entire effective batch, exactly as Trainer + # does, then retain RS on every accumulation backward. + loss_ctxs = engine.model.build_loss_ctx_batch( + [{"seq_ctx": seq_ctx, "shifted_labels": (ids + 1) % 256} for seq_ctx in seq_ctxs] + ) + items = [ + ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx) + for seq_ctx, loss_ctx in zip(seq_ctxs, loss_ctxs, strict=True) + ] + for start in range(0, 8, width): + engine.model.set_is_last_backward(start + width == 8) + engine.train_step(items[start : start + width]) + grad_norm = engine.clip_grad_norm(do_clip=False) + gradients = self._selected_training_tensors(engine, gradients=True) + # Empty local experts are valid; require routed contributions + # globally instead of requiring every EP owner to receive work. + routed_nonzero = torch.stack( + [torch.count_nonzero(value) for name, value in gradients.items() if ".experts." in name] + ).sum() + dist.all_reduce(routed_nonzero) + assert routed_nonzero > 0 + engine.step_optimizer(grad_norm) + observed.append((gradients, self._selected_training_tensors(engine, gradients=False))) + if reference is None: + reference = observed + else: + for step, (actual_step, expected_step) in enumerate(zip(observed, reference, strict=True)): + for kind, (actual, expected) in enumerate(zip(actual_step, expected_step, strict=True)): + assert actual.keys() == expected.keys() + for name in actual: + torch.testing.assert_close( + actual[name], + expected[name], + rtol=1e-2, + atol=1e-3, + msg=lambda message: f"width={width}, step={step}, kind={kind}, {name}: {message}", + ) + # Only clean up successful, rank-coordinated runs. A finally + # barrier would mask rank-local assertions from the test runner. + torch.cuda.synchronize() + dist.barrier() + engine.close() + + def test_micro2_accumulation_requires_sync_on_root_and_layers(self) -> None: + self.create_pg("cuda") + torch.manual_seed(20260805) + engine = TrainEngine( + model_cfg=_tiny_config("qwen", "moonep", compile=False), + optim_cfg=AdamWConfig(foreach=False), + fsdp_cfg=FSDPConfig(ep_size=4, recompute_ratio=1.0, torch_compile=False), + intra_layer_micro_batch=2, + ) + engine.init_model_weights() + try: + units = [module for module in engine.model.modules() if isinstance(module, FSDPModule)] + assert len(units) > 1 + for module in units: + for recurse in (False, True): + with self.assertRaisesRegex(ValueError, "requires gradient ReduceScatter on every backward"): + module.set_requires_gradient_sync(False, recurse=recurse) + module.set_requires_gradient_sync(True, recurse=recurse) + + # Reentrant recomputation and state retention still allow ordinary + # accumulation in sharded gradients after each synchronized RS. + first_gradients = {} + for last in (False, True): + engine.model.set_is_last_backward(last) + engine.train_step([self._model_training_item(engine, offset=offset) for offset in (0, 16)]) + torch.cuda.synchronize() + gradients = self._selected_training_tensors(engine, gradients=True) + if not last: + first_gradients = gradients + else: + assert gradients.keys() == first_gradients.keys() + for name, gradient in gradients.items(): + torch.testing.assert_close(gradient, 2 * first_gradients[name], rtol=1e-2, atol=1e-3) + finally: + torch.cuda.synchronize() + dist.barrier() + engine.close() + + @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) + } + parent = event.cpu_parent + gradient_staging = False + while parent is not None: + gradient_staging |= parent.name == "MoonEP::gradient_handoff" + parent = parent.cpu_parent + if event.name == "aten::copy_" and tensor_shapes & full_home_shapes and not gradient_staging: + full_weight_copies += 1 + # Scheme A deliberately allocates dW via this public op. + # Gradient suffix staging has the same shape as weights, + # so it must not be misclassified as a weight landing copy. + if event.name == "moe::k_grouped_gemm": + full_dw_materializations += 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/model/test_fsdp_model.py b/tests/model/test_fsdp_model.py index bcf67f76f9..2b65ff64b9 100644 --- a/tests/model/test_fsdp_model.py +++ b/tests/model/test_fsdp_model.py @@ -6,6 +6,7 @@ import torch.distributed as dist from safetensors import safe_open from torch import nn +from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DTensor from xtuner._testing.testcase import DeterministicDDPTestCase @@ -85,6 +86,22 @@ class TestFSDPModel(DeterministicDDPTestCase): def world_size(self) -> int: return 4 + def test_cannot_disable_gradient_sync(self): + self.create_pg("cuda") + config = ToyModelConfig(compile_cfg=False) + model = config.build().cuda() + model.fully_shard(FSDPConfig(torch_compile=False)) + + for recurse in (False, True): + with self.assertRaisesRegex(ValueError, "requires gradient ReduceScatter on every backward"): + model.set_requires_gradient_sync(False, recurse=recurse) + model.set_requires_gradient_sync(True, recurse=recurse) + + # The restriction belongs to XTuner instances, not PyTorch's class. + native = fully_shard(nn.Linear(4, 4).cuda()) + native.set_requires_gradient_sync(False) + native.set_requires_gradient_sync(True) + def test_model_forward_backward(self): self.create_pg("cuda") 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..6472b181b5 --- /dev/null +++ b/tests/module/dispatcher/test_moonep_contract.py @@ -0,0 +1,295 @@ +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_capability_check_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_capability import check_config + 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) + + config = _moe_config(dispatcher="moonep", ep_size=4, n_routed_experts=8) + if valid: + check_config(config) + else: + with pytest.raises(RuntimeError, match="grouped_gemm requires GROUPED_GEMM_USE_CUTLASS=1"): + check_config(config) + + +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..fa70017f1b --- /dev/null +++ b/tests/module/dispatcher/test_moonep_dispatcher.py @@ -0,0 +1,436 @@ +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, + _MoonEPExpertGradBridge, + _MoonEPLayerGradJoin, +) + + +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 generation_for(self, ordinal): + return ordinal % 2 + + def landing(self, generation): + return self._landings[generation] + + def prefetch_weights(self, *, buffer, plan, generation): + 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 + + def local_compute_view(self, *, hidden_nvsh, cu_seqlens): + b = self._landings[0][0].shape[0] + counts = torch.tensor([hidden_nvsh.shape[0]] + [0] * (2 * b - 1), dtype=torch.int32) + return hidden_nvsh, counts + + def return_expert_gradients(self, *, buffer, plan, gradients, grad_slot, initialize): + del buffer, plan + b = self._landings[0][0].shape[0] + targets = self._slots[grad_slot] + for target, gradient in zip(targets, gradients, strict=True): + if initialize: + target[:b].zero_() + target[:b].add_(gradient[:b]) + target[b:].copy_(gradient[b:]) + return targets[0][:b], targets[1][:b] + + 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( + moonep_integration._MoonEPResources, + "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, + ep_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, execution_order=["layers.0.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_inputs, layer_states = dispatcher.prepare_layer_inputs([hidden_states]) + layer_input, layer_state = layer_inputs[0], layer_states[0] + 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.resources.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, torch.Tensor) for weight in post["expert_weight_layout"].trainable_weights) + assert post["expert_weight_layout"].trainable_weights[0].shape == (4, 256, 128) + assert torch.equal(result["hidden_states"], hidden_states * 0.5) + assert not result["hidden_states"].requires_grad + assert runtime.resources.buffer.num_sms == 64 + + call_state_ref = weakref.ref(layer_state) + del layer_state, layer_states, pre, dispatched, post, pre_combined, combined + assert call_state_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_inputs, invalid_states = dispatcher.prepare_layer_inputs([torch.randn(4, 128, dtype=torch.bfloat16)]) + dispatcher.dispatch_preprocess( + hidden_states=invalid_inputs[0], + 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_states[0], + ) + + +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, + ep_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, execution_order=["layers.0.experts"]) + + assert _Workspace.allocated[-1].destroyed + + +def test_gradient_reduce_start_hands_dw_to_the_workspace_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 + + gradients = tuple(torch.arange(24, dtype=torch.bfloat16).view(4, 2, 3) + projection for projection in range(2)) + calls: list[tuple] = [] + event = _CompletionEvent() + + class _Workspace: + def return_expert_gradients(self, *, buffer, plan, gradients, grad_slot, initialize): + del buffer, plan + calls.append((gradients, grad_slot, initialize)) + return gradients[0][:2], gradients[1][:2] + + def enqueue(operation, inputs=()): + del inputs + return operation(), event + + resources = SimpleNamespace(workspace=_Workspace(), buffer=object(), enqueue=enqueue) + layer = moonep_integration._MoonEPLayer( + fqn="layers.0.experts", projections=(nn.Linear(3, 3), nn.Linear(3, 3)), ordinal=0 + ) + call_state = moonep_integration._MoonEPLayerCallState( + resources=resources, + layer=layer, + generation=0, + grad_slot=0, + layer_gradients=moonep_integration._MoonEPLayerGradients(), + ) + call_state.plan = object() + home_parameters = ( + nn.Parameter(torch.zeros_like(gradients[0][:2])), + nn.Parameter(torch.zeros_like(gradients[1][:2])), + ) + call_state.home_parameters = home_parameters + + moonep_integration.start_gradient_completion(call_state, gradients) + + assert len(calls) == 1 + assert calls[0][1] == 0 and calls[0][2] is True + assert call_state.layer_gradients.initialized is True + assert event.waits == 0 + + parameters, home_grads = moonep_integration.finish_gradient_completion(call_state) + assert event.waits == 1 + assert parameters is home_parameters + for parameter, actual, expected in zip(parameters, home_grads, gradients, strict=True): + assert parameter.grad is None # Only the layer Join publishes H. + torch.testing.assert_close(actual, expected[:2]) + assert call_state.gradient_completion is None + + +def test_gradient_reduce_start_and_join_preserve_device_order(monkeypatch) -> None: + from xtuner.v1.module.dispatcher import moonep as moonep_integration + + events: list[str] = [] + targets = (torch.zeros(4), torch.zeros(4)) + + def fake_start(call_state, gradients) -> None: + del call_state + for target, gradient in zip(targets, gradients, strict=True): + target.copy_(gradient) + events.append("start") + + def fake_finish(call_state): + del call_state + events.append("finish") + return (nn.Parameter(torch.zeros(4)), nn.Parameter(torch.zeros(4))), targets + + monkeypatch.setattr(moonep_integration, "start_gradient_completion", fake_start) + monkeypatch.setattr(moonep_integration, "finish_gradient_completion", fake_finish) + + class _WriteWGrad(torch.autograd.Function): + @staticmethod + def forward(ctx, value, weight, projection): + ctx.projection = projection + return value + + @staticmethod + def backward(ctx, grad): + events.append(f"projection-{ctx.projection}") + return grad, torch.full_like(grad, ctx.projection + 1), None + + call_state = object() + source = torch.ones(4, requires_grad=True) + + (joined,) = _MoonEPLayerGradJoin.apply((call_state,), source) + started, w0, w1 = _MoonEPExpertGradBridge.apply( + joined, nn.Parameter(torch.ones(4)), nn.Parameter(torch.ones(4)), call_state + ) + projection_0 = _WriteWGrad.apply(started, w0, 0) + projection_1 = _WriteWGrad.apply(projection_0, w1, 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..2ddf894070 --- /dev/null +++ b/tests/module/dispatcher/test_moonep_workspace.py @@ -0,0 +1,119 @@ +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) + hidden_nvsh, _, cu_seqlens, plan = buffer.dispatch( + hidden, + topk_experts_sk=topk_ids, + tokens_per_expert=tokens_per_expert, + ) + + local_weights = workspace.prefetch_weights(buffer=buffer, plan=plan, generation=0) + + # One grouped GEMM receives exactly one contiguous [B+B] segment + # plus its device counts; the home prefix aliases the current + # FSDP landing. + compute_hidden, local_counts = workspace.local_compute_view( + hidden_nvsh=hidden_nvsh, cu_seqlens=cu_seqlens + ) + assert local_counts.shape == (2 * experts_per_rank,) + assert compute_hidden.shape == hidden_nvsh.shape + assert int(local_counts.sum()) == hidden_nvsh.shape[0] + 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], + ) + + # ``return_expert_gradients`` owns the [2B] split, the shared + # home accumulator (zero-or-add), and the EP exact sum. The + # first producer initializes; a later slot adds into the same + # home prefix. + def dw(value: float) -> tuple[torch.Tensor, torch.Tensor]: + return tuple( + torch.full((2 * experts_per_rank, *shape), value, dtype=torch.bfloat16, device=device) + for shape in ((512, 1024), (1024, 512)) + ) + + home_first = workspace.return_expert_gradients( + buffer=buffer, plan=plan, gradients=dw(1.0), grad_slot=0, initialize=True + ) + home_second = workspace.return_expert_gradients( + buffer=buffer, plan=plan, gradients=dw(2.0), grad_slot=1, initialize=False + ) + for first, second in zip(home_first, home_second, strict=True): + assert first.shape[0] == experts_per_rank + assert torch.isfinite(first).all() + # Every slot's home prefix maps the one shared home chunk + # through its own VA; a write is visible through the other. + first.fill_(9.0) + assert torch.all(second == 9.0) + + copied = torch.count_nonzero(plan.experts_to_copy >= 0) + dist.all_reduce(copied, group=ep_group) + assert copied > 0 + hot_gradient = home_first[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..bb4390d908 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) + + assert output.shape == (4, 128) + assert torch.isfinite(output).all() + with pytest.raises(NotImplementedError, match="trainable weight override"): + layer(hidden_states, counts, trainable_weight=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/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..f5a6149e9b 100644 --- a/xtuner/v1/float8/float8_gmm_tile_wise.py +++ b/xtuner/v1/float8/float8_gmm_tile_wise.py @@ -346,7 +346,16 @@ 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, + ) -> torch.Tensor: + if trainable_weight is not None: + raise NotImplementedError("FP8 grouped linear does not support a trainable weight override yet") + 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..b8b5ad0d24 100644 --- a/xtuner/v1/model/base.py +++ b/xtuner/v1/model/base.py @@ -12,6 +12,7 @@ from math import prod from pathlib import Path from shutil import copy, copytree, rmtree +from types import MethodType from typing import Annotated, Any, Generator, Iterable, Literal, Mapping, NamedTuple, Sequence, cast import torch @@ -48,6 +49,7 @@ from xtuner.v1.module.rope import RopeParametersConfig, RopeScalingConfig from xtuner.v1.utils import get_device, get_logger, get_torch_device_module, log_rank0, profile_time_and_memory from xtuner.v1.utils.compile import MaybeCompile, is_compiled_function, maybe_compile +from xtuner.v1.utils.fsdp import set_requires_gradient_sync from xtuner.v1.utils.load_spec import ( HFSavePlan, LoadSpec, @@ -610,6 +612,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() @@ -728,6 +734,10 @@ def traverse(module): offload_policy=offload_policy, ignored_params=ignored_params if ignored_params else None, ) + # Apply the policy to every FSDP unit, including direct layer callers. + # Delayed RS retains full gradients for all layers and can also retain + # MoonEP's reusable VMM slots. Keep PyTorch's global class untouched. + setattr(target, "set_requires_gradient_sync", MethodType(set_requires_gradient_sync, target)) def save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16, safetensors_prefix: str = "model"): # Save may be called without `fully_shard`; refresh from the current runtime layout. diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 15dc0d628e..39de612b8e 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -73,6 +73,8 @@ MoEDecoderLayerOutput, MoEGate, ) +from xtuner.v1.module.dispatcher import EPExecutionRuntime, build_ep_execution_runtime +from xtuner.v1.module.dispatcher.moonep_capability import check_config, check_fsdp_policy from xtuner.v1.module.mtp import MTPBlock, MTPConfig, MTPLayer from xtuner.v1.utils import ( get_device, @@ -159,7 +161,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 @@ -208,7 +221,9 @@ class MoE(BaseModel): def __init__(self, config: MoEConfig): # Concrete MoE configs override build(), so validate dispatcher support - # at the shared model-construction boundary. + # at the shared model-construction boundary. MoonEP's config-only + # capability checks live in one entry point. + check_config(config) if config.dispatcher == "agrs": if config.expert_tp_size > 1: raise NotImplementedError("AGRS with ExpertTP is not supported") @@ -254,6 +269,11 @@ def __init__(self, config: MoEConfig): self.expert_tp_mesh = None self.ep_tp_mesh = None + # Optional model-scoped EP execution runtime (MoonEP today, a no-op + # Adapter otherwise). Its four lifecycle boundaries are called + # unconditionally below. + self._ep_runtime: EPExecutionRuntime = build_ep_execution_runtime(config, self.ep_mesh) + 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) @@ -522,6 +542,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") @@ -1156,6 +1183,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, + ep_runtime=self._ep_runtime, + layer_fqn=f"layers.{layer_idx}.experts", ) if self.config.freeze_routers: layers[str(layer_idx)].gate.requires_grad_(False) @@ -1223,6 +1252,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, + ep_runtime=self._ep_runtime, + layer_fqn=f"mtp_block.layers.{i}.decoder_layer.experts", ) # Wrap decoder layer in MTPLayer @@ -1263,6 +1294,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") + check_fsdp_policy(self.config, fsdp_config) + self._ep_runtime.validate_before_fsdp(fsdp_config) self.fsdp_config = fsdp_config assert self.fsdp_config.ep_size == self.config.ep_size @@ -1410,8 +1443,31 @@ def fully_shard( self._init_load_spec() self._to_empty_meta() + self._ep_runtime.install_after_fsdp( + fsdp_root=self, + execution_order=self.expert_bearing_layers_in_execution_order(), + ) return self + def expert_bearing_layers_in_execution_order(self) -> list[str]: + """FQNs of the routed-expert modules in FSDP execution order. + + This is the single source the MoonEP install cross-checks against its + construction-order registration: main decoder layers first (skipping + the dense prefix), then the MTP physical layers. It is derived from the + module structure, not a stored list, so a drift from construction order + (or a single-layer + MTP model) fails loudly here. + """ + order = [ + f"layers.{layer_idx}.experts" + for layer_idx, layer in self.layers.items() + if isinstance(getattr(layer, "_checkpoint_wrapped_module", layer), MoEDecoderLayer) + ] + if self.mtp_block is not None and self.config.mtp_config is not None: + num_physical = 1 if self.config.mtp_config.share_weights else self.config.mtp_config.num_layers + order += [f"mtp_block.layers.{i}.decoder_layer.experts" for i in range(num_physical)] + return order + @property @override def default_compile_cfg(self) -> dict[str, TorchCompileOption]: @@ -1419,6 +1475,10 @@ def default_compile_cfg(self) -> dict[str, TorchCompileOption]: return MOE_EP_COMPILE_CFG return MOE_NON_EP_COMPILE_CFG + def close_ep_runtime(self) -> None: + """Release optional dynamic-EP resources before PG teardown.""" + self._ep_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 64dbb2f37a..7566a38f1f 100644 --- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py +++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py @@ -28,6 +28,8 @@ from xtuner.v1.module.dispatcher import ( CombineResult, DispatchResult, + EPExecutionRuntime, + ExpertWeightLayout, PostDispatchResult, PreCombineResult, PreDispatchResult, @@ -215,10 +217,17 @@ 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) + gate_up_out = self.fused_w1w3(x, tokens_per_expert, trainable_weight=trainable[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]) return res @@ -251,10 +260,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, + ep_runtime: EPExecutionRuntime | None = None, + layer_fqn: str | None = None, ): super().__init__() self.ep_mesh = ep_mesh @@ -319,14 +330,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", + ep_runtime=ep_runtime, + layer_fqn=layer_fqn, + projections=(self.experts.fused_w1w3, self.experts.fused_w2), ) def forward( @@ -426,6 +441,10 @@ def _forward( position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_kwargs: dict[str, object] | None = None, ) -> MoEDecoderLayerOutput: + # MoonEP uses this identity seam to place Join before attention and + # carries its opaque invocation token into dispatch phase 1. + layer_inputs, layer_states = self.dispatcher.prepare_layer_inputs([hidden_states]) + hidden_states, layer_state = layer_inputs[0], layer_states[0] residual, hidden_states, router_results, attn_outputs = self._pre_moe_forward( hidden_states=hidden_states, seq_ctx=seq_ctx, @@ -444,6 +463,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, @@ -469,7 +490,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, @@ -482,6 +503,7 @@ def _forward( pre_dispatched=pre_dispatched, dispatched=dispatched, post_dispatched=post_dispatched, + async_op=self._async_combine, decoding=False, ) @@ -490,14 +512,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) @@ -507,11 +539,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, @@ -561,14 +588,20 @@ def _micro_batch_forward( dispatched_list: list[DispatchResult] = [] pre_moe_forward_out_list: list[torch.Tensor] = [] + # A single multi-output Join must precede every branch (including + # residuals), so shared expert gradients reach FSDP exactly once. + hidden_states_list, layer_states = self.dispatcher.prepare_layer_inputs(hidden_states_list) + # Attention + gate + pre-dispatch for ( hidden_states, + layer_state, attention_kwargs, seq_ctx, position_embeddings, ) in zip( hidden_states_list, + layer_states, attention_kwargs_list, seq_ctx_list, position_embeddings_list, @@ -586,6 +619,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) @@ -620,7 +655,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..3711f55d42 100644 --- a/xtuner/v1/module/dispatcher/__init__.py +++ b/xtuner/v1/module/dispatcher/__init__.py @@ -1,10 +1,14 @@ +from __future__ import annotations + import os -from typing import Literal +from typing import Any, Literal, Protocol XTUNER_DISPATCHER_DEBUG = os.getenv("XTUNER_DISPATCHER_DEBUG", "0") == "1" import torch.distributed as dist +from torch import nn +from torch.distributed.device_mesh import DeviceMesh from xtuner.v1.utils import get_logger, log_rank0 @@ -13,6 +17,7 @@ CombineResult, DispacherInterface, DispatchResult, + ExpertWeightLayout, NaiveDispatcher, PostCombineResult, PostDispatchResult, @@ -25,17 +30,88 @@ logger = get_logger() +class EPExecutionRuntime(Protocol): + """Model-scoped EP execution lifecycle: four boundaries ``MoE`` calls + unconditionally, with no shared implementation (MoonEP, or the no-op + Adapter below). + + ``bind_layer`` returns this backend's per-layer Dispatcher, or ``None`` to + let ``build_dispatcher`` fall back to a generic Adapter. + """ + + def bind_layer(self, *, layer_fqn: str, projections: tuple[nn.Module, nn.Module]) -> Any: ... + + def validate_before_fsdp(self, fsdp_config: object) -> None: ... + + def install_after_fsdp(self, *, fsdp_root: nn.Module, execution_order: list[str]) -> None: ... + + def close(self) -> None: ... + + +class NoEPExecutionRuntime: + """The "no model-scoped EP execution runtime" Adapter (not ``None``). + + Every EP-runtime lifecycle boundary on ``MoE`` is called unconditionally; + for a backend without one, each boundary is a no-op and ``bind_layer`` + returns ``None`` so ``build_dispatcher`` falls back to a generic Adapter. + """ + + def bind_layer(self, *, layer_fqn: str, projections: tuple[nn.Module, nn.Module]) -> None: + del layer_fqn, projections + return None + + def validate_before_fsdp(self, fsdp_config: object) -> None: + del fsdp_config + + def install_after_fsdp(self, *, fsdp_root: nn.Module, execution_order: list[str]) -> None: + del fsdp_root, execution_order + + def close(self) -> None: + return + + +def build_ep_execution_runtime(config: Any, ep_mesh: DeviceMesh | None) -> EPExecutionRuntime: + """Single place a model-scoped EP execution backend is selected. + + ``config`` is a ``MoEConfig``; it is typed loosely to avoid a model-layer + import cycle. + """ + if config.dispatcher == "moonep": + from .moonep import MoonEPModelRuntime + + assert ep_mesh is not None + return MoonEPModelRuntime( + ep_group=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, + ) + return NoEPExecutionRuntime() + + # 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", + *, + ep_runtime: EPExecutionRuntime | None = None, + layer_fqn: str | None = None, + projections: tuple[nn.Module, nn.Module] | None = None, ) -> DispacherInterface: + # A model-scoped EP execution runtime binds its own per-layer Dispatcher. + if ep_runtime is not None and layer_fqn is not None and projections is not None: + bound = ep_runtime.bind_layer(layer_fqn=layer_fqn, projections=projections) + if bound is not None: + return bound # 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 +119,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 +146,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 +153,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__ = [ @@ -101,9 +171,13 @@ def build_dispatcher( "NaiveDispatcher", "TorchAll2AllDispatcher", "MoEAGRSDispatcher", + "EPExecutionRuntime", + "NoEPExecutionRuntime", "build_dispatcher", + "build_ep_execution_runtime", "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..e5e144aa8d 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,18 @@ HiddenStates: TypeAlias = torch.Tensor +ProjectionPair: TypeAlias = tuple[torch.Tensor, torch.Tensor] + + +class ExpertWeightLayout(NamedTuple): + """Call-local expert weight ownership at the dispatcher/MLP seam. + + A dynamic-EP backend may hand ``MoEBlock`` a call-local weight alias whose + dW still returns through autograd. Direct-output WGrad and external + (two-segment) storage are not part of the first-version contract. + """ + + trainable_weights: ProjectionPair | None = None def _get_backward_pre_hook(backward_previous_event: torch.cuda.Event): @@ -59,6 +71,7 @@ class PostDispatchResult(TypedDict): # TODO: hidden_states: torch.Tensor tokens_per_expert: torch.Tensor + expert_weight_layout: ExpertWeightLayout class PreCombineResult(TypedDict): @@ -102,13 +115,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_inputs( + self, + layer_inputs: list[torch.Tensor], + ) -> tuple[list[torch.Tensor], list[object | None]]: + """Prepare all inputs of one FSDP layer call before branching. + + 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_inputs, [None] * len(layer_inputs) @abstractmethod def dispatch( @@ -136,6 +157,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 +262,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 +280,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 +433,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..eb86bad673 --- /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 + + # MoonEP's layer Join publishes one completed home sum per FSDP + # call. Assignment retains its VMM alias until native copy-in; + # it does not allocate persistent full-model gradient storage. + 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..cf8be807b5 --- /dev/null +++ b/xtuner/v1/module/dispatcher/moonep.py @@ -0,0 +1,1045 @@ +"""MoonEP's model-scoped XTuner integration. + +The backend import remains lazy so unrelated dispatchers do not require +MoonEP. ``MoonEPModelRuntime`` owns model resources and ``MoonEPDispatcher`` +owns one routed layer's static policy. One dispatch/combine call is a pure +data ``_MoonEPLayerCallState`` record advanced by the module-level transaction +functions (``dispatch_forward``, ``prepare_experts``, ``combine_forward``, +``combine_backward``, gradient completion, and the layer-Join handoff). The +private VMM workspace remains the deep module for physical expert layout. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +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 + + +@dataclass(frozen=True) +class _MoonEPLayer: + """One physical routed layer's identity, stored once. + + ``ordinal`` is the FSDP execution-order position; the home generation is + issued from it by the workspace, so it is not stored here. + """ + + fqn: str + projections: tuple[nn.Module, nn.Module] + ordinal: int + + +@dataclass(frozen=True) +class _MoonEPResources: + """The only non-``None`` state after ``install_after_fsdp``. + + The call state borrows this record instead of the runtime; ``buffer_for`` + and ``enqueue`` are its two methods. ``_buffer_box`` is a one-slot mutable + box holding ``(Buffer, Fixed-S)`` because the activation Buffer is still + built lazily on the first forward once Fixed-S is known. + """ + + workspace: _ExpertVMMWorkspace + landing: ExpertLandingAdapter + comm_stream: torch.cuda.Stream + ep_group: dist.ProcessGroup + num_experts: int + experts_per_rank: int + top_k: int + hidden_size: int + gradient_slots: int + num_sms: int + _buffer_box: list + + def buffer_for(self, tokens_per_rank: int) -> Any: + """Dispatch entry: build the Fixed-S Buffer once, then check S.""" + if not self._buffer_box: + assert _moonep_backend is not None + 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, + ) + self._buffer_box.append((buffer, tokens_per_rank)) + buffer, fixed_s = self._buffer_box[0] + if tokens_per_rank != fixed_s: + raise RuntimeError(f"MoonEP fixed S changed: {fixed_s} -> {tokens_per_rank}") + return buffer + + @property + def buffer(self) -> Any: + """Combine/backward entry: the Buffer this call's dispatch built.""" + if not self._buffer_box: + raise RuntimeError("MoonEP Buffer must be created by dispatch first") + return self._buffer_box[0][0] + + def expect_tokens_per_rank(self, tokens_per_rank: int) -> None: + """Reject a changed Fixed-S before the first Dispatcher/VMM op.""" + if self._buffer_box and tokens_per_rank != self._buffer_box[0][1]: + raise RuntimeError(f"MoonEP fixed S changed: {self._buffer_box[0][1]} -> {tokens_per_rank}") + + def enqueue( + self, + operation: Callable[[], Any], + *, + inputs: tuple[torch.Tensor | None, ...] = (), + ) -> tuple[Any, torch.cuda.Event]: + """Run one MoonEP transaction on XTuner's stream and return ``(result, + done event)``.""" + caller_stream = torch.cuda.current_stream() + self.comm_stream.wait_event(caller_stream.record_event()) + for tensor in inputs: + if tensor is not None: + tensor.record_stream(self.comm_stream) + with torch.cuda.stream(self.comm_stream): + result = operation() + done = self.comm_stream.record_event() + return result, done + + def home_generation(self, layer: _MoonEPLayer) -> int: + """The workspace owns the physical chunks and issues the generation.""" + return self.workspace.generation_for(layer.ordinal) + + +class ExpertLandingAdapter: + """Make one generation's home weights ready and hand back its Parameters. + + ``DirectVMMLanding`` is the production path (FSDP unpacks straight into VMM + home rows); ``StagingReferenceLanding`` is the numerical-reference path + (FSDP lands normal storage, then copies into the VMM home rows). Both are + the two Adapters of one Seam, so the runtime and transaction functions + only call ``prepare`` and never branch on a ``staging_reference`` flag. + """ + + def install( + self, *, fsdp_root: nn.Module, workspace: _ExpertVMMWorkspace, layers: tuple[_MoonEPLayer, ...] + ) -> None: + raise NotImplementedError + + def prepare(self, *, layer: _MoonEPLayer, generation: int) -> tuple[nn.Parameter, nn.Parameter]: + raise NotImplementedError + + def uninstall(self) -> None: + raise NotImplementedError + + +class DirectVMMLanding(ExpertLandingAdapter): + """FSDP's final per-parameter unpack lands directly in the VMM home + rows.""" + + def __init__(self) -> None: + self._fsdp_params: tuple[Any, ...] = () + + @override + def install( + self, *, fsdp_root: nn.Module, workspace: _ExpertVMMWorkspace, layers: tuple[_MoonEPLayer, ...] + ) -> None: + self._fsdp_params = install_fsdp_vmm_landing( + fsdp_root=fsdp_root, + targets=tuple( + (layer.fqn, layer.projections, workspace.landing(workspace.generation_for(layer.ordinal))) + for layer in layers + ), + ) + + @override + def prepare(self, *, layer: _MoonEPLayer, generation: int) -> tuple[nn.Parameter, nn.Parameter]: + # "Ready" is one check: FSDP has already materialized the weight in + # the VMM landing, so there is no copy. + del generation + return fsdp_current_unsharded_expert_parameters(layer.projections) + + @override + def uninstall(self) -> None: + if self._fsdp_params: + uninstall_fsdp_vmm_landing(self._fsdp_params) + self._fsdp_params = () + + +class StagingReferenceLanding(ExpertLandingAdapter): + """FSDP lands normal storage first, then this copies it into the VMM home + rows.""" + + def __init__(self) -> None: + 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." + ) + self._workspace: _ExpertVMMWorkspace | None = None + + @override + def install( + self, *, fsdp_root: nn.Module, workspace: _ExpertVMMWorkspace, layers: tuple[_MoonEPLayer, ...] + ) -> None: + # No FSDP binding is installed: that is what distinguishes the two + # Adapters. The copy happens in ``prepare``. + del fsdp_root, layers + self._workspace = workspace + + @override + def prepare(self, *, layer: _MoonEPLayer, generation: int) -> tuple[nn.Parameter, nn.Parameter]: + assert self._workspace is not None + parameters: list[nn.Parameter] = [] + for linear, landing in zip(layer.projections, self._workspace.landing(generation), strict=True): + weight = cast(torch.Tensor, linear.weight) + if not isinstance(weight, nn.Parameter): + raise RuntimeError(f"{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"{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] + + @override + def uninstall(self) -> None: + self._workspace = None + + +def build_landing_adapter(staging_reference: bool) -> ExpertLandingAdapter: + if staging_reference: + return StagingReferenceLanding() + return DirectVMMLanding() + + +class MoonEPModelRuntime: + """Own the lifecycle and ordered layer registry for one model/EP group. + + Construction takes no CUDA resources. ``build_dispatcher`` registers one + physical routed layer per call in construction order; ``install_after_fsdp`` + allocates the workspace, cross-checks registration order against FSDP + execution order, and installs the landing Adapter. + """ + + 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: + # Config-level capability validation (backend version, EP geometry, + # dtype, grouped-GEMM backend, ...) lives in ``moonep_capability`` and + # runs at meta model build. Keep only the optional-backend gate here so + # direct construction still fails fast. + require_moonep_backend() + + 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._num_sms = num_sms + self._gradient_slots = intra_layer_micro_batch + self._landing = build_landing_adapter(staging_reference) + + # Physical routed layers in registration (construction) order. + self._layers: list[_MoonEPLayer] = [] + self._resources: _MoonEPResources | None = None + self._closed = False + + def bind_layer( + self, + *, + layer_fqn: str, + projections: tuple[nn.Module, nn.Module], + ) -> MoonEPDispatcher: + """Register one physical routed layer and return its Dispatcher.""" + if any(layer.fqn == layer_fqn for layer in self._layers): + raise ValueError(f"duplicate MoonEP routed layer: {layer_fqn}") + layer = _MoonEPLayer(fqn=layer_fqn, projections=projections, ordinal=len(self._layers)) + self._layers.append(layer) + return MoonEPDispatcher(runtime=self, layer=layer) + + def validate_before_fsdp(self, fsdp_config: Any) -> None: + # The build-time FSDP policy checks moved to ``moonep_capability``. + # This boundary stays because the Protocol needs it and a future + # backend may have its own FSDP preconditions. + del fsdp_config + + def install_after_fsdp(self, *, fsdp_root: nn.Module, execution_order: list[str]) -> None: + """Allocate execution resources after native FSDP has been + installed.""" + if self._resources 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") + + # Registration order vs FSDP execution order, checked once in the only + # place that can see both. ``moe.py`` hands over the ordered list + # rather than an adapter reading FSDP private structure. + registered = [layer.fqn for layer in self._layers] + if registered != execution_order: + raise RuntimeError( + f"MoonEP registration order does not match FSDP execution order: {registered} != {execution_order}" + ) + + 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._gradient_slots, + home_generations=2, + ) + # 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() + try: + self._landing.install(fsdp_root=fsdp_root, workspace=workspace, layers=tuple(self._layers)) + except Exception: + workspace.destroy() + raise + self._resources = _MoonEPResources( + workspace=workspace, + landing=self._landing, + comm_stream=comm_stream, + ep_group=self._ep_group, + num_experts=self._num_experts, + experts_per_rank=self._num_experts // self._ep_group.size(), + top_k=self._top_k, + hidden_size=self._hidden_size, + gradient_slots=self._gradient_slots, + num_sms=self._num_sms, + _buffer_box=[], + ) + + @property + def resources(self) -> _MoonEPResources: + """One place decides "is MoonEP installed".""" + if self._closed: + raise RuntimeError("MoonEP runtime was closed") + if self._resources is None: + raise RuntimeError("MoonEP FSDP resources must be installed before forward") + return self._resources + + def close(self) -> None: + """Release Buffer before VMM workspace at a coordinated boundary.""" + if self._closed: + return + if self._resources is not None: + resources = self._resources + resources.comm_stream.synchronize() + for buffer, _ in resources._buffer_box: + buffer.destroy() + resources.landing.uninstall() + resources.workspace.destroy() + self._resources = 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 call state is opaque control state for the remaining five stages; + # it never crosses into the compiled tensor-only expert block. + _moonep_call: _MoonEPLayerCallState + + +class MoonEPDispatchResult(TypedDict): + """Stage 2: global dispatch outputs. + + The call state travels on ``MoonEPPreDispatchResult`` only; every later + stage already receives ``pre_dispatched`` and reads ``_moonep_call`` there. + """ + + 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 + + +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 the + projection-paired ``[2B, O_p, I_p]`` call-local weight aliases. + """ + + +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]. + + +@dataclass +class _MoonEPLayerGradients: + """One FSDP call's initialization state, never shared across replay calls. + + Storage belongs to the workspace. Only the first backward producer clears it; forward/checkpoint replay must not + touch a different call's live H. + """ + + initialized: bool = False + + +@dataclass(eq=False) +class _MoonEPLayerCallState: + """One routed-layer call's pure lifecycle state for the transaction + functions. + + The record borrows the installed ``_MoonEPResources`` and the ``_MoonEPLayer`` + identity, plus the call-local plan/event/weight/gradient handles that the + module-level transaction functions read and advance. It owns no behavior + and never references ``MoonEPDispatcher``. Identity, not field equality, + distinguishes two calls, so instances stay hashable by ``id``. + """ + + resources: _MoonEPResources + layer: _MoonEPLayer + generation: int + grad_slot: int + layer_gradients: _MoonEPLayerGradients + + # One MoonEP communication plan and its device-side dependency chain. Each + # event is recorded once its named producer has been enqueued. + plan: Any | None = None + dispatch_done: Any | None = None + weights_ready: Any | None = None + combine_done: Any | None = None + + # Borrowed local [2B, O_p, I_p] weight aliases for this call. + local_weights: ProjectionPair | None = None + + # Current FSDP unsharded home Parameters [B, O_p, I_p] receive the returned + # BF16 home gradients after both local projections complete. + home_parameters: tuple[nn.Parameter, nn.Parameter] | None = None + # Completed home views and the event covering the pair reduction. + gradient_completion: tuple[ProjectionPair, Any] | None = None + + +# --- Transaction functions -------------------------------------------------- +# +# Each function covers one complete device-side sequence for a single call and +# advances ``_MoonEPLayerCallState`` in place. The autograd Functions and the +# Dispatcher are the only callers; the pair backward must reuse the forward +# plan and gradient slot recorded on the call state. + + +def finish_combine(state: _MoonEPLayerCallState, combined: torch.Tensor, *, async_op: bool) -> torch.Tensor: + """Establish the final device dependency for an async combine.""" + if async_op: + assert state.combine_done is not None + state.combine_done.wait() + return combined + + +def prepare_experts(state: _MoonEPLayerCallState, dispatched: MoonEPDispatchResult) -> MoonEPPostDispatchResult: + """Wait at the first weight consumer and expose the tensor-only layout.""" + resources = state.resources + assert state.weights_ready is not None + + with torch.profiler.record_function("MoonEP::prepare_experts"): + # This inserts a device dependency; it never waits on the host. + state.weights_ready.wait() + # Both halves of the ``[E+B] -> [2B]`` contract now live in the + # workspace: the caller cannot receive counts it must still fix up. + hidden_states, local_counts = resources.workspace.local_compute_view( + hidden_nvsh=dispatched["hidden_states"], + cu_seqlens=dispatched["cu_seqlens"], + ) + + local_weights = state.local_weights + state.local_weights = None + assert local_weights is not None + # Join activation and both weight edges so staging precedes upstream + # activation backward; a weight-only hook cannot establish that dependency. + # The bridge already makes the aliases require grad, so grouped GEMM + # returns their dW without a leaf ``nn.Parameter`` wrapper. + hidden_states, w13, w2 = _MoonEPExpertGradBridge.apply(hidden_states, local_weights[0], local_weights[1], state) + return MoonEPPostDispatchResult( + hidden_states=hidden_states, + tokens_per_expert=local_counts, + expert_weight_layout=ExpertWeightLayout( + trainable_weights=(w13, w2), + ), + ) + + +def dispatch_forward( + state: _MoonEPLayerCallState, + 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]: + """Dispatch on a fresh plan and start both projection weight prefetches.""" + resources = state.resources + buffer = resources.buffer_for(source_hidden.shape[0]) + + def dispatch_and_prefetch() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # Landing.prepare makes this generation's home weights ready: a check + # for the direct Adapter, the staging copy for the reference Adapter. + # Its call position (before dispatch's device barrier) is part of this + # sequence, not an implicit precondition of a mode branch. + state.home_parameters = resources.landing.prepare(layer=state.layer, generation=state.generation) + 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 + state.plan = plan + state.dispatch_done = torch.cuda.current_stream().record_event() + # Only local weights are returned now; the gradient slot views are + # created inside ``return_expert_gradients`` when they are needed. + state.local_weights = resources.workspace.prefetch_weights( + buffer=buffer, + plan=plan, + generation=state.generation, + ) + return hidden_nvsh, route_weights_nvs, cu_seqlens + + with torch.profiler.record_function("MoonEP::dispatch_forward"): + result, state.weights_ready = resources.enqueue( + dispatch_and_prefetch, + inputs=(source_hidden, source_route_weights, topk_ids, tokens_per_expert), + ) + assert state.dispatch_done is not None + if not async_op: + state.dispatch_done.wait() + return result + + +def dispatch_backward( + state: _MoonEPLayerCallState, + grad_hidden_nvsh: torch.Tensor, + grad_route_weights_nvs: torch.Tensor, +) -> ProjectionPair: + """Combine activation and route-weight gradients on the saved plan.""" + resources = state.resources + buffer = resources.buffer + grad_hidden_nvsh = grad_hidden_nvsh.contiguous() + grad_route_weights_nvs = grad_route_weights_nvs.contiguous() + + def combine_gradients() -> tuple[torch.Tensor, torch.Tensor]: + grad_hidden, grad_route_weights, no_event = buffer.combine( + plan=state.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 = resources.enqueue( + combine_gradients, + inputs=(grad_hidden_nvsh, grad_route_weights_nvs), + ) + done.wait() + return result + + +def combine_forward( + state: _MoonEPLayerCallState, + expert_output: torch.Tensor, + route_weights: torch.Tensor, + *, + async_op: bool, +) -> torch.Tensor: + """Fuse route scaling into the combine boundary on the saved plan.""" + resources = state.resources + buffer = resources.buffer + + def combine_output() -> torch.Tensor: + output, gathered_weights, no_event = buffer.combine( + plan=state.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, state.combine_done = resources.enqueue( + combine_output, + inputs=(expert_output, route_weights), + ) + if not async_op: + state.combine_done.wait() + return output + + +def combine_backward(state: _MoonEPLayerCallState, grad_output: torch.Tensor) -> tuple[torch.Tensor, Any]: + """Replay duplicated weights on the saved plan and return weighted grad.""" + resources = state.resources + buffer = resources.buffer + grad_output = grad_output.contiguous() + + def dispatch_gradient_and_prefetch() -> tuple[torch.Tensor, torch.cuda.Event]: + # FSDP pre-backward has restored this generation; same Adapter call, + # same sequence. + replay_home_parameters = resources.landing.prepare(layer=state.layer, generation=state.generation) + if state.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, state.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=state.plan, + async_finish=False, + zero_copy=False, + ) + assert no_weights is None and no_cu is None and reused_plan is state.plan + gradient_dispatch_done = torch.cuda.current_stream().record_event() + resources.workspace.prefetch_weights(buffer=buffer, plan=state.plan, generation=state.generation) + return grad_weighted, gradient_dispatch_done + + with torch.profiler.record_function("MoonEP::combine_backward"): + (grad_weighted, gradient_dispatch_done), replay_done = resources.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(state: _MoonEPLayerCallState, gradients: ProjectionPair) -> None: + """Hand the allocation-return dW to the workspace for the home return.""" + if state.gradient_completion is not None: + raise RuntimeError("MoonEP gradient completion was started twice") + resources = state.resources + + with torch.profiler.record_function("MoonEP::gradient_handoff"): + # The workspace owns the ``B`` split, the home-prefix zero-or-add, the + # duplicate-suffix copy, and the EP exact-sum reduction. ``initialize`` + # is the call-local flag the Dispatcher owns (ADR-0027). + home_grads, done = resources.enqueue( + lambda: resources.workspace.return_expert_gradients( + buffer=resources.buffer, + plan=state.plan, + gradients=gradients, + grad_slot=state.grad_slot, + initialize=not state.layer_gradients.initialized, + ), + inputs=gradients, + ) + state.layer_gradients.initialized = True + state.gradient_completion = (home_grads, done) + + +def finish_gradient_completion( + state: _MoonEPLayerCallState, +) -> tuple[tuple[nn.Parameter, nn.Parameter], ProjectionPair]: + """Wait on the device event; the layer Join owns the single handoff.""" + completion = state.gradient_completion + if completion is None: + raise RuntimeError("MoonEP gradient completion was not started") + home_grads, done = completion + done.wait() + + if state.home_parameters is None: + raise RuntimeError("MoonEP gradient completion has no home Parameters") + home_parameters = state.home_parameters + + state.home_parameters = None + state.gradient_completion = None + return home_parameters, home_grads + + +class _MoonEPExpertGradBridge(torch.autograd.Function): + """Consume both dWs before releasing the expert activation gradient.""" + + @staticmethod + def forward( + ctx: Any, + hidden_states: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + call_state: _MoonEPLayerCallState, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ctx.call_state = call_state + return hidden_states, w13, w2 + + @staticmethod + def backward( + ctx: Any, grad_hidden: torch.Tensor, dw13: torch.Tensor, dw2: torch.Tensor + ) -> tuple[torch.Tensor, None, None, None]: + start_gradient_completion(cast(_MoonEPLayerCallState, ctx.call_state), (dw13, dw2)) + # dW is now owned by MoonEP; do not also accumulate it on anchor leaves. + return grad_hidden, None, None, None + + +class _MoonEPLayerGradJoin(torch.autograd.Function): + """Join every microbatch before the native FSDP input backward hook.""" + + @staticmethod + def forward( + ctx: Any, + call_states: tuple[_MoonEPLayerCallState, ...], + *layer_inputs: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + ctx.call_states = call_states + return layer_inputs + + @staticmethod + def backward(ctx: Any, *grad_inputs: 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. Every call state's + # returned views alias the same home H, so completing all of them and + # publishing the last pair once is correct; publishing each would count + # the whole sum repeatedly. Native FSDP consumes it once via copy-in + # before upstream calls may reuse H; RS itself may remain asynchronous. + home_parameters: tuple[nn.Parameter, nn.Parameter] | None = None + home_grads: ProjectionPair | None = None + for call_state in ctx.call_states: + home_parameters, home_grads = finish_gradient_completion(call_state) + assert home_parameters is not None and home_grads is not None + accumulate_fsdp_unsharded_expert_gradients(home_parameters, home_grads) + return (None, *grad_inputs) + + +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, + call_state: _MoonEPLayerCallState, + async_op: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ctx.call_state = call_state + hidden_nvsh, route_weights_nvs, cu_seqlens = dispatch_forward( + call_state, + 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 = dispatch_backward( + cast(_MoonEPLayerCallState, ctx.call_state), + 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, + call_state: _MoonEPLayerCallState, + async_op: bool, + ) -> torch.Tensor: + ctx.call_state = call_state + ctx.save_for_backward(expert_output, route_weights) + return combine_forward(call_state, 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 = combine_backward(cast(_MoonEPLayerCallState, ctx.call_state), 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 + ``_MoonEPLayerCallState`` that the module-level transaction functions + advance with the call's plan, event, weight, and gradient state. + """ + + def __init__( + self, + *, + runtime: MoonEPModelRuntime, + layer: _MoonEPLayer, + ) -> None: + super().__init__( + n_routed_experts=runtime._num_experts, + process_group=runtime._ep_group, + ) + self._runtime = runtime + self._layer = layer + self._next_gradient_slot = 0 + + def _new_call_state(self, layer_gradients: _MoonEPLayerGradients) -> _MoonEPLayerCallState: + """Allocate the next call-local slot and call-state token.""" + resources = self._runtime.resources + grad_slot = self._next_gradient_slot + self._next_gradient_slot = (grad_slot + 1) % resources.gradient_slots + return _MoonEPLayerCallState( + resources=resources, + layer=self._layer, + generation=resources.home_generation(self._layer), + grad_slot=grad_slot, + layer_gradients=layer_gradients, + ) + + @override + def prepare_layer_inputs( + self, + layer_inputs: list[torch.Tensor], + ) -> tuple[list[torch.Tensor], list[object | None]]: + """Create one call-local Join and one plan/duplicate slot per + branch.""" + gradients = _MoonEPLayerGradients() + call_states = tuple(self._new_call_state(gradients) for _ in layer_inputs) + if len(call_states) > self._runtime.resources.gradient_slots: + raise ValueError("MoonEP layer width exceeds the gradient slot ring") + # No-grad original forwards build no backward node and never clear H. + return list(_MoonEPLayerGradJoin.apply(call_states, *layer_inputs)), list(call_states) + + @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_inputs") + if not isinstance(layer_state, _MoonEPLayerCallState): + raise TypeError("MoonEP layer_state must be a _MoonEPLayerCallState") + layer_state.resources.expect_tokens_per_rank(hidden_states.shape[0]) + 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_call=layer_state, + ) + + @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") + # Create the activation autograd edge and start weight prefetch. + hidden_nvsh, topk_weights_nvs, cu_seqlens = _DispatchAutograd.apply( + pre_dispatched["hidden_states"], + pre_dispatched["topk_ids"], + pre_dispatched["tokens_per_expert"], + topk_weights.to(dtype=torch.float32).contiguous(), + pre_dispatched["_moonep_call"], + async_op, + ) + return MoonEPDispatchResult( + hidden_states=hidden_nvsh, + topk_weights=topk_weights_nvs, + cu_seqlens=cu_seqlens, + ) + + @override + def dispatch_postprocess( + self, + *, + pre_dispatched: MoonEPPreDispatchResult, + dispatched: MoonEPDispatchResult, + async_op: bool = False, + ) -> MoonEPPostDispatchResult: + del async_op + return prepare_experts(pre_dispatched["_moonep_call"], 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 post_dispatched, decoding + # Create the fused route-scaled combine autograd edge. + return MoonEPCombineResult( + hidden_states=_CombineAutograd.apply( + pre_combined["hidden_states"], + dispatched["topk_weights"], + pre_dispatched["_moonep_call"], + 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 post_dispatched, pre_combined + return MoonEPPostCombineResult( + hidden_states=finish_combine( + pre_dispatched["_moonep_call"], + 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_capability.py b/xtuner/v1/module/dispatcher/moonep_capability.py new file mode 100644 index 0000000000..6a8eae4cff --- /dev/null +++ b/xtuner/v1/module/dispatcher/moonep_capability.py @@ -0,0 +1,106 @@ +"""MoonEP capability checks at the two points where they become decidable. + +``check_config`` runs during meta model build, when only ``MoEConfig`` is +visible; ``check_fsdp_policy`` runs just before ``fully_shard``, when +``FSDPConfig`` is visible. Neither creates a resource. Two things stay out of +this module by design: the workspace still owns the preconditions for safely +calling VMM (node-local, peer access, chunk alignment), and +``fsdp_vmm_landing`` still owns the torch-version / FSDP-ABI fail-fast. +""" + +from __future__ import annotations + +import os +from typing import Any + +import torch + + +def check_config(config: Any) -> None: + """Point 1: meta model build. Only ``MoEConfig`` is visible. + + Adding a new MoonEP capability constraint means adding one check here, and + it cannot be missed from a second call site. + """ + if config.dispatcher != "moonep": + return + + # Config-shape checks first, so they surface even where the optional + # backend is not installed. + float8_cfg = getattr(config, "float8_cfg", None) + if float8_cfg is not None and float8_cfg.enable_float8: + raise ValueError("MoonEP currently requires BF16 expert compute; FP8 is not supported") + if config.ep_size <= 1: + raise ValueError("MoonEP requires expert parallelism") + if config.ep_size not in (2, 4, 8): + raise ValueError("MoonEP requires ep_size in {2, 4, 8}") + if config.n_routed_experts % config.ep_size: + raise ValueError("MoonEP requires n_routed_experts divisible by ep_size") + if config.moe_bias: + raise ValueError("MoonEP does not support routed-expert linear bias") + if config.expert_tp_size > 1: + raise ValueError("MoonEP first version is TP1 only") + if config.intra_layer_micro_batch < 1: + raise ValueError("intra_layer_micro_batch must be positive") + + from .moonep import require_moonep_backend + + require_moonep_backend() + _require_cutlass_grouped_gemm_when_selected() + + +def check_fsdp_policy(config: Any, fsdp_config: Any) -> None: + """Point 2: before ``fully_shard``. ``FSDPConfig`` is now visible. + + The two points cannot merge: ``FSDPConfig`` does not exist yet at + ``MoE.__init__``. + """ + if config.dispatcher != "moonep": + return + + 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") + + # A blocking rule from the domain model that had no code before: a + # checkpoint-wrapped MTP physical layer must use reentrant checkpointing. + if _any_mtp_layer_checkpoint_wrapped(config, fsdp_config) and not fsdp_config.mtp_checkpoint_use_reentrant: + raise ValueError( + "a checkpoint-wrapped MTP physical layer requires FSDPConfig.mtp_checkpoint_use_reentrant=True with MoonEP" + ) + + +def _require_cutlass_grouped_gemm_when_selected() -> None: + # MoonEP keeps token counts device-resident. Triton already satisfies that + # contract; grouped_gemm does so only with its CUTLASS backend. + 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" + ) + + +def _any_mtp_layer_checkpoint_wrapped(config: Any, fsdp_config: Any) -> bool: + mtp_config = getattr(config, "mtp_config", None) + if mtp_config is None: + return False + if mtp_config.share_weights: + return True + # Non-shared: a non-terminal MTP layer selected by the recompute ratio is + # checkpoint-wrapped. Mirror ``MoE._should_recompute``'s global-index rule. + total_layers = config.num_hidden_layers + mtp_config.num_layers + num_recompute = int(total_layers * (getattr(fsdp_config, "recompute_ratio", 0.0) or 0.0)) + return any((config.num_hidden_layers + mtp_idx) < num_recompute for mtp_idx in range(mtp_config.num_layers - 1)) + + +__all__ = ["check_config", "check_fsdp_policy"] diff --git a/xtuner/v1/module/dispatcher/moonep_workspace.py b/xtuner/v1/module/dispatcher/moonep_workspace.py new file mode 100644 index 0000000000..9f5bc5fb28 --- /dev/null +++ b/xtuner/v1/module/dispatcher/moonep_workspace.py @@ -0,0 +1,539 @@ +"""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 Any, 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]: one home accumulator shared by every invocation of the active call. + home_gradients: 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]: return views [shared home, slot-local duplicate]. + 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, + home_generations: int, + ) -> None: + # One field. Every view's shape, indexing, and storage ownership stay + # centralized on ``_WorkspaceLayout``; adding a view touches only the + # layout and its accessor, and the rank-divergence quarantine keeps + # the whole graph by referencing this one object. + self._layout: _WorkspaceLayout | None = layout + self._ep_group = ep_group + self._ep_rank = ep_rank + self._num_experts = num_experts + self._experts_per_rank = experts_per_rank + # The only place the two-generation ``2`` is stored; ``generation_for`` + # is the single issuer for every consumer of a home generation. + self._home_generations = home_generations + 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, + home_generations: int = 2, + ) -> _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, + home_generations=home_generations, + 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, + home_generations=home_generations, + ) + + @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, + home_generations: 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[_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(home_generations))) + home_gradients.append(allocate_chunk(chunk_shape)) + 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][1] + duplicate_fd = allocations["duplicate_gradients"][projection][slot][1] + # Different virtual views share the physical home accumulator H. + # Only the duplicate suffix is slot-local. These are return + # views, not safe overwrite-output targets for grouped GEMM. + 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 + + @property + def _views(self) -> _WorkspaceLayout: + if self._layout is None: + raise RuntimeError("MoonEP workspace has been destroyed") + return self._layout + + def generation_for(self, execution_ordinal: int) -> int: + """Issue the home generation for one physical layer in execution order. + + Adjacent expert-bearing layers alternate over the two-generation home + ring, so no consumer ever re-validates ``generation in (0, 1)``. + """ + if self._destroyed: + raise RuntimeError("MoonEP workspace has been destroyed") + return execution_ordinal % self._home_generations + + def landing(self, generation: int) -> tuple[torch.Tensor, torch.Tensor]: + """Return projection-paired FSDP targets, each ``[B, O_p, I_p]``.""" + return cast(tuple[torch.Tensor, torch.Tensor], self._views["landings"][generation]) + + def local_compute_view( + self, *, hidden_nvsh: torch.Tensor, cu_seqlens: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Turn ``[NvS, H]`` dispatch output and ``[E+B]`` endpoints into a + grouped-GEMM-ready ``(hidden, [2B] counts)`` pair. + + Three layout-coupled steps stay in the owning module: fold ``[E+B]`` + into this rank's home ``B`` and duplicate ``B`` counts; zero the NvS + padding tail so uninitialized rows never reach the GEMM; and record + ``NvS - sum(counts)`` on the last group so the GEMM still walks the + full ``[NvS, H]``. Only small device metadata is computed here. + """ + _ = self._views + counts = self._local_token_counts(cu_seqlens) + covered = counts.sum() + row_is_covered = torch.arange(hidden_nvsh.shape[0], device=hidden_nvsh.device) < covered + hidden = hidden_nvsh * row_is_covered.unsqueeze(-1) + counts = torch.cat((counts[:-1], counts[-1:] + hidden_nvsh.shape[0] - covered)) + return hidden, counts + + def prefetch_weights(self, *, buffer: Any, plan: Any, generation: int) -> tuple[torch.Tensor, torch.Tensor]: + """Prefetch the global ``[E+B]`` weights and return this generation's + local ``[2B]`` compute aliases. + + ``buffer`` is the MoonEP ``Buffer`` and ``plan`` its opaque plan + object; both stay untyped MoonEP-owned values. + """ + buffer.prefetch_weight( + plan=plan, + projections=self._views["global_weights"][generation], + async_finish=False, + ) + return cast(tuple[torch.Tensor, torch.Tensor], self._views["local_weights"][generation]) + + def return_expert_gradients( + self, + *, + buffer: Any, + plan: Any, + gradients: tuple[torch.Tensor, torch.Tensor], + grad_slot: int, + initialize: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Complete the home-expert gradient boundary in the owning module. + + ``gradients`` is the allocation-return ``[2B]`` dW pair. The home + prefix aliases one accumulator shared by every Domino microbatch of a + single FSDP call, so it is zeroed only on the first producer + (``initialize``, a call-local flag owned by the Dispatcher) and added + to otherwise; the duplicate suffix is slot-local. ``reduce_grad_bf16`` + then sums the EP partials without dividing. + """ + if not 0 <= grad_slot < self._gradient_slots: + raise ValueError(f"gradient slot out of range: {grad_slot}") + b = self._experts_per_rank + # A slot is reused sequentially across physical layers. A fresh + # TensorImpl/version counter over the same VMM storage avoids both + # payload allocation and AOT version clashes. + targets = tuple( + target.new_empty(0).set_( + target.untyped_storage(), + target.storage_offset(), + target.shape, + target.stride(), + ) + for target in self._views["local_grad_outputs"][grad_slot] + ) + for target, gradient in zip(targets, gradients, strict=True): + if initialize: + target[:b].zero_() + target[:b].add_(gradient[:b]) + target[b:].copy_(gradient[b:]) + buffer.reduce_grad_bf16( + plan=plan, + local_grads=targets, + distributed_duplicate_grads=self._views["distributed_duplicate_grads"][grad_slot], + async_finish=False, + ) + return targets[0][:b], targets[1][:b] + + def _local_token_counts(self, cu_seqlens: torch.Tensor) -> torch.Tensor: + 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") + 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 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._layout = None + self._destroyed = True + + def __del__(self) -> None: + if getattr(self, "_destroyed", True) or getattr(self, "_layout", None) is None: + return + warnings.warn( + "MoonEP workspace was not destroyed explicitly; resources may leak.", + ResourceWarning, + ) + # Keep the whole view graph alive by referencing the one layout object + # instead of tearing CUDA/VMM state down after distributed ranks may + # have diverged during interpreter shutdown. + _UNDISPOSED_WORKSPACE_TENSORS.append(self._layout) + self._layout = None 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..0af0addc3c 100644 --- a/xtuner/v1/module/grouped_linear/moe_group_linear.py +++ b/xtuner/v1/module/grouped_linear/moe_group_linear.py @@ -159,9 +159,20 @@ 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) + def forward( + self, + x: torch.Tensor, + tokens_per_expert: torch.Tensor, + *, + trainable_weight: torch.Tensor | None = None, + ): + # 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 out = group_gemm(x, weight, tokens_per_expert) if self.moe_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/fsdp.py b/xtuner/v1/utils/fsdp.py index 8f1096b714..5cc64b5b2e 100644 --- a/xtuner/v1/utils/fsdp.py +++ b/xtuner/v1/utils/fsdp.py @@ -1,4 +1,5 @@ import torch +from torch.distributed.fsdp import FSDPModule from xtuner.v1.utils.device import get_torch_device_module @@ -6,6 +7,16 @@ DEVICE_MODULE = get_torch_device_module() +def set_requires_gradient_sync(module: FSDPModule, requires_gradient_sync: bool, *, recurse: bool = True) -> None: + """Keep XTuner accumulation in FSDP-owned sharded gradients.""" + if not requires_gradient_sync: + raise ValueError( + "XTuner requires gradient ReduceScatter on every backward; " + "set_requires_gradient_sync(False) is not supported." + ) + FSDPModule.set_requires_gradient_sync(module, requires_gradient_sync, recurse=recurse) + + def release_deferred_fsdp_all_gathers(model: torch.nn.Module) -> tuple[int, int]: """Release FSDP2 all-gather buffers that were prefetched but not consumed.""" From 8f3d0aab89c531940dc3fd6ea1f0e6a49e0e194a Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 9 Sep 2026 17:24:00 +0000 Subject: [PATCH 3/3] [Test] Add Qwen3.5 MoonEP/DeepEP acceptance gate Reproducible 20-step DeepEP/MoonEP comparison for the formal Qwen3.5-35B-A3B FSDP2 + EP4 workload, with a config-lock test and a throughput/curve comparator. XTUNER_USE_FA3 and XTUNER_DETERMINISTIC are overridable. --- .../run_qwen35_moonep_acceptance.sh | 90 ++++++ .../sft_qwen35_moonep_acceptance.py | 114 +++++++ tests/engine/test_moonep_acceptance.py | 151 +++++++++ xtuner/_testing/moonep_acceptance.py | 288 ++++++++++++++++++ 4 files changed, 643 insertions(+) create mode 100755 tests/acceptance/run_qwen35_moonep_acceptance.sh create mode 100644 tests/acceptance/sft_qwen35_moonep_acceptance.py create mode 100644 tests/engine/test_moonep_acceptance.py create mode 100644 xtuner/_testing/moonep_acceptance.py diff --git a/tests/acceptance/run_qwen35_moonep_acceptance.sh b/tests/acceptance/run_qwen35_moonep_acceptance.sh new file mode 100755 index 0000000000..aa1f0db531 --- /dev/null +++ b/tests/acceptance/run_qwen35_moonep_acceptance.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -lt 3 || $# -gt 4 ]]; then + echo "Usage: $0 [acceptance-root]" >&2 + exit 2 +fi + +backend=$1 +mtp=$2 +pack_length=$3 +acceptance_root=${4:-work_dirs/moonep_qwen35_acceptance} +# Resolve the checkout that owns this script so commit-level comparisons can +# run in isolated worktrees without importing the developer's dirty checkout. +repo_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +gpu_lock=/mnt/shared-storage-user/zhaopenghao/github/xtuner/zdev/gpu_lock.sh + +if [[ $backend != "deepep" && $backend != "moonep" ]]; then + echo "backend must be deepep or moonep" >&2 + exit 2 +fi +if [[ $mtp != "0" && $mtp != "1" ]]; then + echo "mtp must be 0 or 1" >&2 + exit 2 +fi +if ! [[ $pack_length =~ ^[1-9][0-9]*$ ]]; then + echo "pack-length must be a positive integer" >&2 + exit 2 +fi + +# Re-enter the exact same command while holding the repository-wide 8-GPU +# lock. The marker prevents recursively acquiring the non-reentrant lock. +if [[ ${MOONEP_ACCEPTANCE_LOCK_HELD:-0} != "1" ]]; then + exec "$gpu_lock" env MOONEP_ACCEPTANCE_LOCK_HELD=1 "$0" "$@" +fi + +source /mnt/shared-storage-user/zhaopenghao/miniconda3/etc/profile.d/conda.sh +conda activate pt212_cu132 +cd "$repo_root" + +run_dir="$acceptance_root/${backend}_mtp${mtp}_pack${pack_length}" +if [[ ${MOONEP_ACCEPTANCE_MICRO_BATCH:-1} != "1" ]]; then + run_dir="${run_dir}_micro${MOONEP_ACCEPTANCE_MICRO_BATCH}" +fi +if [[ -e $run_dir ]]; then + echo "refusing to mix acceptance attempts in existing directory: $run_dir" >&2 + exit 2 +fi +mkdir -p "$run_dir" + +export PYTHONPATH="$repo_root" +export MOONEP_ACCEPTANCE_BACKEND=$backend +export MOONEP_ACCEPTANCE_MTP=$mtp +export MOONEP_ACCEPTANCE_PACK_LENGTH=$pack_length +export MOONEP_ACCEPTANCE_WORK_DIR=$run_dir +export MOONEP_ACCEPTANCE_MODEL_PATH=${MOONEP_ACCEPTANCE_MODEL_PATH:-/mnt/shared-storage-user/llmrazor-share/model/Qwen3.5-35B-A3B} +export MOONEP_ACCEPTANCE_DATA_PATH=${MOONEP_ACCEPTANCE_DATA_PATH:-/mnt/shared-storage-user/llmrazor-share/data/alpaca} +export MODEL_COMPILE=1 +export XTUNER_USE_FA3="${XTUNER_USE_FA3:-0}" +# FA3's Hopper backward has no deterministic path for head_dim 256 (Qwen3.5), +# so the FA3 comparison runs non-deterministically; keep the default otherwise. +export XTUNER_DETERMINISTIC="${XTUNER_DETERMINISTIC:-true}" +export XTUNER_ACTIVATION_OFFLOAD=0 +export XTUNER_COMPILE_NO_INPLACE_BUFFERS=1 +export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export PROFILE_RANKS=${PROFILE_RANKS:-0} +export TRITON_CACHE_DIR="$run_dir/triton_cache" +export TORCHINDUCTOR_CACHE_DIR="$run_dir/torchinductor_cache" +unset XTUNER_USE_CUTLASS_GROUP_GEMM +unset GROUPED_GEMM_USE_CUTLASS + +config=tests/acceptance/sft_qwen35_moonep_acceptance.py +python -m xtuner._testing.moonep_acceptance capture \ + --config "$config" \ + --output "$run_dir/acceptance_manifest.json" + +# VMM/NCCL allocations are not all visible to PyTorch's allocator counters. +# Sample all eight devices while the same GPU lock covers the training job. +nvidia-smi --query-gpu=timestamp,index,memory.used --format=csv,noheader,nounits \ + --loop-ms=500 > "$run_dir/device_memory.csv" & +memory_monitor_pid=$! +trap 'kill "$memory_monitor_pid" 2>/dev/null || true; wait "$memory_monitor_pid" 2>/dev/null || true' EXIT + +torchrun \ + --nproc-per-node 8 \ + --master-port "${MOONEP_ACCEPTANCE_MASTER_PORT:-29618}" \ + xtuner/v1/train/cli/sft.py \ + --config "$config" \ + 2>&1 | tee "$run_dir/stdout.log" diff --git a/tests/acceptance/sft_qwen35_moonep_acceptance.py b/tests/acceptance/sft_qwen35_moonep_acceptance.py new file mode 100644 index 0000000000..149fc33e4b --- /dev/null +++ b/tests/acceptance/sft_qwen35_moonep_acceptance.py @@ -0,0 +1,114 @@ +"""Matched Qwen3.5 configuration for the MoonEP/DeepEP 20-step gate. + +The dispatcher, MTP switch, fixed pack length and output directory are the +only run-varying inputs. All other workload choices are deliberately shared. +""" + +import os + +import torch + +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig +from xtuner.v1.datasets import FTDPTokenizeFnConfig +from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model import Qwen3_5_VLMoE35BA3Config +from xtuner.v1.model.moe.qwen3_5_text import MOE_EP_COMPILE_CFG +from xtuner.v1.module.mtp import MTPConfig +from xtuner.v1.train import TrainerConfig + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise ValueError(f"{name} must be set") + return value + + +backend = _required_env("MOONEP_ACCEPTANCE_BACKEND") +if backend not in {"deepep", "moonep"}: + raise ValueError(f"MOONEP_ACCEPTANCE_BACKEND must be deepep or moonep, got {backend!r}") + +mtp_enabled = bool(int(_required_env("MOONEP_ACCEPTANCE_MTP"))) +pack_length = int(_required_env("MOONEP_ACCEPTANCE_PACK_LENGTH")) +micro_batch = int(os.environ.get("MOONEP_ACCEPTANCE_MICRO_BATCH", "1")) +if micro_batch not in (1, 2): + raise ValueError("MOONEP_ACCEPTANCE_MICRO_BATCH must be 1 or 2") +if pack_length <= 0: + raise ValueError("MOONEP_ACCEPTANCE_PACK_LENGTH must be positive") +if os.environ.get("XTUNER_ACTIVATION_OFFLOAD", "0") != "0": + raise ValueError("formal MoonEP acceptance runs require XTUNER_ACTIVATION_OFFLOAD=0") +if os.environ.get("XTUNER_USE_CUTLASS_GROUP_GEMM", "0") == "1": + raise ValueError("formal MoonEP acceptance runs require the Triton grouped-GEMM backend") + +model_cfg = Qwen3_5_VLMoE35BA3Config(only_llm_forward=True) +text_cfg = model_cfg.text_config +text_cfg.ep_size = 4 +text_cfg.dispatcher = backend +text_cfg.moonep_staging_reference = False +text_cfg.router_compute_dtype = "float32" +text_cfg.router_async_offload = False +# FA2's package metadata has no importable extension in pt212_cu132, so the +# default keeps both dispatcher runs on flex attention. Set XTUNER_USE_FA3=1 to +# compare against FlashAttention 3, whose varlen path skips BlockMask +# construction entirely (no per-packed-document Dynamo/SymPy mask subgraph). +text_cfg.attention.attn_impl = ( + "flash_attention" if os.environ.get("XTUNER_USE_FA3", "0") == "1" else "flex_attention" +) +# FlexAttention intentionally compiles behind a graph break so its BlockMask +# tensors become fixed-layout inputs to the kernel graph. Keep all default +# Qwen3.5 compile targets, but let MHA form that one required boundary. +text_cfg.compile_cfg = MOE_EP_COMPILE_CFG | { + "xtuner.v1.module.attention.mha.MultiHeadAttention.forward": {"fullgraph": False} +} +text_cfg.mtp_config = MTPConfig(num_layers=1) if mtp_enabled else None + +dataset_cfg = [ + { + "dataset": DatasetConfig( + name="alpaca", + anno_path=_required_env("MOONEP_ACCEPTANCE_DATA_PATH"), + sample_ratio=1.0, + ), + "tokenize_fn": FTDPTokenizeFnConfig(max_length=262144), + } +] +dataloader_cfg = DataloaderConfig( + dataset_config_list=dataset_cfg, + pack_to_max_length=True, + pack_max_length=pack_length, + pack_level="hard", +) + +profile_step_env = os.environ.get("MOONEP_ACCEPTANCE_PROFILE_STEP") +profile_step = int(profile_step_env) if profile_step_env else None + +trainer = TrainerConfig( + load_from=_required_env("MOONEP_ACCEPTANCE_MODEL_PATH"), + tokenizer_path=_required_env("MOONEP_ACCEPTANCE_MODEL_PATH"), + model_cfg=model_cfg, + optim_cfg=AdamWConfig(lr=6e-5, foreach=False), + lr_cfg=LRConfig(lr_type="cosine", lr_min=1e-6), + loss_cfg=CELossConfig(mode="chunk", chunk_size=1024), + fsdp_cfg=FSDPConfig( + ep_size=4, + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + torch_compile=True, + cpu_offload=False, + ), + dataloader_cfg=dataloader_cfg, + global_batch_size=8 * micro_batch, + intra_layer_micro_batch=micro_batch, + sp_size=1, + total_step=20, + work_dir=_required_env("MOONEP_ACCEPTANCE_WORK_DIR"), + seed=0, + strict_load=False, + auto_resume=False, + debug_skip_save=True, + exp_tracker="jsonl", + profile_step=profile_step, + profile_time=profile_step is not None, + profile_memory=False, +) diff --git a/tests/engine/test_moonep_acceptance.py b/tests/engine/test_moonep_acceptance.py new file mode 100644 index 0000000000..c7fcff36ab --- /dev/null +++ b/tests/engine/test_moonep_acceptance.py @@ -0,0 +1,151 @@ +import json +import runpy + +import pytest + +from xtuner._testing.moonep_acceptance import AcceptanceRun, compare_runs, summarize_memory + + +def _write_tracker(path, *, tgs_scale: float, mtp: bool) -> None: + path.parent.mkdir(parents=True) + with path.open("w", encoding="utf-8") as output: + for step in range(1, 21): + record = { + "step": step, + "runtime_info/text_tokens": 65536, + "runtime_info/tgs": (1000 + step) * tgs_scale, + "loss/reduced_llm_loss": 2.0 - step / 100, + "loss/reduced_balancing_loss": 0.01 + step / 10000, + "loss/local_loss": 2.01 - step / 100 + step / 10000, + "grad_norm": 0.5 + step / 1000, + } + if mtp: + record["loss/reduced_mtp_loss"] = 0.2 - step / 1000 + record["loss/local_loss"] += record["loss/reduced_mtp_loss"] + output.write(json.dumps(record) + "\n") + + +@pytest.mark.parametrize("backend", ["deepep", "moonep"]) +@pytest.mark.parametrize("mtp", [False, True]) +@pytest.mark.parametrize("micro_batch", [1, 2]) +def test_qwen35_acceptance_config_locks_the_formal_workload(monkeypatch, tmp_path, backend, mtp, micro_batch) -> None: + monkeypatch.setenv("MOONEP_ACCEPTANCE_BACKEND", backend) + monkeypatch.setenv("MOONEP_ACCEPTANCE_MTP", str(int(mtp))) + monkeypatch.setenv("MOONEP_ACCEPTANCE_MICRO_BATCH", str(micro_batch)) + monkeypatch.setenv("MOONEP_ACCEPTANCE_PACK_LENGTH", "65536") + monkeypatch.setenv("MOONEP_ACCEPTANCE_WORK_DIR", str(tmp_path / "run")) + monkeypatch.setenv("MOONEP_ACCEPTANCE_MODEL_PATH", "/model") + monkeypatch.setenv("MOONEP_ACCEPTANCE_DATA_PATH", "/data") + + trainer = runpy.run_path("tests/acceptance/sft_qwen35_moonep_acceptance.py")["trainer"] + model = trainer.model_cfg + + assert trainer.total_step == 20 + assert trainer.global_batch_size == 8 * micro_batch + assert trainer.intra_layer_micro_batch == micro_batch + assert trainer.sp_size == 1 + assert trainer.debug_skip_save is True + assert trainer.dataloader_cfg.pack_to_max_length is True + assert trainer.dataloader_cfg.pack_max_length == 65536 + assert trainer.fsdp_cfg.ep_size == 4 + assert trainer.fsdp_cfg.param_dtype.__str__() == "torch.bfloat16" + assert trainer.fsdp_cfg.reduce_dtype.__str__() == "torch.bfloat16" + assert trainer.fsdp_cfg.torch_compile is True + assert trainer.fsdp_cfg.cpu_offload is False + assert model.only_llm_forward is True + assert model.text_config.ep_size == 4 + assert model.text_config.dispatcher == backend + assert model.text_config.moonep_staging_reference is False + assert model.text_config.moonep_num_sms == 64 + assert model.text_config.router_async_offload is False + assert model.text_config.router_compute_dtype == "float32" + assert ( + model.text_config.compile_cfg["xtuner.v1.module.attention.mha.MultiHeadAttention.forward"]["fullgraph"] + is False + ) + assert (model.text_config.mtp_config is not None) is mtp + if mtp: + assert model.text_config.mtp_config.num_layers == 1 + assert model.text_config.mtp_config.share_weights is False + + +def test_acceptance_report_compares_all_steps_and_warm_throughput(tmp_path) -> None: + deepep_tracker = tmp_path / "deepep" / "tracker.jsonl" + moonep_tracker = tmp_path / "moonep" / "tracker.jsonl" + _write_tracker(deepep_tracker, tgs_scale=1.0, mtp=True) + _write_tracker(moonep_tracker, tgs_scale=0.96, mtp=True) + records = [json.loads(line) for line in moonep_tracker.read_text().splitlines()] + for record in records: + for name in tuple(record): + if name.startswith("loss/reduced_") and name.endswith("loss"): + record[name] *= 1.02 + record["grad_norm"] *= 1.04 + moonep_tracker.write_text("".join(f"{json.dumps(record)}\n" for record in records)) + + deepep = AcceptanceRun.from_tracker(deepep_tracker, backend="deepep", mtp=True, pack_length=65536) + moonep = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=True, pack_length=65536) + result = compare_runs(deepep, moonep) + + assert result.passed + assert result.throughput_ratio == pytest.approx(0.96) + assert result.throughput_steps == list(range(6, 21)) + assert set(result.curves) == { + "reduced_llm_loss", + "reduced_mtp_loss", + "reduced_balancing_loss", + "total_loss", + "grad_norm", + } + assert all(curve.cosine_similarity >= 0.99 for curve in result.curves.values()) + assert all(curve.mean_relative_difference < 0.03 for name, curve in result.curves.items() if name != "grad_norm") + assert result.curves["grad_norm"].mean_relative_difference == pytest.approx(0.04) + + +def test_acceptance_report_rejects_incomplete_or_mismatched_runs(tmp_path) -> None: + deepep_tracker = tmp_path / "deepep" / "tracker.jsonl" + moonep_tracker = tmp_path / "moonep" / "tracker.jsonl" + _write_tracker(deepep_tracker, tgs_scale=1.0, mtp=False) + _write_tracker(moonep_tracker, tgs_scale=0.94, mtp=False) + + deepep = AcceptanceRun.from_tracker(deepep_tracker, backend="deepep", mtp=False, pack_length=65536) + slow = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=65536) + assert not compare_runs(deepep, slow).passed + + mismatched = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=32768) + with pytest.raises(ValueError, match="pack_length"): + compare_runs(deepep, mismatched) + + micro2 = AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=65536, micro_batch=2) + with pytest.raises(ValueError, match="micro_batch"): + compare_runs(deepep, micro2) + + lines = moonep_tracker.read_text().splitlines() + moonep_tracker.write_text("\n".join(lines[:-1]) + "\n") + with pytest.raises(ValueError, match="exactly steps 1..20"): + AcceptanceRun.from_tracker(moonep_tracker, backend="moonep", mtp=False, pack_length=65536) + + +def test_memory_report_includes_all_ranks_and_external_device_allocations(tmp_path) -> None: + for rank in range(2): + tracker = tmp_path / "logs" / "exp_tracking" / f"rank{rank}" / "tracker.jsonl" + tracker.parent.mkdir(parents=True) + tracker.write_text( + "\n".join( + json.dumps( + { + "step": step, + "memory/max_memory_GB": (10 if step == 1 else 6) + rank, + "memory/reserved_memory_GB": (12 if step == 1 else 8) + rank, + } + ) + for step in range(1, 21) + ) + ) + (tmp_path / "device_memory.csv").write_text( + "2026/09/07 00:00:00.000, 0, 15360\n2026/09/07 00:00:00.000, 1, 16384\n2026/09/07 00:00:00.500, 0, 14336\n" + ) + result = summarize_memory(tmp_path) + assert result["recorded_ranks"] == 2 + assert result["all_steps"] == {"allocated_gib": 11, "reserved_gib": 13} + assert result["steps_6_20"] == {"allocated_gib": 7, "reserved_gib": 9} + assert result["sampled_device_peak_gib"] == {"0": 15, "1": 16} diff --git a/xtuner/_testing/moonep_acceptance.py b/xtuner/_testing/moonep_acceptance.py new file mode 100644 index 0000000000..9ce8b04eb7 --- /dev/null +++ b/xtuner/_testing/moonep_acceptance.py @@ -0,0 +1,288 @@ +"""Parse and judge matched MoonEP/DeepEP Qwen3.5 acceptance runs.""" + +from __future__ import annotations + +import argparse +import csv +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], ...] + micro_batch: int = 1 + + @classmethod + def from_tracker( + cls, + tracker: str | Path, + *, + backend: str, + mtp: bool, + pack_length: int, + micro_batch: int = 1, + ) -> "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, micro_batch=micro_batch) + + @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"], + micro_batch=manifest.get("micro_batch", 1), + ) + + @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", "micro_batch"): + 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"]), + "micro_batch": trainer.intra_layer_micro_batch, + "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 summarize_memory(work_dir: str | Path) -> dict[str, Any]: + """Report allocator peaks separately from sampled device use (incl. VMM). + + Trainer resets its allocator peak after each step. Aggregate every logged + rank, and retain both whole-run and steady-state maxima, not averages. + Device samples cover the whole process, including startup/compilation. + """ + work_dir = Path(work_dir) + trackers = sorted(work_dir.glob("**/exp_tracking/rank*/tracker.jsonl")) + records = [json.loads(line) for path in trackers for line in path.read_text().splitlines() if line] + result: dict[str, Any] = {"recorded_ranks": len(trackers)} + for phase, selected in ( + ("all_steps", records), + ("steps_6_20", [record for record in records if 6 <= record["step"] <= 20]), + ): + result[phase] = { + label: max((record[key] for record in selected if key in record), default=None) + for label, key in ( + ("allocated_gib", "memory/max_memory_GB"), + ("reserved_gib", "memory/reserved_memory_GB"), + ) + } + device_peaks: dict[str, float] = {} + samples = work_dir / "device_memory.csv" + if samples.exists(): + for _, device, used_mib in csv.reader(samples.read_text().splitlines()): + device = device.strip() + device_peaks[device] = max(device_peaks.get(device, 0.0), float(used_mib) / 1024) + result["sampled_device_peak_gib"] = device_peaks + return result + + +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) + report = result.to_dict() + report["memory"] = {"deepep": summarize_memory(args.deepep), "moonep": summarize_memory(args.moonep)} + args.output.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + return 0 if result.passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main())