# Load\_state() with custom objects and scheduler

**URL:** <https://discuss.huggingface.co/t/load-state-with-custom-objects-and-scheduler/169759>\
**Category:** 🤗Accelerate\
**Created:** [October 31, 2025, 3:12pm UTC](https://discuss.huggingface.co/t/load-state-with-custom-objects-and-scheduler/169759 "2025-10-31T15:12:14Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![rogercc](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/rogercc/32/39761_2.png) [@rogercc](https://discuss.huggingface.co/u/rogercc)\
**Post date:** [October 31, 2025, 3:12pm UTC](https://discuss.huggingface.co/t/load-state-with-custom-objects-and-scheduler/169759/1 "2025-10-31T15:12:14Z")

</div>

Hi!

I am having some issues understanding why I can’t resume training corectly with this code:

```auto
optimizer = optim.AdamW(params, lr=args.learning_rate)

scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(args.num_iterations * args.lr_warmup_fraction), num_training_steps=args.num_iterations)

# custom class which only saves 2 ints
# and implements load_state_dict() and state_dict()
training_state = TrainingStateTracker(iteration=1, global_step=0)
accelerator.register_for_checkpointing(training_state)

agent, optimizer, scheduler = accelerator.prepare(agent, optimizer, scheduler)

if args.checkpoint_dir != "":
    print(f"Loading checkpoint from {args.checkpoint_dir}")
    accelerator.load_state(args.checkpoint_dir)
    print(f"Checkpoint loaded from {args.checkpoint_dir}")

accelerator.wait_for_everyone()

# these will be either the initialized values or the loaded values from the checkpoint
start_iteration = training_state["iteration"]
global_step = training_state["global_step"]

print(f"Starting training from iteration {start_iteration} and global step {global_step}")

```

1. I see `scheduler.bin` being stores when I call `.save_state()` so I don’t have to explicitly register the shceduler for checkpointing , right? Still, when loading state, the first LR that I get when calling `.step()` is not the continuation of the last one before it was stored….

2. the last print statement: `print(f"Starting training from iteration {start_iteration} and global step {global_step}")` always prints starting from iteration 1 and global\_step 0, so even though I see the `custom_checkpoint_0.pkl` being saved (since I registered it for checkpoint), the 2 fields are not loaded correctly…

---

<div class="post-metadata">

**Author:** ![system](https://us1.discourse-cdn.com/hellohellohello/original/2X/d/de4155eb4aa4108ecb32a1389d7cc37ae69f88b7.png) [@system](https://discuss.huggingface.co/u/system)\
**Post date:** [October 31, 2025, 3:12pm UTC](https://discuss.huggingface.co/t/load-state-with-custom-objects-and-scheduler/169759/2 "2025-10-31T15:12:15Z")

</div>



---

<div class="post-metadata">

**Author:** ![system](https://us1.discourse-cdn.com/hellohellohello/original/2X/d/de4155eb4aa4108ecb32a1389d7cc37ae69f88b7.png) [@system](https://discuss.huggingface.co/u/system)\
**Post date:** [October 31, 2025, 11:49pm UTC](https://discuss.huggingface.co/t/load-state-with-custom-objects-and-scheduler/169759/3 "2025-10-31T23:49:44Z")

</div>



---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [November 1, 2025, 12:34am UTC](https://discuss.huggingface.co/t/load-state-with-custom-objects-and-scheduler/169759/4 "2025-11-01T00:34:39Z")

</div>

This is complicated…😅

* * *

Short answer:

- You don’t need to register the scheduler if you passed it to `accelerator.prepare(...)`. Accelerate will checkpoint it as `scheduler` (file `scheduler*`), together with model, optimizer, RNG, scaler. ([huggingface.co](https://huggingface.co/docs/accelerate/en/usage_guides/checkpoint "Checkpointing"))

- Your LR jump and your two integers not restoring come from two distinct causes:

Below is the exact, minimal fix.

# Fix the LR scheduler resume

1. Choose one stepping method and keep it the same across save and load:

2. Rebuild the scheduler with the **same hyperparameters** as when you saved it. `get_linear_schedule_with_warmup` bakes `num_warmup_steps` and `num_training_steps` into the schedule. If those differ, the resumed LR curve will diverge even when the state loads. ([huggingface.co](https://huggingface.co/docs/transformers/en/main_classes/optimizer_schedules "Optimization"))

3. Load after `prepare`. The documented sequence is: build objects → `prepare(...)` → `load_state(...)`. ([huggingface.co](https://huggingface.co/docs/accelerate/en/usage_guides/checkpoint "Checkpointing"))

4. Sanity check right after load:

```python
# do NOT step here; just inspect
print("resumed_lr", scheduler.get_last_lr())
# ... after a real optimizer step ...
optimizer.step()
# if and only if you set step_scheduler_with_optimizer=False:
# scheduler.step()
print("post_step_lr", scheduler.get_last_lr())

```

This verifies the first LR is continuous and that you’re not double-stepping. The Accelerate scheduler wrapper is designed to step with the optimizer only when a real step happened. ([huggingface.co](https://huggingface.co/docs/accelerate/package_reference/torch_wrappers "DataLoaders, Optimizers, and Schedulers"))

# Fix the custom `TrainingStateTracker` not restoring

The file `custom_checkpoint_0.pkl` can exist and still be skipped if the **number of registered custom objects** at load time doesn’t match the number of `custom_checkpoint_*.pkl*` files in the folder. Accelerate logs a warning and skips loading the customs, which leaves your fields at constructor defaults. Clean any extras (e.g., `.orig`) and register the same objects in the same order before `load_state`. ([GitHub](https://github.com/huggingface/accelerate/issues/1563 "LR resume is somewhat buggy · Issue #1563 · huggingface/accelerate · GitHub"))

Concrete checklist:

- Register every run, **before** `load_state`:

- Ensure the checkpoint dir has exactly one matching file (for your single custom object): `custom_checkpoint_0.pkl` and nothing like `custom_checkpoint_0.pkl.orig`. If counts differ, Accelerate prints: “Found checkpoints: N / Registered objects: M — Skipping.” ([GitHub](https://github.com/huggingface/accelerate/issues/1563 "LR resume is somewhat buggy · Issue #1563 · huggingface/accelerate · GitHub"))

# Robust resume template (drop-in)

```python
from accelerate import Accelerator
from transformers import get_linear_schedule_with_warmup

# 1) keep this consistent across runs
accelerator = Accelerator(step_scheduler_with_optimizer=True)

# 2) rebuild exactly as before
optimizer = torch.optim.AdamW(params, lr=args.learning_rate)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=int(args.num_iterations * args.lr_warmup_fraction),
    num_training_steps=args.num_iterations,
)

# 3) register customs BEFORE load_state
training_state = TrainingStateTracker(iteration=1, global_step=0)
accelerator.register_for_checkpointing(training_state)

# 4) prepare, then load
agent, optimizer, scheduler = accelerator.prepare(agent, optimizer, scheduler)
if args.checkpoint_dir:
    accelerator.load_state(args.checkpoint_dir)

# 5) read restored counters
start_iteration = training_state["iteration"]
global_step = training_state["global_step"]

```

Why this works:

- `save_state/load_state` handle model, optimizer, RNG, scaler, and the scheduler automatically when it went through `prepare(...)`. The constants include `SCHEDULER_NAME="scheduler"` (your `scheduler.bin`). ([huggingface.co](https://huggingface.co/docs/accelerate/en/usage_guides/checkpoint "Checkpointing"))
- Custom objects load only if the registered count/order equals the found files. ([GitHub](https://github.com/huggingface/accelerate/issues/1563 "LR resume is somewhat buggy · Issue #1563 · huggingface/accelerate · GitHub"))

# Edge conditions that also break LR continuity

- Manually calling `scheduler.step()` **and** leaving `step_scheduler_with_optimizer=True` → LR steps twice per update on multi-GPU. Users reported this exact drift; the fix is to pick one stepping path. ([GitHub](https://github.com/huggingface/accelerate/issues/2142 "lr\_scheduler.step have undesired result when used on ..."))
- Changing DeepSpeed config to include a scheduler while also passing a real scheduler in code is rejected; you must not specify both. Accelerate enforces this. ([gemfury.com](https://gemfury.com/emaballarin/python%3Aaccelerate/accelerate-1.12.0.dev0-py3-none-any.whl/content/accelerator.py "emaballarin / accelerate v1.12.0.dev0 - python package"))

# Similar cases for reference

- Accelerate skips custom object restore if checkpoint counts don’t match; log shows “Found checkpoints … Registered objects … Skipping.” Resolution: clean extra `custom_checkpoint_*.pkl*` and re-register before load. ([GitHub](https://github.com/huggingface/accelerate/issues/1563 "LR resume is somewhat buggy · Issue #1563 · huggingface/accelerate · GitHub"))
- Re current stepping behavior and defaults: `step_scheduler_with_optimizer` is `True` by default; set it to `False` if you want manual stepping. ([huggingface.co](https://huggingface.co/docs/accelerate/en/package_reference/accelerator "Accelerator"))
- PyTorch call order rule: call `optimizer.step()` before `scheduler.step()` to avoid skipping the first schedule value. ([PyTorch Forums](https://discuss.pytorch.org/t/correct-way-to-use-order-scheduler-in-pytorch-1-6/97229 "Correct way to use(order) scheduler in pytorch 1.6"))

Bottom line:

- Keep `step_scheduler_with_optimizer` and your call pattern aligned.
- Recreate the scheduler with the same `num_warmup_steps` and `num_training_steps`.
- Register the custom tracker every run before `load_state` and ensure the checkpoint dir has a **single** matching `custom_checkpoint_0.pkl`.
- Always `prepare(...)` then `load_state(...)`. ([huggingface.co](https://huggingface.co/docs/accelerate/en/usage_guides/checkpoint "Checkpointing"))