-
Notifications
You must be signed in to change notification settings - Fork 25.5k
[fx graph cache] Refactor FxGraphCachePickler #138682
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Summary: In an upcoming change, we need to modify how FxGraphCachePickler pickles tensor constants. Specifically, we'll pickle differently depending on whether parameters have been frozen. To do that, it will be convenient to change FxGraphCachePickler into a regular object instead of a collection of classmethods. Test Plan: unit tests [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138682
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit dbb15b8 with merge base 8aedc64 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pickler.dumps(details3), | ||
) | ||
|
||
def test_bypass_unsupported(self): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Noticed this code path wasn't being tested.
# Sometimes inductor configs are unpickleable and can fail | ||
raise BypassAOTAutogradCache from e | ||
|
||
def debug_lines(self) -> List[str]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Couldn't figure out why we need this since we can just do AOTAutogradCachePickler.debug_lines(details)
. Same in codecache.py
dispatch_table[torch.Tensor] = _reduce_tensor | ||
def __init__(self): | ||
super().__init__() | ||
self.dispatch_table: Dict |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
So mypy doesn't complain.
|
||
@classmethod | ||
def dumps(cls, obj: Any) -> bytes: | ||
# Run with pickler.fast so it doesn't intern strings, making the hash result more predictable |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This didn't have to move to __init__
, but seemed like the more natural place for it
Summary: In an upcoming change, we need to modify FxGraphCachePickler to behave differently depending on whether the graph has frozen parameters (whether or not we have frozen parameters). To do that, it will be convenient to change FxGraphCachePickler into a regular object instead of a collection of classmethods. Test Plan: unit tests cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov [ghstack-poisoned]
@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: Move all the custom `_reduce_*` functions inside the FxGraphCachePickler class. This is mostly a cosmetic change since they're conceptually members of FxGraphCachePickler. But also in an upcoming diff, I'll add a member variable to the class to control how we handle constant tensors, so it will be convenient to be able to query that setting via `self`. I made the analogous changes to AOTAutogradCachePickler for consistency. Test Plan: unit tests Pull Request resolved: #138683 Approved by: https://github.com/eellison ghstack dependencies: #138681, #138682
Summary: In an upcoming change, we need to modify FxGraphCachePickler to behave differently depending on whether the graph has frozen parameters (whether or not we have frozen parameters). To do that, it will be convenient to change FxGraphCachePickler into a regular object instead of a collection of classmethods. Test Plan: unit tests Pull Request resolved: pytorch#138682 Approved by: https://github.com/eellison ghstack dependencies: pytorch#138681
Summary: Move all the custom `_reduce_*` functions inside the FxGraphCachePickler class. This is mostly a cosmetic change since they're conceptually members of FxGraphCachePickler. But also in an upcoming diff, I'll add a member variable to the class to control how we handle constant tensors, so it will be convenient to be able to query that setting via `self`. I made the analogous changes to AOTAutogradCachePickler for consistency. Test Plan: unit tests Pull Request resolved: pytorch#138683 Approved by: https://github.com/eellison ghstack dependencies: pytorch#138681, pytorch#138682
Stack from ghstack (oldest at bottom):
Summary: In an upcoming change, we need to modify FxGraphCachePickler to behave differently depending on whether the graph has frozen parameters (whether or not we have frozen parameters). To do that, it will be convenient to change FxGraphCachePickler into a regular object instead of a collection of classmethods.
Test Plan: unit tests
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov