'MPS' training Issue(s) with NanoGPT: -Inf, NaN's
Labels: triaged, module: mps

### 🐛 Describe the bug

When I train Karpathy's NanoGPT model, setting device='mps,', after ~2500 iterations the weight's get corrupted with -Inf, NaNs.  If I reduce the block-size to 4, the model trains without errors.
The model trains without errors when device='cpu' and block-size doesn't matter.

https://github.com/karpathy/nanoGPT

This is running on MacOS 13.5 on a 2021 iMac 27" with an AMD Radeon Pro 5700 XT GPU.
The -Inf seem to start in this layer:
```
self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
```
model.py:        torch.autograd.set_detect_anomaly(True) #dbl 3/14/23
Is there a way to configure 'MPS' to generate trace logs of resource usage to see it's it's an 'MPS' resource exhaustion issue?

I reported this previously on the Pytorch forum.
https://discuss.pytorch.org/t/training-issues-in-mps-inf-and-nan/175244

```
/Users/davidlaxer/anaconda3/envs/AI-Feynman/bin/python /Users/davidlaxer/nanoGPT/train.py --dataset=covid --n_layer=16 --n_head=8 --n_embd=64 --device='mps' --compile=False --eval_iters=1 --block_size=1024 --batch_size=4 --max_iters=50000 
Overriding: dataset = covid
Overriding: n_layer = 16
Overriding: n_head = 8
Overriding: n_embd = 64
Overriding: device = mps
Overriding: compile = False
Overriding: eval_iters = 1
Overriding: block_size = 1024
Overriding: batch_size = 4
Overriding: max_iters = 50000
iter 2666: loss 10.0208, time 13285.35ms, mfu 0.02%
self.transformer.ln_f(x):   tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[-1.0961,  1.0357,  0.9074,  ...,  0.9867,  1.0914,  0.8752],
         [-1.1047,  1.0156,  1.0849,  ...,  0.8794,  1.0162,  0.7908],
         [-0.9748,  1.0310,  0.9873,  ...,  1.0442,  1.0484,  1.0544],
         ...,
         [-1.1226,  0.9061,  0.9063,  ...,  1.0940,  1.2296,  0.7446],
         [-1.1805,  0.9367,  0.8593,  ...,  0.6749,  0.5979,  0.7332],
         [-1.0284,  0.8252,  1.1170,  ...,  0.9860,  1.1770,  1.1917]]],
       device='mps:0', grad_fn=<NativeLayerNormBackward0>)
x:  tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[-1.0961,  1.0357,  0.9074,  ...,  0.9867,  1.0914,  0.8752],
         [-1.1047,  1.0156,  1.0849,  ...,  0.8794,  1.0162,  0.7908],
         [-0.9748,  1.0310,  0.9873,  ...,  1.0442,  1.0484,  1.0544],
         ...,
         [-1.1226,  0.9061,  0.9063,  ...,  1.0940,  1.2296,  0.7446],
         [-1.1805,  0.9367,  0.8593,  ...,  0.6749,  0.5979,  0.7332],
         [-1.0284,  0.8252,  1.1170,  ...,  0.9860,  1.1770,  1.1917]]],
       device='mps:0', grad_fn=<NativeLayerNormBackward0>)
logits = self.lm_head(x):  tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[-0.3500,  0.3934, -0.3290,  ..., -0.6095, -0.3002, -0.5246],
         [-0.3567,  0.3950, -0.3281,  ..., -0.6431, -0.3298, -0.5635],
         [-0.3567,  0.3357, -0.3592,  ..., -0.6269, -0.3543, -0.5235],
         ...,
         [-0.3366,  0.3330, -0.3424,  ..., -0.6295, -0.3873, -0.5330],
         [-0.3458,  0.4383, -0.3182,  ..., -0.6497, -0.3405, -0.5221],
         [-0.3601,  0.3462, -0.3220,  ..., -0.6174, -0.2938, -0.4807]]],
       device='mps:0', grad_fn=<LinearBackward0>)
loss:  tensor(nan, device='mps:0', grad_fn=<NllLossBackward0>)
X:  tensor([[  284,  1478,   309,  ...,   290,   281,   287],
        [  393,  4633,   329,  ...,   734, 27887,   355],
        [  257,  1598, 10690,  ...,    11,   290,  3298],
        [22443,   378,   357,  ..., 21966,  3047,    11]], device='mps:0')
Y:  tensor([[ 1478,   309,   423,  ...,   281,   287,    12],
        [ 4633,   329,  9471,  ..., 27887,   355,   880],
        [ 1598, 10690,   357,  ...,   290,  3298,  1143],
        [  378,   357,    50,  ...,  3047,    11, 11791]], device='mps:0')
Loss: tensor(nan, device='mps:0', grad_fn=<NllLossBackward0>)
X:  tensor([[  284,  1478,   309,  ...,   290,   281,   287],
        [  393,  4633,   329,  ...,   734, 27887,   355],
        [  257,  1598, 10690,  ...,    11,   290,  3298],
        [22443,   378,   357,  ..., 21966,  3047,    11]], device='mps:0')
Y:  tensor([[ 1478,   309,   423,  ...,   281,   287,    12],
        [ 4633,   329,  9471,  ..., 27887,   355,   880],
        [ 1598, 10690,   357,  ...,   290,  3298,  1143],
        [  378,   357,    50,  ...,  3047,    11, 11791]], device='mps:0')
logits: tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]],

        [[-0.3500,  0.3934, -0.3290,  ..., -0.6095, -0.3002, -0.5246],
         [-0.3567,  0.3950, -0.3281,  ..., -0.6431, -0.3298, -0.5635],
         [-0.3567,  0.3357, -0.3592,  ..., -0.6269, -0.3543, -0.5235],
         ...,
         [-0.3366,  0.3330, -0.3424,  ..., -0.6295, -0.3873, -0.5330],
         [-0.3458,  0.4383, -0.3182,  ..., -0.6497, -0.3405, -0.5221],
         [-0.3601,  0.3462, -0.3220,  ..., -0.6174, -0.2938, -0.4807]]],
       device='mps:0', grad_fn=<LinearBackward0>)
/Users/davidlaxer/anaconda3/envs/AI-Feynman/lib/python3.10/site-packages/torch/autograd/__init__.py:251: UserWarning: Error detected in LogSoftmaxBackward0. Traceback of forward call that caused the error:
  File "/Users/davidlaxer/nanoGPT/train.py", line 345, in <module>
    logits, loss = model(X, Y)
  File "/Users/davidlaxer/anaconda3/envs/AI-Feynman/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/Users/davidlaxer/anaconda3/envs/AI-Feynman/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/davidlaxer/nanoGPT/model.py", line 267, in forward
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
  File "/Users/davidlaxer/anaconda3/envs/AI-Feynman/lib/python3.10/site-packages/torch/nn/functional.py", line 3055, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
 (Triggered internally at /Users/davidlaxer/pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:119.)
  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
Traceback (most recent call last):
  File "/Users/davidlaxer/nanoGPT/train.py", line 381, in <module>
    scaler.scale(loss).backward() # dbl 3/14/23
  File "/Users/davidlaxer/anaconda3/envs/AI-Feynman/lib/python3.10/site-packages/torch/_tensor.py", line 491, in backward
    torch.autograd.backward(
  File "/Users/davidlaxer/anaconda3/envs/AI-Feynman/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Function 'LogSoftmaxBackward0' returned nan values in its 0th output.

Process finished with exit code 1

```

