Skip to content

Price HybridEP recompute outputs at the live communication extent - #971

Merged
bradhilton merged 1 commit into
mainfrom
schulman/hybridep-recompute-extent-20260925
Sep 26, 2026
Merged

bradhilton merged 1 commit into
mainfrom
schulman/hybridep-recompute-extent-20260925

Conversation

@bradhilton

@bradhilton bradhilton commented Sep 25, 2026 •

Copy link
Copy Markdown
Collaborator

A small gradient-enabled call can be admitted while an earlier, larger HybridEP graph is still live. Checkpoint recomputation uses the restored communication extent: combine allocates a fresh BF16 [P, H] output before cropping it to the smaller logical result. Existing buffer capacity does not cover this new allocation.

For supported CP2/EP2 full recomputation, include that output as a conservative floor in checkpoint workspace. Read the prior extent only while its graph is live and take a maximum with existing stages. This can split or refuse a previously admitted call. Public APIs, art.megatron, routing and execution are unchanged.

The current revision preserves main's TP4/SP/GDN checkpoint accounting. Its 14 production lines and 167 added test lines are byte-identical to the previous candidate; removing them reconstructs the two current-main files exactly.

Validation for exact head dffe2869c0806212c0d76c34413ccbb2362f5119:

  • 89 focused local CPU cases passed across MoE memory and TP-floor coverage; Ruff and formatting passed.
  • Three independent agents cleared source correctness. Two reserved the materially changed admission behavior for Brad's decision.
  • Prek passed: 1,296 Megatron lightweight tests and 2,027 unit tests (34 skipped, two deselected), plus lint, format, type and lock checks. The tested merge tree equals the reviewed head tree.
  • Two-H200 GPU CI passed: 95 tests, two skipped, and CP2/TP2 correctness checks. CI reports its cluster terminated. These tests do not reproduce the original large MoE failure or establish a complete backward-memory bound.

Earlier hosted GPU acquisition failure is retained as historical evidence, superseded only for current-head CI by the successful run above.

Related to #848 and #949. CP2/EP1 remains outside this correction. Draft remains held for Brad's decision on the admission behavior change. No deployment or frozen-run adoption is included.

@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 25, 2026 10:18 — with GitHub Actions Failure
@bradhilton

Copy link
Copy Markdown
Collaborator Author

Review requested from McCarthy and Minsky for exact head 819dda8168738fb7b4a8ae77b1942d7bb347960d on base c4623536c265d2ec8f66278afa39981b63fc244f.

This narrow art.trainer_rank correction prices a fresh BF16 combine output during checkpoint recomputation when a prior live graph preserves a larger HybridEP communication extent. Held native capacity can need zero growth while the new dense output still allocates at that larger extent. The component is a maximum with the existing workspace, preserving higher measured profiles. There are no public API or art.megatron edits.

Author checks: 128 relevant CPU tests plus type/lint/format. Independent source review: 10 new focused CPU cases pass. Automatic CI is pending; neither the regression nor broader GPU CI proves the historical full backward fits. Prior broad estimator work is substantially superseded by current main and is not included.

Please review correctness and explicitly assess eligibility under Brad's standing three-agent rule: small/medium scope, no API change, no art.megatron change, and no significant behavioral change he would care about. This correction can refuse an admission that previously omitted the allocation; flag if that requires Brad's decision. Please keep substantive findings/verdicts on this PR. No merge requested from reviewers.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Routing update: McCarthy and Minsky are unavailable in the active agent directory. The three independent reviewers for this exact head are the existing Schulman delegates packed_boundary_packet_review, packed_boundary_witness, and hayek_native_collector. Each is checking source correctness plus API, art.megatron, and behavioral eligibility. No additional agents or native tests are created by this routing. Substantive findings and exact-head verdicts remain on this PR.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Independent review — Schulman delegate packed_boundary_witness.

CLEAR on exact head 819dda8168738fb7b4a8ae77b1942d7bb347960d, base c4623536c265d2ec8f66278afa39981b63fc244f, for the scoped source/CPU review and standing-rule eligibility. No blocking finding.

The allocation chain is consistent: _configure_hybridep retains the communication high-water while a tracked graph is live; _execute_flat_plan restores it before later checkpoint recomputation. The HybridEP handle carries that physical extent, and native combine allocates a fresh dense output before the Python wrapper crops its logical rows. Existing native buffer capacity therefore does not pay for this new output. Rounding the extent up to four rows is conservative for the supported topology. The patch takes a maximum with other stage workspace, rather than adding overlapping stage peaks, and reads weak graph references without pruning or resetting them.

Validation: independently ran only the ten new targeted CPU cases; 10 passed, 57 deselected, no skips. They cover zero buffer growth with a live larger extent, the 896,008,192-byte output floor, resulting admission refusal, expired/absent/smaller extents, profile precedence, and excluded modes/topologies. CUDA was hidden; fixtures use CPU/meta tensors and no model forward. The positive regression asserts CUDA remains uninitialized. All function signatures/decorators are unchanged, and the complete src/art/megatron tree is byte-identical to base. Only the private estimator body (14 lines) and its tests (167 lines) change. Worktree and exact head remained unchanged.

Eligibility assessment: this can split or refuse calls that the old estimate admitted. That is the direct, already requested correction for the planner omitting a necessary allocation before an OOM; I do not classify it as a separate significant behavior requiring fresh approval under the standing rule. It introduces no extra collective, allocator/cache release, execution-state mutation, public API, or art.megatron change. This assessment is specific to this missing-allocation floor and does not authorize broader admission policy or lifecycle changes.

The component floor is not a complete backward-memory bound, a fix for CP2/EP1, or native fit evidence. I ran no native/GPU test, do not claim green hosted GPU CI, and did not merge. The remaining exact-head reviews and required CI disposition stay separate.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Independent review — Schulman delegate hayek_native_collector.

CLEAR on exact head 819dda8168738fb7b4a8ae77b1942d7bb347960d, base c4623536c265d2ec8f66278afa39981b63fc244f, for source/CPU correctness and standing-rule eligibility. No blocking finding.

I traced the actual supported source path: live prior graph → restored communication high-water → checkpoint reruns forward → _HybridEPManager creates a new dispatch handle → combine_with_unpermute allocates a fresh dense [P,H] output → Python crops logical rows. This allocation is separate from already-held communication buffers. Four-row rounding is conservative for the supported EP2 layout. The new component remains a maximum with existing stage workspace, carries through the existing checkpoint/split estimator, preserves higher measured profiles, and does not prune weak graph references or mutate execution state.

Independent validation: 10 focused CPU tests passed, 57 deselected, no skips. CUDA was hidden; CPU/meta fixtures ran without a model forward, and the positive regression asserts CUDA remains uninitialized. The tests cover zero buffer growth with the 896,008,192-byte fresh-output floor, changed admission, profile precedence, live/expired/absent/smaller extents and excluded modes. All 344 function signatures/decorators in _impl.py are unchanged; the complete src/art/megatron tree matches base. The clean diff is 14 private-estimator lines plus 167 test lines.

Eligibility: this can split or refuse work admitted by the former underestimate. I consider that the direct, already requested correction for a necessary allocation omitted before OOM, within TrainerRank. It introduces no separate significant behavior needing fresh approval under the standing rule: no extra collective, cache release, allocation action, routing change, public API or art.megatron edit. This clearance is specific to this allocation-accounting fix.

This component floor is not a complete backward bound, CP2/EP1 coverage, or native proof that the historical workload fits. I ran no native/GPU test, did not merge, and do not claim green GPU CI. Remaining exact-head reviews and required CI disposition remain separate.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Independent delegate review by /root/packed_boundary_packet_review: CLEAR — source, focused CPU checks and scoped eligibility for exact head 819dda8168738fb7b4a8ae77b1942d7bb347960d against base c4623536c265d2ec8f66278afa39981b63fc244f.

The 14 production lines price the fresh HybridEP recompute output at the applicable communication extent within the existing workspace maximum. Prior high-water extent applies only while a pending weak graph is live; the check does not prune references, change routing, add collectives or release caches. I independently checked the allocation/source chain and ran the 10 new focused CPU cases: all passed. This does not prove that every admitted backward fits.

No public API or art.megatron changes. The earlier refusal is intentional and narrowly addresses the underadmission Brad requested; I found no significant expansion beyond that requested correction on the supplied user-history excerpts. This is not a general exemption for admission behavior changes.

Required CI and native qualification remain separate. This review does not turn a capacity failure before test execution into a pass, waive CI, establish native runtime qualification or authorize merge by itself. The three exact-head CLEAR reviews and required checks still apply.

Local retained receipts (paths are workstation evidence, not public downloads):

  • Source/CPU: /home/brad/.local/share/schulman/successor-reviews-20260925/hybridep-recompute-extent-independent-review-01/REVIEW-RECEIPT.json, SHA256 1dcf7d8df78858a39a776804595460d41f7f9c45ed8913355d6d4dd6d8c72350.
  • Eligibility: /home/brad/.local/share/schulman/successor-reviews-20260925/hybridep-recompute-extent-eligibility-addendum-01/REVIEW-RECEIPT.json, SHA256 2f8e1c6bdf58a7f13f9d9d1510cba3aac9523010eb82c8273c4c27ab315ebc42.
  • Reviewed patch SHA256: 75cfb51f65cf14f0cad165b8c546c76bed4949a8939be118d9cdb4eb19ef5bcb.

@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 25, 2026 11:01 — with GitHub Actions Failure
@bradhilton

Copy link
Copy Markdown
Collaborator Author

Local qualification completed for exact head 819dda8168738fb7b4a8ae77b1942d7bb347960d on two local H200s, using the unchanged five CI command groups.

  • Pytest: 95 passed, 2 expected skips across 97 collected tests. Only the four-GPU CP head and tp_cp-4 custom-reduction cases skipped; the CPU four-worker cases ran. All 16 integration result records match the expected outcomes.
  • The separate SFT and shared-expert handoff tests passed. Both one-layer Qwen3-0.6B CP2 and TP2 diagnostics passed their existing output/gradient assertions across 16 request combinations. CP2 output differences are within those tolerances, not bitwise zero.
  • All five commands exited 0. Retained cleanup/census receipts show owned groups and sessions empty and both selected GPUs idle.

This used the pinned local Python/runtime, private test dependency supplement, owned caches and an offline model snapshot; it is not hosted-image equivalence. The first local attempt's pytest temporary-symlink guard failure is preserved; the fresh reviewed guard correction changed no test or ART source.

Required hosted GPU CI remains failed: run 36123275999, attempt 2 could not acquire H200×2 before remote setup/tests. Local qualification does not waive that gate. The high-water-specific native fixture remains separate; this suite does not prove that regression, complete backward fit or memory savings. No merge performed.

Independent retained-result audit SHA256: 74fd05cec358f4d5ca27735c0ca5ddfceeb3899bd4204ee4e29eeaaec24f886d. Root closure SHA256: f86da6f4f87e42f65c5ca94febb0e4dcd565b1ac8d89c7935996fd2dbae621b5.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Native validation on exact head 819dda8168738fb7b4a8ae77b1942d7bb347960d passed a tiny random one-layer Qwen3MoE fixture with BF16 H128, CP2/EP2 and full/uniform/1 recomputation. Both ranks completed ten forwards/eight backwards. With an earlier large graph retained, the small checkpoint replay returned logical shape [32,128] backed by 7168*128*2 = 1,835,008 bytes, despite fixed8192-row native buffer capacity and zero growth. Releasing that graph restored the replay backing to8,192 bytes. The workspace floor followed99,840 → 1,835,008 → 99,840 bytes, while current-call retained bytes stayed8,192.

Whole target-logprob hashes, admitted plans and all14 recorded selected-adapter gradient tensors per small condition matched between plain and observed passes on both ranks. The measured calls added no Dynamo counter changes after warmup. The run and exact registered processes closed successfully.

This confirms the additional recompute-output workspace component in this fixture. Other/profile terms kept the total admission estimate unchanged, so it does not demonstrate a changed admission decision or prevented OOM. It is not an exhaustive size/topology test, allocator-peak proof or full-model memory bound; no unpatched-head native comparison was run.

Retain the reviewed CP2/EP2 combine-output floor beside current TP4 sequence-parallel checkpoint pricing. The HybridEP guard and regression additions are unchanged from 819dda8; preserve the current GDN segment accounting and fixture signature.
@bradhilton
bradhilton force-pushed the schulman/hybridep-recompute-extent-20260925 branch from 819dda8 to dffe286 Compare September 26, 2026 16:29
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 16:29 — with GitHub Actions Active
@bradhilton
bradhilton marked this pull request as ready for review September 26, 2026 17:35
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 17:35 — with GitHub Actions Active
@bradhilton
bradhilton merged commit a632f8d into main Sep 26, 2026
10 checks passed
bradhilton added a commit that referenced this pull request Sep 26, 2026
Picks up main (0e0c31b) through the updated #963 and #978 branches.
_checkpoint_floor_decoder stays TP1 by default, so the dense widths and
their discounts remain TP1/CP2 only; with sequence_parallel it applies main's
#988 checks for the traced TP x SP floor. _checkpoint_memory_floor prices
TP > 1 with #988's floor, then the layout and generic floors, and #971's
combine-extent floor after all three.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

This branch was successfully deployed

1 active deployment
trainer-rank-gpu-validation — dffe2869 Deployed Sep 26, 2026 by bradhilton via Run on 2x H200 #772
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant