Repository: pytorch
Issue: DISABLED test_flash_attention_vs_math_ref_grads_batch_size_1_seq_len_q_1024_seq_len_k_1024_head_dim_16_is_causal_False_dropout_p_0_48_bfloat16_scale0_cuda_bfloat16 (__main__.TestSDPACudaOnlyCUDA)
Labels: module: rocm,triaged,module: flaky-tests,skipped

Platforms: rocm

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/flakytest?name=test_flash_attention_vs_math_ref_grads_batch_size_1_seq_len_q_1024_seq_len_k_1024_head_dim_16_is_causal_False_dropout_p_0_48_bfloat16_scale0_cuda_bfloat16&suite=TestSDPACudaOnlyCUDA&limit=100) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/27649622705).

Over the past 3 hours, it has been determined flaky in 4 workflow(s) with 4 failures and 4 successes.

**Debugging instructions (after clicking on the recent samples link):**
DO NOT ASSUME THINGS ARE OKAY IF THE CI IS GREEN. We now shield flaky tests from developers so CI will thus be green but it will be harder to parse the logs.
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Click on the Test step of the job so that it is expanded. Otherwise, the grepping will not work.
3. Grep for `test_flash_attention_vs_math_ref_grads_batch_size_1_seq_len_q_1024_seq_len_k_1024_head_dim_16_is_causal_False_dropout_p_0_48_bfloat16_scale0_cuda_bfloat16`
4. There should be several instances run (as flaky tests are rerun in CI) from which you can study the logs.



<details><summary>Sample error message</summary>

