Back Home

GitHub Repo

PyTorch Community Reports Fused SGD Initialization Gap That May Corrupt Momentum State After AMP Skips the First Step

A PyTorch community report on October 21 says fused SGD may leave its momentum buffer uninitialized when a gradient overflow skips the first update, affecting the next valid update. Another participant reported reproducing the issue on CPU, but maintainers have not confirmed it and no official fix has been released.

Pytorch Deepdream (https://github.com/gordicaleksa/pytorch-deepdream) by gordicaleksa (https://github.com/gordicaleksa/pytorch-deepdream/commits?author=gordicaleksa) · MIT · Image source
zh-Hant

On October 21, PyTorch received a training-correctness report: when torch.optim.SGD is configured with both momentum>0 and fused=True, a GradScaler overflow on the first update may leave the optimizer with uninitialized momentum state, which is then read during the next valid update. The reporter provided a minimal reproduction and said the issue occurred on both CPU and CUDA with version 2.14.1 and the October 20 nightly build. Community report

This interaction involves the safeguards used in mixed-precision training. The official AMP examples explain that GradScaler first unscales gradients and skips the update when it encounters inf or NaN, preventing abnormal gradients from changing parameters. So when investigating this kind of issue, it is necessary to check not only whether parameters remain unchanged at the time, but also whether the optimizer retains state that could affect later operations. AMP documentation

Another participant said they reproduced the issue on CPU and traced it to the fused path: the code allocates the momentum buffer, then the kernel returns early when it encounters the overflow flag, while the outer code still saves the buffer. The next iteration therefore treats it as existing momentum. In the report’s deterministic mode, uninitialized memory is filled with NaN, causing parameters to become NaN afterward. This is a test condition that exposes the gap; it should not be interpreted to mean that all training runs will necessarily produce the same result. Reproduction and source analysis

The official SGD documentation specifies that the initial momentum is set directly to the current gradient, without dampening on the first step. The fused version improves performance by combining operations and is not currently the default. This means a fix must avoid reading uninitialized values while preserving the mathematical semantics of the first valid update. SGD documentation

For engineering teams using this configuration, a reasonable validation approach is to deliberately trigger overflow on the first step, then compare parameters and momentum state with fused=True and fused=False. Follow-up items include maintainer confirmation, regression tests covering nonzero dampening, and identifying the release that will include the fix. As of the check, discussion remained at the community reproduction and proposed-fix stage, with no linked fix PR. Issue status

Sources

  1. Fused SGD keeps an uninitialized momentum buffer when GradScaler skips the first step
  2. SGD — PyTorch 2.14 documentation
  3. Automatic Mixed Precision examples — PyTorch 2.14 documentation