Reach through the DeepSpeed optimizer wrapper in AcceleratedOptimizer.eval()

[vineethsaivs]

What this fixes

AcceleratedOptimizer.train() has a branch for the case where DeepSpeed wraps the user optimizer one level deeper:

elif (
    hasattr(self.optimizer, "optimizer")
    and hasattr(self.optimizer.optimizer, "train")
    and callable(self.optimizer.optimizer.train)
):
    # the deepspeed optimizer further wraps the optimizer
    self.optimizer.optimizer.train()

eval() carries the same docstring but never got the matching branch, so under DeepSpeed a schedule-free optimizer is put into train mode by train() and never taken back out by eval(). Nothing raises; evaluation just silently runs with train-mode weights, which for a schedule-free optimizer means the wrong parameters.

The history suggests this was an oversight rather than a deliberate asymmetry. #2631 introduced train()/eval() as a mirrored pair, #3055 changed both together in one diff, and #3266 ("support for wrapped schedulefree optimizer when using deepspeed") added the unwrap to train() only, with its patch hunk ending on the line def eval(self):.

The fix

Mirror the train() branch in eval(), comment included, so the two stay symmetric.

Tests

CPUOptimizerTester had no train()/eval() coverage at all, so this adds three CPU tests covering the DeepSpeed-style double-wrapped optimizer, the plain schedule-free optimizer, and an optimizer with no train/eval support (which should stay a no-op).

The double-wrap test is the regression test for this bug; it fails on main with the optimizer still reporting train after eval() and passes with the fix. The other two pass either way and are there so the unwrapped paths cannot regress. Reproduces on CPU with no GPU or DeepSpeed install needed.

$ pytest tests/test_optimizer.py -q
4 passed, 1 skipped

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[verma8076]

This looks correct and worth merging — it's the exact symmetric counterpart to the DeepSpeed wrapper reach-through that #3266 added to train(), using the same one-level self.optimizer.optimizer pattern. Tests cover the plain case, the wrapped case, and the no-mode-support case, so no regressions there.

Real-world impact isn't cosmetic: without this, optimizer.eval() silently no-ops under DeepSpeed, so a schedule-free optimizer's checkpoint keeps train_mode=True and accelerator.save_state() captures the wrong (train-view) fp32 master weights with no error raised (see the linked issue for a concrete repro of that failure mode).

It's been open over a month with no review and just got the stale-bot comment. @muellerzr @BenjaminBossan — worth a look before it gets auto-closed?

[vineethsaivs]

Still applies, and thank you @verma8076 for the read: AcceleratedOptimizer.eval() on current main still has only the direct branch, so the DeepSpeed wrapper case is unchanged.

if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval):
    self.optimizer.eval()

Re-verified today against main: with the fix reverted, test_accelerated_optimizer_train_eval_with_wrapped_optimizer fails (- eval / + train, the mode never leaves train) and the other three pass; with it, all four pass. The @require_cpu class is skipped on an accelerator-backed machine, so I forced torch_device == "cpu" locally to collect it.

I also gave the third test something to assert. It was the no-mode-support control and it only called train()/eval(), so it proved nothing beyond "did not raise". It now runs over the wrapper shape as well, which is the negative case for the branch this PR adds, and asserts the inner optimizer never grows a mode it did not have.

Not stale from my side. Happy to rebase whenever a maintainer has a moment.

[SunMarc]

A minor fix but worth merging i guess

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.