[nn-module] Use standard dict for _parameters, _modules and _buffers #129164
Conversation
[ghstack-poisoned]
anijain2305
requested review from
albanD,
jbschlosser and
mikaylagawarecki
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/129164
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit ae8e301 with merge base f2f4dde ( FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
ghstack-source-id: 410b138b7d88377f36a116fdb914b04445a19805 Pull Request resolved: #129164
…d _buffers" [ghstack-poisoned]
ghstack-source-id: 62efbe2258cefd704b8a7eb7b9c82b292994013e Pull Request resolved: #129164
|
@albanD I modified the description to answer the questions to best of my knowledge. |
|
@albanD A gentle ping. Let me know if you have some experiments in mind to check perf impact. Or any other suggestions to move on this PR. |
There was a problem hiding this comment.
From discussion with @albanD, we are probably ok with this one
The only slight concern might be BC for an nn.Module that was torch.save-ed, but I think that should not be an issue
|
@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 |
|
@pytorchbot revert -m 'Sorry for reverting your change but it seems to break some internal dper3 tests' -c ghfirst From the imported diff D59091876, the error is |
|
@pytorchbot successfully started a revert job. Check the current status here. |
…buffers (#129164)" This reverts commit f2840bb. Reverted #129164 on behalf of https://github.com/huydhn due to Sorry for reverting your change but it seems to break some internal dper3 tests ([comment](#129164 (comment)))
|
@anijain2305 your PR has been successfully reverted. |
|
It was a silly impl in internal - fixed in https://www.internalfb.com/diff/D59166026 |
|
@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 |
Stack from ghstack (oldest at bottom):
TorchDynamo guard mechanism guards on the key order on the dictionaries if the user iterates over the dictionary. For standard dict, we can write a fast C++ implementation by using PyDict_Next. But with OrderedDict, we have to rely on
keysPython API to get the key ordering. This makes guard evaluation slow.With Dynamo inlining into inbuilt nn modules, I am seeing many guards over the OrderedDict on
_modules,_parameters. From reading the code, I don't see any reason to not use standard dicts. I think OrderedDict was preferred over dict because of the ordering, but dicts are now ordered. With this PR, I am observing ~20% reduction in guard overhead of a HF model.Functionality impact
move_to_endmethod for OrderedDict (link). But the changes here are internal to nn module, and we do not usemove_to_endfor_parameters,_modulesand_buffers. We usemove_to_endfor hooks but this PR keeps the OrderedDict for hooks untouched (we should still followup with hooks but in a separate PR).Perf impact
dictis completely implemented in C. OrderedDict is Python wrapper over dict with only few method overridden (link).Typing impact
_modulesetc. We have iterators likenamed_parameterswhich return an Iterator of Parameter. So, no typing changes required.