sdpa for bert casues nan when using bfloat16 with padding.

[Leoyzen]

System Info

• transformers: transformers==4.41.1
• pytorch: 2.3.0+cu121
• python: 3.10

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

In [1]: import torch.nn.functional as F

In [2]: import torch

In [3]: data = torch.load("reproduce_data.pt", map_location='cuda')

In [4]: with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
   ...:     print(F.scaled_dot_product_attention(data['q'], data['k'], data['v'], data['attn_mask'], data['dropout_p'], data['is_causal']).isnan().any((1
   ...: ,2)))
   ...:
tensor([[False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],               # <----- without pandding
        ...,
        [ True,  True,  True,  ...,  True,  True,  True],               # <----- with pandding
        [ True,  True,  True,  ...,  True,  True,  True],               # <----- with pandding
        [ True,  True,  True,  ...,  True,  True,  True]], device='cuda:0')

In [5]: with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=True):
   ...:     print(F.scaled_dot_product_attention(data['q'], data['k'], data['v'], data['attn_mask'], data['dropout_p'], data['is_causal']).isnan().any((1
   ...: ,2)))
   ...:
tensor([[False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        ...,
        [ True,  True,  True,  ...,  True,  True,  True],
        [ True,  True,  True,  ...,  True,  True,  True],
        [ True,  True,  True,  ...,  True,  True,  True]], device='cuda:0')

# it's okay with offical math kernel
In [6]: with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=False):
   ...:     print(F.scaled_dot_product_attention(data['q'], data['k'], data['v'], data['attn_mask'], data['dropout_p'], data['is_causal']).isnan().any((1
   ...: ,2)))
   ...:
