Price the adapter gradients a recompute backward holds at its peak

[bradhilton]

A full-recompute backward allocates each layer's adapter gradients as it passes, and the step's optimizer frees them. While it recomputes layer i it still holds the saved boundaries of layers 0..i. So a short wave peaks near layer 0, with almost every layer's gradients live. The checkpoint floor priced only the other end (all boundaries), so short first waves were admitted below their peaks.

Change:

• The floor adds the largest excess of pending gradients over released boundaries, while a gradient group's slot still has unallocated adapter gradients (the step's first gradient wave).
  • It is computed over the real per-layer sizes, with no uniform-layer assumption.
  • A parameter shared by several layers counts once, at its highest layer; only local shards count.
  • Parameters outside the decoder, including any a head shares with a layer, and a checkpoint's other trainable parameters count as live throughout.
• Gradient groups run their backward one after another, not layer by layer together. Autograd drains the last-forwarded group first, and separate backward calls can come in either order. While one group runs, groups already run hold their gradients and groups not yet run hold their boundaries. The price is the worst order, in closed form. A single group prices exactly as before.
• Split children training the same slots charge it once; children with different slots each charge their own.
• Unprofiled gradient waves also carry 64 MiB of measured first-execution transients. The non-full-recompute path already has an equivalent allowance.
• The cheap width estimate defers to the exact plan while a slot has pending gradients, since only slot metadata can price them.

Before/after (Qwen3.6-35B-A3B, 40 layers, CP2, one real sequence per wave; raw = estimate ÷ 1.1 over the peak, admitted = estimate over the peak):

	before raw / admitted	after raw / admitted
EP1 cold, 2,047 tokens	-19.5% / -11.5%	+18.7% / +30.6%
EP1 cold, 4,095 tokens	-10.9% / -2.0%	+3.8% / +14.2%
EP1 cold, 8,337 tokens	+7.0% / +17.7%	+10.0% / +21.0%
EP1 cold, 20,245 tokens	+14.9% / +26.4%	+16.3% / +27.9%
EP2 cold, 2,047 tokens	+11.3% / +22.4%	+28.1% / +40.9%

Testing: 21 new CPU tests, including brute-force checks over random per-layer sizes and every backward order of up to six gradient groups, mutation-checked. The trainer-rank suite passes. The GPU runs above are single-group waves, whose pricing the sequential-groups change leaves unchanged.

CI status: GPU validation passed on this exact head in run 36399223939 ('Run on 2x H200' in 20m58s). The later run 36404599822 was started automatically by the ready-for-review event, and I cancelled it before it provisioned a GPU to keep shared H200 jobs one at a time. Its failed status is that cancellation, not a code failure.

Limits:

• Measured on one MoE model at CP2 (EP1 and EP2). Dense models, TP>1, and waves with two or more trainable slots (priced one group's backward after another) are covered by CPU tests only.
• Short waves now over-estimate more once their gradients are pending: a warm 2k EP1 wave goes from +17% to +34%, and a 2k EP2 wave from +35% to +50%. Waves above ~8k tokens change by 1–3 points.
  • Allocator snapshots put most of the remaining EP2 excess in main's own terms: the MoE workspace's 1.5x routing allowance and the full-copy input gradient.
  • Pricing each backward layer from only its own live parts would tighten this, but on main it would under-admit EP1 4k (its live MoE stage and mixer exceed the priced workspace), so that belongs with a more complete workspace model.
• The 64 MiB allowance is the transients measured on this model; other models' first executions weren't measured.
• The cost components gain two fields, so reports recorded before this change replay (with --allow-source-drift) as mismatched.

🤖 Generated with Claude Code