Skip to content

fix: release GDN LoRA components before output addition - #1048

Merged
bradhilton merged 2 commits into
mainfrom
schulman/gdn-lora-output-lifetime-candidate-20260929
Sep 30, 2026
Merged

bradhilton merged 2 commits into
mainfrom
schulman/gdn-lora-output-lifetime-candidate-20260929

Conversation

@bradhilton

@bradhilton bradhilton commented Sep 29, 2026 •

Copy link
Copy Markdown
Collaborator

GDN LoRA projection keeps qkv, z, beta and alpha tensors alive after concatenating them. Release those four Python owners before the final output addition to reduce temporary component storage; tensor math and the out-of-place output stay unchanged.

A real TE component comparison on one H200, with fixed synthetic BF16 values at the 64,902-row Qwen geometry, measured a 1,603,339,264-byte (1.493 GiB) reduction in peak allocated memory across all 12 pairs with equal starting allocations. All 18 output/input-gradient/four-LoRA-gradient comparisons were bitwise exact, and the designated tensor state was unchanged. Peak reserved memory and synchronized physical-free readings were unchanged. This is component evidence; it establishes no whole-model savings or planner allowance.

Validation: changed-file Ruff lint and formatting checks pass. The maintained regression uses the real GDN fixture in no_grad and gradient paths. On the same H200, its two cases fail on the unmodified runtime at the intended storage-lifetime assertion, then both pass with this change, including input and all four LoRA gradient checks. Setup and teardown pass in both arms; normal pytest artifacts, JUnit and process exit statuses agree.

Draft for review. Merge of this art.megatron change is held for Brad's decision.

@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 29, 2026 21:54 — with GitHub Actions Error
@bradhilton

Copy link
Copy Markdown
Collaborator Author

Agent validation summary for exact head fb322a874fc62e4761916fc967166eed6a6f4407:

  • Two independent Astra source reviews and one Opus review cleared this head. The maintained regression produced both intended storage-lifetime failures on the stock runtime, then both cases passed on this head, including setup, teardown and gradient checks.
  • The real TE component completed 36 calls and 18 bitwise output/gradient comparisons. All 12 pairs had equal starting allocations and 1.493 GiB lower peak allocated memory. Designated tensor state was unchanged. Reserved memory and synchronized physical-free readings were unchanged; no whole-model savings or planner allowance is claimed.
  • Ruff lint/format passed. CI found a test-only tuple-annotation issue; a one-line correction is prepared for focused review. GPU CI is still running.

The PR remains draft. Merge of the art.megatron change requires Brad's decision.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Exact-head addendum for c007fcd809c92e10486340f67999ae85ab7ec6c3: the only change from the validated predecessor is the test helper's args: tuple annotation, fixing the three CI type errors. Two independent Astra reviews and focused Opus review clear this exact commit. The runtime file is byte-identical; the test AST after removing that annotation and the dispatch method code object are unchanged. Targeted ty and changed-file Ruff checks pass; local whole-tree ty retains the same 18 unrelated environment-dependent diagnostics seen on the predecessor.

The earlier maintained old-red/new-green and component results remain applicable: 1.493 GiB lower allocated peak, 18 bitwise comparisons, and no improvement in reserved peak or synchronized physical free memory. There is no whole-model or planner claim. No local GPU tests were repeated for the annotation. Normal CI is running on the new head; this PR remains draft until applicable checks pass, and merging the runtime change remains held for Brad.

@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 29, 2026 22:13 — with GitHub Actions Active
@bradhilton
bradhilton marked this pull request as ready for review September 29, 2026 22:34
@bradhilton

Copy link
Copy Markdown
Collaborator Author

Ready for review at exact head c007fcd809c92e10486340f67999ae85ab7ec6c3. Quality CI and 2x H200 validation passed. The ready-for-review gate reused that exact successful GPU run, with no second GPU job. Both independent Astra reviews and Opus clear this exact head; the retained old-red/new-green and bitwise component evidence still apply.

