Adding an option to only retain optimizer states for the last N checkpoints
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days. This issue was closed because it has been inactive for 7 days since being marked as stale. PR is merged: https://github.com/NVIDIA/NeMo/pull/10098
**Is your feature request related to a problem? Please describe.**

Currently when training LLMs, the checkpoints are saved with all of the optimizer states, which generally increases the checkpoint size about 6 times. When the LLM itself is very large (e.g. > 30B parameters) and we are running many training jobs at the same time, this requires having a very large storage on the cluster. In case of SFT experiments where we almost never need to restart training from the middle of the job, saving optimizer states for intermediate checkpoints is really wasteful and it would be great to have some control over this.

**Describe the solution you'd like**

Have a config parameter "save_optimizer_states_for_last_checkpoints" (probably a less verbose name is needed). I'd default it to 2 since we always need a last checkpoint state in case job fails or is preempted and it's good to keep another one in case the last state is corrupted for some reason. Then we can set it to -1 to keep the current behavior of saving all optimizer states, set it to 0 in case we don't want to save any, or anything in between. 

The way I'd implement it is by modifying PTL/megatron logic in checkpoint saving and whenever we save a new checkpoint, we would then simply remove the optimizer states from one of the old checkpoints based on the config parameter. 

**Describe alternatives you've considered**

A current alternative that we have to use is basically to manually remove the optimizer states when we run out of space, which is of course very inconvenient. Another potential way to solve this is by having a script that runs alongside the training jobs and monitors the training folder and whenever a new checkpoint is saved, it does the cleaning. This way we don't have to modify nemo code, but it does not feel like a good solution.

**Additional context**

n/a