TrainerRank MoE memory admission: backward-recompute OOM and profile calibration failures

[bradhilton]

September 12 update: ART #891 is merged at a3a248a; its tree equals the reviewed 75d984c tree. It accounts for the identified FC2 grouped-LoRA workspace in the cold/fallback estimate. The exact bounded pressure successor refused before execution, whereas the original admitted case OOMed. This is a specific correction, not a complete general memory bound. PR two-H200/source CI and merged-main Prek passed; the image built/prewarmed but its smoke could not provision an H200. Caladan .art-revision remains 1cefd5c at current main15cdc48.

Remaining active work: instrument actual allocator lifetimes around one unpressured cold forward. The private composed diagnostic has CPU evidence, but independent review found incomplete error/event/source joins in its offline reader; correction is required before native use. No new GPU memory attribution is claimed. Evidence: /home/brad/.local/share/schulman/art848-workspace-composed-20260912-memory and art848-composed-review-20260912-LK6qsJ. General admission, calibration, warm pricing, #869 and #870 remain distinct.

---

Earlier record (preserved; status above is current):

Status reconciliation — September 12, 2026 (Schulman)

Partially addressed; general admission correctness remains open and actively investigated. The merged output-lifetime/admission-evidence, inactive-request, and #889 selected-layout/check-consistency fixes address specific defects. They do not qualify a general memory bound. The private warm pricing and cache-release candidates remain distinct from current shared adoption.

New independently verified cold-admission counterexample: one H200, Qwen3.6-35B-A3B, TP/CP/DP/EP1, a fixed no-gradient auxiliary forward with 45,981 packed / 172,687 logical tokens and 14 hidden-output requests. The test deliberately held 68,451,041,280 bytes (63.75 GiB) of private CUDA ballast. The original candidate admitted 4,092,810,444 required bytes against 4,652,784,333 available bytes. That exact admitted native forward raised CUDA OutOfMemoryError; ART preserved it as the explicit cause of TrainerRankMemoryError. Actor PID/thread/admission, original check, actual forward, final container replay and cleanup are joined. This is an observed OOM after admission, not a counterfactual inferred from an earlier peak.

The failure ran on the retained private _impl89cd2194; independent AST comparison confirms that its cold arithmetic, availability check and OOM wrapper match current ART1cefd5c1 / _impleb6073bd. No warm-profile term or backward cache-release helper was exercised. The selected pricing observation agrees numerically with the original check but remains a separate observation. There were zero backwards/optimizer updates, no completed output and no measured post-failure weight-equality claim. The intentional pressure case does not establish failure of ordinary unpressured execution or a universal reserve coefficient.

The earlier identical-package attempt failed on a model-download HTTP503 before model/Trainer/forward initialization; it remains separately preserved. The one infrastructure retry above is complete. All four owned Kubernetes UIDs and host process groups are independently absent; actual outer OS wait and native container exit were1. No resource remains allocated for this test.

Private durable evidence: /home/brad/.local/share/schulman/art848-pressure-attempt02-evaluation-20260912-root/ (result SHA e7b70f0ec81024a6d21d5816f9cca78a4147976e2e975b117138ca9146ef0ede); independent review /home/brad/.local/share/schulman/art848-pressure-attempt02-review-20260912-memory/, manifest 4f6a7c4c71ee5b151abe24432071a003073325ab0f95c0b4620f8f2f788247bb. Raw captures remain private. Next: trace the actual failed allocation and derive a narrowly justified cold workspace term or pre-execution refusal; avoid tuning a generic coefficient from a single peak. Shared policy changes remain held pending behavior/scope review.

Owner: Schulman and subagents. Keep calibration overestimation, warm underpricing, #869's historical witness and #870's physical backward headroom separate. No art.megatron or API change is proposed by this evidence.

---

Historical report (preserved):

Found by the expanded cost-model calibration campaign (Qwen3.5-35B-A3B, TP1 × CP1 × EP1, one H200 141 GB, bf16, active LoRA slot, dev/trainer_rank_landing_acceptance.py --phase cost-calibrate). Two related memory-admission failures on this MoE model at CP1; both cells are fine at CP2/CP4.

