Is it possible to get save_state to save step as well?

[tshmak]

Is it possible to get save_state to save accelerator.step as well? Or is there another way to save and load the current iteration/step?

[muellerzr]

Sure, that seems reasonable certainly :) Would you like to make a PR doing so? If not let me know and I can look into that this week!

[tshmak]

I'm not confident in doing it myself. Thanks for looking into it!

[bipinKrishnan]

Hey @muellerzr , are you planning to work on this? If not, may be I could take a look into this if you could provide some context into the feature :)

[muellerzr]

Hey @bipinKrishnan! Sure. The saving logic resides in src/accelerate/checkpointing.py. We would essentially need to add the step in IMO part of the states file for now. (As it definitely is part of RNG related resuming).

On restoring it should just set Accelerator.step.

This will likely need some modifications to get step working (such as having load_accelerator_state return the items we want to override in the Accelerator itself as a dict perhaps) and then loading that part in after the call to load_accelerator_state during Accelerator.load_state

[bipinKrishnan]

Hi @muellerzr , just a quick reminder on the PR created to close this issue. Let me know if any further changes are needed :)