refactor(moe): unify runtime and checkpoint layout by S1ro1 · Pull Request #3411 · PrimeIntellect-ai/prime-rl

S1ro1 · 2026-08-27T00:37:17Z

Summary

Model and checkpoint contract

Replace the separate standard, GPT-OSS, and Nemotron expert/router implementations with one canonical MoE tree.
Store routed expert weights as separate gate_proj, up_proj, and down_proj parameters. Non-gated experts omit gate_proj.
Nest every custom router at mlp.router.gate, routed experts at mlp.experts, and optional shared experts inside mlp.
Store the single selection-only router bias at mlp.router.selection_bias; model-specific checkpoint names are converted only at the HF boundary.
Keep shared expert projections canonical. Qwen3.5 alone nests its scalar sigmoid gate at mlp.shared_expert.output_gate, with conversion mapping the standalone Hugging Face key at the boundary.
Use the same activation protocol and FeedForward implementation for dense MLPs, shared experts, and grouped experts.
Make shared-expert construction exclusively parent-owned: model layers pass an explicit FeedForward/subclass or None, while MoEArgs and MoE never infer or construct one.
Keep existing HF source layouts at the conversion boundary. GPT-OSS alone adds the structural split/join required for its interleaved source tensor.
Require .prime-v1 in reusable <conversion_dir>/prime caches so legacy Prime layouts fail before loading.

MoE runtime contract

Use one routed execution sequence for every backend: dispatch tokens, run canonical grouped experts, combine tokens.
Select grouped expert compute through explicit BF16, DeepGEMM FP8, or MXFP8 implementations. MoEArgs remains architecture-only.
Select local/Torch or DeepEP token dispatch once during model setup. Dispatchers own routing metadata, communication, padding, chunking, and synchronization.
Keep TokenDispatcher as the structural interface while a shared base owns the canonical run/synchronize sequence; local, Torch, and DeepEP implementations own only their dispatch/combine details.
Keep DeepEP combine events, concatenate streams, and flattened dispatch state owned by each dispatcher; chunking uses the equal input shape shared by EP ranks and does not synthesize empty chunks.
Ignore dispatch configuration when expert parallelism resolves to one rank; local dispatch remains active and DeepEP has no optimizer side effects.
Keep grouped expert execution directly in GroupedExperts.forward and router confidence accounting directly in TokenChoiceTopKRouter.forward.
Shard every canonical expert parameter with one backend-independent Shard(0) placement.
Keep dense linear quantization, routed-expert compute, and EP transport independent.
Remove the TorchTitan module hooks, expert decorator, dynamically attached process group, backend attributes, tensor-subclass interception, and process-global MXFP8 patches from the reachable path.
Vendor the narrow MXFP8 grouped-GEMM and Torch EP transport implementation in prime-kernels PR #1, preserving the pinned torchao revision and BSD license.

Public contracts

Checkpoint tree:

mlp.router.gate: nn.Linear
mlp.router.selection_bias: optional persistent selection-only buffer
mlp.experts.gate_proj: stacked nn.Parameter for gated experts, absent for non-gated experts
mlp.experts.up_proj: stacked nn.Parameter
mlp.experts.down_proj: stacked nn.Parameter
mlp.shared_expert.{gate_proj,up_proj,down_proj}: FeedForward projections when present
mlp.shared_expert.output_gate: Qwen3.5-only scalar nn.Linear
expert matrices: [num_experts, output_features, input_features]

Runtime configuration:

[trainer.model.moe.compute]
type = "bf16"  # bf16 | deepgemm_fp8 | mxfp8

[trainer.model.moe.dispatch]
type = "torch"  # torch | deepep
transport = "bf16"  # bf16 | mxfp8 for torch dispatch

DeepEP owns its tuning fields:

[trainer.model.moe.dispatch]
type = "deepep"
num_sms = 20
token_chunk_size = 4096

This is intentionally breaking. enable_grouped_gemm, enable_a2a, ep_comm_backend, deepep_num_sms, and deepep_token_chunk_size are removed without aliases.

Validation completed

Foundation conversion, detector, marker, state-loading, checkpoint, NIXL graph, and wire-format tests passed.
Focused CPU configuration and MoE contracts: 143 passed.
Focused MoE runtime configuration parametrizations after the EP=1 no-op change: 6 passed.
Affected custom-model CUDA suite: 24 passed, 3 skipped across AFMoE, GLM, GPT-OSS, Nemotron, Qwen3/Qwen3.5, and state loading.
Final H200 model suite after execution-path cleanup: 81 passed, 3 skipped (Slurm 2597).
Torch EP on two H200 GPUs passed with BF16 and DeepGEMM FP8, top-k 2, both score locations, shared experts, gated/non-gated/bias variants, uneven routing, a rank receiving zero routed tokens, and input/router/expert gradients.
Final Torch EP BF16 and DeepGEMM FP8 matrix passed after the cleanup (Slurm 2596).
DeepEP on two H200 GPUs passed with BF16 and DeepGEMM FP8, chunking on/off, both score locations, shared experts, a rank receiving zero routed tokens, and gradients.
Final DeepEP BF16 and DeepGEMM FP8 matrix passed after the cleanup (Slurm 2595).
Parent-owned shared-expert construction: 16 focused CPU tests and 64 H200 model tests passed, with 3 expected skips (Slurm 2600).
Full-graph MoE and shared FeedForward compile validation passed after the ownership cleanup (Slurm 2601).
Local DeepGEMM FP8 passed across all supported topology, activation, and bias combinations.
torch.compile(fullgraph=True) forward/backward passed for the full MoE and shared FeedForward across all supported topology/activation combinations.
Single-rank DCP and Torch EP + FSDP DCP model/optimizer resume passed.
The mxfp8_moe wheel build includes its public module, manifest, and license; the module imports against the pinned torchao build.
Unsupported H200 MXFP8 selection fails before model mutation.
Review audit: every queued DeepEP combine event is now synchronized, including chunked execution (two-H200 Slurm 2604).
Dispatcher review cleanup passed on two H200s: DeepEP BF16/DeepGEMM with chunking and a zero-receive rank (Slurm 2847), plus Torch EP topology/bias/zero-receive coverage (Slurm 2848).
The exact reverse_text_moe integration passed end-to-end on two H200s with an isolated conversion cache (Slurm 2607).
Conversion ranks now make the create decision on master and validate .prime-v1 only after the final all-rank barrier.
Canonical router selection bias: 12 CPU MoE/conversion tests passed; CUDA checkpoint loading, Nemotron, and GLM model tests passed, with the existing GLM BF16 gradient comparison passing on isolated rerun after one threshold-edge miss (Slurm 2851/2852).
Ruff, Ruff formatting, compileall, and git diff --check passed.

Deliberate exclusions

MultiLoRA is not adapted to the new runtime and may fail.
Fused-MoE trainer integration is removed and will be reconsidered only after the canonical runtime is finalized.
Routed-expert selective activation checkpointing is removed; the remaining selective targets are unchanged.

Remaining before marking ready

Run numerical MXFP8 grouped-GEMM and MXFP8 Torch EP transport forward/backward on SM100 hardware. The available Slurm fleet currently has only H200 nodes.
Merge prime-kernels PR Use hf model #1, publish wheels containing mxfp8_moe, and update the pyproject.toml/lock wheel pins. The currently pinned v0.8.0 wheel does not contain this module.

Note

High Risk
This changes the MoE forward/backward path, expert-parallel dispatch, checkpoint key layout, and trainer TOML in breaking ways across all custom MoE models.

Overview
Breaking trainer config: MoE is configured via [trainer.model.moe.compute] (bf16, deepgemm_fp8, mxfp8) and [trainer.model.moe.dispatch] (torch with transport, or deepep with num_sms / token_chunk_size). Dense [trainer.model.quantization] no longer drives expert GEMMs or EP transport. Removed flags include ep_comm_backend, deepep_*, moe_use_grouped_mm, moe_fused_kernel, and quantization enable_grouped_gemm / enable_a2a.

Unified MoE execution: All custom models share one MoE path—TokenChoiceTopKRouter, stacked GroupedExperts (gate_proj / up_proj / down_proj), and a pluggable TokenDispatcher (LocalTokenDispatcher, TorchTokenDispatcher, MXFP8TorchTokenDispatcher, DeepEPTokenDispatcher). configure_moe_runtime wires grouped GEMM backends and dispatch at setup; EP weight sharding is only ExpertWeightParallel. DeepEP dispatch/combine, grouped-GEMM permutation, and chunk pipelining live in the new dispatcher layer instead of ad hoc MoE hooks.

Checkpoint and model surface: PrimeRL checkpoints rename routed weights from w1/w2/w3 to gate/down/up_proj, move selection bias to mlp.router.selection_bias, and add GPT-OSS HF↔prime conversion. LatentMoE / Nemotron-specific expert classes are folded into the canonical tree; dense layers use FeedForward. Reusable prime conversion caches must include a .prime-v1 marker.

Removed / docs: Fused flash_moe trainer integration and bench_fused_moe.py are dropped; routed_experts selective activation checkpointing is removed. Docs and examples (e.g. GLM-5) are updated for the new tables.

^{Reviewed by Cursor Bugbot for commit ffd5be7. Bugbot is set up for automated code reviews on this repo. Configure here.}

garrett361 · 2026-08-28T19:23:05Z

+    for event in _pending_combine_events:
+        event.current_stream_wait()
+    _pending_combine_events.clear()


For my own understanding, why did we move from syncing via a single global event to a list of events?

Also a little unclear to me why we have a global _pending_combine_events list rather than it being an attribute on each DeepEPTokenDispatcher instance. I see we export sync_combine, but don't seem to actually consume it outside of this file, so I'm missing why it's so decoupled from the dispatcher class

garrett361 · 2026-08-28T19:37:50Z

    score_before_experts: bool = True,
 ) -> _PendingDispatchState:
+    num_input_tokens = hidden_states.shape[0]
+    if num_input_tokens == 0:


when are we dispatching zero tokens? Looks like a debugging path?

garrett361 · 2026-08-28T19:42:48Z

    hidden_states = _unpermute_tokens(hidden_states, state.permuted_indices, state.num_recv_tokens)
-    return _DeepEPCombine.apply(hidden_states, state.handle_id)
+    combined = _DeepEPCombine.apply(hidden_states, state.handle_id)
+    return combined[: state.num_input_tokens]


Why is the slice needed? Would have expected num_input_tokens are automatically returned from the combine

garrett361 · 2026-08-28T19:56:59Z

+
+@dataclass(frozen=True)
+class DeepEPDispatchState:
+    backend: _DispatchState


I find this config a bit confusing. DeepEPDispatchState owns a _DispatchState attr (rather than being a subclass of _DispatchState?) and the owned state is called backend?

garrett361 · 2026-08-28T21:08:43Z

+            if not chunk_ranges:
+                chunk_ranges = [(0, 0)]
+
+        def dispatch_chunk(start: int, end: int) -> _PendingDispatchState:


We are redefining this dispatch_chunk function and run_pending_chunk below every time run is called. Refactor to avoid this?

garrett361 · 2026-08-28T21:28:34Z

+class LocalDispatchState:
+    num_tokens: int
+    token_indices_experts_sorted: torch.Tensor
+    scores_after_experts: torch.Tensor | None


In future work, we can probably simplify this to scores and drop all score_before_experts = True support. More of a note for myself.

garrett361 · 2026-08-28T21:30:36Z

+    token_indices_experts_sorted: torch.Tensor,
+    scores_after_experts: torch.Tensor | None,
+) -> torch.Tensor:
+    if scores_after_experts is not None:


Another optimization note for future work: would like to see if dispatching the scores and fusing the scoring into the swiglu is more benficial than just doing the scoring locally in the combine.

garrett361 · 2026-08-28T21:34:48Z

+        return None
+
+
+class TorchTokenDispatcher(LocalTokenDispatcher):


nit: some odd inheritance patterns here. True comms-dispatching classes being subclasses of LocalTokenDispatcher is a bit off semantically, but then also we end up overriding many of the base class's methods entirely in the subclasses and have trivial methods elsewhere (e.g. LocalTokenDispatcher.sync. Seems like more of an ABC pattern is expected here

garrett361 · 2026-08-28T21:39:35Z

+import torch
+
+
+class GroupedGemm(Protocol):


just noting that I like the use of Protocols here and throughout, think it's a nice pattern

garrett361 · 2026-08-28T21:54:49Z

-        computed = iter(torch.autograd.grad(out, wanted, grad_out) if wanted else ())
-        grads = {p: (next(computed) if leaf.requires_grad else None) for leaf, p in zip(leaves, grad_poses)}
-        return grads[0], grads[1], grads[2], grads[3], None, grads[5], None, None
+    bias = torch.repeat_interleave(bias, num_tokens_per_expert.to(torch.int64), dim=0)


