consistent checkpoint filenaming

[stas00]

Currently we get this for the checkpoint filenames:

zero_pp_rank_0_mp_rank_00_model_states.pt
zero_pp_rank_0_mp_rank_00optim_states.pt
zero_pp_rank_1_mp_rank_00_model_states.pt
zero_pp_rank_1_mp_rank_00optim_states.pt

this PR fixes optim to be consistent with model, leading to:

zero_pp_rank_0_mp_rank_00_model_states.pt
zero_pp_rank_0_mp_rank_00_optim_states.pt
zero_pp_rank_1_mp_rank_00_model_states.pt
zero_pp_rank_1_mp_rank_00_optim_states.pt

Thanks.

[jeffra]

Thanks @stas00, I agree this naming looks a lot better to me. I do have one concern though, if someone checkpointed a model previously with the old naming scheme I don't think they'd be able to load it after this PR :(

[stas00]

Ah, makes sense. I didn't think of that!

How do you do a deprecation cycle and I will code a transition to the new way.

edit: well actually there is nothing to deprecate here other than perhaps telling the user how to rename the old checkpoint - it's not user's responsibility. I will just code it to handle both and say, in a year the old one can be removed if desired ;)

If I'm not mistaken _get_all_zero_checkpoints has a flaw, that it checks if the checkpoint exists and if it doesn't it warns but still tries to load it, so I can fix it while I'm at it.

edit: I missed it has an early return - it's coded perfectly

[stas00]

how about this backward compatible rename (last commit)?

I'm yet to test it.

I'm just writing a test that verifies that all the checkpoint files are in place and it's always easier to compose the multiple similar filenames on the fly than hardcode them.

I would totally understand if you won't want to clutter the code and would rather leave the old filename, since most people will never see it. I will adapt the test to use the current naming then. It's no problem.

[stas00]

seems to work well - I saved a checkpoint with the new way, renamed it to the old way and then was able to load it with this PR.