Release HybridEP dispatcher state after combine

[bradhilton]

Part of #949. #858 and #861 stopped MCore's EP1 all-to-all dispatcher from keeping each MoE layer's checkpoint graph alive after backward. The flex dispatcher that ART uses at EP>1 with HybridEP has the same retention and was left out.

MCore's _HybridEPManager keeps routing_map, token_probs and dispatched_probs after combine, until that layer's next setup_metadata and dispatch. The dispatched probabilities are a differentiable output of the dispatch, so their graph keeps the following alive:

• the layer's finished checkpoint graph;
• the recomputed layer input (a checkpoint leaf) and its .grad;
• the router probabilities;
• the dispatch handle.

TrainerRank now adapts exact MoEFlexTokenDispatcher instances whose manager is exactly _HybridEPManager, in the model it is given, the same way #861 adapts combine_postprocess. After the upstream combine_postprocess returns, it clears those three manager fields. Installation also clears any state already held. Backward keeps what it needs through its own graph, and the next setup and dispatch recreate these fields. Not adapted:

• CUDA graph capture, which reads the fields back from the manager;
• instances with their own combine_postprocess;
• DeepEP managers;
• other dispatcher types.

Evidence

Qwen3.6-35B-A3B with random weights, 8 layers, CP2/EP2, 194,753 tokens, on two local H200s with #955. The allocator traces are the ones from #963's requalification:

• Cold call. At the peak on rank 1 (35.24 GB), 4.0 GB belongs to other layers: 1.92 GB of router, probability and handle state across seven layers, and 2.10 GB of retained input gradients from the five layers already backwarded. The gradients are allocated in MCore's checkpoint backward (random.py:630, autograd input-buffer accumulation).
• Warm call. The working set equals the cold one (35.01 GB), but 9.28 GB of the previous call's state is freed during the call. Of that, 7.37 GB is freed when each layer's dispatch replaces dispatched_probs and handle, and 1.03 GB of router state in setup_metadata. That is about 1 GB per MoE layer at this size, scaling with layers and tokens. Admission counts it as used memory.

Validation on two H200s

Three EP2/CP2 runs of 3 calls each, with random seeded weights. The weight and input fingerprints are identical across runs. Two base runs used #963's head 4b3d1e059; the third added this change. All figures are for rank 1 (105,153 tokens), except the rank-0 bullet.

	base	with this change
cold peak, including memory held before the call	44.44 GB	40.71 GB
warm peak, including memory held before the call	45.03 GB	41.31 GB
memory held between calls	19.30 GB	10.91 GB
call time, cold / warm / warm	22.1 / 4.5 / 3.2 s	22.7 / 5.1 / 3.3 s

The second base run's warm call also took 5.1 s. That is ordinary variance: collection time alone varies by 0.6 s between the two base runs.

• Rank 0 drops by 4.4 GB at both peaks and by 7.78 GB between calls.
• Traces. The allocator traces with this change show no other-layer HybridEP state or retained input gradients at the cold peak. The warm call frees nothing from the previous call except the harness's own objects (0.89 GB).
• Numerics. Losses are bitwise identical between runs for each corresponding call. Gradients at EP2 are not bitwise deterministic run to run even without this change: the two base runs differ by 1.30–1.32e-2 relative L2, in the active LoRA slot's gradients of 7 of 8 layers. This change differs from each base run by the same 1.30–1.32e-2, with matching worst-tensor ratios, so it is indistinguishable from run-to-run noise at that level. It only drops Python references to tensors the graph keeps anyway, so by construction it changes no computation. The measurement can bound an effect only above about 1e-2.
• Planner. Admission sees the smaller held state directly. Predictions do not change here. With Price the recomputed layer's attention or GDN activations in the checkpoint floor #963's planner, the EP2 raw prediction's margin over the measured peak grows, because the retained state had been filling part of the floor's allowances.

Schulman's independent audit of a full-depth (40 MoE layer) diagnostic run found the same retention: after a final GC, 40 managers still held these three fields, 9.64–10.23 GB of CUDA tensors directly, all with HybridEPDispatchBackward.

Full-depth run (Schulman, independently audited). This change alone was added to a frozen ART tree, and a synthetic 19-history step was run: 1,073,093 tokens in groups of 10 and 9, on two H200s. Everything else in that setup stayed as it was: the original callback, probe, lane and compiler-GC settings, and the selective FC1 offload.

• The step completed: both forwards and backwards, and one finite update.
• First group, allocated peak per rank: 123.38 / 128.30 GiB before, 72.50 / 81.72 GiB after.
• After its backward: 79.65 / 85.98 GiB before, 5.84 GiB after.
• Second group: the baseline refused it for lack of memory. With this change, normal admission accepted it, needing 32.66 GiB with 40.80 GiB usable and no cache recovery.

The before and after runs used different physical GPU pairs and fresh allocator and compiler state. So these are a historical comparison, not an exclusive attribution. The audit makes no claim about numerical parity, one-GPU fit or production throughput. The remaining first-group estimator miss it reports belongs to that frozen tree's planner, not to current main or #963.

Tests

tests/integration/megatron/model_support/test_dispatcher_graph_retention.py is in the Megatron lightweight CI lane. It drives the pinned MCore flex dispatcher and _HybridEPManager methods (dispatch_preprocess, setup_metadata, dispatch, combine, the post-processing steps) through CheckpointFunction. Only HybridEP's fused dispatch and combine kernels are replaced, by CPU equivalents. The tests cover:

• all four checkpoint inputs retained after backward before adaptation, none after, with bitwise-equal loss and input and router gradients. This runs eager and compiled with aot_eager: the compiled case installs into an already warmed model without a compiler reset, and the adaptation is also installed with the checkpoint backward still pending;
• installation dropping state that an earlier, unadapted forward left behind;
• the manager's fields cleared after every combine while backward still produces finite gradients;
• two outstanding forwards with repeated retain_graph backward, eager and compiled, with and without checkpointing, matching the unadapted model bitwise;
• idempotent installation, with the class method untouched;
• no adaptation for DeepEP managers, CUDA graph capture, instance combine_postprocess overrides, or dispatcher and manager subclasses;
• pickle and deepcopy preserving the adaptation.

The whole file passes locally (38 tests, run with GPUs hidden). Deleting either release makes its tests fail: without the release after combine, 11 fail; without the release at installation, the installation test fails. ruff and ty are clean on the changed files.

🤖 Generated with Claude Code