Skip to content

Fix eval mode for wrapped optimizers - #4161

Open
uuzzrm wants to merge 1 commit into
huggingface:mainfrom
uuzzrm:codex/fix-accelerated-optimizer-eval
Open

Fix eval mode for wrapped optimizers#4161
uuzzrm wants to merge 1 commit into
huggingface:mainfrom
uuzzrm:codex/fix-accelerated-optimizer-eval

Conversation

@uuzzrm

@uuzzrm uuzzrm commented Aug 12, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes #4159.

AcceleratedOptimizer.train() already reaches through the extra optimizer wrapper used by DeepSpeed. eval() only checked the outer wrapper, so schedule-free optimizers could remain in training mode when an evaluation checkpoint was saved.

This adds the matching one-level fallback and a regression test using a small schedule-free-like optimizer wrapped in a DeepSpeed-like object.

Testing

  • PYTHONPATH=src uv run --with torch --with pytest --with numpy --with packaging --with psutil --with pyyaml --with huggingface_hub --with safetensors --with '.[test_prod]' pytest tests/test_optimizer.py -q
  • uv run --with ruff ruff format --check tests/test_optimizer.py src/accelerate/optimizer.py
  • git diff --check

AI writing disclosure

  • AI-assisted: Codex helped investigate the issue, draft the focused patch and regression test, and run the checks. I reviewed the changed lines, reproduced the failure before the fix, and verified the passing test result.
  • AI-generated
  • No AI usage

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • I read the contributor guideline and searched existing issues and PRs before starting.
  • I linked the issue addressed by this change.
  • Documentation needs updating.
  • I added a regression test for the changed behavior.

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.

AcceleratedOptimizer.eval() does not reach through DeepSpeed wrapper, silently saving train-view schedule-free weights

1 participant