Skip to content

[CP] Fix cuDNN CP LSE dimension bug#163231

Closed
fegin wants to merge 2 commits intogh/fegin/322/basefrom
gh/fegin/322/head
Closed

[CP] Fix cuDNN CP LSE dimension bug#163231
fegin wants to merge 2 commits intogh/fegin/322/basefrom
gh/fegin/322/head

Conversation

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Sep 18, 2025

🔗 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 (image):

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.

[ghstack-poisoned]
@fegin fegin changed the title [CP] Fix a cuDNN CP bug [CP] Fix cuDNN CP LSE dimension bug Sep 18, 2025
@ezyang
Copy link
Contributor

ezyang commented Sep 18, 2025

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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@eqy what was the reason again for adding the extra dim to LSE?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@fegin
Copy link
Contributor Author

fegin commented Sep 18, 2025

@ezyang This PR fixes an existing test. I updated the PR summary.

@fegin
Copy link
Contributor Author

fegin commented Sep 20, 2025

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 20, 2025
@pytorchmergebot
Copy link
Collaborator

Merge started

Your 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

Advanced Debugging
Check the merge workflow status
here

mansiag05 pushed a commit to mansiag05/pytorch that referenced this pull request Sep 22, 2025
cleonard530 pushed a commit to cleonard530/pytorch that referenced this pull request Sep 22, 2025
dsashidh pushed a commit to dsashidh/pytorch that referenced this pull request Sep 26, 2025
@github-actions github-actions bot deleted the gh/fegin/322/head branch October 21, 2025 02:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request Merged module: context parallel PyTorch Context Parallel oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: context parallel

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants