Make memory logging great again #817
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/817
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit ddc0197 with merge base 41341fd ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@kartikayk Don't we want the recipe to control all logging? there could be use cases where you just want the peak memory stats without logging it. I purposely did not add it for that reason, otherwise I agree |
|
We already have logging done in many of our utils (eg: seed.py), so I think its ok? Plus this function sets expectation that it'll log. In a previous version of this, I was actually using this same function in setup_model as well to log stats after model initi. So it would be good to just do it in the function otherwise we'd just be copying all of this many time? |
|
Sorry missed the convo on here before stamping and just looked at the code. I'm fine with moving this to a util in memory.py, the one thing I would ask is that we keep the current version of |
What does this mean? :) |
This is dependent on the changes in this pytorch stack: pytorch/pytorch#146217 Add support for running `ZBVZeroBubbleSchedule` and v-shaped CSV schedules in torchtitan Fixes pytorch/torchtitan#774 --------- Co-authored-by: tianyu-l <150487191+tianyu-l@users.noreply.github.com>
Better format the memory stats log output
Ran recipes to confirm logs