Account for dense CP rank loads in retained-activation memory

[bradhilton]

Follow-up to Shannon’s review of #915. Dense selective and disabled-recompute requests charged every CP rank for the global token count, so requests that fit could be split or refused. The 27B 4k pair now estimates 69.5 GiB against 62.9 GiB observed and completes backward; its previous 117.4 GiB estimate exceeded the available 81.9 GiB budget.

Price retained token activations using the busiest rank’s actual attention/GDN layout, including TP padding. CP ownership is uneven: that pair has 6,964 packed tokens but 4,096 on its busiest rank, so simply dividing by CP would underprice it. Keep recurrent segment states, gathered outputs, learned-profile floors, and the cold allowance unchanged. CP width selection uses the existing exact-plan fallback, at the cost of additional CPU planning. Full recompute, no-grad, and MoE retain their previous estimates; expert dispatch needs separate CP calibration before a discount.

Native H200 CP2 measurements (incremental allocated GiB, maxima across both ranks and repetitions):

Model / pair length	Recompute	Previous estimate	New estimate	Observed forward
27B / 2k	selective	69.123	34.954	31.518
27B / 4k	selective	117.374	69.524	62.918
4B / 4k	selective	41.478	20.922	18.727
4B / 4k	none	41.478	20.922	19.356
4B / 4k	selective + mlp	25.760	13.063	11.581
1.7B / 4k	selective	21.002	15.477	13.932

Calibration report · CP CSV

All 13 GPU cells (52 rank-samples) cover the measured forward peak, with at least 8.1% headroom. Every backward completes with finite losses and adapter gradients. This calibrates eager dense TP1/CP2, bf16, random weights, and rank-1 LoRA; larger CP, combined TP/CP, and MoE CP remain uncalibrated. The report records a short-input backward peak above its forward estimate; the estimate is for forward admission. Local GPUs are released; no calibration clusters were needed.

Also records the disabled-recompute behavior change requested in review: None, including the none override and EP-overlap trainers, now receives the retained-activation floor from #915. Previously admitted workloads can split or be refused; balanced MoE can be priced about 3× above its peak. This note was added to #915 before merging it.

Validation: 637 TrainerRank tests pass (2 skipped), including uneven CP admission, unchanged GDN-state costs, the measured 27B peak, and collective propagation of CP planning failures. Lint and type checking pass. CI is green on 44ea923bf: quality checks passed 783 Megatron tests and 1,561 unit tests (73 skipped); native GPU validation passed, including CP/TP public-API backward checks. The CI cluster was terminated successfully.

[bradhilton]

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

[bradhilton]

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

[bradhilton]

Consolidated review record (automated agent review traffic removed 2026-09-17)

Change. Follow-up to Shannon's review of #915: dense selective and disabled-recompute CP requests charged every rank for the global token count, so fitting requests could be split or refused. Retained activations are now priced from the busiest rank's actual attention/GDN layout with TP padding; the 27B 4k pair estimates 69.5 GiB (62.9 observed) instead of 117.4.

Review. Minsky and McCarthy independently reviewed the full diff at 44ea923; both found source-correctness CLEAR with no changes requested, noting the admission change is intentional and confined to dense CP selective/disabled-recompute paths, with full recompute, no-grad, MoE and CP1 pricing unchanged. Minsky ran 188 TrainerRank tests with a base negative control; McCarthy reproduced all 52 calibration samples across 13 cells plus 156 unchanged-path controls. CI was green: 783 Megatron and 1,561 unit tests plus native GPU validation.

Deferred / follow-ups.

• Larger CP, combined TP/CP and MoE CP remain uncalibrated; the estimate covers forward admission only (a short-input backward peak exceeded it).