DeepSpeed
0d0aa30b - Clear stale offload scratch before restoring low-precision gradients (#8728)

Commit
4 days ago
Clear stale offload scratch before restoring low-precision gradients (#8728) ## Problem With unmanaged gradient accumulation, ZeRO-2 CPU offload and BF16/FP16 master weights, a genuinely overflowing step can cause the next entirely finite step to be skipped. `_restore_cpu_offload_grad_to_gpu` copies only the locally owned fragment into a shared scratch tensor for low-precision masters. The norm and optimizer-buffer consumers read that fragment, but the overflow tracker reads the entire restored gradient. NaNs left in the unowned region by another parameter group can therefore set `local_overflow` for a finite window. ## Change Clear the destination view before the partial copy. The FP32-master full-copy path is unchanged. The regression uses the public Engine training loop with two parameter groups (12 and 8 elements) and DP=2. It first produces a real NaN step, then requires two finite steps to update both parameters. All parameters participate in every step; the test does not fabricate optimizer state or poison scratch storage directly. It covers BF16/FP16 with both low-precision and FP32 masters. ## Validation - Historical device-validation base: `c687df79653459d08c0c6fada6d124971d966662`. The device results below belong to that original revision, not the new merge head. - Two NVIDIA B300 GPUs; Python 3.12.13, Torch 2.11.0+cu130, CUDA 13.0; actual CPUAdam extension. - On unmodified master, the added BF16 and FP16 low-precision-master cases fail at `finite step did not update large`; both FP32-master controls pass. - With this change, all four registered distributed regression cases pass. - Official `pre-commit run --files` passes with the unchanged repository configuration and pinned hooks on Linux/Python 3.12. - The original module's two CPU-accumulator cases pass. Its other eight cases require the no-offload GPU FusedAdam comparison and could not run successfully in this environment: extension compilation fails on missing `cusparse.h`. They are not claimed as passing. This is a correctness fix, not a performance or convergence claim. **Merge refresh (2026-11-03):** merged upstream master `dc046b15d7a58f6c8dacba435508b52b1ebc620c` as `8efca4cd5bfb0af39822ed1c69272743eecdb787`. The test-file conflict was resolved by retaining the upstream request-boundary tests and the original overflow-recovery test plus its `CPUAdamBuilder` import. All upstream test definitions and the original overflow test retain identical ASTs; the production fix is unchanged. Relative to the merged master, the PR remains two files with 63 added lines. Exact repository changed-file pre-commit checks passed on Linux/Python 3.12.3 for the committed tree, with no hook edits; syntax and `git diff --check` passed. Distributed/GPU tests were not rerun for this merge, and hosted CI for the new head is tracked separately from the historical results above. AI assistance was used for this conflict refresh and validation; the submitting human reviewed the merge diff before the signed-off push. **Current-head hosted CI:** the DCO app check and Modal test collection succeeded; three Modal execution/smoke checks were skipped. Five GitHub Actions workflows (`Formatting`, `python`, `cpu-torch-latest`, `nv-pre-compile-ops`, and `DCO / required`) currently report `action_required` with zero jobs, pending upstream approval. These are admission gates, not executed test failures; no retry has been requested. Signed-off-by: Xucheng Zhou <aden1350@outlook.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Author
Parents
Loading