Repository: pytorch
Issue: [MPS] softmax returns NaN attention probabilities for large tensors, in float16 and float32.
Labels: triaged,module: mps

### 🐛 Describe the bug

I've been making large images in stable-diffusion in float16 on my 4090 for a while.  
I tried on my M1 Max, generating an image at a size which _is_ possible in float16 on CUDA.

Using a recent PyTorch nightly (`2.1.0.dev20230310`), macOS 13.3 public beta.  
fails in float32 also.

```python
from torch import full, softmax, FloatTensor
import torch

# fill value is based on mean() from some real attention scores I had
attention_scores: FloatTensor = full([10, 12416, 12416], 0.0598, device='mps', dtype=torch.float16)

attention_probs: FloatTensor = attention_scores.softmax(dim=-1)

assert not attention_probs.isnan().any().item(), "found NaN in attention_probs"
# interestingly, attention_scores *also* becomes NaN. does softmax() mutate the input?
```

to pinpoint _where_ in the softmax it fails: here's a decomposed softmax.

```python
from torch import full, FloatTensor
import torch

attention_scores: FloatTensor = full([10, 12416, 12416], 0.0598, device='mps', dtype=torch.float16)

def softmax(x: FloatTensor, dim=1) -> FloatTensor:
  maxes = torch.max(x, dim, keepdim=True).values
  assert not maxes.isnan().any().item(), "found NaN in maxes"
  diffs = x-maxes
  assert not diffs.isnan().any().item(), "found NaN in diffs" # this assertion gets tripped
  x_exp = torch.exp(diffs)
  assert not x_exp.isnan().any().item(), "found NaN in x_exp"
  x_exp_sum = torch.sum(x_exp, dim, keepdim=True)
  assert not x_exp_sum.isnan().any().item(), "found NaN in x_exp_sum"
  quotient = x_exp/x_exp_sum
  assert not quotient.isnan().any().item(), "found NaN in quotient"
  return quotient

attention_probs: FloatTensor = softmax(attention_scores, dim=-1)
assert not attention_probs.isnan().any().item(), "found NaN in attention_probs"
```

interestingly, the NaNs are introduced when computing `diffs = x-maxes`. even though this really should just be computing lots of `0.0598-0.0598=0`s.

I have exported a realistic `attention_scores` tensor, but it's 3GB. not sure how to share a file that large. tell me if it'd be useful to have.

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0.dev20230310
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.3 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: version 3.22.1
Libc version: N/A

Python version: 3.11.0 | packaged by conda-forge | (main, Jan 14 2023, 12:26:40) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.3-arm64-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Max

Versions of relevant libraries:
[pip3] clip-anytorch==2.5.2
[pip3] numpy==1.24.2
[pip3] torch==2.1.0.dev20230310
[pip3] torchdiffeq==0.2.3
[pip3] torchsde==0.2.5
[pip3] torchvision==0.15.0.dev20230310
[conda] clip-anytorch             2.5.2                    pypi_0    pypi
[conda] numpy                     1.24.2                   pypi_0    pypi
[conda] pytorch                   2.1.0.dev20230310        py3.11_0    pytorch-nightly
[conda] sympy                     1.11.1                  py311_0    pytorch-nightly
[conda] torchdiffeq               0.2.3                    pypi_0    pypi
[conda] torchsde                  0.2.5                    pypi_0    pypi
[conda] torchvision               0.15.0.dev20230310          pypi_0    pypi
```

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev