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.
• 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.)

Evidence of correct number of checkpoints being saved

(joe-torchtune) [jrcummings@devvm4767.pnb0 ~/projects/joe-torchtune (impl-step-based-ckpt)]$ ls /tmp/torchtune/llama3_2_1B/full_single_device/
step_100  step_125  step_150  step_175  step_200  step_25  step_50  step_75  torchtune_config.yaml

Evidence of correct resuming from ckpt mid-epoch

Evidence of correct resuming from ckpt at epoch boundary

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.

❌ 1 New Failure, 2 Cancelled Jobs

As of commit 755a7ac with merge base fa92c96 ():

NEW FAILURE - The following job has failed:

• GPU tests / gpu_test (3.11, stable) (gh)
  tests/recipes/test_qat_lora_finetune_distributed.py::TestQATLoRAFinetuneDistributedRecipe::test_save_and_load_merged_weights[llama3/8B_qat_lora-llama3-tune]

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

• GPU tests / gpu_test (3.10, stable) (gh)
  tests/recipes/test_qat_lora_finetune_distributed.py::TestQATLoRAFinetuneDistributedRecipe::test_save_and_load_merged_weights[llama3/8B_qat_lora-llama3-tune]
• GPU tests / gpu_test (3.9, stable) (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?..... 🥺