fix: stream-ordering crash in SignSGD under FSDP2 - #3373
Closed
mikasenghaas wants to merge 2 commits into
Closed
Conversation
The per-param loop launched thousands of tiny kernels per step; that long asynchronous tail overlapped the gradient frees and collectives that follow the step and hit a stream-ordering race under FSDP2 + CP (async CUDA illegal memory access, deterministic after step 1, on both the plain and optim_cpu_offload paths; full offload was immune because it runs no GPU optimizer kernels). CUDA_LAUNCH_BLOCKING=1 makes the crash vanish, confirming ordering rather than indexing. Foreach batching collapses the step into a few fused launches per (device, dtype) bucket — the same kernel pattern AdamW uses, which has never triggered the race. Update semantics are unchanged: decoupled decay p *= 1 - lr*wd (algebraically identical to the previous self-aliased add) then p -= lr * sign(g).
empty_cache returns cached blocks to the driver, so any kernel still in flight that references a cached block faults with an illegal memory access once the block is freed under it. The pre-broadcast empty_cache (added for FP8 gather headroom) ran right after optimizer.step(); optimizers with long kernel tails (SignSGD's per-param loop) deterministically crashed after step 1 under FSDP2 + CP, while AdamW's compact foreach step happened to finish in time. Drain all pending work before releasing blocks.
Member
Author
mikasenghaas
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - One-line fix for the deterministic `CUDA error: an illegal memory access` that killed every SignSGD run without full optimizer offload since #3234. - `torch.cuda.empty_cache()` returns cached blocks to the driver; a kernel still in flight that references a cached block faults once its block is freed under it. The pre-broadcast `empty_cache` (added in #3234 for FP8 gather headroom) runs immediately after `optimizer.step()` — SignSGD's per-param loop leaves thousands of tiny kernels still running at that point, while AdamW's compact foreach step happens to finish in time, which is why only SignSGD crashed. Draining pending work with `torch.cuda.synchronize()` before releasing blocks closes the race for every optimizer, at ~ms cost once per step. ## Verification All on the same GLM-4.5-Air scaleswe config (2 trainer + 2 inference nodes, `sign_sgd` + `optim_cpu_offload`, no full offload): | state | outcome | |---|---| | main | illegal memory access right after step 1, 2/2 slurm attempts | | main + `CUDA_LAUNCH_BLOCKING=1` | no crash — ordering, not indexing | | **main + this one line** | [runs clean](https://wandb.ai/primeintellect/swe-ablations?nw=nwuser&search=sgd-nooffload): 11+ steps and counting, healthy grad norms, no restarts | Supersedes #3373 — its foreach SignSGD rewrite turned out to be unnecessary for correctness (verified by running the untouched optimizer with only this line). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Touches the per-step RL training hot path and GPU memory lifecycle before weight broadcast; the change is narrow and matches existing synchronize-before-empty_cache patterns elsewhere, but any extra sync can affect step timing. > > **Overview** > Fixes a **CUDA illegal memory access** that could occur right after `optimizer.step()` when the trainer frees GPU cache before broadcasting weights. > > The pre-broadcast `torch.cuda.empty_cache()` (for FP8 gather headroom) can return cached blocks to the driver while **optimizer kernels are still in flight**—especially with SignSGD’s many small per-parameter updates. The change **drains pending GPU work** via `torch.cuda.synchronize()` immediately before `empty_cache()`, with expanded comments explaining the race. > > This is a one-line behavioral fix in the RL training loop’s weight-broadcast path; it does not change broadcast semantics, only ordering so memory is not reclaimed under active kernels. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 8947aa4. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY -->
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes the deterministic
CUDA error: an illegal memory accessthat killed every SignSGD run without full optimizer offload (crash right after step 1, all attempts, plain andoptim_cpu_offloadpaths alike).torch.cuda.empty_cache()(added in feat: overlap CPU optimizer with backward #3234 for FP8 gather headroom) runs immediately afteroptimizer.step()and returns cached blocks to the driver. A kernel still in flight that references a cached block faults once the block is freed under it. SignSGD's per-param loop launches thousands of tiny kernels whose async tail is still running at that point; AdamW's compact foreach step happens to finish in time, which is why only SignSGD crashed.CUDA_LAUNCH_BLOCKING=1makes the crash vanish (ordering, not indexing), and the failing stack points atempty_cacheitself.torch.cuda.synchronize()before thatempty_cache()— drain pending work before releasing blocks. Costs ~ms once per step, protects every optimizer.SignSGD.stepwith foreach ops over local shards, bucketed per (device, dtype) — a few fused launches instead of thousands (aten._foreach_signhas no DTensor sharding strategy, henceto_local(); every op is pointwise and grads share the param's placement at step time, so the update commutes with sharding). Update semantics are unchanged and verified against the previous formula (decoupled decayp *= 1 - lr*wd, thenp -= lr * sign(g)).Verification
Before/after on GLM-4.5-Air scaleswe, 2 trainer + 2 inference nodes,
sign_sgd+optim_cpu_offload, no full offload:CUDA_LAUNCH_BLOCKING=1Also verified on the fake-data debug config (
torchrun --nproc-per-node=2 -m prime_rl.entrypoints.trainer @ configs/debug/fake/rl.toml --optim.type sign_sgd): 5/5 steps.🤖 Generated with Claude Code