Repository: pytorch
Issue: RuntimeError: philox_cuda_state for an unexpected CUDA generator used during capture
Labels: triaged,module: random,module: cuda graphs

### 🐛 Describe the bug

I have created a custom GraphCudaCapture Module but it raised the error below. I wonder why the generator needs to be defined and fixed. 

```python
@dataclass
class CudaGraphRecord:

    graph = None
    args = None
    kwargs = None
    output = None
    cuda_graph_created: bool = False
    enable_cuda_graph: bool = True


class CudaGraphBatchRecord(dict):

    def __init__(self, enable_cuda_graph):
        super().__init__()
        self.enable_cuda_graph = enable_cuda_graph

    def __getitem__(self, key):
        if key not in self:
            self[key] = CudaGraphRecord(enable_cuda_graph=self.enable_cuda_graph)
        for item, value in self.items():
            if item == key:
                return value
        raise Exception()


class CudaGraphInferenceModule(torch.nn.Module):

    # Inspired from https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/model_implementations/diffusers/vae.py

    inference_methods = ["forward"]

    def __init__(self, module, enable_cuda_graph = True):
        super().__init__()
        self.module = module
        self.module.requires_grad_(requires_grad=False)
        self.module.to(memory_format=torch.channels_last)
        self.cuda_graph_records = {}

        for method_name in self.inference_methods:
            fn = getattr(self, f"_{method_name}")
            assert fn
            self.cuda_graph_records[method_name] = CudaGraphBatchRecord(enable_cuda_graph=enable_cuda_graph)
            setattr(self, method_name, partial(self._apply_fn, fn=fn, graph_record=self.cuda_graph_records[method_name]))

    def __getattr__(self, key):
        if hasattr(self._modules["module"], key) and key not in self.__dict__:
            return getattr(self._modules["module"], key)
        if key == "module":
            return self._modules["module"]
        return object.__getattribute__(self, key)

    def _graph_replay(self, graph_record, *args, **kwargs):
        for i in range(len(args)):
            if torch.is_tensor(args[i]):
                graph_record.args[i].copy_(args[i])
        for k in kwargs:
            if torch.is_tensor(kwargs[k]):
                graph_record.kwargs[k].copy_(kwargs[k])
        graph_record.graph.replay()
        return graph_record.output

    def extract_batch_size(self, *args, **kwargs) -> int:
        raise NotImplementedError

    def _apply_fn(self, *args, fn=None, graph_record=None,  **kwargs):
        batch_size = self.extract_batch_size(*args, **kwargs)
        if graph_record[batch_size].enable_cuda_graph:
            if graph_record[batch_size].cuda_graph_created:
                outputs = self._graph_replay(graph_record[batch_size], *args, **kwargs)
            else:
                self._create_cuda_graph(fn, graph_record[batch_size], *args, **kwargs)
                outputs = self._graph_replay(graph_record[batch_size], *args, **kwargs)
            return outputs
        else:
            return self._forward(*args, **kwargs)

    def _create_cuda_graph(self, fn, graph_record, *args, **kwargs):
        # Warmup to create the workspace and cublas handle
        cuda_stream = torch.cuda.Stream()
        cuda_stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(cuda_stream):
            for _ in range(3):
                fn(*args, **kwargs)
        torch.cuda.current_stream().wait_stream(cuda_stream)

        # Capture inputs to the graph
        graph_record.graph = torch.cuda.CUDAGraph()
        graph_record.args = args
        graph_record.kwargs = kwargs

        with torch.cuda.graph(graph_record.graph):
            # Store output
            graph_record.output = fn(*graph_record.args, **graph_record.kwargs)

        graph_record.cuda_graph_created = True
```

