[Pytorch][Attention][CP][CI] Make 3221 changes robust - #3438
[Pytorch][Attention][CP][CI] Make 3221 changes robust#3438sudhakarsingh27 wants to merge 7 commits into
Conversation
FlashAttention 2 backward needs zeroed gradient outputs for deterministic grouped-query attention in this path. Permit unpadded FlashAttention backends because inter-sequence padding requires no backend-specific handling, and align the autograd return arity with the forward inputs. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis follow-up broadens unpadded no-load-balance THD context parallelism to supported FlashAttention backends, initializes FA2 backward accumulation buffers, and corrects the custom autograd return arity.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains. No blocking failure remains. Important Files Changed
Reviews (6): Last reviewed commit: "Rely on existing L1 attention runs" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| assert ( | ||
| use_fused_attention or use_flash_attn_3 | ||
| ), "No-load-balance THD partitioning requires FusedAttention or FlashAttention 3." | ||
| assert not ( | ||
| use_flash_attn_3 and pad_between_seqs | ||
| ), "No-load-balance THD partitioning with FlashAttention 3 does not support padding yet." | ||
| use_fused_attention or not pad_between_seqs | ||
| ), "No-load-balance THD partitioning only supports padding with FusedAttention." |
There was a problem hiding this comment.
This relaxed assertion enables FlashAttention 2 for unpadded no-load-balance THD all-gather context parallelism, but the existing test covers only FlashAttention 3. Add FA2 forward and backward coverage so regressions in output, dQ/dK/dV initialization, backend argument compatibility, and autograd behavior are detected.
Knowledge Base Used: PyTorch attention execution
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
Addressed. The no-load-balance CP test now queries production backend selection for the BF16 THD all-gather GQA configuration and validates output, dQ, dK, and dV. In L1, FA4 is disabled and FA3 is not installed, so the existing deterministic and non-deterministic suite runs exercise FA2 without a separate forced invocation. The focused path was also validated with FA2 on SM90 and SM100 and with FA3/FA4 on supported hardware.
| use_flash_attn_3 and pad_between_seqs | ||
| ), "No-load-balance THD partitioning with FlashAttention 3 does not support padding yet." | ||
| use_fused_attention or not pad_between_seqs | ||
| ), "No-load-balance THD partitioning only supports padding with FusedAttention." |
There was a problem hiding this comment.
Is this change because FA3 and pad_between_seqs are incompatible? I think it's better to keep them separate right?
There was a problem hiding this comment.
This is not a general incompatibility between FA3 and pad_between_seqs. The limitation is specific to the no-load-balance partitioning strategy. Backend selection remains separate: any supported FlashAttention version can run the unpadded path, while inter-sequence padding in this path remains FusedAttention-only. The assertion and comments now reflect that separation.
Document why FA2 initializes backward gradient buffers and why inter-sequence padding remains limited to FusedAttention. This makes the supported backend matrix explicit without changing behavior. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Reuse the focused CP test across supported FlashAttention generations and force FA2 from L1 in a fresh process. Backend controls are read at import time, so a separate invocation covers deterministic FA2 GQA backward without adding runner plumbing. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Query the existing backend selector for the exact THD context-parallel configuration instead of duplicating architecture and version gates in pytest. This keeps capability decisions in one place while allowing version-forced test invocations to exercise production selection logic. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Document why the focused deterministic case runs in a fresh process and how the version flags force FA2, while the full runs retain normal backend selection. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
The general context-parallel test is already collected by both L1 attention invocations. Since this suite selects FA2, a separate forced deterministic invocation duplicates coverage without exercising a distinct path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
sudhakarsingh27
changed the title
sudhakarsingh27
changed the title
KshitijLakhani
left a comment
KshitijLakhani
left a comment
There was a problem hiding this comment.
LGTM for merge after CI passes successfully !
3 participants
Description
Follow-up to #3221 that makes no-load-balance THD all-gather context parallelism robust for unpadded FlashAttention backends and fixes a bug in CP
Validation
git diff --checkqa/L1_pytorch_distributed_unittest/test.sh: all tests passed