fix(trainer-rank): account for recomputed mixer memory and head stages

[bradhilton]

Agent: Dalinar

Full one-layer recompute keeps a layer’s mixer live beside its MoE stage. Price that overlap and previously hidden residual, norm, GDN, shared-expert, routing and first-call TE workspaces. Where every MoE layer’s recompute is covered, charge one incoming gradient instead of one per boundary. HybridEP prices one dispatched copy on the EP group’s balanced rows when EP and CP groups coincide.

Price eligible warmed CP2 target-only fused head statistics as a separate backward stage; unsupported heads retain the existing floor. EP routing allowances become 1.4/1.6/2.0 for EP2/4/8, based on measured routing. Reports freeze these facts for CPU replay. This branch includes #1002; review that first. #978 adds per-rank layouts and #981 adds observations.

Brad’s decision remains required for strict fused statistics: a wave admitted using the staged head price raises if its fused kernel fails instead of using a wider FP32 fallback. Its CP peer can wait at the next collective, as with a rank-local OOM.

Recorded trainer-rank/replay CPU tests and local H200 runs passed. On the full stack with #1002, Qwen3.6-35B-A3B CP2 single-sequence cold errors were +0.9–2.0% at EP1 and +3.0–13.4% at EP2; warm errors were +2.8–24.9%. No measured case was under. Staged-head evidence covers 29 CP2 EP1/EP2 waves. These measurements do not establish a general peak bound.

EP2 still exceeds the 10% accuracy target, largely from the routing allowance. Routing measurements cover sampled workloads only; EP4/EP8 have no memory runs, and a small EP8 sample approached its 2.0 allowance. PyTorch peaks omit the HybridEP buffer and other external memory: one cold EP2 case grew 2.76 GB outside PyTorch while the estimate included only its 1.05 GB HybridEP buffer. A warm routing discount is separate work.