```python
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/uvicorn/protocols/http/httptools_impl.py", line 419, in run_asgi
    result = await app(  # type: ignore[func-returns-value]
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/uvicorn/middleware/proxy_headers.py", line 78, in __call__
    return await self.app(scope, receive, send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/fastapi/applications.py", line 270, in __call__
    await super().__call__(scope, receive, send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/applications.py", line 124, in __call__
    await self.middleware_stack(scope, receive, send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/middleware/errors.py", line 184, in __call__
    raise exc
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/middleware/errors.py", line 162, in __call__
    await self.app(scope, receive, _send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/middleware/exceptions.py", line 79, in __call__
    raise exc
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/middleware/exceptions.py", line 68, in __call__
    await self.app(scope, receive, sender)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/fastapi/middleware/asyncexitstack.py", line 21, in __call__
    raise e
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/fastapi/middleware/asyncexitstack.py", line 18, in __call__
    await self.app(scope, receive, send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/routing.py", line 706, in __call__
    await route.handle(scope, receive, send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/routing.py", line 276, in handle
    await self.app(scope, receive, send)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/routing.py", line 66, in app
    response = await func(request)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/fastapi/routing.py", line 235, in app
    raw_response = await run_endpoint_function(
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/fastapi/routing.py", line 163, in run_endpoint_function
    return await run_in_threadpool(dependant.call, **values)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/starlette/concurrency.py", line 41, in run_in_threadpool
    return await anyio.to_thread.run_sync(func, *args)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/anyio/to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/anyio/_backends/_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/anyio/_backends/_asyncio.py", line 867, in run
    result = context.run(func, *args)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/lightning/app/components/serve/python_server.py", line 253, in predict_fn
    return self.predict(request)
INFO:   File "app.py", line 34, in predict
    images = self._model.predict_step(prompts=texts, batch_idx=0)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/lightning.py", line 92, in predict_step
    samples_ddim, _ = self.sampler.sample(
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/sampler.py", line 85, in sample
    x = dpm_solver.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=2, lower_order_final=True)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/dpm_solver.py", line 1057, in sample
    model_prev_list.append(self.model_fn(x, vec_t))
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/dpm_solver.py", line 372, in model_fn
    return self.data_prediction_fn(x, t)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/dpm_solver.py", line 356, in data_prediction_fn
    noise = self.noise_prediction_fn(x, t)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/dpm_solver.py", line 350, in noise_prediction_fn
    return self.model(x, t)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/dpm_solver.py", line 311, in model_fn
    noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/dpm_solver.py", line 264, in noise_pred_fn
    output = model(x, t_input, cond, **model_kwargs)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/dpm_solver/sampler.py", line 75, in <lambda>
    lambda x, t, c: self.model.apply_model(x, t, c),
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/models/diffusion/ddpm.py", line 858, in apply_model
    x_recon = self.model(x_noisy, t, **cond)
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/deepspeed_replace.py", line 83, in _apply_fn
    outputs = self._graph_replay(graph_record[batch_size], *args, **kwargs)
INFO:   File "/content/venv/lib/python3.8/site-packages/ldm/deepspeed_replace.py", line 73, in _graph_replay
    graph_record.graph.replay()
INFO:   File "/home/zeus/.local/lib/python3.8/site-packages/torch/cuda/graphs.py", line 87, in replay
    super(CUDAGraph, self).replay()
INFO: RuntimeError: philox_cuda_state for an unexpected CUDA generator used during capture. In regions captured by CUDA graphs, you may only use the default CUDA RNG generator on the device that's current when capture begins. If you need a non-default (user-supplied) generator, or a generator on another device, please file an issue.
ERROR:    Exception in ASGI application
```

### Versions

```bash
Collecting environment information...
PyTorch version: 1.13.1+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 14 2022, 12:59:47)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-135-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090
GPU 2: NVIDIA GeForce RTX 3090
GPU 3: NVIDIA GeForce RTX 3090
GPU 4: NVIDIA GeForce RTX 3090
GPU 5: NVIDIA GeForce RTX 3090
GPU 6: NVIDIA GeForce RTX 3090

Nvidia driver version: 520.61.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.24.0
[pip3] open-clip-torch==2.7.0
[pip3] pytorch-lightning==1.8.6
[pip3] torch==1.13.1+cu116
[pip3] torchmetrics==0.11.0
[pip3] torchvision==0.14.1+cu116
[conda] Could not collect
```

cc @pbelevich @mcarilli @ezyang