### Versions

```
% python collect_env.py
Collecting environment information...
PyTorch version: 2.1.0a0+git8472c24
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.5 (x86_64)
GCC version: Could not collect
Clang version: 14.0.6
CMake version: version 3.22.1
Libc version: N/A

Python version: 3.10.12 (main, Jul  5 2023, 15:58:50) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-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:
Intel(R) Core(TM) i7-10700K CPU @ 3.80GHz

Versions of relevant libraries:
[pip3] audiolm-pytorch==0.0.1
[pip3] configmypy==0.1.0
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.4
[pip3] pytorch-transformers==1.1.0
[pip3] tensorly-torch==0.4.0
[pip3] torch==2.0.1
[pip3] torch-cluster==1.6.1
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torch-spline-conv==1.2.2
[pip3] torch-struct==0.5
[pip3] torch-summary==1.4.5
[pip3] torch-utils==0.1.2
[pip3] torchaudio==2.1.0.dev20230630
[pip3] torchtraining-nightly==1604016577
[pip3] torchvision==0.16.0.dev20230630
[pip3] vector-quantize-pytorch==0.9.2
[conda] nomkl                     3.0                           0  
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpy-base                1.25.2          py310hd8f4981_0  
[conda] pytorch-transformers      1.1.0                    pypi_0    pypi
[conda] tensorly-torch            0.4.0                    pypi_0    pypi
[conda] torch                     2.1.0a0+git28720ad          pypi_0    pypi
[conda] torch-cluster             1.6.1                    pypi_0    pypi
[conda] torch-scatter             2.1.1                    pypi_0    pypi
[conda] torch-sparse              0.6.17                   pypi_0    pypi
[conda] torch-spline-conv         1.2.2                    pypi_0    pypi
[conda] torch-struct              0.5                      pypi_0    pypi
[conda] torch-summary             1.4.5                    pypi_0    pypi
[conda] torch-utils               0.1.2                    pypi_0    pypi
[conda] torchaudio                2.1.0.dev20230630          pypi_0    pypi
[conda] torchtraining-nightly     1604016577               pypi_0    pypi
[conda] torchvision               0.16.0.dev20230630          pypi_0    pypi
[conda] vector-quantize-pytorch   0.9.2                    pypi_0    pypi
```

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