Add UUID-based cache key support for pre-grad custom passes#177403
Closed
frgossen wants to merge 5 commits intogh/frgossen/10/basefrom
Closed
Add UUID-based cache key support for pre-grad custom passes#177403frgossen wants to merge 5 commits intogh/frgossen/10/basefrom
frgossen wants to merge 5 commits intogh/frgossen/10/basefrom
Conversation
pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/177403
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 6361bb2 with merge base 6a461fe ( 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. |
frgossen
added a commit
that referenced
this pull request
Mar 13, 2026
pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. ghstack-source-id: 09ddb53 Pull Request resolved: #177403
This PR needs a
|
This was referenced Mar 13, 2026
aorenste
approved these changes
Mar 16, 2026
pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo Lucaskabela [ghstack-poisoned]
pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo Lucaskabela [ghstack-poisoned]
pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo Lucaskabela [ghstack-poisoned]
pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy kadeng muchulee8 amjames chauhang aakhundov coconutruben jataylo Lucaskabela [ghstack-poisoned]
This was referenced Mar 17, 2026
zou3519
approved these changes
Mar 17, 2026
Collaborator
|
Starting merge as part of PR stack under #177428 |
This was referenced Mar 18, 2026
Collaborator
|
Starting merge as part of PR stack under #177428 |
Collaborator
|
Starting merge as part of PR stack under #177428 |
2 similar comments
Collaborator
|
Starting merge as part of PR stack under #177428 |
Collaborator
|
Starting merge as part of PR stack under #177428 |
pytorchmergebot
pushed a commit
that referenced
this pull request
Mar 18, 2026
Pull Request resolved: #177428 Approved by: https://github.com/mlazos, https://github.com/zou3519 ghstack dependencies: #177397, #177403
pytorchmergebot
pushed a commit
that referenced
this pull request
Mar 18, 2026
…77429) Add a pre_grad_pass_timing config ("early", "late", or "default") that controls when pre-grad passes run relative to the AOT autograd cache lookup. - "early": passes run before cache lookup, so they execute on every compile (including cache hits) and the cache key reflects the already-transformed graph. - "late": passes run after cache lookup (only on cache miss); requires custom passes to provide a UUID for the cache key. - "default": automatically resolves to "late" when possible (no custom pass, or a custom pass with a UUID), and falls back to "early" when the custom pass has no UUID. Explicitly setting "late" with a UUID-less custom pass now raises a RuntimeError instead of silently bypassing the cache. The existing test_pre_grad_passes_called_on_cache_miss_only test is renamed and pinned to "late" timing, and new tests cover early timing, both default timing branches, and the error case. Pull Request resolved: #177429 Approved by: https://github.com/aorenste, https://github.com/zou3519 ghstack dependencies: #177397, #177403, #177428
ryanzhang22
pushed a commit
to ryanzhang22/pytorch
that referenced
this pull request
Mar 19, 2026
…177403) pre_grad_custom_pass was the only custom pass config without UUID-based cache key integration. It was excluded from config serialization but not handled specially via UUID extraction, so its effect was only captured indirectly through the resulting FX graph. This meant two different passes producing the same graph could incorrectly share a cache entry. Align pre_grad_custom_pass with post-grad and joint passes: change its type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so the UUID is extracted explicitly), include it in FxGraphHashDetails, and validate it in _check_can_cache. Pull Request resolved: pytorch#177403 Approved by: https://github.com/aorenste, https://github.com/zou3519 ghstack dependencies: pytorch#177397
ryanzhang22
pushed a commit
to ryanzhang22/pytorch
that referenced
this pull request
Mar 19, 2026
Pull Request resolved: pytorch#177428 Approved by: https://github.com/mlazos, https://github.com/zou3519 ghstack dependencies: pytorch#177397, pytorch#177403
ryanzhang22
pushed a commit
to ryanzhang22/pytorch
that referenced
this pull request
Mar 19, 2026
…torch#177429) Add a pre_grad_pass_timing config ("early", "late", or "default") that controls when pre-grad passes run relative to the AOT autograd cache lookup. - "early": passes run before cache lookup, so they execute on every compile (including cache hits) and the cache key reflects the already-transformed graph. - "late": passes run after cache lookup (only on cache miss); requires custom passes to provide a UUID for the cache key. - "default": automatically resolves to "late" when possible (no custom pass, or a custom pass with a UUID), and falls back to "early" when the custom pass has no UUID. Explicitly setting "late" with a UUID-less custom pass now raises a RuntimeError instead of silently bypassing the cache. The existing test_pre_grad_passes_called_on_cache_miss_only test is renamed and pinned to "late" timing, and new tests cover early timing, both default timing branches, and the error case. Pull Request resolved: pytorch#177429 Approved by: https://github.com/aorenste, https://github.com/zou3519 ghstack dependencies: pytorch#177397, pytorch#177403, pytorch#177428
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
pre_grad_custom_pass was the only custom pass config without UUID-based
cache key integration. It was excluded from config serialization but not
handled specially via UUID extraction, so its effect was only captured
indirectly through the resulting FX graph. This meant two different
passes producing the same graph could incorrectly share a cache entry.
Align pre_grad_custom_pass with post-grad and joint passes: change its
type to CustomGraphPassType, add it to _cache_config_ignore_prefix (so
the UUID is extracted explicitly), include it in FxGraphHashDetails, and
validate it in _check_can_cache.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo @Lucaskabela