Repository navigation
Fix forward chunk tracking for resident activation groups - #8243
WEIXILIN0307 wants to merge 12 commits into
Conversation
…t copy `ChunkOffloadHandler.reload()` enqueues its host-to-device copy on the caller's current stream -- which on the `tensor_pop` path is the compute stream -- with no ordering edge on the group's `_offload_event`, recorded on `d2h_stream`. Only `bulk_reload_group` waits on that event; `tensor_pop` calls `reload()` inline and bypasses the wait. A backward that reaches a group before any prefetch ran can therefore read the pinned staging buffer while the device-to-host copy that fills it is still in flight. The saved activation silently becomes stale buffer contents and the gradients computed from it are wrong. Nothing raises, and the loss curve still descends. This contradicts the feature's own documented design: "Before the group's backward, tensors are reloaded from CPU to GPU on `h2d_stream`, and the compute stream waits for the transfer to complete." The inline path is not rare. The prefetch trigger `on_group_start_backward` is the backward of the group-start marker, which autograd runs at the end of that group's backward, after every `tensor_pop` in it, so the reload LIFO has to be primed one group ahead by `pre_reload_last_layer`. During warmup nothing primes it: `should_bulk_offload` returns True unconditionally while `is_warmup`, and `tensor_pop` drained a group's tensors while leaving the group object on `_groups_to_reload`, so every `bulk_reload_group` re-targeted an already-drained tail and the priming branch was never reached. Measured on a 4-layer GPTModel with `--offload-modules core_attn`, the warmup iteration took the unordered inline path for all 24 of its reloads, with 0 prefetched. Add a `reload_pending` latch to `OffloadTensorGroup`, set when the group's device-to-host copy is enqueued and cleared by whichever reload consumes it. On the inline path, `tensor_pop` waits on the group's offload event on the compute stream and drops the group from the LIFO so the next prefetch advances. Keying the latch on group identity is necessary rather than waiting unconditionally: `_offload_event` is a single event per group, reused every iteration, so a group that was not offloaded this iteration would otherwise wait on a stale edge. `reset()` clears the latch for groups offloaded but never reloaded, and `bulk_reload_group` clears it too so a prefetched group never re-serializes compute behind the device-to-host copy. With no artificial stall at all -- a 256 MiB device-to-host copy and nothing else -- 20 trials per arm corrupted 19/20 on B300 and 20/20 on A40 before this change, and 0/20 on both after it. Add a regression test that stalls `d2h_stream` with `torch.cuda._sleep` before the commit to make the race deterministic, asserts the stall took effect, and then checks both the ordering and the reloaded values. Reverting the latch turns it red. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: GarlGuo <wg247@cornell.edu>
The single-group test stayed green when `_groups_to_reload.remove(group)` was dropped from `tensor_pop`. Use two forward-ordered groups, so the inline pop of the second must release the LIFO and let the next prefetch advance to the first. Raised in review on NVIDIA#7399. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: GarlGuo <wg247@cornell.edu>
fc9e737 to
af95692
Compare
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
svcnvidia-nemo-ci
added
waiting-on-customer
Signed-off-by: Pinoeer-kingxi <13022943007@163.com>
af95692 to
fb271eb
Compare
WEIXILIN0307
changed the title
Signed-off-by: Pinoeer-kingxi <13022943007@163.com>
fb271eb to
d8be77f
Compare
svcnvidia-nemo-ci
removed
the
waiting-on-customer
WEIXILIN0307
changed the title
|
Thanks for tracing the scheduler and pool lifetimes so carefully. I’ve removed the extra H2D-to-D2H wait and the test that bypassed normal warmup. I’ve also made #7399 a prerequisite and removed my duplicate fallback wait. Reload ordering and queue bookkeeping now come from that PR. I took another look at the remaining changes and removed the evaluation-first reset handling, the tensor-growth special case, and the new scheduler argument. The evaluation-first test only showed an internal state difference, without demonstrating a later training failure. Tensor growth also goes beyond what this PR needs. The fix I’ve kept is in finish_all_groups(): empty transfer queues don’t mean that the current microbatch has finished its forward groups. With fraction zero, that shortcut can advance to the next chunk too early and cause Chunk mismatch during backward. The new regression uses fixed-length training from the start: sequence 128, two microbatches and three SGD steps. The old code fails at fraction zero on both ranks, while fractions 0.5 and 1 pass. The updated code passes all three. I also reran the dependent GDN tests and complete Qwen TP2+SP+PP2 training on four A6000s. The logs and commands are here, including the old-code failures and exact commits. #7399 still needs to land first. |
|
/ok to test 1336c9d |
|
/review |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
LGTM. No actionable findings in this light review. Inspected chunk progression, inline-reload ordering and queue bookkeeping, and the regression tests. Static review only; GPU tests were not run.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex. This was a full review. Request a full review with /review reset=true.
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
/ok to test f9d0b7a |
svcnvidia-nemo-ci
added
waiting-on-customer
Signed-off-by: Pinoeer-kingxi <13022943007@163.com>
…PT training Signed-off-by: Pinoeer-kingxi <13022943007@163.com>
Signed-off-by: Pinoeer-kingxi <13022943007@163.com>
Signed-off-by: Pinoeer-kingxi <13022943007@163.com>
svcnvidia-nemo-ci
added
waiting-on-maintainers
What does this PR do?
Fix a
Chunk mismatchassertion during fixed-length GPU training when fine-grained activation offloading is enabled with--activation-offload-fraction 0. The failure was reproduced in a native four-layer GPT training test with synthetic tokens and two microbatches.finish_all_groups()treats empty offload/reload queues as a completed forward chunk. At fraction zero, those queues can be empty while later groups of the same scope still remain. Advancing the handler at that point gives backward the wrong microbatch's chunk.Remove the queue-empty shortcut and use the existing remaining-group lookup and cursor. The regression runs the real forward/backward schedule and optimizer steps, starting with the normal offload warmup lifecycle. Evaluation-first initialization and changes to tensor sizes after warmup are outside this PR.
Dependencies and scope
Requires #7399 (
73ac08c254a0ff8997a0ff559799a4d7f073248b) for fallback reload ordering and reload-LIFO bookkeeping. Its original commits are preserved; merge #7399 first. #7908 depends on this PR.The original change after that prerequisite contains the chunk-tracking fix and its regression.
Validation
Current head:
5bc9e72bda70129f088dd78328d4064f1d838c62, including a merge frommainat3ddfe0fe07efc73dbdc40d4a02acbc1c1f8a9f08.On one RTX PRO 6000, seven selected tests pass: the native fixed-length training regression at fractions 0/.5/1, three TE CUDA graph cases, and the upstream checkpoint-without-output storage regression. Commands and rank logs are available.
At
d8be77f56f4d38713cb0306c5c4893aa3cb2331a, the same fixed-length regression fails withChunk mismatchon the unchanged prerequisite at fraction zero and passes with the fix. Both positive-fraction controls pass. The pinned original evidence retains commands, per-rank logs, state comparisons and two-/four-rank results. Those multi-rank executions are historical; the current-head rerun uses one GPU. Upstream GPU CI is pending.Issue tracking
Related to #6757.
Code provenance
No code copied or adapted from external repositories.
Contribution process
Pre-checks