Uh oh!
There was an error while loading. Please reload this page.
[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training - #2056
[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056jayhenry wants to merge 4 commits into
Conversation
jayhenry
commented
Sep 1, 2026
@claude review |
Uh oh!
There was an error while loading. Please reload this page.
| 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" | ||
| ) | ||
| if self._moonep_runtime is not None and len(seq_ctx) > self.config.intra_layer_micro_batch: |
There was a problem hiding this comment.
改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch
| staging_reference: bool, | ||
| num_sms: int = 64, | ||
| ) -> None: | ||
| self._backend = require_moonep_backend() |
There was a problem hiding this comment.
直接检查是否安装 moonep,不做过度检查。
另外,在文件顶层 try import moonep 即可,后续引用也直接用 moonep,去掉 self._backend
| # moe forward | ||
| # (e, ) | ||
| tokens_per_expert = torch.histc(topk_ids, bins=self.n_routed_experts, min=0, max=self.n_routed_experts) | ||
| tokens_per_expert = torch.bincount(topk_ids.flatten(), minlength=self.n_routed_experts) |
There was a problem hiding this comment.
torch.bincount 存在强制 Host-Device 同步,改回原来的 torch.histc
| n_routed_experts=n_routed_experts, | ||
| ep_group=process_group, | ||
| tp_group=tp_group, | ||
| ep_tp_group=ep_tp_group, |
There was a problem hiding this comment.
先hardcode bf16,去掉这个参数,assert 不支持 fp8
| min=0, | ||
| max=self.n_routed_experts, | ||
| ) # .view(self.ep_mesh.size(), -1) | ||
| tokens_per_expert = torch.bincount(topk_ids.flatten(), minlength=self.n_routed_experts) |
There was a problem hiding this comment.
同样要使用不host sync的方式,比如 histc
| from torch import Tensor | ||
| try: |
There was a problem hiding this comment.
保留 try except 的import 方式
| 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. |
There was a problem hiding this comment.
将各种检查封装到一个私有函数,保证顶层的线性步骤
| self._fixed_tokens_per_rank: int | None = None | ||
| self._closed = False | ||
| def bind_dispatcher( |
There was a problem hiding this comment.
改名为 build_dispatcher 是否更合适?
| """Own one fresh MoonEP plan and all of its device-side completion | ||
| edges.""" | ||
| def __init__(self, runtime: MoonEPRuntime, *, layer_id: int, grad_slot: int) -> None: |
| topk_ids, | ||
| tokens_per_expert, | ||
| topk_weights, | ||
| self, |
There was a problem hiding this comment.
将 invocation 整个传参是否参数过宽?
| # completed home gradients are handed back before FSDP post-backward. | ||
| with torch.profiler.record_function("MoonEP::prepare_experts"): | ||
| home_parameters = self._current_home_parameters() | ||
| local_weights, gradient_targets, weights_ready = workspace.materialize( |
There was a problem hiding this comment.
这个函数会 prefetch weight 产生GPU通信,如果想要 domino ep 正常工作,是不是应该将它放到 dispatcher.dispatch 函数中?
| raise NotImplementedError("MoonEP fixed-S training dispatch does not implement decoding") | ||
| grad_slot = self._next_gradient_slot | ||
| self._next_gradient_slot = (grad_slot + 1) % self._runtime._intra_layer_micro_batch | ||
| return _MoonEPInvocation(self._runtime, layer_id=self._layer_id, grad_slot=grad_slot).dispatch( |
There was a problem hiding this comment.
现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?
Summary
This PR adds an optional
dispatcher="moonep"path for node-local BF16 MoE training. Native FSDP2 remains the sole owner of expert parameters, optimizer state, and checkpoint identity; MoonEP owns only the communication/VMM execution workspace.torch.compile._StridedShardrecognition in distributed grad-norm calculation.Design
MoonEP is imported lazily only when selected. The integration validates a versioned backend capability before allocating resources, so DeepEP, All2All, AGRS, and non-EP configurations do not acquire a MoonEP dependency.
The hot path is ordered with CUDA events and GPU-side EP barriers. The profiler regression gate verifies:
FSDP post-all-gather hooks install the VMM landing tensors while preserving the original DTensor parameters. The backward path returns BF16 home gradients to those parameter edges; existing FSDP communication then reduce-scatters them and produces FP32 sharded gradients for the FP32 optimizer update.
Supported scope
2.12.1+cu132.Not claimed in this first version: Expert TP, FP8 experts, cross-node MoonEP, FSDP
no_sync, pipeline parallelism, or decoding.External MoonEP dependency
This XTuner branch expects MoonEP XTuner integration API v2. The validated companion implementation is currently commit
c14bd43001efd8233950bb99e8eac9b1bafbdcc4on the local MoonEPxtuner-integrationbranch.The companion MoonEP patch still needs to be published separately before this PR can be merged or reproduced outside the development environment. The XTuner package does not add MoonEP as a mandatory dependency.
Validation
The formal workload is Qwen3.5-35B-A3B on one node with 8 H200 GPUs, FSDP2 + EP4, BF16 parameters/reduction, Direct VMM landing, Triton grouped GEMM, and 20 training steps. Each comparison changes only the dispatcher.
For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.
Additional GPU regressions:
See the acceptance report for the full configuration, per-step metrics, profiler evidence, and reproduction commands.
Result
The first-version objective is met for single-node BF16 FSDP2 + EP training: MoonEP matches DeepEP numerically and reaches equivalent steady-state throughput while preserving FSDP parameter/checkpoint ownership and a host-sync-free communication hot path.