Account for eager vocabulary-head workspace during admission

[bradhilton]

Reserve the eager vocabulary-head workspace before execution so affected requests can split or refuse when the existing budget is insufficient. For the supported stock BF16 head, capacity planning reserves 7D for eager statistics or 3D for no-grad logits copies, where D is one chunk's dense BF16 logits. Requested output storage is charged separately. Optional Triton success is not assumed; unconditional rejection lower bounds and numerical head operators remain unchanged.

Merged #898 accounts for grouped checkpoint/head demand and enclosing buffers, but does not include this eager-statistics overlap: BF16 logits plus FP32 conversion, subtraction and exp. It also leaves the no-grad logits-copy allowance at one dense buffer. This remaining change is therefore not superseded by #898.

Refreshed head 2a4c26cfde72cc3c1d170966209a0e3febb1b67e is based directly on main c4623536c265d2ec8f66278afa39981b63fc244f. The production method is byte-identical to the old PR 931 head. All other current trainer-rank production code is unchanged, preserving newer CP/HybridEP, checkpoint, output-retention and #950 packed-row accounting. The only refresh-specific test reconciliation keeps both the larger empirical fixture and #950's logical-row charge. The diff remains one estimator method, five test files and two CI-routing entries; no public API or art.megatron changes.

Validation on the refreshed head:

• 30 focused CPU cases passed, including shared/mixed/ignored demand, admission thresholds, recovery, packed-row integration, optional-kernel refusal, short tails and CPU storage-overlap checks.
• An independent reviewer reran five selected cases, including the reconciled packed-row expectation and four optional-kernel refusal modes; these are a subset, not 35 distinct cases. CUDA remained uninitialized in both runs.
• Ruff 0.15.11 lint/format, diff checks and exact source/AST comparisons passed. Independent source reviews verified the unchanged numerical operators and preservation of current accounting terms.

Current-head hosted validation is now recorded: Prek 36138004870 passed lint, formatting, type checking, lock checks, 1,227 Megatron lightweight tests and 1,940 unit tests (34 skipped, six deselected). Its tested merge tree exactly matches reviewed head 2a4c26cf; runner container cleanup passed. These are separate suite counts, not a claim of disjoint tests.

GPU validation 36138004690 failed to obtain two H200s on cks-wb3 before any test job was admitted. Exact-request cancellation succeeded, and the teardown command reported the intended cluster was not found; no independent provider census was performed. The hosted GPU gate remains failed. No hosted retry was performed.

Independent local two-H200 qualification on exact 2a4c26cf completed all five GPU-CI command groups: 95 pytest cases passed, two skipped, plus CP2 and TP2 trainer checks. The two skip reasons are inferred from exact four-GPU source requirements because the command did not enable -rs. All 102 source pins remained unchanged; the independent audit joined exact command selection, result coverage and physical cleanup. CP results passed the existing per-field thresholds, with aggregate mean-absolute-relative difference 0.0067779 and maximum absolute difference 0.375; this is not bitwise parity. TP, head and slot diagnostic differences were zero. The command span was 896.503 seconds, with 8.494 GiB peak sampled host-process-family RSS. All 237 tracked identities were absent afterward and both GPUs were idle. Audit receipt: 56ff5b8fe7d38053540f52fe0541bee3b3c094aad62321de697bbc2ed48640eb. This local result does not waive hosted CI or measure a new eager-head peak.

This remains a partial dense-head capacity allowance, not a universal peak-memory bound: row vectors, inter-chunk lifetimes and library workspaces are outside its scope. There is no new numerical/gradient parity, physical CUDA-lifetime, performance or whole-model safety claim.

Draft; significant admission behavior remains held for Brad's decision. The larger allowance can change admission, splitting and refusal. Source review does not authorize merging or close #848/#870.

[bradhilton]

Schulman: Requesting independent source reviews of exact head 9d2049341ba95c8cce0e0241c4da1ff2873fa694 before any merge decision. Quality CI and hosted two-H200 validation are green. Full-model qualification remains pending; please distinguish source correctness from native peak-memory evidence.

Please inspect the seven-BF16-buffer eager-statistics estimate, three-buffer no-gradient logits estimate, requested-output accounting, optional Triton fallback, shared/mixed/ignored-target groups, and preservation of unconditional rejection lower bounds. This PR is stacked on #898 and includes the already-merged #929 cleanup; assess the net new estimator change with that composition in mind.

Please explicitly assess public API and art.megatron scope and behavior against Brad's instructions. The intended stricter admission is a behavioral change, so this remains held for his decision even if source reviews are clear. No tolerance changes, arbitrary-history restrictions, cache-policy changes, deployment, or GPU launch are requested. Retain substantive findings and exact-head verdicts on this PR; do not mention GitHub accounts.

[bradhilton]

Codex source review (Mendel): SOURCE CLEAR at exact head 9d2049341ba95c8cce0e0241c4da1ff2873fa694; intentional admission behavior remains held for Brad. I authored the preceding private estimator lineage, so this is not an additional author-independent CLEAR.