```
Traceback (most recent call last):
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/testing/_comparison.py", line 1228, in not_close_error_metas
    pair.compare()
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/testing/_comparison.py", line 710, in compare
    self._compare_values(actual, expected)
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/testing/_comparison.py", line 837, in _compare_values
    compare_fn(
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/testing/_comparison.py", line 1016, in _compare_regular_values_close
    matches = torch.isclose(
RuntimeError: atol must be greater than or equal to zero, but got nan

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_transformers.py", line 3041, in test_flash_attention_vs_math_ref_grads
    self.assertEqual(key.grad, key_ref.grad.to(key.grad.dtype),
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/testing/_internal/common_utils.py", line 3682, in assertEqual
    error_metas = not_close_error_metas(
  File "/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/testing/_comparison.py", line 1234, in not_close_error_metas
    raise RuntimeError(
RuntimeError: Comparing

TensorOrArrayPair(
    id=(),
    actual=tensor([[[[-0.0177, -0.0280, -0.0278,  ..., -0.0242, -0.0325, -0.0272],
          [-0.1079, -0.1235, -0.1245,  ..., -0.1172, -0.1235, -0.1309],
          [-0.0055,  0.0045, -0.0220,  ..., -0.0106, -0.0143, -0.0194],
          ...,
          [ 0.0098,  0.0270,  0.0244,  ...,  0.0166,  0.0371,  0.0283],
          [-0.0608, -0.0522, -0.0559,  ..., -0.0515, -0.0503, -0.0457],
          [ 0.0139,  0.0117,  0.0220,  ...,  0.0160,  0.0168,  0.0135]],

         [[-0.1177, -0.1230, -0.1147,  ..., -0.1230, -0.0947, -0.1260],
          [ 0.0947,  0.0703,  0.0747,  ...,  0.0444,  0.0957,  0.0757],
          [-0.0767, -0.0840, -0.0752,  ..., -0.0786, -0.0505, -0.0713],
          ...,
          [ 0.1196,  0.1436,  0.1191,  ...,  0.1201,  0.1196,  0.1143],
          [-0.0703, -0.0835, -0.0806,  ..., -0.0640, -0.0664, -0.0608],
          [ 0.0923,  0.1094,  0.1143,  ...,  0.1011,  0.0996,  0.0972]],

         [[ 0.0723,  0.0801,  0.0791,  ...,  0.0835,  0.0840,  0.0752],
          [ 0.1484,  0.1738,  0.1660,  ...,  0.1768,  0.1611,  0.1660],
          [-0.0089, -0.0066,  0.0037,  ...,  0.0058,  0.0021,  0.0005],
          ...,
          [-0.0203, -0.0157, -0.0330,  ..., -0.0231, -0.0060, -0.0066],
          [ 0.0806,  0.0830,  0.0640,  ...,  0.0845,  0.0747,  0.0649],
          [-0.1040, -0.1250, -0.1299,  ..., -0.1152, -0.1118, -0.1348]],

         [[-0.0859, -0.0850, -0.0806,  ..., -0.0933, -0.0850, -0.0850],
          [ 0.1426,  0.1475,  0.1553,  ...,  0.1562,  0.1504,  0.1387],
          [-0.0466, -0.0408, -0.0327,  ..., -0.0280, -0.0405, -0.0535],
          ...,
          [ 0.0168,  0.0265,  0.0092,  ...,  0.0123,  0.0182,  0.0036],
          [ 0.0216,  0.0173,  0.0302,  ...,  0.0153,  0.0281,  0.0413],
          [-0.0177, -0.0270, -0.0291,  ..., -0.0199, -0.0344, -0.0305]]]],
       device='cuda:0', dtype=torch.bfloat16),
    expected=tensor([[[[ 2.8038e-04, -2.0409e-04,  1.2112e-04,  ...,  6.4373e-05,
           -1.8883e-04,  4.3392e-05],
          [-4.6968e-05, -1.3161e-04,  2.0142e-03,  ..., -1.0681e-04,
            2.0027e-05,  8.9645e-04],
          [-2.1458e-04,  6.5804e-05,  5.1498e-04,  ...,  4.4632e-04,
            2.7847e-04,  2.9755e-04],
          ...,
          [-1.8311e-04, -5.5695e-04, -9.5749e-04,  ..., -9.9182e-04,
            1.9169e-04, -3.2616e-04],
          [ 5.0354e-04,  4.4346e-05,  7.4005e-04,  ..., -7.6675e-04,
           -1.9741e-04,  3.4523e-04],
          [-2.5368e-04,  4.4060e-04, -2.5749e-04,  ...,  4.9591e-04,
           -1.9264e-04, -4.5300e-05]],

         [[-1.6880e-04,  5.3024e-04, -1.0633e-04,  ..., -4.1199e-04,
            2.6584e-05, -7.0953e-04],
          [ 2.5558e-04, -2.0504e-04,  7.1716e-04,  ...,  6.0272e-04,
            8.3447e-05,  3.0708e-04],
          [-2.3270e-04, -3.4094e-05, -4.7302e-04,  ...,  1.6308e-04,
           -3.4332e-04,  5.7459e-05],
          ...,
          [-3.6430e-04,  1.3733e-04,  4.1962e-04,  ...,  8.4686e-04,
            8.3923e-04,  1.3275e-03],
          [ 2.5558e-04, -1.0788e-05, -4.7493e-04,  ..., -2.8419e-04,
            2.9922e-05, -6.1417e-04],
          [-9.5749e-04,  5.9128e-05,  6.8665e-04,  ...,  7.7438e-04,
           -5.5313e-04,  1.2131e-03]],

         [[-3.9339e-05, -3.5667e-04,  8.6308e-05,  ...,  1.9169e-04,
            4.2725e-04,  7.0095e-05],
          [-5.9509e-04, -1.4591e-04, -8.4877e-05,  ..., -4.4441e-04,
            1.9989e-03, -4.9973e-04],
          [-2.2507e-04, -8.2850e-06, -4.9973e-04,  ..., -3.0136e-04,
           -5.6839e-04,  2.9922e-05],
          ...,
          [ 3.0899e-04,  3.2043e-04,  4.3297e-04,  ...,  1.1396e-04,
           -2.7275e-04,  3.2997e-04],
          [ 2.7657e-04, -7.4387e-04,  2.9755e-04,  ..., -1.4782e-04,
            1.4114e-03, -6.9046e-04],
          [-4.7386e-06,  6.6757e-04,  2.3460e-04,  ...,  3.5763e-05,
           -2.5635e-03,  2.7657e-04]],

         [[-1.1520e-03, -3.0518e-04, -3.4094e-05,  ...,  3.5286e-04,
           -6.8283e-04,  9.9945e-04],
          [ 2.5177e-03,  2.7275e-04,  9.8228e-05,  ...,  4.6349e-04,
            1.5030e-03, -1.2817e-03],
          [-3.6716e-05,  7.8201e-04, -9.3460e-05,  ...,  1.2207e-04,
           -3.7384e-04, -9.5367e-05],
          ...,
          [ 1.1587e-04, -3.0899e-04,  2.0218e-04,  ...,  1.4877e-04,
           -2.7275e-04, -1.4019e-04],
          [-3.3379e-04,  3.0041e-05, -1.5545e-04,  ..., -4.4060e-04,
            2.4986e-04,  1.0300e-04],
          [-5.8746e-04, -6.4468e-04,  1.0252e-04,  ...,  8.0109e-04,
            2.7847e-04,  2.3460e-04]]]], device='cuda:0', dtype=torch.bfloat16),
    rtol=339973952.0,
    atol=nan,
    equal_nan=True,
    check_device=False,
    check_dtype=True,
    check_layout=False,
    check_stride=False,
)

resulted in the unexpected exception above. If you are a user and see this message during normal operation please file an issue at https://github.com/pytorch/pytorch/issues. If you are a developer and working on the comparison functions, please except the previous error and raise an expressive `ErrorMeta` instead.

To execute this test, run the following from the base repo dir:
    PYTORCH_TEST_WITH_ROCM=1 python test/test_transformers.py -k TestSDPACudaOnlyCUDA.test_flash_attention_vs_math_ref_grads_batch_size_1_seq_len_q_1024_seq_len_k_1024_head_dim_16_is_causal_False_dropout_p_0_48_bfloat16_scale0_cuda_bfloat16

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0
```

</details>


Test file path: `test_transformers.py`

ConnectionTimeoutError: Connect timeout for 5000ms, GET https://raw.githubusercontent.com/pytorch/pytorch/main/test/test_transformers.py -2 (connected: false, keepalive socket: false, socketHandledRequests: 1, socketHandledResponses: 0)
headers: {}

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @clee2000