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>