trainer_rank: preserve memory estimates and account for overlapping expert buffers

[bradhilton]

TrainerRank can underprice an expert forward by omitting simultaneously live dispatch and FC1 buffers. It can also discard a learned peak beyond its 8× calibration window, so increasing a request from 80 to 81 tokens can turn a refusal into admission. This change addresses those two specific defects in #848.

• Extend the known MoE working-set floor for the observed single-rank, fused gate/up, no-offload expert path. Three complete intervals in a retained H200 allocation trace show seven distinct overlapping buffers; the previous floor covered four. For 45,981 packed rows, the additional component is 3,766,763,520 bytes. Unsupported paths retain their existing behavior.
• Keep the empirical estimate beyond the calibration window while retaining the independent calibration/retained-memory trust gates. Remove the now-obsolete search branch for a falling post-window estimate.

Public APIs, art.megatron, loss/gradient rules, output accounting, and allocator policy are unchanged. Corrected estimates can select smaller internal groups or refuse requests previously admitted. One synthetic case now needs four subforwards instead of two at the same 20 GiB budget; tests verify complete ordered outputs and every subforward's memory requirement. Different grouping can affect throughput and floating-point rounding.

Validation: 191 combined CPU tests pass, including gradient/no-gradient direct-forward regressions, MoE guards, split coverage/budgets, caller-peak observations and weird shapes. The five new calibration regressions fail on the parent. Ten allocation-trace integrity controls and an independent raw-event replay establish the three complete overlap intervals. Ruff, format and diff checks pass.

This is a partial fix, not closure of #848 or #870. The corrected component still falls below the retained whole-forward peak. The history ring's beginning is missing, so it cannot establish the full allocation history. These changes have not yet completed native candidate/warm/backward qualification. The separately held cache-release PR #888 is not included.

Current review and CI: McCarthy, Minsky and Taravangian each cleared source correctness at fe975382158cd8fe368045375999187311868bd8, with an explicit behavioral-adoption/merge hold. Prek and the automatic two-H200 workflow passed (GPU run). This does not replace the targeted cold/warm/backward qualification above.

Private review evidence:

• /home/brad/.local/share/schulman/art848-forward-lane-20260914/REPORT.md
• /home/brad/.local/share/schulman/art848-backward-trust-20260914-evidence/REPORT.md
• /home/brad/.local/share/schulman/memory-resolution-20260914-root/independent-tail-review.json
• /home/brad/.local/share/schulman/memory-resolution-20260914-root/integration-02-ownership.json

[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]

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. Fixes two TrainerRank memory-estimate defects from #848: the MoE expert floor now prices seven live buffers instead of four on the observed fused single-rank path, and learned per-token peaks persist beyond the 8x calibration window rather than discarded. Corrected estimates can split further or refuse previously admitted requests; loss/gradient rules and public APIs are unchanged.

Review. McCarthy, Minsky and Taravangian each reviewed the diff at fe97538: source correctness CLEAR, while flagging the stricter admission as an intentional behavior change needing an explicit adoption decision. Minsky ran the changed test files (152 passed; nine expected parent failures), Taravangian replayed the 4,096-event allocation trace confirming three overlap intervals, and McCarthy audited the sealed evidence and reported 191 CPU passes. Prek and the two-H200 GPU workflow passed.

Deferred / follow-ups.

• Partial fix; TrainerRank MoE memory admission: backward-recompute OOM and profile calibration failures #848 and TrainerRank reusable-cache admission can leave insufficient physical CUDA memory for backward library allocations #870 remain open pending native candidate/warm/backward qualification.
• Preserve captured conditioning when retaining sampled lineage #885 and cache-release PR Share budgeted CUDA cache recovery with gradient handoff #888 held separately.
• Non-blocking (Minsky): FC1-branch comment should read "compiled path observed once".