[PyTorch] Disable Flash Attention backend in Userbuffers tests

[timmoon10]

Description

We have experienced some Userbuffers test failures on A100s, apparently because the Flash Attention backward pass introduces numerical errors. This test is primarily intended to test the linear layers and not attention, so as a quick fix I've just disabled the Flash Attention backend.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Disable Flash Attention backend in Userbuffers tests

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[ksivaman]

LGTM as a fix for green CI, @cyanguwa we should document + try and root cause this

[greptile-apps]

Greptile Summary

• Disables Flash Attention backend in Userbuffers layer tests by setting NVTE_FLASH_ATTN=0 to prevent numerical errors on A100s
• The fix is appropriately scoped to _run_layer_with_overlap function only, affecting tests for Linear, LayerNormLinear, LayerNormMLP, MultiheadAttention, and TransformerLayer

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The change is minimal and surgical - it only adds an environment variable to disable Flash Attention in specific tests that were experiencing numerical errors. The fix properly sets and unsets the environment variable, maintaining clean test isolation. The scope is appropriately limited to layer tests (not low-level GEMM tests), which aligns with the PR description stating this is a quick fix for Userbuffers tests that are primarily intended to test linear layers, not attention.
• No files require special attention

Important Files Changed

Filename	Overview
tests/pytorch/distributed/test_comm_gemm_overlap.py	Added NVTE_FLASH_ATTN=0 environment variable to disable Flash Attention backend in layer overlap tests to avoid numerical errors

Sequence Diagram

sequenceDiagram
    participant Test as test_layers_with_overlap_*
    participant Helper as _run_layer_with_overlap
    participant Env as Environment
    participant Subprocess as run_layer_with_overlap.py
    
    Test->>Helper: Call with layer parameters
    Helper->>Env: Set "NVTE_FLASH_ATTN=0"
    Helper->>Env: Set "PYTORCH_JIT=0"
    Helper->>Env: Set "NVTE_TORCH_COMPILE=0"
    Helper->>Env: Set "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0"
    Helper->>Subprocess: Run test with modified environment
    Subprocess-->>Helper: Return test result
    Helper->>Env: Unset "NVTE_FLASH_ATTN"
    Helper->>Env: Unset "PYTORCH_JIT"
    Helper->>Env: Unset "NVTE_TORCH_COMPILE"
    Helper->>Env: Unset "NVTE_ALLOW_NONDETERMINISTIC_ALGO"
    Helper-->>Test: Return success/failure

Loading

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[timmoon10]

/te-ci pytorch L1

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[KshitijLakhani]

LGTM
Agree with @ksivaman on documenting it. On the JAX side, we've mostly been sprinkling TODO for such things - I wonder if @cyanguwa would prefer something like that or a different approach ?