Explicit check for step when loading the state

[muellerzr]

What does this PR do?

To not break prior releases, checks if step exists in our checkpoint to try and load it in.

Fixes #2923

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@BenjaminBossan

[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.

[BenjaminBossan]

For my understanding, if a user loads an older checkpoint, they won't have step, so override_attributes will just be an empty dict, right? Would this now fail?

https://github.com/huggingface/accelerate/pull/2765/files#diff-2d7515874eaecac2687c7fc1a9c720be53f802bf14b4c3dcebe14ad443d075dcR3114

[muellerzr]

Yes indeed, thanks

[BenjaminBossan]

Thanks for fixing this.

[alexanderswerdlow]

Just wanted to comment on this since I believe this PR has a bug [or can cause a bug/unexpected behavior in some cases]. I tried this exact fix on a local fork and found that this caused issue where the master rank had the step saved but other ranks didn't. I have a suspicion it is because I was gating my state_state to only occur on the main process which was fine for DDP, but not in this case, see my comment here. In my case, a gather worked to make sure all ranks agreed, but this took a few hours of debugging to figure out this was needed:

            starting_steps = gather_object([self.accelerator.step])
            print(f"Starting steps: {starting_steps}")
            if not all([x > 0 for x in starting_steps]):
                print(f"Not all ranks have >0 step, setting to: {starting_steps[0]}")
                self.accelerator.step = starting_steps[0]

[muellerzr]

Oh that's very interesting. Would you like to make a follow-up with this solution @alexanderswerdlow ? (Interesting they got out of sync like that). Otherwise I can do so shortly

[alexanderswerdlow]

@muellerzr Apologies, I don't have a ton of bandwidth atm, but happy to test it. It unexpected to say the least and took a sec to find because steps being out of sync doesn't always matter!

Not sure how problematic a silent internal gather is inside accelerate [or if this is common in accelerate], but I think a warning would also suffice if the step isn't found [since this gather behavior may not be the most robust in all cases].

[simonhessner]

Has this PR already been released in a new version? I am using 0.33.0 and I get KeyError: 'step'

[ReppiksProductions]

@muellerzr How exactly do PR work if I need this fix? Since this was a simple change in two scripts I changed them manually. Is there a way to automatically do this from the command line?

I'm on a Windows 11 system