huggingface / huggingface/accelerate
[Bug] Fused AdamW AMP overflow leaves step_was_skipped false and advances the scheduler
- Dominant language
- Python
- Stars
- 9.9k
- Forks
- 1.5k
- Avg merge
- 5d 2h
- Merged PRs (30d)
- 27
Description
### Summary
With FP16 and `AdamW(fused=True)`, a gradient overflow skips the parameter update, but `AcceleratedOptimizer.step_was_skipped` stays `False`. `AcceleratedScheduler` then advances the learning-rate schedule even though no optimizer update took place.
The unfused AdamW path reports the same skip correctly.
### System Info
**Environment:** Accelerate 1.14.0, commit [beb0672a](https://github.com/huggingface/accelerate/blob/beb0672aa8444ea7647aee056f624effe5996346/src/accelerate/optimizer.py), PyTorch 2.13.0+cu130, CUDA.
The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in [f13f7c13](https://github.com/huggingface/accelerate/blob/f13f7c13b64c10b6eb7e2d73171ea1b94c748701/src/accelerate/optimizer.py); the full experiment has not been rerun on that commit.
A separate current-source regression was run on Ubuntu 26.04 LTS (x86_64), Python 3.11.15, NumPy 2.4.6, and PyTorch 2.13.0+cpu, using source baseline `f13f7c13` and PR #4243. It constructs `GradScaler` and the optimizer/scheduler wrappers directly; no Accelerate launch configuration is used for that CPU test. The pinned CUDA reproduction uses FP16 mixed precision and fused AdamW.
### Information
- [ ] The official example scripts
- [x] My own modified scripts
### Tasks
- [x] My own task or dataset (give details below)
Synthetic gradients with an injected overflow, using the native optimizer and scheduler wrappers.
### Reproduction
The test prepares a model, a fused AdamW optimizer, and a scheduler through an FP16 `Accelerator`. After `backward()`, it overwrites one gradient element with `torch.inf`, then calls the prepared optimizer's `step()` followed by the scheduler's `step()`.
**Expected:** parameters and the optimizer step counter unchanged, `step_was_skipped == True`, scheduler not advanced.
| Case | Parameter update | `step_was_skipped` | Scheduler steps taken |
| --- | --- | --- | ---: |
| Fused AdamW, overflow | Skipped | `False` | 1 |
| Unfused AdamW, same overflow | Skipped | `True` | 0 |
| Fused AdamW, overflow, with scale-backoff repair | Skipped | `True` | 0 |
| Fused AdamW, finite gradient (control) | Applied | `False` | 1 |
In the affected case the optimizer's own state step is also unchanged, so the update really is skipped. The mismatch is in the wrapper's report and the scheduler step that follows from it. Reproduced on a second host.
The regression code is in [PR #4243](https://github.com/huggingface/accelerate/blob/d28f872bdf1f59ef7f9e3c52d7e48050ac6738de/tests/test_optimizer.py). From a checkout of that PR with its test dependencies installed:
```bash
python -m pytest -q tests/test_optimizer.py tests/test_scheduler.py
```
The CPU run passed 10 tests and skipped 2 accelerator-only tests. The new fused-AdamW overflow case fails on the unmodified baseline.
### Expected behavior
On overflow, parameters and the optimizer state step should remain unchanged, `step_was_skipped` should be `True`, and the scheduler should not advance. A subsequent finite update should advance the optimizer and scheduler together.
### Root cause
`AcceleratedOptimizer` detects a skipped update by checking whether its patched Python `optimizer.step()` was invoked.
For fused AdamW, `GradScaler` always invokes that method and passes the non-finite-gradient information down to the fused kernel, which skips the update internally. The Python call therefore happens in both the successful and the skipped case, so it cannot distinguish overflow on this path.
### Proposed fix
Record `scaler.get_scale()` before `scaler.step()` and compare it with the scale after `scaler.update()`. With the tested `GradScaler` behavior, a decreased scale indicates overflow, and the wrapper should set `_is_overflow = True`.
This restores the skip flag and the scheduler behavior in the reproduction. Regression coverage should include the fused-overflow, unfused-overflow, and finite-gradient cases and check the parameter update, optimizer step counter, skip flag, and scheduler clock together.
Contributor guide
Research direction
Start in src/accelerate/optimizer.py at AcceleratedOptimizer and inspect the scaler handling around optimizer.step(). Run python -m pytest -q tests/test_optimizer.py tests/test_scheduler.py, including the regression coverage in PR #4243. Done means fused and unfused overflow cases report step_was_skipped correctly, leave parameters and optimizer state unchanged, and do not advance the scheduler, while finite updates advance both.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 38/100