[DTensor] single_dim fix symint + _create_expanded_strategy#172421
[DTensor] single_dim fix symint + _create_expanded_strategy#172421wconstab wants to merge 2 commits intogh/wconstab/498/basefrom
Conversation
Claude's summary
torch/distributed/tensor/_ops/single_dim_strategy.py:242-257 - Modified the caching logic to handle SymInts from dynamic shapes by:
- Renaming the cached function to _create_expanded_strategy_impl
- Creating a wrapper _create_expanded_strategy that tries the cached version first
- Catching TypeError (from unhashable SymInts) and falling back to the uncached implementation
The test confirms that the caching logic in _create_expanded_strategy gracefully falls back to uncached execution when output_tensor_meta contains unhashable SymInts.
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/172421
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 24427b0 with merge base e09f7ac ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pianpwk
left a comment
There was a problem hiding this comment.
Would it make sense to work towards a hash key for symints? maybe something like (id(ShapeEnv), ShapeEnv version, symbol name, symbol hint) could be a conservative start.
Or just have a symint-aware cache that forces cache miss when they're present - I feel like we hit this error a lot.
|
would it be possible to |
|
@pianpwk do you think we can make it so 1. We hash/cache successfully on symint including its size hint or its set of constraints if present, and 2. If we use the symint during shard prop we are allowed to use it via its size hint only or via its constraints, but it causes a hard error if directly accessed as int? I guess if we do that,I'm still not sure how we handle error cases where we try to access an unbacked symint. |
|
@weifengpy sure, I can land this early. I'll do it later this morning or you can just rebase underneath your work right now before I get to it. |
Claude's summary
torch/distributed/tensor/_ops/single_dim_strategy.py:242-257 - Modified the caching logic to handle SymInts from dynamic shapes by:
- Renaming the cached function to _create_expanded_strategy_impl
- Creating a wrapper _create_expanded_strategy that tries the cached version first
- Catching TypeError (from unhashable SymInts) and falling back to the uncached implementation
The test confirms that the caching logic in _create_expanded_strategy gracefully falls back to uncached execution when output_tensor_meta contains unhashable SymInts.
[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 |
Merge failedReason: 1 jobs have failed, first few of them are: trunk / before-test / target-determination Details for Dev Infra teamRaised by workflow job |
|
@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 |
…172421) Claude's summary torch/distributed/tensor/_ops/single_dim_strategy.py:242-257 - Modified the caching logic to handle SymInts from dynamic shapes by: - Renaming the cached function to _create_expanded_strategy_impl - Creating a wrapper _create_expanded_strategy that tries the cached version first - Catching TypeError (from unhashable SymInts) and falling back to the uncached implementation The test confirms that the caching logic in _create_expanded_strategy gracefully falls back to uncached execution when output_tensor_meta contains unhashable SymInts. Pull Request resolved: pytorch#172421 Approved by: https://github.com/pianpwk, https://github.com/weifengpy
Claude's summary
torch/distributed/tensor/_ops/single_dim_strategy.py:242-257 - Modified the caching logic to handle SymInts from dynamic shapes by:
- Renaming the cached function to _create_expanded_strategy_impl
- Creating a wrapper _create_expanded_strategy that tries the cached version first
- Catching TypeError (from unhashable SymInts) and falling back to the uncached implementation
The test confirms that the caching logic in _create_expanded_strategy gracefully falls back to uncached execution when output_tensor_meta contains unhashable SymInts.
ghstack-source-id: d8abf38
Pull Request resolved: pytorch/pytorch#172421
Stack from ghstack (oldest at bottom):
Claude's summary
torch/distributed/tensor/_ops/single_dim_strategy.py:242-257 - Modified the caching logic to handle SymInts from dynamic shapes by:
- Renaming the cached function to _create_expanded_strategy_impl
- Creating a wrapper _create_expanded_strategy that tries the cached version first
- Catching TypeError (from unhashable SymInts) and falling back to the uncached implementation
The test confirms that the caching logic in _create_expanded_strategy gracefully falls back to uncached execution when output_tensor_meta contains unhashable SymInts.