[dynamo] use proxies to nn.Module in dynamo generated GraphModules #120756
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/120756
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 7bb33b1 with merge base 61ff41f ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
…hModules" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
…hModules" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
…hModules" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
|
|
||
| # Super hacky - returns a proxy object for a nn.Module that references | ||
| # to the original nn.Module, but is a distinct object. | ||
| def nn_module_proxy(mod): |
There was a problem hiding this comment.
could you name this wrapper instead, due to the other uses of proxy?
torch/_dynamo/utils.py
Outdated
| return mod | ||
| try: | ||
| # best effort shallow copy | ||
| proxy = torch.nn.Module.__new__(torch.nn.Module) |
There was a problem hiding this comment.
Could you use a WeakRef here instead? and do some routing to the underlying weakref'd object?
Then we could use the finalizer to cleanup the graph/invalidate anything else.
There was a problem hiding this comment.
I attempted using weakrefs but discovered that a big difficulty in using weakrefs is that some downstream code (HOPs, export, backends) may require the actual module itself.
torch/_dynamo/utils.py
Outdated
| # best effort shallow copy | ||
| proxy = torch.nn.Module.__new__(torch.nn.Module) | ||
| proxy.__class__ = mod.__class__ | ||
| proxy.__dict__ = mod.__dict__ |
There was a problem hiding this comment.
Won't the parameters will still be held alive by this?
There was a problem hiding this comment.
Parameters are lifted as inputs to GraphModules, so AFAIK a GraphModule should never have a Parameter as an attribute.
The point of this nn.Module wrapper/proxy is so that when the user deletes all their references to the original module, dynamo cache invalidation will work because the compiled GraphModule doesn't hold a reference to the original module.
The parameters will be freed when cache invalidation runs.
…hModules" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
|
@jansel do you think it would be better if we explicitly write out which attributes of |
|
Yeah we could prune things down. |
…hModules" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
|
@anijain2305 and I talked offline - we can keep the |
|
@pytorchbot merge |
Merge startedYour 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 |
Summary: Fixes remaining refleaks found when debugging pytorch/pytorch#119607, tests added in pytorch/pytorch#120657. Also fixes some tests that xfail: pytorch/pytorch#120631 (not entirely sure why), but introduced tests now fail. X-link: pytorch/pytorch#120756 Approved by: https://github.com/jansel Reviewed By: huydhn Differential Revision: D55225127 Pulled By: williamwen42 fbshipit-source-id: 0d402deb4246d981304985761b94f637970fabce
…126332) * [dynamo] use proxies to nn.Module in dynamo generated GraphModules (#120756) Fixes remaining refleaks found when debugging #119607, tests added in #120657. Also fixes some tests that xfail: #120631 (not entirely sure why), but introduced tests now fail. Pull Request resolved: #120756 Approved by: https://github.com/jansel * [dynamo] use proxies to nn.Module in dynamo generated GraphModules (#120756) Fixes remaining refleaks found when debugging #119607, tests added in #120657. Also fixes some tests that xfail: #120631 (not entirely sure why), but introduced tests now fail. Pull Request resolved: #120756 Approved by: https://github.com/jansel
Stack from ghstack (oldest at bottom):
Fixes remaining refleaks found when debugging #119607, tests added in #120657.
Also fixes some tests that xfail: #120631 (not entirely sure why), but introduced tests now fail.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @aakhundov