Repository navigation
Conversation
added 4 commits
October 19, 2026 09:57
Apply router scores to the expert activations before the down (W2) projection instead of multiplying the expert output in combine, so the combine reduces to a plain sum: - Token dispatchers expose absorb_router_scores. When set, dispatch returns routed_scores_R in its metadata and combine skips its score multiply. Covered backends: standard (local + all-to-all), TorchAO (EP=1 padded and EP>1), and DeepEP/HybridEP (opt-in). - GroupedLinear.forward takes routed_scores_R and scales the projection input in float32, restoring the input dtype before the grouped GEMM. RoutedExperts passes it to W2 only, and only when a score is actually absorbed, so W2 modules that do not know about scores are unaffected. - GptOssGroupedLinear scales its post-W2 bias rows by the same scores, so the result stays score * (h @ W2 + b2). - DeepEP/HybridEP dispatchers take the per-row scores the backend already tracks and clear them from the dispatch state, which is what makes the backend combine skip its fused score multiply. - make_routed_experts_config leaves the standard backend on auto and the DeepEP family off. Auto resolves when RoutedExperts is built: the score is absorbed unless output_postprocess owns the expert outputs, because absorbing scales W2's input while the postprocess runs after W2, and a scale-invariant postprocess (e.g. post-expert RMSNorm) would drop the score. Asking for both explicitly raises a ValueError. - swap_token_dispatcher propagates the flag through MoE quantization. Replay: with router scores absorbed, combine no longer scales expert outputs, so the DeepEP compact combine would read the W2 output from a bare _unpermute_tokens on the replay path. Fold the score multiply, unpermute and combine kernel into one never-replayed region, which keeps the W2 and combine outputs off the saved-tensor set. On the standard all-to-all path the absorbed scores are an output of the saved dispatch region, so mark them before W2's bare scaling; that marker is a no-op on the DeepEP/HybridEP paths and without an active checkpoint. Squashed from the router-score-pre-w2-review series and rebased onto current upstream main. Absorbed and post-combine paths agree to bf16 rounding; loss convergence on a representative dataset is still pending.
Give the local, all-to-all and TorchAO dispatchers the combine structure DeepEP got with router-score absorption: one never-replayed "<fqn>.combine" region whose body reads the w2 output, instead of bare ops guarded by recompute_needs_tensor markers. - Local and TorchAO EP=1 ran the scatter-add (and the padded unpermute) as bare ops, so replay needed the w2 output persisted: the marker pinned the full expert-output tensor (R x D) per MoE layer even with absorbed scores, where no backward consumer needs it (the scatter-add saves only indices). - The EP>1 combine region was policy-replayable. Replaying it re-ran the combine all-to-all and required the w2 output real on the replay path (pinning it under policies that save w2), while saving it retains only the small permutation indices and scores, so no policy benefits from replaying it. Dispatch stays policy-driven everywhere, like DeepEP's dispatch. Replay: the combine body runs only on the original forward; replay skips it, so nothing on the replay path reads the w2 or combine outputs, and with absorbed scores the w2 output is freed after the forward. Also fix the stale DeepEP docs and AC test left by the absorption commit: the DeepEP combine region is never replayed, so only dispatch call counts grow under the full policy. Tests: - New tests/unit_tests/cpu/models/common/test_token_dispatcher.py: the combine body runs exactly once under every policy while the w2 projection's replay still follows the policy; with absorbed scores and w2 saved, the w2 output storage is freed after the forward; loss and grads match the un-checkpointed run bitwise. It fails on the previous code. - test_remat_all_to_all: the combine never replays, so the full policy replays 2 collectives (was 3) and saving dispatch alone stops all replay; the w2-retention test now saves only w2.grouped_mm, pinning that the combine is always saved. Docs: REMAT.md and the activation-remat config docstrings now state that the dispatcher combines are always saved and only dispatch follows the policy. Numerics: remat-structure change only. On CPU, loss and gradients are bitwise identical to the un-checkpointed run (and to the previous code) for the local and EP=1 paths and for the EP>1 path through a 1-rank EP mesh, absorbed and not, under the full, default and save-w2-only policies.
…hers Add two named debug recipes so the token-dispatcher backend matrix is runnable on DeepSeek V3, alongside the existing HybridEP config: - deepseek_v3_debugmodel_alltoall: the stock dispatcher at EP=2, so routed tokens actually cross ranks. The plain deepseek_v3_debugmodel (EP=1) already runs the same dispatcher's local fallback, which covers the local backend without a duplicate config. - deepseek_v3_debugmodel_deepep: DeepEP v2 at EP=2 via TokenDispatcherTransform, following qwen3_moe_deepep: the compact, host-synced training path, buffer capacity from the microbatch shape, CUDA graphs off like the stock all-to-all, plus the NVLink-only env notes for local devgpu. Both keep the default activation-remat policy, under which every dispatcher's combine is saved (never replayed) and the standard backend absorbs router scores before W2, so these exercise the paths changed by the never-replayed-combine work. Registered in the debug-config defaults test. Validated on CPU: both configs build, DeepEPTokenDispatcher.Config gets hidden_dim=256 (from w13.in_features) and num_max_tokens_per_rank=16384 (the 8 x 2048 microbatch over CP x TP), and the base EP=1 config keeps CUDA graphs on while both EP=2 configs disable them.
One fp64 parity test per dispatcher family: with rounding gone, the absorbed and post-combine paths compute the same function, so a dropped, doubled or misaligned router score is an O(1) error. Both tests assert the loss and every input/parameter gradient at the float32 scale W2's input scaling applies. - Standard dispatchers (GPU, world 2): local, all-to-all at EP=1 (local fallback) and EP=2, TorchAO at EP=1 (padded permute, zero-sentinel scores) and EP=2, through the existing _grouped_mm override seam. - DeepEP (CPU, gloo, world 2): stub deep_ep's ElasticBuffer with a pure-PyTorch implementation (all_gather + local selection, deterministic layout, autograd-correct dispatch/combine backward), so the dispatcher, the torchtitan deepep wrapper, its custom ops and their autograd run in float64 without the bf16-only CUDA kernels. The stub module is injected only when deep_ep is absent, and the CPU op impls reuse the wrapper's CUDA impls. Replaces the bf16-tolerance and activation-replay variants and the HybridEP-only file. Verified the DeepEP test fails on double-applied and on misaligned (rolled) scores.
pytorch-bot
Bot
added
ciflow/fake-pg
Mystri
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
October 19, 2026 16:35
meta-cla
Bot
added
the
CLA Signed
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
1 similar comment
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
|
Hi @Mystri — I think the @ tianyu-liu mention in the PR description was intended for @ tianyu-l, who is listed as a requested reviewer. Could you please update that mention to @ tianyu-l? Thanks! |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[MoE] Apply router scores before W2 projection
Supersedes #4095: same motivation, reworked so the combine of every dispatcher is one never-replayed remat region and the router scores are absorbed before W2.
Motivation
Current
RoutedExpertsapplies router scores to each routed expert output after the W2 projection inTokenDispatcher.combine:This multiplication causes the post-W2
routed_output_RDto retain for backward. Because W2 is linear, the router score can be applied earlier to the SwiGLU activation, immediately before W2:See the diagram below:
Expected memory gain in full model:
Combine output:
T * D * 2 = 4096 * 7168 * 2 = 56.0 MiBper layer,* 58 layers = 3.17 GiBW2 output:
(T * top_k) * D * 2 = 4096 * 8 * 7168 * 2 = 448.0 MiBper layer,* 58 layers = 25.4 GiBTotal per GPU:
25.4 + 3.17 = 28.6 GiBExpected recomputation gain in FullAC
None added: the w2 grouped GEMM still replays exactly once per MoE layer, and the combine is a never-replayed region whose body (score multiply, unpermute, combine communication) runs only on the original forward. Measured on the debug model under FullAC, the replay path also rebuilds fewer saves (15 -> 5 per forward at EP=1), because the absorbed combine is a pure reduction that saves fewer tensors.
Implementation
The dispatcher returns the per-row router scores in its dispatch metadata,
GroupedLinear.forward(..., routed_scores_R=...)scales W2's input in float32 before the grouped GEMM, and combine reduces to a plain sum.Config:
make_routed_experts_config(absorb_router_scores: bool | None = None)forwards the flag to thestandard
AllToAllTokenDispatcher.Config.None(auto) resolves at build time inRoutedExperts.__init__to "absorb unless the experts carry anoutput_postprocess".TokenDispatcherTransform(dispatcher=..., kwargs={"absorb_router_scores": True});their fused combines apply scores themselves, so they default to
False.swap_token_dispatcher(quantization) carries the flag across dispatcher swaps.Dispatcher:
all-to-all splits), gathers them expert-major with the same
permuted_indices, and returns themas
routed_scores_Rin the metadata (Nonewhen not absorbing --torch_remat.regionsupportsNonetuple leaves). TorchAO overrides the gather to zero-fill padded rows.combineskips the score multiply when the scores were absorbed (_score_and_scatter_addtakestopk_scores_N: Tensor | None).dispatchtakes the tracked scores out of the backend state and clears them, sothe fused combine applies no score; the scores ride
EPDispatchMetadata.routed_scores_R.W2:
GroupedLinear.forward(..., routed_scores_R=None)scales the input in float32 and casts backbefore the grouped GEMM;
GptOssGroupedLinearscales its post-W2 bias rows the same way.RoutedExpertsrejectsoutput_postprocess+ absorption (post-W2 postprocess must see unscaledoutputs) and resolves auto at build time.
Remat region improvement
Problem:
routed_output_RDfrom the producing w2 region's durable output slots, sorouted_output_RDisforce-required: the
recompute_needs_tensor(routed_output_RD)markers pack the full(T * top_k) x Dexpert-output tensor onto the tape as SAVE_OUTPUT.With the scores absorbed no backward consumer is left (the scatter-add saves only indices), and the marker
would still pin the tensor for the replay. FullAC additionally re-run the combine all-to-all.
Solution: the unpermute (with the score multiply and the combine communication) becomes one
never-replayed
<fqn>.combineregion (recompute=False, like the routing decision and thetrailing adds). The body runs only on the original forward and reads
routed_output_RDonlythere, so
routed_output_RDis not saved: nothing on the replay path reads it anymore, and withabsorbed scores no backward consumer is left either, so it is freed right after the forward
(under a replaying policy, w2's replay rebuilds only what w2's own backward needs, without the
unpermute re-running on top of the output). The
recompute_needs_tensor(routed_output_RD)markers are gone. Saving the combine retains only what its backward needs -- the permutation
indices, plus the scores (and the w2 output they multiply) in the baseline, indices only with
absorption -- so no policy benefits from replaying it; dispatch keeps following the policy
everywhere, like DeepEP's dispatch.
Measured memory/recompute improvment
Measured by adding:
in
def _non_pp_forward_backward_microbatchintraining_engine.py.Local dispatcher
deepseek_v3_debugmodelwithLocalTokenDispatcher, FullAC (T=16384, top_k=3). Under fullrecompute nothing MoE-related stays resident, so the difference is what the replay path must
rebuild:
All2All dispatcher
The config used for the FullAC A/B (the all-to-all dispatcher falls back to its local methods at
EP=1, so the numbers match the local row above):
Baseline output (absorb only removes rebuilt saves here, same as local):
Absorb output:
At EP=2 with RegionAC save-all, the baseline combine region retains the same tensors as DeepEP
(per MoE layer; the 24.00 MiB bf16 row is the w2 output):
With absorbed scores the combine retains only the permutation indices, as measured for DeepEP
below.
DeepEP dispatcher
deepseek_v3_debugmodel_deepep(T=16384, top_k=3, EP=2, RegionAC save-all),ep_communication.combineregion per MoE layer.
Baseline:
Absorb:
Debug model A/B (
deepseek_v3_debugmodel_deepep, T=16384, top_k=3, EP=2, RegionAC save-all):ep_communication.combineretainedWhole-run resident: 2.52 GiB -> 2.43 GiB. Nothing else in the report changes.
671B (T=4096, top_k=8, D=7168, 58 MoE layers, per GPU): w2 output
R*D*2= 448 MiB/layerT*D*2=56 MiB/layer (independent marker cleanup, 3.17 GiB) = 28.6 GiB/GPU avoided retention, with no
recompute cost in exchange (w2 GEMM and combine are not replayed).
HybridEP dispatcher
Same mechanism as DeepEP: dispatch takes the per-row scores the HybridEP backend tracks and
clears them from the dispatch state, so
hybridep.combine_tokensskips its fused multiply and thescores ride
EPDispatchMetadata.routed_scores_R. Opt-in like DeepEP. No separate A/B run: theHybridEP buffer is GB200/NVLink72-only.
Tests and verifications
tests/unit_tests/gpu/distributed/test_standard_dispatchers_absorb_router_scores.py,GPU, world 2): local, all-to-all at EP=1 (local fallback) and EP=2, TorchAO at EP=1 (padded
permute, zero-sentinel scores) and EP=2. With rounding gone, the absorbed and post-combine
paths compute the same function, so a dropped, doubled or misaligned score is an O(1) error;
loss and every input/parameter gradient are asserted at the float32 scale W2's input scaling
applies (rtol 1e-5).
tests/unit_tests/cpu/distributed/test_deepep_absorb_router_scores.py,CPU, gloo, world 2): the deep_ep kernels are bf16-only, so the test stubs
ElasticBufferwitha pure-PyTorch implementation (all_gather + local selection, deterministic layout,
autograd-correct dispatch/combine backward) and runs the dispatcher, the torchtitan deepep
wrapper, its custom ops and their autograd in float64. Same assertions; verified to fail on
double-applied and on misaligned (rolled) scores.
tests/unit_tests/cpu/models/common/test_token_dispatcher.py): the combine bodyruns exactly once under every policy while the w2 projection's replay still follows the policy;
with absorbed scores and w2 saved, the w2 output storage is freed after the forward; loss and
grads match the un-checkpointed run bitwise.
Verification reports
torch_remat+RegionAC): absorb vs post-combine scoring are numericallyequivalent (rel diff ~0.25% = bf16 rounding); gradients flow to the scores in the absorbed path.
dispatch/combineregions: backward clean,w2.grouped_mmretains nothing with bf16 params, combine retains indices only, backward runsexactly one grouped GEMM (the w13 replay).
format_saved_tensors_report.DeepEP through a pure-PyTorch
deep_epstub (CPU) agree in loss and all gradients at thefloat32 scale the absorb path applies (~1e-7 relative; asserted at rtol 1e-5). The DeepEP test
fails on double-applied and on misaligned scores.
FusedMOE
The fused MoE path is an optional, SM100+-only backend that only activates when explicitly selected, so on every path we actually run (all-to-all / DeepEP / HybridEP) this optimization applies unchanged — and even when selected, it replaces the module wholesale and is unaffected rather than broken.