perf(moe): reuse the DeepEP dispatch layout on activation-checkpoint recompute - #3684
perf(moe): reuse the DeepEP dispatch layout on activation-checkpoint recompute#3684akoumpa wants to merge 3 commits into
Conversation
akoumpa
changed the title
…ation checkpointing
Combining dispatcher="hybridep" with activation checkpointing dies in the first
backward with an opaque error:
CheckpointError: Recomputed values for the following tensors have different
metadata than during the forward pass.
saved: torch.Size([2791, 2816])
recomputed: torch.Size([2701, 2816])
The message points at the checkpoint machinery, which sends readers to the
router: the documented cause of a shape change like this is non-deterministic
re-routing, and ignore_router_for_ac exists to prevent it. That is not what is
happening here.
Instrumenting Gemma4Gate on 8xH100 with DiffusionGemma-26B-A4B shows the
routing is reproduced exactly. Comparing the encoder pass against its recompute,
the top-k index checksum matches on every rank:
rank 0: fwd 282623 == rec 282623 rank 4: fwd 425381 == rec 425381
rank 1: fwd 275769 == rec 275769 rank 5: fwd 279545 == rec 279545
rank 2: fwd 276947 == rec 276947 rank 6: fwd 278018 == rec 278018
rank 3: fwd 273901 == rec 273901 rank 7: fwd 262048 == rec 262048
and all 30 decoder layers match as well (0/30 mismatches per rank). Despite
identical routing on every rank, hybrid_ep_dispatch returned 2791 rows in the
forward and 2701 on replay. The drift is inside the dispatch, so no router-side
mitigation can prevent it; per-layer dispatch managers
(dispatcher_share_token_dispatcher=False) do not help either, and pad_multiple
is set once at construction and never reassigned.
DeepEP's dispatch replays reproducibly and is unaffected.
Warn rather than raise: the failure depends on whether replayed token counts
happen to drift for a given model and shape, and other recipes run HybridEP with
activation checkpointing today. The warning names both working configurations --
dispatcher="deepep" with checkpointing, or HybridEP with checkpointing disabled
(measured the fastest of the two: 5.46 s/step vs 5.83 s/step on 8xH100).
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
…recompute Activation checkpointing replays a block's forward during backward, and the replayed MoE dispatch recomputes its routing layout from scratch -- get_dispatch_layout() plus DeepEP's notify_dispatch -- even though the routing is identical to the forward's. The AC policy pins the router projection and top-k precisely so that recompute routes the same way, so recomputing the layout from that identical routing is redundant work. DeepEP already exposes the shortcut: handing dispatch() the handle from a previous dispatch skips the layout exchange. FusedDispatch.backward already does this, and the profile shows what it is worth -- on 8xH100 with DiffusionGemma-26B-A4B, per step: forward dispatch (computes layout) 150 launches 4420 ms backward dispatch (reuses handle) 60 launches 3.96 ms notify_dispatch (forward) 150 launches 4844 ms cached_notify_dispatch 60 launches 797 ms Of the 150 forward dispatches the model only needs 90 (30 layers x 3 passes: causal encoder, no_grad self-conditioning decode, real decode). The other 60 are activation checkpointing replaying the encoder and decode-2 passes. A recorder is created per checkpointed call, so the forward logs each dispatch's handle and routing metadata in call order and the recompute consumes them in the same order -- a block called once per pass keeps its passes separate. Cached-mode dispatch returns only recv_x, so the routing metadata comes from the record; recv_x itself is deliberately not retained, so the activation memory that checkpointing saves is preserved (only the handle and the per-token routing metadata are held, ~1.7% of the dispatched activation). If a replay ever outruns its log the recorder falls back to a full dispatch rather than replaying a mismatched layout. Scoped to the ignore_router branch on purpose: that is the path that pins routing across recompute. On the non-ignore_router path the router is recomputed and may route differently, where a replayed layout would silently mis-route instead of failing loudly. Measured on eos, 8xH100, 8 steps, DiffusionGemma-26B-A4B SFT with dispatcher=deepep: before 10.6 s/step after 9.77 s/step (-7.8%) Loss, dllm_loss and grad_norm are bit-identical to the unpatched run for all 8 steps, as expected for a change that only avoids recomputing a layout that was already determined by the (pinned) routing. Two DeepEP configuration knobs were measured on the same workload and made no difference, so they are not used here: dispatcher_async_dispatch=True (10.66 s/step) and dispatcher_num_sms=32 (10.61 s/step). Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
e2ae813 to
6644561
Compare
|
/ok to test 6644561 |
|
Re-ran the benchmark against the exact code in this PR (the numbers in the description came from a functionally-equivalent earlier form of the context manager, so this closes that gap): Loss,
Matching |
The HybridEP activation-checkpointing warning added in this PR walked the
model with model.modules(), which made apply_ac require more of its argument
than the checkpointing itself does and broke all 19 apply_ac unit tests
("AttributeError: 'DummyModel' object has no attribute 'modules'"). The walk
only gates a warning, so skip it when the model does not expose modules().
Also fix two test doubles that the new code exposed:
- create_selective_checkpoint_contexts stubs returned a bare "CTX" sentinel,
but the real torch API returns (forward_ctx, recompute_ctx) and the block
context wrappers unpack it. Return a pair instead.
- fused_a2a is imported lazily, so test_parallelizer only passed when some
earlier test had already imported it under the real torch. Stub it in
_import_parallelizer_with_stubs so the module stands alone.
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
|
/ok to test ee94057 |
nemo-ci validation on AC + EP recipes: no regression foundRan this branch against 10 nemo-ci recipes that combine activation checkpointing with expert parallelism, paired against a matched baseline at this PR's exact merge-base (
The set was chosen to cover the path this PR actually changes ( Result: 8/10 pass on both arms; the 2 failures reproduce identically on the baseline
No Why the non-identical rows are not attributable to this PRTwo independent controls, since raw loss deltas alone can't separate a real change from recipe noise:
The two failures
HybridEP warningFires exactly once, only on the Scope of this runThis validates correctness and no-regression, not the perf claim. Measured step-time noise floor from the two same-commit baseline runs is ±3% ( |
1 participant
Two shared-MoE changes found while profiling DiffusionGemma-26B-A4B on 8×H100. Neither is dLLM-specific — both apply to any MoE model using expert parallelism with activation checkpointing.
1. Reuse the DeepEP dispatch layout on recompute (the optimization)
Activation checkpointing replays a block's forward during backward, and the replayed MoE dispatch recomputes its routing layout from scratch —
get_dispatch_layout()plus DeepEP'snotify_dispatch— even though the routing is identical to the forward's. The AC policy pins the router projection and top-k precisely so recompute routes the same way, so recomputing the layout from that identical routing is redundant.DeepEP already exposes the shortcut: handing
dispatch()the handle from a previous dispatch skips the layout exchange.FusedDispatch.backwardalready does exactly this, and the profile shows what it's worth. Per step, rank 0:ctx.handle)notify_dispatch(forward)cached_notify_dispatchOf those 150 forward dispatches the model only needs 90 (30 layers × 3 passes: causal encoder,
no_gradself-conditioning decode, real decode). The other 60 are activation checkpointing replaying the encoder and decode-2 passes.How
A recorder is created per checkpointed call, so the forward logs each dispatch's handle and routing metadata in call order and the recompute consumes them in the same order — a block called once per pass keeps its passes separate, which matters here because the three passes route differently.
Cached-mode
dispatch()returns onlyrecv_x(recv_topk_idx,recv_topk_weightsandnum_recv_tokens_per_expert_listall come backNone), so the routing metadata comes from the record.recv_xitself is deliberately not retained, so the activation memory checkpointing saves is preserved — only the handle and per-token routing metadata are held, ~1.7% of the dispatched activation. If a replay ever outruns its log, the recorder falls back to a full dispatch rather than replaying a mismatched layout.Scoped to the
ignore_routerbranch on purpose: that is the path that pins routing across recompute. On the other branch the router is recomputed and may route differently, where a replayed layout would silently mis-route instead of failing loudly.Measured
eos, 8×H100, 8 steps, DiffusionGemma-26B-A4B SFT with
dispatcher=deepep:Loss,
dllm_lossandgrad_normare bit-identical to the unpatched run for all 8 steps — matchinggrad_normmeans the backward produces identical gradients, which is the correctness bar for this change.Two DeepEP knobs were measured on the same workload and make no difference, so they are not used:
dispatcher_async_dispatch=True(10.66 s/step) anddispatcher_num_sms=32(10.61 s/step).Note (not fixed here)
The
torch.ops.deepep.dispatch/hybridep.dispatchentries inactivation_checkpointing.py's selective-AC save list are inert: dispatch runs throughFusedDispatch.apply, anautograd.Function, which a__torch_dispatch__policy never observes. Those four entries have never matched anything. Worth a separate fix.The dLLM recipe change that motivated this profiling is #3654 (recipe YAMLs only).