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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

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

Open
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep
Open

[Feature] Integrate MoonEP dispatcher for FSDP expert-parallel training#2056
jayhenry wants to merge 4 commits into
InternLM:mainfrom
jayhenry:moonep

Conversation

@jayhenry

Copy link
Copy Markdown
Collaborator

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.

  • Add a model-scoped MoonEP runtime and adapt it to XTuner's six-stage dispatcher API.
  • Land FSDP all-gathered BF16 expert weights directly into MoonEP VMM aliases.
  • Dispatch activations and weights, run the existing grouped GEMM path, and return duplicated expert gradients in BF16 before FSDP reduce-scatter.
  • Support MTP (shared and unshared weights), Domino intra-layer micro-batching, sequence parallelism, activation recompute, and torch.compile.
  • Integrate DCP/HF persistence, activation offload, optimizer swap, Muon, and explicit runtime teardown.
  • Add a reproducible Qwen3.5 20-step DeepEP/MoonEP acceptance gate and detailed validation report.
  • Fix PyTorch 2.12 FSDP _StridedShard recognition in distributed grad-norm calculation.
flowchart LR
A["FSDP all-gather: complete BF16 experts"] --> B["MoonEP VMM direct landing"]
B --> C["Dispatch + grouped GEMM + combine"]
C --> D["BF16 duplicated-gradient return"]
D --> E["FSDP reduce-scatter"]
E --> F["FP32 shard gradient + optimizer state"]
Loading

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:

  • no complete home-expert weight copy in the direct path;
  • no complete local duplicated-gradient temporary;
  • no CUDA device/event/stream host synchronization between planning and duplicated-gradient handoff.

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

  • PyTorch 2.12.1+cu132.
  • Single-node BF16 EP2/EP4/EP8; formal end-to-end acceptance uses FSDP2 + EP4 on 8 GPUs.
  • TP1.
  • MTP, Domino micro-batching, sequence parallelism, activation recompute, and compile.
  • Synchronous/asynchronous DCP, HF export/load, AdamW, swap AdamW, and Muon.

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 c14bd43001efd8233950bb99e8eac9b1bafbdcc4 on the local MoonEP xtuner-integration branch.

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.

GateDeepEP median tokens/sMoonEP median tokens/sMoonEP / DeepEP
MTP disabled, steps 6-205980.9996047.073101.10%
MTP1, steps 6-205357.1645408.172100.95%

For both gates, all loss and grad-norm curves have cosine similarity >= 0.99 and mean relative difference < 1%.

Additional GPU regressions:

  • MoonEP forward/MTP/Domino/SP/compile: 16 passed.
  • DCP/HF/offload/optimizer/lifecycle: 10 passed.
  • Dispatcher and grouped-GEMM regressions: passed.
  • FSDP2 + EP4 nested-shard grad norm: passed on PyTorch 2.12; compatibility regression also passed on PyTorch 2.9.
  • GLM FSDP2 + EP4 training, MoonEP/DeepEP two-step numerical parity, MTP reentrant micro2, and DCP cold resume: passed.

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.

@jayhenry

Copy link
Copy Markdown
CollaboratorAuthor

@claude review

Comment threadxtuner/v1/module/dispatcher/moonep.py
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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改为通用检查 len(seq_ctx) == self.config.intra_layer_micro_batch

staging_reference: bool,
num_sms: int = 64,
) -> None:
self._backend = require_moonep_backend()

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

直接检查是否安装 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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

先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)

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同样要使用不host sync的方式,比如 histc

from torch import Tensor


try:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

保留 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.

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将各种检查封装到一个私有函数,保证顶层的线性步骤

self._fixed_tokens_per_rank: int | None = None
self._closed = False

def bind_dispatcher(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

改名为 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:

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 runtime 整个传参是否参数过宽?

topk_ids,
tokens_per_expert,
topk_weights,
self,

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

将 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(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个函数会 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(

@jayhenryjayhenrySep 2, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 dispatcher太浅,只做简单的转发。将 Invocation 退化成一个对象类,并没有 dispatch, combine 等方法,让 dispatcher 本身有具体的dispatch等逻辑,这样改是否合理?

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@jayhenry