Save checkpoint without optimizer states

[odedovadia]

Really awesome repo, thanks for the great work!

Currently, saving a checkpoint always includes the full optimizer state, which can be quite expensive storage-wise and is often unnecessary.

Do you plan on adding a save_optimizer flag to make this optional?

If this feature aligns with your roadmap I’d be happy to open a quick PR.

Thanks!

[guyueh1]

@odedovadia some clarification questions

1. are you experimenting with the DTensor path or Mcore path, or where do you want this feature to land?
2. what's the expected behavior when resuming, if your checkpoint doesn't contain a optimizer states, how does the resumed experiment initialize optimizer states? Do you expect it to be random initialized?

Skipping the saving can be done easily with small amount of code, but we need to define the behavior of resuming in this mode. if that's resolved, I am happy to help shepherd your PR

[odedovadia]

Thanks @guyueh1

1. I use the Mcore path, but that shouldn't matter since the change is at the algorithm level. All your policy workers already support saving without the optimizer.
2. Resuming with randomly initialized states makes sense IMO (it's also the default behavior in HuggingFace). We could add a warning when when resuming from a checkpoint without an optimizer.

Opened a PR based on this: #1843

[guyueh1]

Left comment in PR