[distributed] fix corrupted gradient problem under ring CP. - #14274
[distributed] fix corrupted gradient problem under ring CP.#14274sayakpaul wants to merge 7 commits into
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Cc: @pthombre could you try this PR and tell if it works for your use case as expected? |
| ctx: torch.autograd.function.FunctionCtx, | ||
| grad_out: torch.Tensor, | ||
| *args, | ||
| query: torch.Tensor | None = None, |
There was a problem hiding this comment.
I think backward ops for varlen flash should also be updated for this?
There was a problem hiding this comment.
Wanted to keep the PR lean and really scoped to what we have tried on the AutoModel side.
| value = value.transpose(1, 2).contiguous() | ||
| out = out.transpose(1, 2).contiguous() | ||
| # Model layout (B, S, H) -> the (B, H, S, 1) shape the cuDNN backward kernel expects. | ||
| lse = lse.transpose(1, 2).unsqueeze(-1).contiguous() |
There was a problem hiding this comment.
Does this apply to torch<=2.9? I believe the trailing dim is expected in later versions?
There was a problem hiding this comment.
Yeah good catch. Should have been addressed now.
| lse = softmax_lse.permute(0, 2, 1).contiguous() if return_lse else None | ||
| # Save LSE in model layout (B, S, H) so the backward pass — including the ring loop, which | ||
| # overrides it with the reduced LSE — always sees the same layout regardless of `return_lse`. | ||
| lse = softmax_lse.permute(0, 2, 1).contiguous() |
There was a problem hiding this comment.
Should this apply to FA2 backends as well?
There was a problem hiding this comment.
Wanted to keep the PR very lean and cater to what we have verified already.
3 participants
Fixes #14265