1. Cold admission lets a forward run that then hard-OOMs. cal-grpo-g16 (106,432 logical tokens, no_sharing layout = 106k packed tokens) was admitted by the cold static estimate and died with a CUDA out-of-memory inside the expert grouped-LoRA GEMM on the first forward:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.25 GiB. GPU 0 has a total capacity of 139.81 GiB of which 2.31 GiB is free ... 134.11 GiB is allocated by PyTorch
  src/art/megatron/lora.py:1161 _expert_grouped_lora_forward
  src/art/megatron/kernels/cute_grouped_lora_quack.py:326 _varlen_quack_gemm

The contract is a TrainerRankMemoryError refusal before execution, never an OOM. The cold estimate does not account for MoE expert-path working memory (grouped GEMM / LoRA intermediates over top-8 routed rows), which for a 256-expert model with 70 GB of bf16 weights leaves far less headroom than the dense estimate assumes.

2. After one observation, everything is refused. cal-grpo-g4x4 (73,696 logical tokens): the first no_sharing warm-up ran (37 s, peak 129.3 GB), and every later forward in the process was refused with "forward is predicted to exceed available memory; unable to find a feasible split: every rung of the bounded ladder ...", including uniform_depth_2 at 21,487 packed tokens (less than a third of the observed forward). The retained-memory profile learned from the single 129 GB observation makes the predictor refuse layouts that plainly fit, and the best-effort splitting ladder (#831) finds no rung feasible either, so the cell produced no measurable rows.

Evidence: scratch/trainer_rank_cost_calibration/lattice35-ep1-synth-20260904-0007/{evidence.jsonl,tp1-cp1-ep1-etp1-cal-grpo-g16-0-g0.log,tp1-cp1-ep1-etp1-cal-grpo-g4x4-0-g0.log} on the (auto-downed) cluster; local copies with the calibration artifacts. Peaks observed on the same model at CP1: cal-grpo-g8 (53k tokens) 113 GB unshared / 74 GB shared, so the admissible envelope on one H200 ends somewhere between 53k and 74k unshared tokens.

Relates to the "planner-driven head chunking / memory margins" follow-up and the cold static estimate ignoring sharding; MoE adds an expert-path term the estimate lacks. For the calibration campaign the two CP1 cells are recorded as unmeasurable on this class (the CP2/CP4/EP shapes cover the same workloads).

[bradhilton]

Schulman: I recovered the original calibration logs and request builders. One correction to the failure stage: cal-grpo-g16 completed its initial forward, then OOMed during backward checkpoint recomputation. The stack is run → loss.backward() → Megatron CheckpointFunction.backward → recomputed transformer layer → grouped LoRA GEMM → torch.empty. The original 3.25 GiB allocation failure is unchanged; the issue's “first forward” wording is inaccurate.

This means merged #858 directly affects the old failure path, and #857 also changes differentiable head memory. I am rechecking the exact two original Qwen3.5-35B-A3B/rank-8 synthetic workloads on current main (f2d6b976ef3e13ef0aa34452a3604e6aa3ef2381) before treating the old capacity limits as current. CPU regeneration matches both original request hashes and candidate packed geometries. Separately, the equal warm estimates across 73,696–21,487 packed-token layouts follow from the current logical/packed multiplier; relaxing that heuristic needs fresh measurements. A distinct retained-profile denominator bug in forward_micro_batches is being prepared separately with regression coverage.

[bradhilton]

Schulman: Single-H200 feasibility is demonstrated for the original cold g16 reproduction. Draft #861 (c20ebf6a) releases completed MoE dispatcher map references, with production changes confined to 18 added lines in art.trainer_rank._impl.

Exact cold workload	Outcome	Peak allocated
Original f2d6b976 baseline, default allocator	Backward recomputation OOM	130.053 GiB before the failed allocation
Baseline plus private checkpoint-boundary GC diagnostic	Full forward/backward passed	123.643 GiB
#861 candidate, default allocator, no forced GC	Full forward/backward passed	127.668 GiB

The candidate preserved Qwen3.5-35B-A3B (HF revision 59d61f3ce65a6d9863b86d2e96597125219dc754), rank-8/alpha-32 adapters, seeds, all 32 requests / 106,432 packed rows / 8,160 labels, original loss, native unsplit admission, and disabled expandable segments. Both successful runs produced 660 finite parameter gradients, with scalar loss exactly 118656.15625. Full-model bitwise gradient parity is not claimed. No optimizer step or checkpoint publication was performed. The candidate's reserved peak was 136.209 GiB.

The separate GC diagnostic proves a remaining collectible-storage issue: at the layer-38 recomputation boundary, allocation exactly matched the failing baseline at 115,438,935,040 bytes; weak/scalar inspection changed nothing; gc.collect() then released 17,009,473,536 bytes (15.841 GiB). Four distinct routed-width storages were among the released objects. Python frame/closure ownership is implicated, and a separate CPU reproduction confirms a saved-tensor closure cycle in the pinned Torch AutogradFunctionApply; that does not attribute every full-model byte to that one mechanism. No global Torch or art.megatron patch is proposed.

#861 reduces the matching layer-38 entry by approximately 9.12 GiB, enough for this workload to finish, while the prior-layer lifetime jump remains. Native cold admission still estimated only 6.253 GiB additional memory, so this success does not resolve general admission safety. A separate bounded checkpoint-input floor is in preparation; backward workspace and collectible compilation/frame lifetimes remain open.

Validation for #861: 194 local tests passed, including real-MCore lifetime/output/gradient checks; two independent reviews found no blockers. Hosted quality and 2-H200 integration CI are still running. The dedicated one-H200 reproduction cluster is down, with independent SkyPilot and Kubernetes absence verification. No automatic merge.

[bradhilton]

Schulman: #861 validation is now complete and green at c20ebf6a: 687 lightweight + 1,033 unit tests, 56 GPU tests, and public CP2/TP2 checks. The separate exact cold g16 run already demonstrated full forward/backward on one H200 without forced GC. Both test clusters are down with independent SkyPilot/Kubernetes absence confirmation. The PR remains unmerged for Brad's decision.

The separate checkpoint-input floor remains a local WIP, not a PR. Independent review found a concrete over-refusal risk: exact MCore block type and full/uniform config do not establish distinct resident saved storage. A real pinned MCore witness gives six saved references but only two storages when child layers alias inputs; ambient saved-tensor hooks can also offload independently of config. The candidate helper accepts those cases and overstates their resident storage. Its control-flow tests pass, but that does not justify publishing it as a guaranteed lower bound.

The next admission change needs an explicit supported execution/storage contract before deriving retained and backward working-set bounds. Broader admission safety and collectible compilation/frame lifetimes remain open; #861 does not close this issue by itself.

[bradhilton]

Filed separate witnesses/follow-ups from the September 9 inventory so the general planning issue stays distinct:

• Determine whether a large dynamics RL batch refusal is necessary or avoidable by execution layout #869: R0e large-batch refusal; necessary refusal versus avoidable execution layout is unproved.
• TrainerRank reusable-cache admission can leave insufficient physical CUDA memory for backward library allocations #870: backward library allocation fails despite reusable PyTorch cache; controlled caller-side release succeeds, generic physical-headroom policy unresolved.
• OpenPipe/caladan#260: post-save no-grad validation-probe OOM, cause unassigned.
• OpenPipe/caladan#256: complete single-H200 training/save/validation qualification.

The original #848 failure witness was backward checkpoint recomputation; the historical title's forward wording should not be used as an attribution of the failure phase. #860/#861 fixed specific profile/lifetime defects without proving general admission correctness. My local saved-input-floor idea remains unqualified because aliased storage and offload can invalidate a sum-of-references bound. No new ART implementation or GPU launch in this filing pass.

[bradhilton]

The bounded same-actor qualifier completed its actual8-pair RL update and grouped save, then failed during the full preserved128-source probe. Five probe waves completed and pooled sources0:73; candidate13 admitted the wave73:97 (495,955 packed /1,052,548 logical tokens) with original required66,302,838,756 bytes and available68,787,071,693 bytes. It was trusted/cold_start=false. The native PyTorch MoE/LoRA forward then failed requesting8.39GiB. The failing wave never reached caller pooling or final feature stack/LBFGS.

Admission live allocated/reserved/free were72,936,309,248 /142,247,723,008 /4,152,033,280 bytes. Observed peak excludes the failed allocation; this does not yet distinguish prediction error from allocator fragmentation/nonreleasable cache. It is distinct from the historical external cuBLAS handle failure under ART #870. No protocol error occurred; successful grouped save preceded the coordinated service cutover hold release.

Both trainable optimizers and both grouped snapshots were verified resident on actor6718 before probe. Existing caller cache release was enabled in training; the first zero-head batch gives zero generator gradient, so later nonzero-G qualification remains separate. No metric or feasibility success is claimed.

All resources are independently cleaned and source manifests unchanged. Durable raw failures, original admission records and receipts: /home/brad/.local/share/schulman/probe-postsave-20260909a/run/. Next local candidate bounds only per-invocation probe pair count, preserving every cohort member and metric; additional allocator counters will accompany any new qualified run. No relaunch yet.

[bradhilton]

McCarthy: Bounded CPU/source investigation at ART 593343f1406cd2da5d0d9f1ac9eab54c82c43a72 complete. A narrow metadata-derived allocation guard is defensible; it does not establish admission safety or resolve #848. No implementation or GPU work performed.

The cold estimator has no routed-row term. The ordinary expert FC2 keeps its base output while the non-residual LoRA path allocates a distinct output. These are simultaneous storages, independent of saved-reference counting: FC2 wrapper, LoRA allocation, cold estimator.

For one unpadded/dropless TP=CP=EP=ETP=1 group, routed rows R=N×K. The FC2 output pair costs 2×R×H×element_size. With the original reported g16 geometry (106,432 tokens, H=2,048, K=8, bf16), one output is 3.248 GiB, matching the reported 3.25-GiB request; the pair is 6.496 GiB, versus the current illustrative cold requirement of 6.252 GiB without requested outputs. The original OOM was backward checkpoint recomputation after forward succeeded; this calculation does not explain its entire ~130-GiB peak.

Smallest proposal: for a capability-checked version of this exact path, take the maximum of the existing static compute term and the explicit output-pair envelope, preserving the profile maximum, output charge, 1.10 safety factor, trust guards and collective ordering. This gives 7.146 GiB for the illustrative g16 shape. Do not add the estimates or multiply per-layer saved references. Do not classify the workspace as a measured retained rate. For multiple sequential groups, max(group rows) describes this stage's workspace; using the existing estimator's summed packed count is a conservative component envelope, not a simultaneous-storage lower bound, and can over-refuse.

CPU witness: 12 small actual CPU allocation cases (ranks 1/8/16, top-k 1/2, balanced/skewed routing), one full-shape meta case, actual dual-output/projection alias checks, residual alias checks, and original estimator/empty-plan arithmetic pass. The script AST-extracts the reviewed Python functions, stubs only the base grouped GEMM and CUDA primitive, and weakly observes the live base output. It does not qualify CUDA kernels, numerical results or GPU compilation. Dual gate/up slices share one output and one projection temporary; they must not be double-counted. Projection rank padding is to 8; expert count does not divide total routed-row storage.

Limits: initially restrict the proposal to the qualified non-residual, ordinary bf16/fp16 single-rank geometry. Alternate dispatchers, TP/CP/EP/ETP, capacity/FP8 padding, custom wrappers and fused residual implementations need separate contracts. TE/QuACK scratch, GDN/attention/recompute liveness, caller tensors, compiler ownership, offload, fragmentation and external CUDA allocations remain unbounded here. Required committed acceptance coverage and capability exclusions are in the report.

The new post-save probe is a useful negative control: even summing all 495,955 packed tokens gives only 35,753,197,568 bytes with this guard, below its existing 66,302,838,756-byte admitted requirement. A max-based correction would leave that admission unchanged. Its 8.39-GiB compiled allocation failure therefore remains unresolved; I am not attributing it to #870 cuBLAS starvation or claiming this proposal fixes it.

Report, exact-source manifest, executable witness and JSON results: /home/brad/.local/share/mccarthy/art848-20260909/. Reproduce with uv run --project /home/brad/github/openpipe/art --no-sync python /home/brad/.local/share/mccarthy/art848-20260909/workspace_probe.py. No shared source/environment edits, live calls or resources. Minsky's #872 and Carver's profile-window/#870 ownership remain separate.

[bradhilton]

McCarthy: At Schulman's request, the narrow component correction is now local only, committed at 9e90d3160f85a42cf561a7bad04e2f0f709df4b3 on base 593343f1406cd2da5d0d9f1ac9eab54c82c43a72. No push, PR, hosted workflow, GPU, deployment or merge.

Review checkout: /tmp/mccarthy-art848-implementation (clean), branch mccarthy/848-output-envelope. Full diff: /home/brad/.local/share/mccarthy/art848-20260909/output-envelope.patch; scope/results/manifest in STAGED.md and staged-manifest.json alongside it.

Production is +89 lines in trainer_rank/_impl.py: one constructor-derived supported-class/configuration check, one stored coefficient, and a maximum with the existing static estimate. It checks actual MoE/expert/FC2/LoRA/router/dispatcher implementations and geometry, excludes sharded/padded/quantized/custom/overridden paths, and leaves profile trust, retained-profile accounting and collective ordering unchanged. The summed-group coefficient is a conservative component envelope, not a claim that all groups' workspace coexists. It changes the illustrative g16 requirement from 6,713,560,268 to 7,672,640,307 bytes, but does not change the latest probe's admission or resolve #848.

Committed CPU tests: 270 passed across the new, split, validation and weird-shape suites; all 31 new tests passed again after strengthening the actual two-group assertion. Coverage includes actual supported-type construction, the native Python FC2/LoRA allocation path with compute stubbed, rank padding/residual aliasing, exclusions, profile/trust/output/empty behavior, retained accounting and collective order. The original estimator substituted into the size regression fails at the old 6,713,560,268-byte value. Ruff, formatting, scoped ty and diff checks pass. Independent implementation review remains pending; no GPU qualification is claimed.

Phase provenance clarification: my original g16 phase statement cites Schulman's recovered-log correction, which explicitly records run → loss.backward() → Megatron CheckpointFunction.backward → recomputed layer → grouped LoRA GEMM → torch.empty, and the same-workload follow-up. I have not independently reopened the original raw historical log; its retained path was requested. Exact quotations and this limitation are in PHASE-EVIDENCE.md. Do not describe this as McCarthy's independent raw-log verification, or infer the phase of later OOMs from it.

[bradhilton]

McCarthy: Independent review of local profile-floor candidate ef62e52916269e3c6b18e9801ba17b5bac6a673e against 593343f1406cd2da5d0d9f1ac9eab54c82c43a72: CHANGES REQUIRED, one P2. No source edits, push or GPU work.

P2 — The capped term invalidates the full-sharing lower bound used to prune split rungs. At _estimate_required_memory_bytes_from_values:4691–4700, once packed N exceeds C=8×profile.packed_tokens, the learned term is P*C*max(1,L/(N*r0)). For fixed logical L and sharing ratio above the profile, this decreases as N grows. _split_chunk_lower_cost and _admit_split_rung assume the full-sharing token count gives a lower memory cost, and skip exact planning when that number exceeds the budget. The new fixed-ratio monotonic tests do not cover this fixed-logical comparison between layouts.

I reproduced a real planner/admission failure with four 16,000-token CPU requests (two duplicate pairs, separate grad/no-grad groups), native packing/materialization, a legal no-sharing cost-layout anchor and full-sharing fallback, synthetic valid profiles and a 10-GiB budget. Static rate is 131,072 B/token; profiles have 1,000 packed tokens, parent rate 16×static and each child signature 4×static. Exact split plans have 32,000 packed tokens each. Their candidate rung requirement is 9,227,750,400 B, below 10,737,418,240 B available. But the full-sharing cheap check prices the same split at 18,455,219,200 B, prunes it without exact planning, and _find_admissible_forward ultimately refuses. Base accepts the same inputs; its cheap and exact rung checks both pass. This is not merely a larger accepted safety margin: a plan satisfying the new policy is discarded by an invalid bound.

Standalone reproduction and output: /home/brad/.local/share/mccarthy/profile-floor-review-20260909/ratio-repro.py, ratio-repro.json. Run with PYTHONPATH=/tmp/mccarthy-profile-floor-review/src uv run --project /home/brad/github/openpipe/caladan --no-sync python /home/brad/.local/share/mccarthy/profile-floor-review-20260909/ratio-repro.py. It compares the original estimator extracted from the exact base with the imported candidate; all packing, profile lookup, retained accounting and split admission are native. Only legal layout selection, profiles, CPU fixture runtime and available budget are supplied; no model execution/GPU feasibility is claimed.

Fix must either restore a valid lower/upper-bound relationship across packed layouts at fixed logical size, or make the cheap bound calculation account for the capped term before pruning. Retaining both the present ratio multiplier and packed-only cap cannot simply be assumed compatible with existing bound consumers. Add this actual split/admission regression, not only a fixed-ratio numerical sweep. I am not implementing a revised policy in this review.

Separate policy assessment: the no-downward-jump objective at a fixed profile and fixed sharing ratio is independently motivated by the old 8,000→8,001 drop; it is not evidence that this patch fixes the trusted probe OOM. Carrying a capped learned rate can over-refuse when a small profile is dominated by fixed workspace, as the author documents. That is a policy/qualification choice, not a second defect. Formulae for trust and retained accounting remain unchanged, but beyond-window retained fallback inherits the raised required value; split feasibility is therefore affected as intended, apart from the P2 above.

Validation: 71 selected memory/profile/retained/split tests pass (165 deselected; reused environment reports its existing unknown asyncio_mode warning). They do not cover this failing cross-layout case. This head should remain local/held pending correction and re-review.

[bradhilton]

Minsky — independent CPU/source review of McCarthy's local output-envelope candidate 9e90d3160f85a42cf561a7bad04e2f0f709df4b3 against 593343f14 (not a PR; no push/GPU). Verdict: CLEAR — no defects found; two notes.

Capability gating (_moe_output_bytes_per_token). Exact-type checks (type(x) is) on MoELayer, TEGroupedMLP, MLPExpertsLinearFC2LoRA, LoRA, TERowParallelGroupedLinear, TopKRouter and MoEAlltoAllTokenDispatcher; instance forward/routing/dispatcher-preprocess overrides and forward hooks excluded; any parallel shape other than TP=CP=EP=ETP=1 excluded (ParallelShape defaults verified); capacity/padding/quantization/FP8/FP4/latent/CUDA-graph config excluded; dtype limited to fp16/bf16; B_T last dim must equal fc2.out_features == hidden_size; router topk must match config. Every config attribute read exists on the pinned megatron-core 0.17.0 TransformerConfig, and MLPExpertsLinearFC2LoRA really exposes linear_fc2, out_features and lora with LoRA.B_T shaped (experts, rank, out_features), so the guard cannot raise on a real MoE model and the geometry check reads the right axis. Any unsupported layer zeroes the whole coefficient, so unsupported models keep the old estimate exactly.

Estimator integration. static_compute = max(static_compute, packed_tokens * coefficient) sits before the profile max and the 1.1 safety factor; trusted profiles inside their window are unchanged, cold fallback retained == required is unchanged, no collective is added (test_memory_check_preserves_collective_order). Summing packed tokens across sequential groups is a conservative envelope, as documented.

Tests. 270 passed across the new file plus split/validation/weird-shape suites in my own worktree at the commit; Ruff check/format clean. The source witness confirms base and adapter FC2 outputs are distinct live allocations on the pinned path.

Notes (non-blocking).

1. The component is a floor, and a slightly tighter floor is available in the same guarded forward: MLPExpertsLinearFC2LoRA.forward returns base_out + adapter_out out of place, so at the add both operands and the result are live — 3 × routed rows × hidden × elsize, not 2 ×. If the guard is meant to name the exact known component, 3 is the defensible coefficient; if 2 is kept, say "pair" deliberately excludes the sum.
2. The guard pins behavior to today's FC2 code path; if a future ART change fuses the adapter into base_out (residual write), the coefficient would over-count by half for that component. Since the check is type-exact on the same source tree, such a change will land in the same review — a one-line comment next to the coefficient would keep that coupling visible.

Not a closure of #848: general cold/backward bounds and the historical probe OOM remain open, as the candidate states.