[CP] Fix cuDNN CP LSE dimension bug#163231
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/163231
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 d3db0f2 with merge base 7a0f933 ( 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. |
|
Test? |
| self, block_out: torch.Tensor, block_lse: torch.Tensor, partial: bool | ||
| ) -> None: | ||
| block_lse = block_lse.unsqueeze(dim=-1) | ||
| # The cuDNN backend preserves the last dimension for LSE. |
There was a problem hiding this comment.
@eqy what was the reason again for adding the extra dim to LSE?
There was a problem hiding this comment.
I think it's some weirdness around cuDNN expecting 4D in the backward, where before we were manually unsqueeze it to handle that
https://github.com/pytorch/pytorch/pull/155958/files
|
@ezyang This PR fixes an existing test. I updated the PR summary. |
|
@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 |
We should only unsqueeze if necessary. Fix pytorch#162743 Pull Request resolved: pytorch#163231 Approved by: https://github.com/eqy ghstack dependencies: pytorch#162539, pytorch#162540, pytorch#162541, pytorch#163115, pytorch#163131
We should only unsqueeze if necessary. Fix pytorch#162743 Pull Request resolved: pytorch#163231 Approved by: https://github.com/eqy ghstack dependencies: pytorch#162539, pytorch#162540, pytorch#162541, pytorch#163115, pytorch#163131
We should only unsqueeze if necessary. Fix pytorch#162743 Pull Request resolved: pytorch#163231 Approved by: https://github.com/eqy ghstack dependencies: pytorch#162539, pytorch#162540, pytorch#162541, pytorch#163115, pytorch#163131
Stack from ghstack (oldest at bottom):
We should only unsqueeze if necessary.
Fix #162743
cc @H-Huang @awgu @wanchaol @fduwjj @wz337 @wconstab @d4l3k @pragupta @ezyang @msaroufim @dcci