Skip to content

trainer_rank: retain completed split-forward memory peaks - #894

Merged
bradhilton merged 6 commits into
mainfrom
schulman/split-memory-floor-20260914
Sep 15, 2026
Merged

bradhilton merged 6 commits into
mainfrom
schulman/split-memory-floor-20260914

Conversation

@bradhilton

Copy link
Copy Markdown
Collaborator

A split forward_micro_batches call records its child forward peaks but drops the later caller/backward peak. A matching later call can therefore be admitted below a previously observed memory requirement. Preserve the first child baseline and earlier child peaks across CUDA counter resets, then record a bounded empirical floor on normal iterator resumption. Matching split admissions use the larger of that floor (with the existing margin) and the existing estimate.

The composition key preserves request mappings, partition, checkpoint/gradient/output mix and layout geometry while ignoring cost-only child reordering. It retains no tensors, graphs or token values. Child retained-memory profiles and collective ordering remain unchanged. Throw, close and partial-execution failure do not establish a completed split profile.

This is a partial fix for #848, stacked on held #893. It does not bound a first unseen backward, different partitions/token routing, compiler/library state, or every execution order. Unsupported keys and a full bounded cache stay unlearned. No public API, art.megatron, dependency, experiment, loss or allocator policy changes. Previously observed matching calls may now split further or refuse; behavioral eligibility and native qualification remain review items. Keep this draft held pending those decisions and the parent.

Validation: 195 CPU tests passed (one native allocation test explicitly excluded), followed by 12 passing focused tests after adding throw/close/partial-failure coverage. Ruff, formatting and diff checks pass. Original runtime AST matches the independently reviewed private candidate. Old-version failure witnesses and the corrected test-double setup failure are preserved. No native GPU qualification is claimed for this patch.

Exact head: 4c8cb6a10fbbc1904cd997e7add8ef4be0e3da79; base: fe975382158cd8fe368045375999187311868bd8.

Durable local evidence: /home/brad/.local/share/schulman/art848-split-integration-evidence-20260914-forward-bpz11vnz/REPORT.md and /home/brad/.local/share/schulman/art848-split-lifecycle-integration-20260914-forward-b3kqc_a8/REPORT.md.

@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 14, 2026 21:15 — with GitHub Actions Active

@bradhilton bradhilton left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

@bradhilton bradhilton left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

@bradhilton bradhilton left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 14, 2026 21:37 — with GitHub Actions Active

@bradhilton bradhilton left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

@bradhilton bradhilton left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

@bradhilton bradhilton left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review text removed on 2026-09-17; the consolidated review record is in the summary comment on this PR.

@bradhilton
bradhilton marked this pull request as ready for review September 15, 2026 00:26
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 15, 2026 00:26 — with GitHub Actions Active
@bradhilton
bradhilton changed the base branch from schulman/848-admission-floor-20260914 to main September 15, 2026 00:44
@bradhilton
bradhilton merged commit 8e06bfd into main Sep 15, 2026
9 checks passed
@bradhilton

Copy link
Copy Markdown
Collaborator Author

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

Change. Split forward_micro_batches calls dropped the later caller/backward peak, so a matching later call could be admitted below an observed requirement. Preserve the first child baseline and earlier child peaks across CUDA counter resets and record a bounded empirical floor that matching split admissions must honor. Partial fix for #848, stacked on #893.

Review. McCarthy, Minsky and Taravangian gave source-correctness CLEAR at 4d9d788 after confirming runtime parity with the previously cleared head. Verified: 195 CPU tests (one native test excluded), 12 focused lifecycle tests, a 98-test run with a parent negative control, scalar counter fixtures, Ruff/format/ty. All flagged the intentional behavior change: matching splits may split further or refuse; no native GPU qualification claimed.

Findings addressed during review.

  • CI caught test-file typing errors; fixed by a test-only commit, runtime unchanged.

Deferred / follow-ups.

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