Cudagraphs support for compiled optimizers

[mlazos]

This removes extra copies introduced by torch.compile when compiling the optimizer. These copies would tank perf and make cudagraphs unusable.

Marks all params/optimizer state as static addresses and a finalizer which cleans up the graph attributes when the optimizer goes out of scope.

**Note: this does not mark grads as static because this will increase memory usage significantly

There are two cases:

1. The upstream graph is cudagraphed - this case will work fine OOTB
2. The upstream graph is not cudagraphed - in this case, there will be a lot of copies introduced from the upstream (to copy the grads) into cudagraphed-owned memory, unless the user explicitly marks the grads as static. If the user does this, this will also require not deallocating the grads in zero_grad() (either the mod or optimizer version) by setting them to zero vs None. There is a PR (Throw error if setting static grads to None in zero_grad() #107853) in flight to throw an error if zero_grad attempts to set static grads to None.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @anijain2305

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

✅ No Failures

As of commit 283cb16 with merge base ad17e5e ():
💚 Looks good so far! There are no failures yet. 💚

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

[eellison]

cool! A few questions:

• should grads be marked as static? I dont know that they will be. if they're not, it will just lead to endless recompilation.
• The cleanup stuff is cool, i dont know if it could be more generic. Also, it wouldnt be sufficient to clean up cudagraph memory usage.. we'll need something where we invalidate a particular compilation within cudagraphs. i can submit that pr separately, not a blocker for this.
  -You might want someone to help with dynamo reviewing that is not me

[jansel]

I will let @eellison handle the review on this one

[mlazos]

@pytorchbot merge

[pytorch-bot]

This PR needs to be approved by an authorized maintainer before merge.

[mlazos]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here