Reduce overhead in CUDAGraph Trees

[eellison]

Stack from ghstack (oldest at bottom):

• Flip Switch Redux #98341
• -> Reduce overhead in CUDAGraph Trees #98529
• Better Handling of Storage Cache #98254
• Make sure we dealloc on recording, not just replay #97440
• Account for forwards which whose corresponding backwards are not invoked #98112

Significantly reduces overhead of constructing Tensors and Storages and checking Storage Liveness. Removes the regression for HF models that I tested and removes 75% of overhead of the extremely overhead bound resnet50 training we have in torchbench. (.91x base commit, 1.02x torchinductor default, 1.16x this PR, 1.25 previous cudagraphs impl).

This PR takes care of all of the lower hanging fruit.

• Computes storage aliasing at record time instead of during at runtime. We no longer need to use a runtime storage cache, and can instead index directly into the existing alias if there is one, or construct a new Storage

• Moves the heavyweight C++ calls into a batch - getting storage weakrefs and constructing tensors

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @desertfire

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 7a9b0a6:
💚 Looks good so far! There are no failures yet. 💚

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

[ngimel]

It's still slower than previous cudagraphs implementation was? Do you have an idea why?
Can we switch between tapes and previous implementation depending on how many graphbreaks there are?

[eellison]

@ngimel it's not really slower on any real model, (same speed for huggingface, etc) but resnet in torchbench the graph.replay() is so quick that just reconstructing the tensors and storages is a noticeable overhead. We also do not have a well tuned batch size for that model.

It's not really about graph breaks - even just with no graph breaks, there is still a substantial memory overhead in the existing implementation because you are not sharing memory between forward and backward.

For inference, there are so few tensors output compared to training that there should be practically no difference. A further training optimization would be to reuse the same TensorImpls for saved tensors for the backward since we know they will be live if the backward hasn't been run, and do not escape the backward.

[eellison]

@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