Price the recomputed layer's attention or GDN activations in the checkpoint floor

[bradhilton]

Part of #949. This makes TrainerRank's cold memory estimate count what is actually live at the backward peak under full one-layer recompute. The expert-parallel routing allowance also now depends on the EP size.

Builds on #1002, which is merged into this branch; review #1002 first. This PR prices #1002's adapter-gradient term against its own one-incoming-gradient floor. On the full stack with #1002 (#963 + #978 + #981; Qwen3.6-35B-A3B, 40 layers, CP2, one real sequence per wave), the raw estimate against the measured peak is:

• EP1 cold first waves (2k to 20k tokens): +0.9% to +2.0%;
• EP2 cold first waves: +3.0% to +13.4%;
• warm waves: +2.8% to +24.9%.

None are under. The table below predates #1002.

What changes:

• Recomputed mixer. The recomputed layer's attention or GDN activations stay live beside its MoE stage, and the estimate now counts them. Leaving them out under-predicted CP2/EP1 by 23%.
• Incoming gradient. When every layer is a MoE layer whose expert stage, FC1 included, the estimate prices, it charges the one incoming gradient live at the peak, not one per layer. The old charge over-counted about 17 GB on the random-weight 40-layer run and 8.4 GB on the real-data one. Other models keep the old charge, which also covers recompute work the estimate doesn't price.
• Items the old charge hid. These are now priced: the recomputed layer's residual and norm rows, a sixth GDN value tensor, the shared expert's FC1/GLU outputs, router and permutation state, and Transformer Engine's cuBLAS workspaces (first call only).
• HybridEP dispatch. HybridEP keeps one dispatched copy of each routed row, not two. It dispatches the whole EP group's tokens, so when the EP group is the CP group, routed rows are now priced on the group's even share, not on the busiest CP rank's rows. For example, on a real-data CP2 batch one rank held 52,480 rows, but each rank's even share of routed tokens was 48,397.
• EP routing allowance. It was a flat 1.5 and is now 1.4 / 1.6 / 2.0 at EP2 / 4 / 8. The values come from a forward-only routing measurement on 3.5M tokens of real retail trajectories. In 200k-token batches the worst layer reached up to 1.22 / 1.40 / 1.62 with pretrained weights and 1.24 / 1.41 / 1.61 with a trained policy; a small rollout sample reached 1.95 at EP8, close to its 2.0 allowance. The old flat 1.5 was too low for EP8. EP4 and EP8 have no memory runs.

Raw estimate vs measured PyTorch peak on the most-loaded rank, Qwen3.6-35B-A3B on two H200s. Random-weight runs use 194,753 tokens. Real-data runs use about 206k tokens of real retail agent trajectories (97k after prefix sharing).

case	before (cold)	after (cold / warm)
CP1	+2.1%	+0.5% / +15.9% ¹
CP2/EP1	−23% (under)	+0.3% / +6.8% ¹
EP2/CP2, random weights, 8 layers	+38.5%	+26.5% / +23.6%
EP2/CP2, random weights, 40 layers	not run	+19.7% / +17.5%
EP2/CP2, pretrained, real data, 40 layers	not run	+18.2% / +19.2%
EP2/CP2, trained policy (059, step 225), real data, 40 layers	not run	+18.3% / +19.4%

¹ Measured on an earlier commit of this PR. The later changes don't alter the estimates for these two configurations.

In the pretrained EP2 cold case, the estimate is 26.81 GB against a 22.67 GB PyTorch peak. 1.05 GB of that estimate is HybridEP's communication buffer, which lives outside PyTorch, so the PyTorch peak can't show it. Leaving it out, the estimate is +13.6%. The GPU's memory outside PyTorch grew 2.76 GB during that call; the estimate prices only the HybridEP buffer part of that.

Limitation. EP2 is still well above the 10% target. With this formula and allowance, on these measured workloads, most of the remaining gap is the routing allowance: 1.4 is priced, while the layer at the peak saw about 1.0. The allowance covers the worst imbalance measured: 1.16 to 1.24 per layer on real data, and about 1.35 inferred from one production run. A cold estimate can't know which layer will be imbalanced, and these routing samples bound only what was measured, not all possible routing. Warm calls could price the routing they have already seen; that is a separate proposal.

Testing: all trainer-rank unit tests pass. The table's runs are on local H200s.

🤖 Generated with Claude Code