New OptimizerInBackward

[joecummings]

OptimizerInBackward

Motivation

Previously, we utilized a process very similar to the one outlined in this blog post. Although it worked, it involved lots of if/else switching in the recipe based on what composed well with optimizer fused in backwards and what did not.

Goal

Simplify our recipes (729 LOC -> 677 LOC)! This PR provides a canonical OptimizerInBackward class that can be used as a drop-in replacement for any PyTorch optimizer. It integrates into the full_finetune_single_device.py recipe and adds tests + updates documentation.

Testing

To-do

• Integrate into full_finetune_distributed.py
• Integrate into rest of recipes
• Deprecate the utility functions used to create the optim in backward hooks previously

[pytorch-bot]

🔗 Helpful Links

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

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

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

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 74.10072% with 36 lines in your changes missing coverage. Please review.

Project coverage is 60.12%. Comparing base (c8e670b) to head (5674f9a).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
recipes/full_finetune_single_device.py	0.00%	35 Missing ⚠️
torchtune/modules/optim.py	97.05%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2719      +/-   ##
==========================================
- Coverage   60.64%   60.12%   -0.52%     
==========================================
  Files         428      432       +4     
  Lines       26091    26520     +429     
==========================================
+ Hits        15823    15946     +123     
- Misses      10268    10574     +306     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[SalmanMohammadi]

I know this was already discussed but looking at it here it's confusing that we have a class that's meant to be a drop in optimizer, but if you actually want to use it in the recipe we aren't using torchtune's regular instantiation-based configuation for this component, but rely on logic inside the recipe itself to construct the optimizer.

I don't have a neat solution that doesn't involve breaking all the configs unfortunately. Maybe we don't advertise this as a separate modular component but as a core recipe feature? Power users can bear the brunt of implementing it in their own recipes

[joecummings]

Fair， I'll work on the wording here.

[SalmanMohammadi]

Worth raising an issue to follow up on this?

[joecummings]

Yeah fair, I might just opt to raise a meta-Issue tracking how we integrate this into the rest of our recipes and any cleanup that needs to happen to deprecate the old material.

[SalmanMohammadi]

Really clean work

[krammnic]

Amazing work! (I've attempted this, but my RFC was rejected, so I'm really glad to see this merged.)