Model retained activations and allow selective recompute in TrainerRank

[bradhilton]

Selective and no-recompute execution retain decoder activations across layers, so the previous full-recompute heuristic admitted #913's long group and then OOMed. This follow-up to #914 replaces the selective refusal with a retained-activation floor that accounts for TP/SP, the actual attention/GDN layer mix, native MLP fusion, replicated QKV, prefix-segment states, and effective MLP/MoE checkpoints. Profiles can only raise the floor.

Native H200 calibration now puts the 27B estimate about 10–13% above measured peaks for the paired 2k-and-larger workloads. The previous TP4 estimate was about 60% high. Values are incremental allocated GiB; observed peaks are maxima across ranks and repetitions.

Qwen3.8-27B workload	Previous estimate	New estimate	Observed forward
TP4, 2k pair	32.974	22.854	20.550
TP4, 4k pair	56.075	38.789	35.077
TP4, 8k pair	112.151	77.273	69.926
TP4, 12k pair	—	115.757	104.780
TP8, original 19,221 + 19,222 pair	153.866 (refused)	126.188	114.665
TP4, 8k pair with MLP recompute	112.151	42.294	37.881

The previously unmeasured long TP8 case now completes forward and backward. TP4 still refuses the original long pair. Additional runs cover 4B/9B models, unfused eager/compiled execution, short batches with eight sequences, and MoE routing concentrated on one expert rank. MoE dispatch retains an imbalance allowance; the report documents its remaining conservatism and a native TE crash when expert peers are completely empty.

Calibration report · CSV with source/driver hashes, peaks, and negative controls

Validation: 45 final GPU cells, 360 measured rank-samples and 4 refused rank-samples. Every measured forward is covered, every backward completes, and all measured losses/gradients are finite. Local TrainerRank tests: 631 passed, 2 skipped. Ruff and type checking pass. Full-recompute and no-grad behavior retain their previous estimates; the report explicitly limits the calibration's scope.

CI is green on b86ad4bdd: quality checks (783 Megatron tests and 1,553 unit tests passed) and native GPU validation. Calibration clusters and the CI cluster were terminated successfully.

Behavior with recompute disabled

The retained-activation floor also applies when recompute_granularity=None, including the none override and EP-overlap MoE trainers. Requests that previously fit the smaller heuristic may now be split or, for indivisible groups, refused. The MoE allowance protects uneven expert routing and can be roughly 3× the measured balanced-routing peak. Full recompute (the production default, including Caladan) and no-grad estimates are unchanged.

Context parallelism is not discounted in this version, so selective recompute on the default multi-GPU dense CP topology may be overestimated. CP scaling and GPU calibration are being addressed in a follow-up.

[bradhilton]

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

Change. Replaces TrainerRank's selective-recompute refusal with a retained-activation floor accounting for TP/SP, attention/GDN layer mix, MLP fusion, replicated QKV, prefix-segment states and MLP/MoE checkpoints, after the full-recompute heuristic admitted #913's long group and OOMed (follow-up to #914). H200 calibration puts 27B estimates 10–13% above measured peaks; the long TP8 pair now fits.

Review. Shannon reviewed and recommended merging at b86ad4b, confirming the floor applies only to grad-enabled non-full recompute (full-recompute and no-grad estimates unchanged), prices None as well as selective, and that profiles can only raise it. Verified: 631 local TrainerRank tests passed (2 skipped), ruff and type checks, green CI (783 Megatron and 1,553 unit tests plus native GPU validation), 45 GPU calibration cells.

Deferred / follow-ups.

• Context parallelism not modeled, so selective recompute on the default multi-GPU CP topology is overestimated; fixed in Account for dense CP rank loads in retained-activation memory #922 with CP calibration.
• Behavior change for recompute-disabled (None) paths, including EP-overlap MoE trainers, documented in the PR description after review.