torch.no_grad() yields NaN values on mps device, 4D attention mask
Author: AmitMY
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.56.0
- Platform: macOS-15.6.1-arm64-arm-64bit
- Python version: 3.12.2
- Huggingface_hub version: 0.34.3
- Safetensors version: 0.5.3
- Accelerate version: 1.10.1
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1 (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

Don't know.

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

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

### Reproduction

on macOS, define a 4D attention mask that is padded, and call a model.
```py
@pytest.mark.parametrize("device", ["cpu", "mps"])
def test_isolated_language_model_mps(device):
    model = AutoModelForCausalLM.from_pretrained("sbintuitions/tiny-lm").to(device)
    input_ids = torch.tensor([[0, 1, 0, 0], [0, 1, 2, 3]], device=device)

    attention_mask = torch.tensor([
        [[[ True, False, False, False],
          [ True,  True, False, False],
          [False, False, False, False],
          [False, False, False, False]]],
        [[[ True, False, False, False],
          [ True,  True, False, False],
          [ True,  True,  True, False],
          [ True,  True,  True,  True]]]], device=device)
    # attention_mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]], dtype=torch.bool, device=device)

    with torch.no_grad():
       outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    assert not torch.isnan(outputs.logits).any(), "Logits contain NaN values"
```

### Expected behavior

not NaN - same as it is on CPU.

Also, if removing `torch.no_grad()` it is also not NaN.

<img width="291" height="75" alt="Image" src="https://github.com/user-attachments/assets/9d99ff28-c3bd-4ca8-8515-29ba4434d9b2" />

Comments:
Might be related to https://github.com/huggingface/transformers/issues/40242#issuecomment-3225888611. TL;DR; with MPS devices when the full row is masked, the attentions scores get NaNs which is a problem stemming from pytorch, As a workaround, you can try to post-process the mask as follows (https://github.com/huggingface/transformers/blob/75d6f17de68f372284ecb5b40db6f83007ffa394/src/transformers/masking_utils.py#L463-L467) because ig torch didn't yet release a new version which contains the fix for MPS devices

Thanks @zucchini-nlp ! 
It might be related, yes.
Strange then, that it does work when there are gradients, but only fails with `torch.no_grad()`



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.

not stale

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.

not stale

I don't think this is relevant to transformers as pointed earlier, seems to be an issue with how pytorch handles masking internally. Would you consider moving the issue to `torch` @AmitMY ? 

Closing in favor of https://github.com/pytorch/pytorch/issues/167515