The measured benefit remains 1.493 GiB lower component allocated peak, with unchanged reserved peak and synchronized physical free memory. There is no whole-model or planner claim. The CI cluster's termination is recorded in its job log. This PR is unmerged; the art.megatron runtime change remains held for Brad's merge decision.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Whole-model qualification update for unchanged head c007fcd809c92e10486340f67999ae85ab7ec6c3 against baseline 1c05be4eb08ba1148776af57492791a4571f64bb:

The genuine 64,902-token A+B forward comparison completed with zero allocated or reserved peak reduction. Both arms reached 94,780,217,856 bytes (88.271 GiB) peak allocated. Corresponding target hashes, starting allocator counters, natural plans, model state, RNG and inputs match. Warm2 was compiler-quiet in both arms. This was gradient-enabled forward only; it did not run backward.

This limits the earlier 1.493 GiB component result: that local saving has not translated into a whole-model forward benefit on this workload. The exact CP1 source route reaches the patched projection under a recursively compiler-disabled call, but no local allocation trace establishes why the global peak is unchanged. A separate full backward comparison is being prepared; no backward, physical-memory or larger-batch admission benefit is claimed.

Both retained operation histories are closed, including the prior candidate plan-serialization refusal; the final candidate reused the completed immutable baseline after the full canonical-plan guard correction. The final 63 recorded process identities/groups are absent and GPU2 is returned. Independent result receipt: gdn-lifetime-candidate-reuse-independent-result-review/CLOSED-JOIN.json SHA256 a7956b18632bfb53776c9af6b6ed29d63db6768a56c4b24245df583a1e26ec3a, under /home/brad/.local/share/schulman/recovery-20260929/.

@bradhilton

Copy link
Copy Markdown
Collaborator Author

Full-model forward/backward qualification is now complete for reviewed head c007fcd809c92e10486340f67999ae85ab7ec6c3. The earlier isolated 1.493 GiB component saving did not reduce the full-model allocated or reserved peak in this workload. Measured savings were zero in all three phases.

Phase Natural execution plan Forward allocated peak Forward + backward allocated peak
Cold 64,902 rows together 88.271 GiB 93.848 GiB
Warm 1 / warm 2 33,706 and 31,196 rows in separate children 82.182 GiB 84.890 GiB

Corresponding baseline/candidate source, input, initial weights, RNG, plans and allocator starts matched. Warm 2 alone had no observed compilation-counter changes. The roughly 9 GiB cold-to-warm reduction is not evidence that warming makes the same execution plan cheaper: the planner changed the execution plan.

All 660 gradients were present and finite. Warm target outputs/losses matched, but gradients varied within unchanged arms as well as across arms: relative L2 was 2.6285% between baseline repeats, 1.1323% between candidate repeats and 2.6255% between the final baseline/candidate passes. The first repeat still compiled. These small numbers of observations establish neither a causal numerical effect of this patch nor numerical equivalence. Cold loss differed by 0.03095%; objective formula, masks and denominator were unchanged.

Both fresh processes completed with actual exit status 0 and no optimizer updates. All recorded owned identities and process groups were absent at closure; the GPU was released. No larger-batch admission, physical-memory peak or production-throughput improvement is claimed.

Independent report: /home/brad/.local/share/schulman/recovery-20260929/gdn-lifetime-backward-independent-result-review-v1/REPORT.md, SHA256 050bbbdafb62734789961de7402cb0668791b762ffd7b9fc2b910f1e50886974. Original comparison SHA256 63af1eee56415d8e7ba71f3c16718329c7222b2d5e58d5f18bec5f0f3cad77b0.

@bradhilton
bradhilton merged commit 6bfa495 into main Sep 30, 2026
10 checks passed

This branch was successfully deployed

1 active deployment
trainer-rank-gpu-validation — c007fcd8 Deployed Sep 29, 2026 by bradhilton via Run on 2x H200 #999
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