Repository: pytorch
Issue: RuntimeError: derivative for aten::_scaled_dot_product_efficient_attention_backward is not implemented
Labels: triaged,module: sdpa

### 🐛 Describe the bug

Hi, 
I run the code of a WGAN-GP model，it report an error when backwarding gradient_penalty. 
I have rarely seen this error before and there is hardly any search result for a solution.

```
    def calculate_gp(self, text, real_x, fake_x, mask):
        alpha = torch.rand(real_x.shape[0], 1, 1, 1).to(self.device)
        x = (alpha * real_x + (1 - alpha) * fake_x).requires_grad_(True)
        d = self.Discriminator(text, x, mask)
        fake = torch.ones_like(d).to(self.device)
        g = torch.autograd.grad(
            outputs=d,
            inputs=x,
            grad_outputs=fake,
            create_graph=True,
            retain_graph=True
        )[0]
        g = g.contiguous().view(real_x.shape[0], -1)
        gp = torch.mean((torch.sqrt(torch.sum(g ** 2, dim=1) + 1e-12) - 1) ** 2)
        return gp

    def training_step(self, batch, batch_idx):
        motion, cond = batch
        cond['y'] = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in cond['y'].items()}
        text = cond['y']['text']
        lengths = cond['y']['lengths']
        mask = self.creat_mask(lengths, self.seq_len, True)
        batch_size = motion.size(0)
        
        # optimizor
        g_opt, d_opt = self.optimizers() # manual optimizers
        g_lr_schedule, d_lr_schedule = self.lr_schedulers()

        # -------------------        
        # Train Discriminator
        # -------------------  
        for i in range(self.critic_iter):
            d_opt.zero_grad()  
            noise = torch.randn(motion.size(), device=self.device)
            gen_motion = self.Generator(text, noise, mask)
            
            real_validity = self.Discriminator(text, motion, mask) # expect 1
            fake_validity = self.Discriminator(text, gen_motion.detach(), mask) # expect 0
            gradient_penalty = self.calculate_gp(text, motion, gen_motion, mask)
            
            fake_labels = self.get_labels(batch_size, True)
            real_labels = self.get_labels(batch_size, False)
            
            d_loss = - (torch.mean(real_validity - real_labels)) + torch.mean(fake_validity - fake_labels) + gradient_penalty
            
            self.manual_backward(d_loss)
            d_opt.step() 

        # ---------------
        # Train Generator
        # ---------------
        g_opt.zero_grad()
           
        fake_validity = self.Discriminator(text, gen_motion.detach(), mask)
        labels = self.get_labels(batch_size, True)
        
        g_loss = -torch.mean(fake_validity - labels)
        
        self.manual_backward(g_loss)
        g_opt.step()
        
        # log and save
        g_lr_schedule.step()
        d_lr_schedule.step()
        self.log_dict({"G_loss": g_loss.item(), "D_loss": d_loss.item()})  
```

output of terminal
```
Create Model: motion_WGAN_GP
Laod Dataset: humanml3D
Loading dataset humanml3D ...
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 23384/23384 [00:14<00:00, 1643.63it/s]
Pointer Pointing at 0
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
You are using a CUDA device ('NVIDIA GeForce RTX 3090') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]

  | Name          | Type          | Params
------------------------------------------------
0 | Generator     | Generator     | 153 M 
1 | Discriminator | Discriminator | 256 M 
------------------------------------------------
107 M     Trainable params
302 M     Non-trainable params
409 M     Total params
1,638.659 Total estimated model params size (MB)
Epoch 0:   0%|                                                                                                                                                                            | 0/255 [00:00<?, ?it/s]Traceback (most recent call last):
  File "trainer.py", line 42, in <module>
    main()
  File "trainer.py", line 38, in main
    trainer.fit(model=model, train_dataloaders=train_loader)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 544, in fit
    call._call_and_handle_interrupt(
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 580, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 989, in _run
    results = self._run_stage()
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1035, in _run_stage
    self.fit_loop.run()
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 202, in run
    self.advance()
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 359, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 136, in run
    self.advance(data_fetcher)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 242, in advance
    batch_output = self.manual_optimization.run(kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/manual.py", line 92, in run
    self.advance(kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/loops/optimization/manual.py", line 112, in advance
    training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 309, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/strategies/strategy.py", line 382, in training_step
    return self.lightning_module.training_step(*args, **kwargs)
  File "/home/lebin/DreamMotion/model/motion_WGAN_GP.py", line 89, in training_step
    self.manual_backward(d_loss)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/core/module.py", line 1059, in manual_backward
    self.trainer.strategy.backward(loss, None, *args, **kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/strategies/strategy.py", line 204, in backward
    self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/plugins/precision/precision.py", line 69, in backward
    model.backward(tensor, *args, **kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/lightning/pytorch/core/module.py", line 1078, in backward
    loss.backward(*args, **kwargs)
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/home/lebin/miniconda3/envs/DreamMotion/lib/python3.8/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: derivative for aten::_scaled_dot_product_efficient_attention_backward is not implemented
```

### Versions

`Collecting environment information...
```
PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.8.18 (default, Sep 11 2023, 13:40:15)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 12.2.140
Nvidia driver version: 535.129.03

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] pytorch-lightning==2.1.3
[pip3] torch==2.1.2
[pip3] torchaudio==2.1.0
[pip3] torchmetrics==1.2.1
[pip3] torchvision==0.16.0
[pip3] triton==2.1.0
[conda] blas                      1.0                         mkl  
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0            py38h5eee18b_1  
[conda] mkl_fft                   1.3.8            py38h5eee18b_0  
[conda] mkl_random                1.2.4            py38hdb19cb5_0  
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] pytorch                   2.1.2           py3.8_cuda12.1_cudnn8.9.2_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_5    pytorch
[conda] pytorch-lightning         2.1.3              pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.1.0                    pypi_0    pypi
[conda] torchmetrics              1.2.1              pyhd8ed1ab_0    conda-forge
[conda] torchtriton               2.1.0                      py38    pytorch
[conda] torchvision               0.16.0                   pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi`
```




cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @mikaylagawarecki @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519