Implement step based checkpointing

[joecummings]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

Closes #2105. This is a widely requested feature that allows users to have greater control over checkpointing frequency in torchtune.

TODO: Add commentary on design decisions. Acknowledge spaghetti code. Beg forgiveness.

Changelog

• Update FullModelHFCheckpointer to accept a step parameter when saving a checkpoint. Use that step to designate the checkpoint folder name. Keep epoch_{} as a fall-back for BC.
• Add helper function prune_surplus_checkpoints that removes extra checkpoints in order to free up space for the user. This is controlled by the keep_n_last_checkpoints.
• Add helper function get_all_checkpoints_in_dir that scans a directory (meant for the output_dir) and returns all checkpoint directories.
• Add tests for corresponding helper functions
• Modify the full_finetune_single_device.py recipe to utilize step-based checkpointing.
• Add tests for `full_finetune_single_device.py`` recipe w/ step-based checkpointing.

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
• include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

Ran single device recipe w/ step-based checkpointing
Evidence of succesful run

Evidence of correct number of checkpoints being saved

Evidence of correct resuming from ckpt

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

• I did not change any public API
• I have added an example to docs or docstrings

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2384

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 6 New Failures, 5 Cancelled Jobs

As of commit 4d11a48 with merge base 2128559 ():

NEW FAILURES - The following jobs have failed:

• Build Docs / build_docs (3.11) (gh)
  Process completed with exit code 2.
• GPU tests / gpu_test (3.11, stable) (gh)
  Process completed with exit code 2.
• Lint / lint (3.10) (gh)
  recipes/full_finetune_single_device.py:10:1: F401 'typing.Tuple' imported but unused
• Recipe Tests / recipe_test (3.10) (gh)
  Process completed with exit code 2.
• Recipe Tests / recipe_test (3.11) (gh)
  Process completed with exit code 2.
• Unit Test / unit_tests (3.10) (gh)
  Process completed with exit code 2.

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• GPU tests / gpu_test (3.10, stable) (gh)
  Process completed with exit code 2.
• GPU tests / gpu_test (3.9, stable) (gh)
  ##[error]The operation was canceled.
• Recipe Tests / recipe_test (3.9) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.11) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.9) (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[bogdansalyp]

• recipe_state is still saved to ${output_dir}, not ${output_dir}/step_XXX
• resume_from_checkpoint logic should be updated
  • RN it looks for ${output_dir} checkpoint, not step_XXX
  • maybe replace top level cfg.resume_from_checkpoint to have cfg.checkpointer.resume_from which is either "latest" (default) or the path to the checkpoint to resume from. Or separate mutually exclusive resume_from: /path/ and resume_from_latest: True
  • offtopic but cfg.resume_from_checkpoint is mentioned in code as deprecated and replaced by should_load_recipe_state but de facto resume_from_checkpoint is mandatory and should_load_recipe_state doesn't work
• recipe_state has proper step and epoch to continue from but the train cycle still starts from 0 -> logs start from 0 & checkpointing start from 0
• lr schedulers aren't synced with the resume step
• maybe save the wandb run?..... 🥺