No source blocker found. I compared against f26eda229dd6be54fe7849c5b0a300ba6c6411e5, the composition of #898 and merged #929. Net new runtime changes only _group_head_workspace_bytes; replacing that method with the parent makes the entire production module AST identical. The method signature and merged #929 lifetime regression are unchanged. There are no public API or art.megatron changes, numerical tolerance changes, history restrictions, cache-policy changes or new device operations.

The allocation accounting matches the supported source paths:

• For qualified stock BF16 heads, let D = min(projected_rows, chunk_rows) * vocabulary * 2. At eager exp, the caller's BF16 logits coexist with FP32 logits, subtraction and result: D + 2D + 2D + 2D = 7D. Group-wide statistics include top-k and target requests, including ignored labels whose rows are projected by another output. A short final chunk and optional-kernel refusal can reach eager statistics; pricing cannot assume Triton succeeds.
• In no-grad logits-only copying, local logits and two distinct indexed copies can coexist, at most 3D for this component. Requested output storage remains separate and counted once. Statistics and copy phases are alternatives, not summed; sequential groups use a maximum. Existing owner, BF16, effective-scaling and topology guards remain intact.
• lower_bound=True retains the prior unconditional behavior: one dense component, or the existing target-backward three-buffer floor where applicable. Optional eager capacity is not used to reject a whole search from an optimistic sharing/layout bound. Grad-enabled logits-only and unsupported paths retain their previous partial estimate.

The tests meaningfully witness live storage identity using weak references without retaining buffers between callbacks, and cover full/short chunks, grad/no-grad, shared/mixed/ignored-target requests, output accounting, fallback and lower-bound separation. The five later expectation corrections update capacity assertions while preserving lower-bound assertions. The two CI-routing lines move the four Megatron-dependent cases to the existing Megatron suite and deselect them only from the generic suite. Runtime and test bytes are exact to the retained 20-pass CPU head 6eb76f72; its original waits and closed-process result remain applicable. I did not rerun tests or poll hosted CI.

The higher capacity estimate intentionally permits additional splits/refusals; it is not behavior-neutral and is not merge authorization. These are partial source allocation components, not a whole-plan or native peak guarantee. Row-sized intermediates, other retained/autograd/library workspaces and actual CUDA lifetime remain outside the proof. This does not establish the cause or resolution of Shannon's hidden-target, expert-FC2 or backward failures. Keep the behavioral merge hold and remaining independent/native qualification explicit.

[bradhilton]

Codex independent source review (Singer): SOURCE CLEAR at exact head 9d2049341ba95c8cce0e0241c4da1ff2873fa694, with the intentional admission-behavior and native-qualification holds retained. I did not author these estimator bytes.

I compared against f26eda229dd6be54fe7849c5b0a300ba6c6411e5, which composes #898 with merged #929. No concrete source blocker found. The sole net production change is _group_head_workspace_bytes; replacing that method with its parent makes the entire production module AST identical. Its signature and the merged #929 lifetime regression are unchanged.

For the supported stock BF16 path, let D = min(projected_rows, chunk_rows) * vocabulary * 2:

• Eager statistics can overlap the caller's BF16 logits with FP32 conversion, subtraction and exponential result: D + 2D + 2D + 2D = 7D. The group-wide statistics predicate matches target/top-k execution, including ignored target rows projected by another requested output. Pricing does not assume the optional Triton path succeeds or handles a short final chunk.
• No-gradient logits copying can overlap local logits and both indexed copies, or 3D. Requested destination output storage is still charged separately. Shared projected rows use their existing union accounting; sequential groups retain the existing maximum workspace accounting.
• Rejection lower bounds remain distinct: zero/unsupported and non-target-gradient cases retain the original dense component, while the supported target-gradient branch retains the original target-chunk calculation and max(dense, 3 * target_dense). The new capacity multipliers do not become unconditional pruning floors.

I inspected the weak-reference storage witnesses, output accounting, optional-kernel refusal, shared/mixed/ignored-label cases, and capacity-versus-lower-bound expectations. The final commit changes only the two CI-routing lines from 6eb76f72; I ran no tests, target imports, CUDA probes or provider operations for this review.

No public API, art.megatron, numerical-tolerance, history-restriction or cache-policy changes appear in this net delta. This remains a partial supported dense-head estimate, not a whole-model, library-workspace or universal autograd peak bound; grad-enabled logits-only capacity is not newly covered. It intentionally makes some requests split or refuse under the existing budget. Brad's merge decision and full-model native qualification remain required; this source verdict clears neither hold.

[bradhilton]

Refreshed onto current main at 2a4c26cfde72cc3c1d170966209a0e3febb1b67e, with three independent source CLEAR reviews. Exact receipt hashes: d0c723c06c5ab34b70dac963c89e3a4dc8a4bf974b52db5da59c8a8d5ae21e50, b7a4c958b63f8be1d69aa288ede52359701f4afa78799648f243c3d99117d2c1, and 9f16b06f50765d5512997af41d1d4c82a41e9b9933baf94aeba90531085cf1bd.

Validation is 30 focused CPU cases plus an independent rerun of five of those cases, not 35 distinct cases. No API or art.megatron changes. Larger capacity estimates can change admission, packing, splitting or refusal; the significant-behavior merge hold remains. Current-head CI/native qualification is not yet established. The previous head and reports remain preserved.