fix: synchronize before the pre-broadcast empty_cache

[mikasenghaas]

Summary

• One-line fix for the deterministic CUDA error: an illegal memory access that killed every SignSGD run without full optimizer offload since feat: overlap CPU optimizer with backward #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 feat: overlap CPU optimizer with backward #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: 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

---

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.

^{Reviewed by Cursor Bugbot for commit 8947aa4. Bugbot is set up for automated code reviews on this repo. Configure here.}