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) [[email protected] ~/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 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• VolumeLimitExceeded Issue for linux.2xlarge and linux.4xlarge

❌ 2 New Failures, 2 Unrelated Failures

As of commit 650d91d with merge base 3d73591 ():

NEW FAILURES - The following jobs have failed:

• GPU tests / gpu_test (3.11, stable) (gh)
  tests/recipes/test_qat_lora_finetune_distributed.py::TestQATLoRAFinetuneDistributedRecipe::test_training_state_on_resume_with_async_checkpointing[llama3/8B_qat_lora-llama3-tune-False]
• Lint / lint (3.10) (gh)
  Process completed with exit code 1.

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• GPU tests / gpu_test (3.10, stable) (gh) (trunk failure)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.9, stable) (gh) (trunk failure)
  ##[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?..... 🥺

[felipemello1]

var "step" is missing in distributed recipe

[bogdansalyp]

rebased to #2869