Repository navigation
fix(gtp): preserve scratch lifetime during full-iteration capture - #7780
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/review |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
The buffer-lifetime changes look reasonable in this light review. Please add regression coverage for the FP32 accumulation scratch path; the new registered-pool tests bypass it. One minor style correction is also noted inline. GPU tests were not run in this read-only environment.
Review generated by codex
596d3a8 to
e7673c8
Compare
e7673c8 to
35dbdb0
Compare
|
/review mode=claude |
|
Invalid |
|
/review mode=strict model=claude |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
Strict follow-up review: both earlier findings are fixed, and I found nothing new to report.
This review checks the two findings from review 5375767897, which was run on 596d3a8e. I could not compare against that earlier tree because it was not available. The PR is now one rebased commit, so I couldn't isolate what changed since the last review. Under the review rules, a new finding would have to be more severe than the highest level already recorded for this PR ("critical"). Nothing can be above critical, so I only re-checked the earlier findings.
Fixed: regression test for the FP32-accumulation scratch path (inline comment 4152528993).
- The new test
test_full_iteration_fp32_accumulation_scratch_lifetimecovers the case the earlier comment asked for:- group size greater than two;
reduce_scatter_with_fp32_accumulation=True,async_reduction=Trueand_FULL_ITERATION=True;- it asserts that the symmetric-memory pool is not registered.
- Under those settings, the branch at L2194-L2225 goes through
_reduce_scatter_fp32_accum. The test'slen(weight._rs_a2a_bufs) == 1assertion confirms the all-to-all scratch was really allocated. _wait_reduce_scatter(finalize_grad=True)then queues the deferred FP32 sum and passesrs_streamto_release_wgrad_scratch. For graphed chains, that function now callsrecord_streamon the scratch (L2076-L2080).- The test makes early reuse of the scratch visible:
- a 5 ms sleep on
rs_stream; - eight overwrite allocations before the main stream waits on
rs_stream; - exact (
atol=0, rtol=0) gradient checks over ten graph replays; calls == [False, True], which shows that one warmup and one capture happened and that replays do not re-enter Python allocation.
- a 5 ms sleep on
- Under the graph-pool allocator (I checked this by reading the source; I did not run it), the scratch is allocated inside the graph (L338-L340). That is the release path the earlier comment said had no test.
Fixed: string quote style in the new test file (inline comment 4152529003). test_gtp_full_iteration_buffers.py now uses double quotes throughout, including "_wgrad_buf_pool" and "cuda", as style-guide.md requires.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0.
Overall risk: low for this follow-up. Both earlier findings are fixed, and none are still open.
I could not run GPU tests here. Everything above comes from reading the source.
Strict review passed — no significant issues found. LGTM
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by claude
35dbdb0 to
65deeb2
Compare
|
/review |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
LGTM for this light follow-up. No unresolved eligible findings remain.
Verified against the current head:
- The FP32-accumulation regression test exercises the non-registered, asynchronous full-iteration path with more than two ranks, confirms all-to-all scratch allocation, and checks exact gradients across ten replays under allocation pressure. The release path passes the RS stream and records it before dropping graphed scratch references.
- The new test file uses double-quoted string literals, resolving the earlier style finding.
Scope was verification of the earlier findings and claimed fixes, not a fresh whole-PR review. The previously reviewed tree at 35dbdb0bd3c317a2d14e7f4f77dfe73d8ee0cad7 is unavailable, so subsequent changes could not be independently isolated. The retained historical impact ceiling is critical; no new finding can exceed it. GPU tests were not executed; verification was source-only.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex
Use internal events for full-iteration chains: the record and wait nodes of an external event carry no ordering edge inside one graph, so a replay could observe the previous iteration's record. Keep captured wgrad and FP32 reduction scratch in graph-owned storage and record the RS stream as last reader before releasing it. The eager scratch pool is never touched while capturing, so a graph cannot bake in an address that eager reuse later overwrites, and capture allocations are never cached for eager reuse. Retire registered-buffer fences before capture and when the graphs are reset: a fence recorded inside a graph is never re-recorded by replay, and an eager wait on it fails. Preserve the registered-buffer drain for repeated backward so capture reuses the warmed slot. Build on the upstream wgrad-reuse fences and dummy-gradient readiness fix. Cover graph-owned scratch, registered-buffer reuse across capture, reset and recapture, and repeated backward with GPU regression tests. Signed-off-by: Jiangfei Duan <jiangfeid@nvidia.com>
65deeb2 to
eb0d91d
Compare
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/39666335651 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/39666399608 |
Port NVIDIA#7780 at 65deeb2, superseding the narrow repeated-buffer capture fix in a8722a9. Use internal full-iteration events, track RS readers before releasing scratch, keep graph allocations out of the eager pool, clear completed registered-buffer fences before capture and on graph reset, and scope the repeated-backward drain to full-iteration capture. Adapt import/docstring context and the regression's async-RS selection to this branch's chain-based API. Retain the upstream regressions in place of the overlapping local test. Four-GB300 validation (4634802): 21 symmetric-memory plus 11 full-iteration/gradient tests per rank. Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/40251550576 |
Five commits landed on main since the last merge (0a2f6ee); one of them conflicts with this branch, NVIDIA#7047 "[DCP V2 3/4] Integrate dynamic CP with standard train and eval schedules" (ea99fdf). The other four (NVIDIA#8381 CI test selection, NVIDIA#8269 PR template, NVIDIA#6502 checkpoint conversion integrity, NVIDIA#7780 gtp scratch lifetime) do not touch this branch's files. Conflicts, all resolved in this branch's structure: - megatron/core/datasets/gpt_dataset.py: imports (warnings for the new deprecation warning; no `field`, token_dtype_code is a property here). NVIDIA#7047's warn-and-rewrite of the deprecated hybrid_context_parallel alias landed in validate_without_tokenizer(), which asserts only: it keeps the warning and the "not both" error and does not rewrite the alias; the SBHD exclusion counts the alias. - pretrain_gpt.py, pretrain_hybrid.py: NVIDIA#7047 edited the legacy factory and rank predicate this branch replaced with argument_utils and megatron.training.datasets.providers; this branch's imports and the removals stay, NVIDIA#7047's get_batch changes (pg_collection, dynamic_cp) merge as they are. - tools/prepare_cache.py: same; the shared factory maps the new field. The renamed field through the data section: dynamic_context_parallel joins CONTAINER_WRITTEN_FIELDS next to the deprecated alias; the container's cross-section check reads either name against model.transformer's dynamic_context_parallel and finalize() writes the canonical field and clears the alias; the standalone factory forwards both from the arguments, which validate_args already normalizes (NVIDIA#7047). The ownership, section, provider and Core config tests follow, with cases for the alias; the frozen legacy factories the provider tests use as their oracle take the same one-line rename NVIDIA#7047 gave the originals. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: asolergi-nv <asolergibert@nvidia.com>
What does this PR do?
Full-iteration CUDA-graph capture with GTP had three independent problems:
external=True, which is right for partial graphs (record in one graph, wait in another) but gives no ordering edge when record and wait are captured into the same graph, so a replay could observe the previous iteration's record. Full-iteration chains now use internal events.handle.wait(), which orders only the RS stream. Under capture the graph pool could hand that storage to the next compute-stream allocation, producing wrong expert gradients. Graphed chains nowrecord_stream(rs_stream)before releasing. The eager scratch pool is also never read or written while capturing, so a graph cannot bake in an address that eager reuse later overwrites (relevant to partial graphs with mixed chains) and capture allocations are never cached for eager reuse.wait_eventrejects withcudaErrorInvalidValue. Fences are now dropped before capture and when the graphs are reset, and the repeated-backward drain also runs during full-iteration capture.Builds on the merged wgrad-reuse fences (#7927) and dummy-gradient readiness fix (#7467).
Memory note (non-registered path only). Scratch that is
record_stream'd during capture stays allocated until the capture ends unlessPYTORCH_CUDA_ALLOC_CONF=graph_capture_record_stream_reuse:Truelets the graph pool reuse it once the RS stream rejoins. The registered path (--gtp-remat-nccl-ub/--gtp-expert-remat-nccl-ub) owns its buffers and is unaffected.Tests
test_full_iteration_fp32_accumulation_scratch_lifetimerecord_streamon releasetest_repeated_wgrad_reuses_one_buffer_in_full_iteration_capture(4 GPUs)test_full_iteration_reuses_registered_scratch_after_warmup(4 GPUs)cudaErrorInvalidValuetest_capture_does_not_borrow_eager_wgrad_storageValidated on 4x GB300 (PyTorch 2.13 nv26.07, TE 2.20, NCCL 2.30.7) with the allocator reuse option on and off, plus
test_gtp_grad_correctness.py;tools/autoformat.shclean againstmain.Issue tracking
Related PRs: #7467, #7927 (both merged; this change relies on their fences and readiness handling).
Contribution process
Pre-checks
jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it intests/unit_tests/determinism/kernels/manifest.py(guide)