/root/.local/share/conda/envs/bytednlp/lib/python3.10/site-packages/torch/backends/cuda/__init__.py:342: FutureWarning: torch.backends.cuda.sdp_kernel() is deprecated. In the future, this context manager will be removed. Please see, torch.nn.attention.sdpa_kernel() for the new context manager, with updated signature.
  warnings.warn(
tensor([[False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        ...,
        [False, False, False,  ..., False, False, False],               # <----- with pandding, which is correct
        [False, False, False,  ..., False, False, False],               # <----- with pandding, which is correct
        [False, False, False,  ..., False, False, False]], device='cuda:0')

In [7]: mask = data['attn_mask']

# we use min / 2 as float('-inf')
In [10]: mask2 = mask.masked_fill(mask.bool(), torch.finfo(mask.dtype).min / 2)

In [11]: with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True, enable_mem_efficient=True):
    ...:     print(F.scaled_dot_product_attention(data['q'], data['k'], data['v'], mask2, data['dropout_p'], data['is_causal']).isnan().any((1,2)))
    ...:
/root/.local/share/conda/envs/bytednlp/lib/python3.10/site-packages/torch/backends/cuda/__init__.py:342: FutureWarning: torch.backends.cuda.sdp_kernel() is deprecated. In the future, this context manager will be removed. Please see, torch.nn.attention.sdpa_kernel() for the new context manager, with updated signature.
  warnings.warn(
tensor([[False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        [False, False, False,  ..., False, False, False],
        ...,
        [False, False, False,  ..., False, False, False],              # <----- with pandding, which the result is correct.
        [False, False, False,  ..., False, False, False],               # <----- with pandding, which the result is correct.
        [False, False, False,  ..., False, False, False]], device='cuda:0')

Expected behavior

the output should without nan when using bfloat16 and sdap enabled.

I think it is safe to use torch.finfo(dtype).min / 2 instead of torch.finfo(dtype.min.

[amyeroberts]

c @ArthurZucker @fxmarty

[ArthurZucker]

yep, this was always there I think and related to the mask creation that probably overflows.
Do you want to open a PR for a fix? 🤗

[fxmarty]

@Leoyzen can you share the repro tensors and/or reproduction with a transformers example?

[Leoyzen]

@Leoyzen can you share the repro tensors and/or reproduction with a transformers example?

The reproduce_ata.pt which dumped from the private code repo is quite large(with (torch.Size([31, 1, 2000, 2000])) and almost 1GB).

We use Bert from transformers and weights from (stella-v2)[https://huggingface.co/infgrad/stella-large-zh-v2] to do some finetuning work.

Training with bert large and bfloat16 should reproduce the bug.

reproduce.zip

[ArthurZucker]

Mmmm if this is still a problem, we need to propagate the changes from #32227 to bert and bert sdpa!

[ArthurZucker]

Leaving it to the community unless I get time !

[SakshamDhawan]

Hi guys picking this up to propagate the changes from the fixes in Gemma and Llama implementations. Does anyone have guidance on how to replicate the bug?

[DuyguA]

hellos @ArthurZucker and all, I think this issue is resolved by the latest PyTorch release. Issue is that sdpa didn't play well with all zero rows, here's a complaining issue and the fix PR:
pytorch/pytorch#103749
pytorch/pytorch#133882

I tested OP's input data along with a minimal example. All tests were done on Colab A100.

OP data:

data = torch.load("reproduce_data.pt", map_location='cuda')

with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
   print(F.scaled_dot_product_attention(data['q'], data['k'], data['v'], data['attn_mask'], data['dropout_p'], data['is_causal']).isnan().any((1,2)))

tensor([[False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False, False, False, False, False, False, False,
         False, False, False, False]], device='cuda:0')

Minimal data:

import torch
import torch.nn.functional as F

# Define sequence lengths
seq_length = 5
max_seq_length = 10

# Initialize random queries, keys, and values
# Shape: (batch_size=1, num_heads=2, target_seq_length=seq_length+1, hidden_dim=8)
q = torch.randn(1, 2, seq_length + 1, 8, dtype=torch.float32)

# Shape: (batch_size=1, num_heads=2, source_seq_length=max_seq_length, hidden_dim=8)
k = torch.randn(1, 2, max_seq_length, 8, dtype=torch.float32)

# Shape: (batch_size=1, num_heads=2, source_seq_length=max_seq_length, hidden_dim=8)
v = torch.randn(1, 2, max_seq_length, 8, dtype=torch.float32)

# Define the attention mask
# Shape: (batch_size=1, num_heads=1, target_seq_length=seq_length+1, source_seq_length=max_seq_length)
mask = torch.tensor(
    [[[
        [ True, False, False, False, False, False, False, False, False, False],  # regular mask
        [ True,  True, False, False, False, False, False, False, False, False],
        [ True,  True,  True, False, False, False, False, False, False, False],
        [ True,  True,  True,  True, False, False, False, False, False, False],
        [ True,  True,  True,  True,  True, False, False, False, False, False],
        [False, False, False, False, False, False, False, False, False, False],  # padding mask
    ]]]
).to(torch.bool)  # Keep as boolean for masking logic

# Convert mask to float (0 for masked positions, 1 for valid positions)
# Replace `False` (masked positions) with `-inf` and `True` with 0 (no penalty in logits)
attention_mask = mask.to(torch.float32)  # Convert boolean to float for compatibility
attention_mask = attention_mask.masked_fill(~mask, torch.finfo(torch.float32).min)  # Replace False with -inf

# Dropout probability
dropout_p = 0.0

# Causal flag
is_causal = False

# Perform scaled dot product attention
try:
    output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p, is_causal)
    print("Output:", output)
    print("Output Shape:", output.shape)
except Exception as e:
    print(f"Error occurred: {e}")

Output: tensor([[[[-1.0444,  0.5993,  1.0736,  1.3209, -0.0323,  0.1462,  1.6096,
            0.9682],
          [-1.3124,  0.1635,  0.6778,  0.6054, -0.0425,  0.0155,  1.3656,
            1.1903],
          [-0.8785, -0.0359,  0.2625,  0.6747, -0.1908, -0.0602,  1.0220,
            1.1996],
          [-0.8473, -0.4083, -0.0684, -0.0791, -0.1161, -0.1803,  0.8802,
            1.0146],
          [-0.8343,  0.3752,  0.7345,  1.1757, -0.1228,  0.0637,  1.3455,
            0.9949],
          [-0.5261,  0.2097,  0.0073, -0.0036, -0.5551, -0.7264,  0.8049,
            0.1770]],

         [[-0.4158, -0.4586, -0.8774,  0.0774, -0.8315, -0.0257, -0.3001,
           -1.1196],
          [-0.6589, -0.0338, -0.3365, -0.7406, -0.5868, -0.3964, -0.0485,
           -0.7849],
          [-0.5010, -0.3385, -0.5654,  0.0210, -0.6496, -0.2028, -0.2179,
           -0.9440],
          [-0.6298, -0.0294, -0.0096,  0.0652, -0.5166, -0.4455, -0.1599,
           -0.4978],
          [-0.7547,  0.1647,  0.0735, -0.2379, -0.2827, -0.5888, -0.1862,
           -0.2329],
          [-0.4864,  0.1173, -0.1487,  0.0667,  0.0135, -0.5921, -0.1294,
            0.0637]]]])
Output Shape: torch.Size([1, 2, 6, 8])

I made a quick minor PR to test the issue anyway.

[ArthurZucker]

hellos @ArthurZucker and all, I think this issue is resolved by the latest PyTorch release. Issue is that sdpa didn't play well with all zero rows, here's a complaining issue and the fix PR:

This is fairly known, and we have had a fix for that in transformers for a while AFAIK

[ArthurZucker]

but will have a look thanks