[simpleqt]

What does this PR do?

GradientAccumulationPlugin.adjust_scheduler is documented to default to True ("Whether to adjust the scheduler steps to account for the number of steps being accumulated. Should be True if the used scheduler was not adjusted for gradient accumulation."), but the setting has silently never taken effect:

Cause. KwargsHandler.to_kwargs() only returns the attributes whose values differ from the dataclass defaults. Since adjust_scheduler's default is True, both the default and any explicit adjust_scheduler=True are stripped when GradientState stores plugin_kwargs. GradientState.adjust_scheduler then falls back to False:

@property
def adjust_scheduler(self) -> bool:
    return self.plugin_kwargs.get("adjust_scheduler", False)  # always taken

So AcceleratedScheduler.step never bumped self.scheduler._step_count during gradient-accumulation micro-steps, i.e. the adjustment introduced in #1187 ("Make the Scheduler adjust the steps taken relative to the gradient accumulation steps") has been dead code for every user since March 2023, and schedulers sized on the total number of batches (e.g. OneCycleLR created without dividing total_steps) are not accounted for. An explicit False is honored (it differs from the default, so to_kwargs() keeps it) — the setting only ever worked in one direction.

Fix. Minimal, three small changes:

1. GradientState.adjust_scheduler falls back to True, matching the declared plugin default. An explicit adjust_scheduler=False is still honored.
2. The _step_count bump in AcceleratedScheduler.step is guarded with hasattr(self.scheduler, "_step_count"), because schedulers without a step counter (e.g. ReduceLROnPlateau) would otherwise raise AttributeError now that the default is active.
3. The OneCycleLR-style guard compares against self.scheduler.last_epoch (actual steps taken) instead of _step_count. Since _step_count now also advances during micro-steps, _step_count <= total_steps would stop such schedulers early; last_epoch < total_steps steps them exactly total_steps times, which is identical to the behavior on main (where the guard 1 + steps_taken <= total_steps also yielded exactly total_steps steps). Schedulers without total_steps are untouched.

Test. New SchedulerAdjustmentTester in tests/test_scheduler.py (CPU only, single process):

• test_adjust_scheduler_reaches_gradient_state: the default gradient_accumulation_steps path and an explicit adjust_scheduler=True yield gradient_state.adjust_scheduler is True; an explicit False yields False. Fails on main (returns False).
• test_accumulation_accounts_for_skipped_scheduler_steps: with gradient_accumulation_steps=2, the wrapped scheduler's _step_count advances on non-sync micro-steps ([2, 3, 4, 5] after 4 micro-steps) while last_epoch shows only the 2 real updates. Fails on main ([1, 2, 2, 3]).
• test_accumulation_with_plateau_scheduler: ReduceLROnPlateau + accumulation keeps working (one LR update per optimizer step) — guards fix point 2.
• test_accumulation_with_one_cycle_scheduler: OneCycleLR(total_steps=4) under accumulation is still stepped exactly 4 times without raising — guards fix point 3.

Results:

commit	SchedulerAdjustmentTester
clean main (3008af9)	2 failed, 2 passed
this PR	4 passed

Run with:

python -m pytest tests/test_scheduler.py::SchedulerAdjustmentTester -v

Note for reviewers: a previous attempt at this fix (#4262) was closed without review in September. This PR takes the same core diagnosis (the False fallback contradicting the plugin default) with the two companion guards that keep the restored default safe for schedulers lacking _step_count and for OneCycleLR-style total_steps bounds. Users who depended on the old effective default can keep the previous behavior explicitly with GradientAccumulationPlugin(num_steps=..., adjust_scheduler=False).

Before submitting

• This PR fixes a bug or is a tiny typo (check)
• Did you read the contributor guidelines?
• Was this discussed/approved via a GitHub issue? (not a duplicate of an open PR)
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review? @SunMarc

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[Michael-RDev]

No blocking issues found in the reviewed paths. The fallback matches the plugin’s declared default while also preserving False. All four new tests passed, along with existing distributed scheduler checks and accumulation factors of 1, 2, and 4 👍

[JonnyFaeth]

Added a checkpoint-continuation test at d8ca02295630ffffb3b87cf0beff1880614d3b35 to complement the accumulation coverage: save optimizer/scheduler state at an accumulation boundary, recreate the Accelerator and wrappers, restore, and compare the remaining OneCycleLR sequence with uninterrupted execution. Both explicit adjustment settings are covered, including last_epoch and bookkeeping counters.

Windows/Python 3.12, PyTorch 2.11 CPU: 2 new cases passed; 4 existing SchedulerAdjustmentTester cases passed / 9 deselected. Reverting scheduler.py and state.py to their parent gives 1 failed, 1 passed (the adjustment-enabled bookkeeping case fails); both files restored. Ruff checked. Uses in-memory state_dict copies; no filesystem Accelerator.save_state, multi-process or GPU claim.

AI disclosure: prepared and tested with Codex assistance.

Checkpoint-continuation regression patch

diff --git a/tests/test_scheduler_checkpoint_adjustment.py b/tests/test_scheduler_checkpoint_adjustment.py
new file mode 100644
index 00000000..41474cca
--- /dev/null
+++ b/tests/test_scheduler_checkpoint_adjustment.py
@@ -0,0 +1,64 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+import pytest
+import torch
+
+from accelerate import Accelerator
+from accelerate.state import AcceleratorState, GradientState
+from accelerate.utils import GradientAccumulationPlugin
+
+
+@pytest.mark.parametrize("adjust", [False, True])
+def test_one_cycle_accumulation_checkpoint_roundtrip(adjust):
+    def make():
+        GradientState._reset_state()
+        AcceleratorState._reset_state(True)
+        accelerator = Accelerator(
+            cpu=True, gradient_accumulation_plugin=GradientAccumulationPlugin(num_steps=2, adjust_scheduler=adjust)
+        )
+        model = torch.nn.Linear(2, 1)
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.1, total_steps=4)
+        return accelerator, *accelerator.prepare(model, optimizer, scheduler)
+
+    def advance(parts, count):
+        accelerator, model, optimizer, scheduler = parts
+        rates = []
+        for _ in range(count):
+            with accelerator.accumulate(model):
+                accelerator.backward(model(torch.ones(1, 2)).sum())
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                rates.append(scheduler.get_last_lr())
+        return rates
+
+    try:
+        parts = make()
+        advance(parts, 4)
+        optimizer_state = copy.deepcopy(parts[2].state_dict())
+        scheduler_state = copy.deepcopy(parts[3].state_dict())
+        expected = advance(parts, 4)
+        resumed = make()
+        resumed[3].load_state_dict(scheduler_state)
+        resumed[2].load_state_dict(optimizer_state)
+        assert advance(resumed, 4) == expected
+        assert resumed[3].scheduler.last_epoch == 4
+        assert resumed[3].scheduler._step_count == (9 if adjust else 5)
+    finally:
+        GradientState._reset_state()
+        AcceleratorState._reset_state(True)

[simpleqt]

Applied as tests/test_scheduler_checkpoint_adjustment.py and pushed in the single commit (457926e1), with the copyright header kept.

I re-ran your claims on Windows / Python 3.11 / torch 2.15.0+cpu:

• with the branch as it was: your 2 new cases pass, and the 4 existing SchedulerAdjustmentTester cases pass (6 passed);
• with scheduler.py and state.py reverted to their parent (tests kept): 1 failed, 1 passed, the failure being the adjustment-enabled _step_count assertion, exactly as you described.

ruff check and ruff format --check pass on the new file. The commit body now names your test so the contribution is visible in the history.