Log memory usage during checkpoint save

[rohan-varma]

Context

• To help debug peak memory usage during checkpoint save, adding logs at the end of save_checkpoint.

Changelog

• ...

Test plan

• ....

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 0b4fbe1 with merge base c9d1cdc ():
💚 Looks good so far! There are no failures yet. 💚

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

[kartikayk]

This seems mostly for debugging right? Why does the user care if this works? I think we should think a bit about logs that are for debugging (mainly for us) and for actual users. Even iof the peak memory is high, I think the % of users who'll understand this is because of the _load_state_dict call is really small. I think we should think about the intent of these logs a bit and see how to help the user on this journey. Going to request changes for now and lets discuss.