change saving logic

[felipemello1]

Context

What is the purpose of this PR? Is it to

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

Changelog

• Instead of saving model files to output_dir/base_model, we save them to output_dir/epoch_{}. That way, every epoch is ready for inference, without requiring an extra step. In the future, we can design this extra step better, e.g. 'tune prepare_for_inference args'
• Have a file size limit of 100MiB. All expected files are configs + tokenizer, so they shouldn't be larger than that.

Test plan

tree -a /tmp/torchtune/llama3_2_3B/full_single_device

testing with HF

resuming from ckpt

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 446e8cc with merge base cdf5ea2 ():
💚 Looks good so far! There are no failures yet. 💚

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

[ebsmothers]

Is there actually a way for someone to modify these values from the CLI/config? Seems like no, right?

[felipemello1]

no, but i dont see it being an issue. It shouldnt happen, and If users really want to come a 100MB file to every epoch, they can do it manually.

[ebsmothers]

Stamping to unblock. Please make sure to test that after the changes to the directory structure both resume from checkpoint and HF from_pretrained functionality are unaffected.