fix overflow for bfloat16 when using cuda kernel.
Author: Leoyzen
Pull request: yes

# What does this PR do?
using `torch.finfo(dtype).min / 2` instead of `torch.finfo(dtype).min` to avoid overflow caculation when using `F.scaled_dot_product_attention`, especially bfloat16.

Fixes # (issue)
* https://github.com/huggingface/transformers/issues/31035#issuecomment-2149675331

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
 @ArthurZucker and @younesbelkada



Comments:
Hi @Leoyzen, thank you.

If this issues stems from fully masked rows in the attention mask, this issue is related to https://github.com/pytorch/pytorch/issues/110213. In that case, the proper solution would be to use https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/modeling_attn_mask_utils.py#L189 as done for other models. However, I don't see in which case we may have fully masked rows for bert.

fyi @hackyon

@fxmarty No, it not the same issue.

It will be nan with any padding length.

You can see the attn_mask from reproduce example.

Note the fully masked rows for bert exists when using packed sequence/padding mask for speedup training/inferece, such as [PackedBert](https://medium.com/graphcore/packed-bert-how-to-accelerate-fine-tuning-and-inference-for-nlp-tasks-with-packing-de16ac4aac1).

We use such code to produce 3D-mask(maybe fully masked with some rows).

```python
# Generate the boolean extended attention mask.
# see https://github.com/graphcore/examples/blob/master/tutorials/blogs_code/finetuning-packedBERT/walkthrough.ipynb # noqa: E501
attention_mask_3d = attention_mask[:, None, :].repeat(1, attention_mask.shape[1], 1)
attention_mask_3d = (attention_mask_3d == attention_mask_3d.transpose(1, 2)) * (attention_mask_3d != 0)
```

@fxmarty you're right. I tested with `_unmask_unattended` and the nan disapears.

```python
def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
    """
    Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
    `(batch_size, key_value_length)`

    Args:
        mask (`torch.Tensor` or `None`):
            A 2D attention mask of shape `(batch_size, key_value_length)`
        dtype (`torch.dtype`):
            The torch dtype the created mask shall have.
        tgt_len (`int`):
            The target length or query length the created mask shall have.
    """
    #  batch_size, key_value_length = mask.shape
    *_, key_value_length = mask.shape
    tgt_len = tgt_len if tgt_len is not None else key_value_length

    # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
    # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
    # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).
    is_tracing = (
        torch.jit.is_tracing()
        or isinstance(mask, torch.fx.Proxy)
        or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling())
    )

    if not is_tracing and torch.all(mask == 1):
        if tgt_len == 1:
            # For query_length == 1, causal attention and bi-directional attention are the same.
            return None
        elif key_value_length == tgt_len:
            return None
        else:
            # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
            # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
            # Reference: https://github.com/pytorch/pytorch/issues/108108
            expanded_4d_mask = AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
    else:
        expanded_4d_mask = AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
        
    # Attend to all tokens in masked rows from the causal_mask, for example the relevant first rows when
    # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
    # Details: https://github.com/pytorch/pytorch/issues/110213    
    if not is_tracing and expanded_4d_mask.device.type == "cuda":
        expanded_4d_mask = AttentionMaskConverter._unmask_unattended(
            expanded_4d_mask, min_dtype=torch.finfo(dtype).min
        )
    return expanded_4d_mask
```

But it's strange why nan only occured for bfloat16 but not for fp32/fp16.



This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.