Skip to content

Fix skipped-step detection for AMP-aware optimizers - #4243

Open
gss10282023 wants to merge 1 commit into
huggingface:mainfrom
gss10282023:fix-fused-adamw-overflow-skip
Open

Fix skipped-step detection for AMP-aware optimizers#4243
gss10282023 wants to merge 1 commit into
huggingface:mainfrom
gss10282023:fix-fused-adamw-overflow-skip

Conversation

@gss10282023

Copy link
Copy Markdown

Summary

Detect skipped AMP-aware optimizer steps from the GradScaler scale backoff. Fused AdamW calls its Python step() even when the kernel skips the update, so the existing call marker reports success and allows the scheduler to advance.

The extra scale reads are limited to optimizers that handle AMP scaling themselves; other optimizers keep the existing detection path.

Fixes #4242.

Validation

  • python -m pytest -q tests/test_optimizer.py tests/test_scheduler.py: 10 passed, 2 skipped on PyTorch 2.13.0+cpu. The two skipped tests require accelerator hardware.
  • The new regression fails on unmodified f13f7c13 for fused AdamW. It checks the parameter update, optimizer step counter, skip flag, and scheduler together, including consecutive overflows, recovery, scale growth, and unfused AdamW.
  • Ruff 0.13.1 lint and format checks, and git diff --check, pass for the changed files.

The CUDA reproduction is documented in the linked issue; CUDA tests were not rerun on this branch.

Signed-off-by: gss10282023 <76646928+gss10282023@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Fused AdamW AMP overflow leaves step_was_skipped false and advances the scheduler

1 participant