[Dev] Recompute EP A2A overlap activations per layer instead of per model chunk - #6311
Open
Wohox wants to merge 4 commits into
Open
[Dev] Recompute EP A2A overlap activations per layer instead of per model chunk#6311Wohox wants to merge 4 commits into
Wohox wants to merge 4 commits into
Conversation
Full activation recompute under --overlap-moe-expert-parallel-comm previously replayed the whole VPP stage (model chunk) as one unit, so every layer's activations were materialized simultaneously during the chunk backward. Split the chunk's layers into RecomputeSegments instead. Only each segment's input tensor is retained across the forward->backward gap, and a segment's forward is replayed right before its own backward, so at most one segment's activations are live at a time. Peak activation goes from N*I + L*A to N*ceil(L/k)*I + k*A (N in-flight microbatches, L layers, A per-layer activation, I layer input, k = recompute_num_layers). recompute_method / recompute_num_layers now mean exactly what they mean without the overlap flag: 'uniform' groups decoder layers by recompute_num_layers and gives each MTP depth its own segment, 'block' recomputes the first recompute_num_layers decoder layers and skips MTP. The per-layer recompute scope is [attn, moe_dispatch, mlp, moe_combine], mirroring the non-overlap checkpoint scope, which leaves MTP _postprocess outside the checkpoint. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
submodule_mtp_attn_forward stores the attention node's detached input in chunk_state.mtp_hidden_states and submodule_mtp_postprocess_forward concatenates it into the layer output, so the post-process graph reaches back across the node boundary. Building that torch.cat during the initial no_grad forward left it anchored on a non-grad-tracking leaf, dropping the decoder's gradient contribution through the MTP branch (observed as a decoder layernorm weight grad mismatch against the non-overlap reference). Recompute the whole layer instead, and move the segment replay ahead of mtp_post_process.backward. Also park the mHC bridge leaf on the chunk state so the replayed MTP post-process clearing chunk_state.mhc_multistream cannot lose the gradient the MTP backward accumulated on it. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
… test
recompute_method='block' leaves the trailing decoder layers, and every MTP
layer, outside the recompute scope, so the gradient has to cross from a
replayed segment into an eagerly-built graph. Add ('block', 3), which
recomputes every decoder layer of the 3-layer model while its MTP layer stays
eager, alongside the existing ('block', 1).
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
Recompute is per layer segment now, not per model chunk, so the comments left behind in ScheduleNode.forward_no_grad, TransformerLayerNode.reset_for_recompute and the token dispatchers' reset_transient_forward_state no longer describe what the code does. Comment-only. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
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
1 participant
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.
Summary
PR #5869 added full activation recompute to the EP A2A overlap path
(
--overlap-moe-expert-parallel-comm) at model-chunk / VPP-stage granularity:the whole stage runs its initial forward under
no_grad, and every layer'sactivations are rebuilt at once at backward time.
This PR moves that recompute to layer granularity. Only each recompute
segment's input tensor survives the forward->backward gap, and a segment is
replayed with grad enabled immediately before its own backward, so at most one
segment's activations are materialized at a time.
The chunk-level path is removed rather than kept alongside — there is a single
recompute path now.
Semantics
recompute_method/recompute_num_layersnow mean the same thing with andwithout
--overlap-moe-expert-parallel-comm, mirroringmegatron.core.recompute.checkpointed_forward(decoder) andMultiTokenPredictionLayer._checkpointed_forward(MTP):uniform: decoder layers in groups ofrecompute_num_layers; one segment perMTP layer (
recompute_num_layersmust be 1 when MTP is present, as upstream).block: the firstrecompute_num_layersdecoder layers, one segment each;MTP layers are not recomputed (warn + skip, as upstream).
The validation exemptions that let the overlap path leave these unset are gone,
so both flags are required exactly as on the non-overlap path.
Why a segment object rather than a
checkpoint()wrapperThe non-overlap path can wrap a layer in
tensor_parallel.random.checkpointbecause a layer is an opaque callable there. Under EP A2A overlap a layer is
decomposed into
ScheduleNodes (attn / dispatch / mlp / combine) that the 1F1Bscheduler interleaves with other microbatches' nodes — which is the whole point
of the overlap path.
checkpoint()needs its region to run as one contiguousautograd call and to backward as one contiguous unit, so the replay has to be
driven by the scheduler at segment-backward time instead.
Two things follow from N segments per chunk instead of one:
input_ids,position_ids,padding_mask,mtp_hidden_states,mhc_multistream). Theseare mutated in place during the forward, so by backward time they hold
end-of-chunk values rather than what the segment saw.
mhc_multistreamis the only cross-layer detach bridge. Backward runs inreverse, so the MTP segment is replayed and backwarded before the decoder
segment that produces the leaf, and that replay installs a freshly detached
tensor. The gradient is carried across the replay explicitly.
mtp_post_processis inside the recompute scope: excluding it breaks thetorch.cat(mtp_hidden_states)graph and silently drops decoder weight gradients.Testing
Unit tests (
tests/unit_tests/a2a_overlap/):test_recompute_segments.py— 15 CPU-only tests over the segmentation rules(uniform / block, MTP present or not, ragged tails, guards).
test_schedule_chunk_1f1b.py::test_1f1b_schedule_model_chunk_full_recomputecompares gradients from
TransformerModelChunkSchedulePlan.runagainst thenon-overlap
gpt_model.forwardreference under the same recompute config,parametrized over
mtp_layers0/1, dispatcher configs,get_valid_fp8_flags()and
(uniform,1) (uniform,2) (block,1) (block,3).Note that
get_valid_fp8_flags()is architecture-gated:blockwiseis Hopperonly and
mxfp8is Blackwell only, so no single CI machine exercises both.E2E, DeepSeek-V3-Lite-deter (the deterministic benchmark recipe from
megatron-moe-scripts: TP1 PP4 EP2 CP1, VPP2, MBS1 GBS16, seq 256, 16 layers,
flex/DeepEP dispatcher,
unfusedattention, 100 iterations, one node). Everycomparison below runs each arm twice and requires both arms to reproduce
themselves before the arms are compared:
The rows above run the recipe's NonMTP branch. Repeating them on its MTP branch
(
--mtp-num-layers 1, 8 layers, every layer MoE,--pipeline-model-parallel-layout Et|(t|)*6tmLin place of VPP) — which is what exercisesmtp_post_processinsidethe recompute scope and the
mhc_multistreamgradient carry-over:Full recompute is itself bitwise neutral in this recipe: on the non-overlap arm,
recompute on and recompute off agree 100/100.
Memory, measured separately on the DeepSeek-V3 proxy config (TP1 PP2 EP2 VPP2,
MBS2 GBS64, blockwise FP8): peak allocated 66,764 MB -> 21,662 MB (-67.6%),
which lands on top of the non-overlap full-recompute figure (21,325 MB).
Throughput 96.44 -> 68.23 TFLOP/s/GPU, retaining +1.8-3.6% over non-overlap full
recompute.
Stripping the debug instrumentation used during development was verified to be
numerically inert: the submitted tree and the instrumented tree agree bitwise on
both arms.
BF16 and blockwise ran on Hopper at TP1 PP4 EP2 CP1; mxfp8 is Blackwell-only and
that node has 4 GPUs, so it ran at TP1 PP2 EP2 CP1. Both arms of each row share a
config, so the within-row comparison is unaffected.
Not covered: the unit tests have not been executed on Blackwell, so the mxfp8
branch of
get_valid_fp8_flags()is exercised only by the e2e row above.Delayed-scaling FP8 remains rejected by validation on this path (inherited from
#5869) — the rationale is in the assertion comment, but the restriction itself
has not been measured.