Repository: pytorch
Issue: CUDA graph capturing fails for nn.Embedding and large batch sizes
Labels: module: cuda,triaged,module: embedding,module: cuda graphs

### 🐛 Describe the bug

Capturing CUDA graphs fails with a somewhat unspecific error when using `nn.Embedding` (and back-propagating through it) with batch sizes larger than 3072. I assume that this is because of an internal optimization in the respective CUDA kernel, which [performs sorting if more than 3072 inputs are used](https://github.com/pytorch/pytorch/blob/1cafb1027f223f2174f842945dd337cfa0fc120e/aten/src/ATen/native/cuda/Embedding.cu#L262). The (truncated) backtrace when this error is encountered with `CUDA_LAUNCH_BLOCKING=1` looks like this:

```
  [...]
  File "[...]/test_graph.py", line 785, in test
    loss.backward()
  File "[...]/lib/python3.10/site-packages/torch/_tensor.py", line 363, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "[...]/lib/python3.10/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: unique_by_key: failed to synchronize: cudaErrorStreamCaptureUnsupported: operation not permitted when stream is capturing
```

For reference, here's a testing script with which I determined the threshold:
```py
import torch as th
from torch import nn, optim
from bisect import bisect_left

model = nn.Embedding(5, 30).cuda()
opt = optim.Adam(model.parameters())
inp = (th.arange(0, 10000) % 5).cuda()

def test(N):
    opt.zero_grad(set_to_none=True)
    out = model(inp[:N])
    loss = out.mean()
    loss.backward()
    return None

def capture(N):
    th.cuda.synchronize()
    s = th.cuda.Stream()
    s.wait_stream(th.cuda.current_stream())
    with th.cuda.stream(s):
        for _ in range(3):
            test(N)
    th.cuda.current_stream().wait_stream(s)

    graph = th.cuda.CUDAGraph()
    with th.cuda.graph(graph):
        res = test(N)

def try_capture(N):
    print(f'capture {N}')
    try:
        capture(N)
    except:
        print(f'failed {N}')
        return 2
    print(f'ok {N}')
    return 0

thres = bisect_left(list(range(inp.shape[0])), 1, key=lambda x: try_capture(x))
print(f'>> threshold {thres}')
```

As far as I understand, the sorting optimization creates dynamically-sized tensors which are indeed [not supported in CUDA graphs](https://pytorch.org/docs/1.11/notes/cuda.html#constraints). I would see several possibilities to address this:
- The optimization could be disabled with an additional argument to `nn.Embedding()` and `F.embedding()`
- An exception could be raised if a CUDA graph capture is underway and the threshold for sorting inputs is reached.
- As a minimum, refer to this (and similar?) optimizations in the constraints section for the CUDA graph docs.

### Versions

```
Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.31

Python version: 3.10.0 (default, Mar  3 2022, 09:58:08) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-81-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration:
GPU 0: Quadro GP100
GPU 1: Quadro GP100

Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Versions of relevant libraries:
[pip3] mypy==0.961
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.0
[pip3] pytorch3d==0.6.2
[pip3] torch==1.11.0
[pip3] torchaudio==0.11.0
[pip3] torchfile==0.1.0
[pip3] torchvision==0.12.0
[pip3] torchviz==0.0.2
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            12_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            12_linux64_mkl    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-fft                   1.3.1                    pypi_0    pypi
[conda] mkl-random                1.2.2                    pypi_0    pypi
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0
[conda] mkl_random                1.2.2           py310h00e6091_0
[conda] numpy                     1.23.0                   pypi_0    pypi
[conda] numpy-base                1.22.3          py310h9585f30_0
[conda] pytorch                   1.11.0          py3.10_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch3d                 0.6.2                    pypi_0    pypi
[conda] torch                     1.11.0                   pypi_0    pypi
[conda] torchaudio                0.11.0                   pypi_0    pypi
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchvision               0.12.0                   pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi
```

cc @ngimel @mcarilli @ezyang