repeat_interleave will cause a cuda sync here, which we should be able to avoid by specifying the output_size. Think that should all be known CPU side at launch time

garrett361 · 2026-08-28T22:03:45Z

+            if self.selection_bias is not None:
+                selection_scores = selection_scores + self.selection_bias
+            if expert_bias is not None:
+                selection_scores = selection_scores + expert_bias
+            _, selected_experts_indices = torch.topk(
+                selection_scores,
+                k=self.top_k,
+                dim=1,
+                sorted=self.topk_sorted,
+            )
+            top_scores = scores.gather(dim=1, index=selected_experts_indices)


This all looks a little weird to me. There are scenarios where we maybe adding both an expert_bias and a self.selection_bias buffer to the selection_scores?

Also, I think selection_bias is an excellent semantic name, but am concerned that this will mess up checkpoint loading for checkpoints which persist the expert bias buffer. Usually the buffer has a different name than this, and I'm not seeing any code that handles a possible necessary fqn conversion.

garrett361 · 2026-08-28T22:18:50Z

+    def prepare_expert_input(self, x: torch.Tensor) -> torch.Tensor:
+        return x

-        shared_output = self.shared_expert(x) if self.shared_expert is not None else None
-        sync_combine()
-        routed_output = routed_outputs[0] if len(routed_outputs) == 1 else torch.cat(routed_outputs, dim=0)
-        return routed_output if shared_output is None else shared_output + routed_output
+    def prepare_expert_output(self, x: torch.Tensor) -> torch.Tensor:
+        return x


This is purely for Nemotron to override?

yes, but I think it'a fairly nice "API surface" to expose, so wouldn't mind it

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit fa8cf1e. Configure here.}

PR #3411 replaced the per-model expert kernels with one canonical MoE tree, so the V4 implementation now composes the shared pieces instead of carrying its own. `DeepseekV4Experts` becomes a `GroupedExperts` subclass. Its stacked `w1`/`w2`/`w3` were already the canonical `[num_experts, out_features, in_features]` shape, so they simply take the canonical names `gate_proj`/`up_proj`/`down_proj`; the conversion chain renames the on-disk per-expert `w1`/`w2`/`w3` onto them via `routed_experts_op`'s `proj_order`. Both `_run_deepseek_v4_experts_*_impl` kernels, the `expert_parallel` decorator, the `ep_comm_backend` plumbing and `set_ep_comm_backend` go with them: `configure_moe_runtime` now picks the grouped GEMM and the token dispatcher, and `ExpertWeightParallel` shards every expert parameter on dim 0 by name-independent placement. V4's clamped SwiGLU stays V4's. Its limit comes from `config.swiglu_limit`, but the shared `ActivationDispatch` holds stateless classes keyed by name, so there is nowhere in it to put a per-model number. `_ClampedSwiglu` carries the limit as instance state and is assigned to `self.activation`, which is the single point `GroupedExperts.forward` reads, so the shared grouped-GEMM forward is inherited rather than copied. `DeepseekV4MLP` reparents onto `FeedForward`, whose `init_weights` already matches what V4's override did. The load-balancing bias moves with the router: `mlp.expert_bias` becomes `mlp.router.selection_bias`, a buffer the router applies to its own selection rather than a forward argument the layer passes down. `DeepseekV4Router` keeps its `forward` override, since `sqrtsoftplus` is outside the shared `ScoreFuncType` and the scoring chain has no extension hook, but it is re-derived from the current base method so it picks up the fp32 gate path and the int64 token counts. `config.use_grouped_mm` is removed, matching every other model's config: the compute backend is now `[trainer.model.moe.compute]`. Tolerances move from the float32 floor to the bf16 one wherever the routed experts are in the comparison. `GroupedExperts.forward` casts to bfloat16 whatever dtype it is handed, while the deleted for-loop path kept float32, which is what these tests ran on. Verified by forcing the experts back to float32 per-expert matmuls, under which every one of these comparisons passes at its old tolerance unchanged. Each new bound records the measured deviation, and the packed-versus-unpacked invariants keep two orders of magnitude of separation from an actual document leak (measured: 5.1e-1 logits and 2.8 gradients when boundaries are deliberately removed, against bounds of 3e-3 and 8e-2). Two test-side consequences worth naming. `_MOE` moves to the real `swiglu_limit=10.0` default: next to a saturating clamp, bf16 rounding flips which entries get clipped, so a clipped entry's gradient jumps between `silu'(gate) * up` and exactly zero, swinging routed-expert gradients 39% against HF instead of 0.66%. `_CLAMPED_MOE` keeps a saturating limit for the shared expert's clamp test, which stays in float32 and compares exactly. And `test_deepseek_v4_dequantize_e2e_hf.py` never imported the autouse `_seed_rng` fixture that every sibling module does, so its deviations moved with whatever ran before it; importing it makes them reproducible. `test_moe_grouped_mm_experts_match_the_for_loop` is deleted. Both implementations it compared now resolve to `GroupedExperts` with `BF16GroupedGemm`, so it would compare a thing to itself. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…spatch PR #3411 removed `ep_comm_backend` without an alias, so the SFT config no longer parses. The backend now lives in its own table, `[model.moe.dispatch] type = "deepep"`. The rationale recorded alongside it is rewritten rather than carried over. The measurement still stands: on a 4-layer single-node run the torch dispatch deadlocked every rank while DeepEP completed 5 steps cleanly. The explanation does not. It cited `DeepEPExpertParallel` in `trainer/distributed/expert_parallel.py` and argued that only DeepEP hoists the token dispatch into `MoE.forward()`, keeping its collectives outside the activation-checkpoint boundary. #3411 deleted that class, routed every backend through one `TokenDispatcher.run()` call inside `MoE.forward()`, and dropped the `routed_experts` selective-AC target, so the asymmetry that explanation rested on no longer exists. Whether the torch dispatch still deadlocks is untested on the new runtime, so DeepEP stays and the comment says so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…el entry PR #3411 left one path through `GroupedExperts.forward`, in bfloat16, where `use_grouped_mm=False` used to select per-expert `torch.matmul` in whatever dtype it was handed. On the mini checkpoint that costs under 1% on the residual stream, which is enough to flip near-tied top-k selections in the score-routed layers: 4, 3 and 12 of 256 selections in layers 2 through 4, moving those blocks' outputs by 28% to 40% and the logits by 1.19 on a scale of 6.7. The two hash-routed layers are the control, with a frozen table, zero flips and textbook bfloat16 error. The sensitivity is generic to top-k routing rather than V4's; what is V4's is noticing it, since its tests were written to float32 exactness while `test_qwen3_moe.py` asserts `atol=1e-0` on logits and `atol=2048` on gradients. The note says so, and says the flip counts were measured on V4 alone. Two follow-ups recorded: `scripts/mini_moe.py --arch deepseek_v4` is left failing its `assert max_diff < 0.1` on purpose, since the bound is a precision one and what breaks it is discrete; and the trainer against vLLM should be expected to disagree on a similar share of selections, which lands on the mismatch-KL and wants measuring on the real checkpoint. The `GptOssGroupedExperts` entry goes: it described the `expert_parallel` decorator's fixed signature clashing with a 6-argument call under `use_grouped_mm=False`, and #3411 deleted the decorator, that class and the flag. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

The EP note's conclusion survives #3411 but its reasoning does not. It said EP worked because `DeepseekV4Experts` held literal `w1`/`w2`/`w3` params, matching the names torchtitan's `ExpertParallel._partition_fn` sharded by. That class is gone. `ExpertWeightParallel._partition_fn` shards every `named_parameters(recurse=False)` on `Shard(0)` without looking at names, so the stacked `gate_proj`/`up_proj`/`down_proj` inherited from `GroupedExperts` are covered by construction rather than by coincidence of naming. Re-verified rather than assumed: an `ep=8` SFT run on `sft-mini-ep-check.toml` against the mini checkpoint gives finite losses (12.66, 12.24, 13.01), nonzero varying grad norms, no NaNs and 12.3 GiB peak, which also exercises `configure_moe_runtime` discovering `DeepseekV4MoE` and installing the token dispatcher. Recorded that this run uses the default torch dispatch and sets no `[model.ac]`, so it says nothing about the DeepEP-versus-torch deadlock `sft.toml` documents under `ac.mode = "full"`. The config's own header is updated too: it described a fused `gate_up_proj`/`down_proj` layout that `f9573853a` un-fused and #3411 then renamed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>