Repository: pytorch
Issue: torch.nn.functional.scaled_dot_product_attention returns NaN values after backward pass.
Labels: needs reproduction,module: numerical-stability,triaged,module: sdpa

### 🐛 Describe the bug

When using `torch.nn.functional.scaled_dot_product_attention` with autograd a tensor filled with NaN values are returned after a few backward passes.  `Using torch.autograd.set_detect_anomaly(True)` returned this error message for `loss.backward()`:

```RuntimeError: Function 'MseLossBackward0' returned nan values in its 0th output.```

This issue has also been brought up by another user in the [forum](https://discuss.pytorch.org/t/nan-backward-for-scaled-dot-product-attention/193584).

Edit:
Reproducing code with attention class and forward function:
``` Python
class QKVAttention(nn.Module):
    def __init__(self,hidden_size,attn_mask=None,):
        super().__init__()
        self.hidden = hidden_size
        self.Q_linear = nn.Linear(hidden_size,hidden_size)
        self.K_linear = nn.Linear(hidden_size,hidden_size)
        self.V_linear = nn.Linear(hidden_size,hidden_size)
        self.attn_mask = attn_mask.unsqueeze(-2)

    def forward(self,x):
        x = x.unsqueeze(-2)
        
        Q = self.Q_linear(x)
        K = self.K_linear(x)
        V = self.V_linear(x)
      
        #scaled_dot_product_attention() is beta and subject to change as of 2024-05-17
        #problem was produced here:
        attn_out = nn.functional.scaled_dot_product_attention(Q,K,V,attn_mask=self.attn_mask,dropout_p=0.1,is_causal=False)
        
        return attn_out

class MainNet(nn.module):                             
    #init function was skipped
     def forward(self,input):
          h0 = torch.zeros(self.num_layers, self.batch_size,self.hidden_size, requires_grad=True).to(device)
          c0 = torch.zeros(self.num_layers,self.batch_size,self.hidden_size,requires_grad=True).to(device)
          
          packed_lstm_out,(_,_) = self.lstm1(input,(h0,c0))
          
          padded_lstm_out,_ = rnn.pad_packed_sequence(packed_lstm_out,padding_value=44444,batch_first=True)
          
          #creates a row-mask for the attention layer:
          positive_mask = padded_lstm_out == 44444
          positive_row_mask = torch.any(positive_mask,dim=-1,keepdim=True)
          attn_mask = ~positive_row_mask

          attention_layer = QKVAttention(self.hidden_size,attn_mask=attn_mask)
          context = attention_layer(padded_lstm_out)
          squuezed_context = context.squeeze(-2)

          final_output = self.fco(squuezed_context[:,-1,:])
        
          return final_output

```
Training  loop:
```  
optimizer  = torch.optim.AdamW(params=lstm_model.parameters(recurse=True),weight_decay=1e-2,lr=1e-4)
loss_fn = torch.nn.MSELoss()

def training_loop():
    model.train()
    torch.autograd.set_detect_anomaly(True)
    for batch_n, (X,y) in enumerate(train_dataloader):
        pred = model(X)
        
        #calculate loss & backward pass
        loss = loss_fn(pred,y)
        loss.backward()
        
        #adjust params
        optimizer.step()
        optimizer.zero_grad()
        
        #print loss
        if batch_n % report_freq == 0:
            current_values = batch_n*batch_size + len(X)
            print(f"Train loss: {running_loss:>7f} [{current_values:>5d}/{size:>5d}]")

```


### Versions

PyTorch version: 2.3.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.24.4
Libc version: N/A

Python version: 3.12.2 (main, Mar  1 2024, 19:22:10) [Clang 15.0.0 (clang-1500.1.0.2.5)] (64-bit runtime)
Python platform: macOS-14.1-arm64-arm-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

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.3.0
[pip3] torch-tb-profiler==0.4.3
[pip3] torchviz==0.0.2
[conda] Could not collect

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki