pytorch / pytorch/pytorch

Mixing cuda graph and non cuda graph code works poorly in the presence of autograd.

Open
#164,302 2 comments 0 reactions 0 assignees View on GitHub
module: autograd module: cuda graphs triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

Consider this code snippet:

```
import torch

def test():
def func(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
z = x + y
return z * x

x = torch.randn([2, 2], device="cuda", requires_grad=True)
y = torch.randn_like(x, requires_grad=True)

side_stream = False
if side_stream: # success
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
y = y + 1
torch.cuda.current_stream().wait_stream(s)
else: # failure
y = y + 1

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, capture_error_mode="global"):
z = func(x, y)
z.sum().backward()

if __name__ == "__main__":
test()
```

if side_stream is False, this will fail with the following message:

```
cudaErrorStreamCaptureImplicit(906): operation would make the legacy stream depend on a capturing blocking stream
```

Setting side_stream=True will alleviate the issue by doing the relevant work in the non default stream.

As I understand it, this problem comes from the fact that, in autograd, Nodes have the current stream saved when they are constructed. when torch.Tensor.backward() is called, Autograd engine will then [retrieve this saved stream](https://github.com/pytorch/pytorch/blob/99e28ffab3b301980b8517e58a772b70d7ac539e/torch/csrc/autograd/engine.cpp#L1068-L1069) and ["wait on it"](https://github.com/pytorch/pytorch/blob/99e28ffab3b301980b8517e58a772b70d7ac539e/torch/csrc/autograd/engine.cpp#L1069) in order to make sure that the relevant backward computations are ordered after the relevant forward computations.

This normally makes sense from a program correctness point of view. However, if you have code in your forward pass that happened outside of a cuda graph, and this code was run on the "NULL stream", then you will end up joining the null stream against a capturing stream, which is disallowed in CUDA because of the weird semantics that that null stream synchronizes with every other stream that was not created with cudaStreamNonBlocking:

> Stream capture can be used on any CUDA stream except cudaStreamLegacy (the “NULL stream”).

Source: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#creating-a-graph-using-stream-capture

Note that the default stream in pytorch is in fact the null stream: https://github.com/pytorch/pytorch/blob/1f1de20ba965221713a5736dba356539fd6a7217/c10/cuda/CUDAStream.cpp#L239

https://github.com/pytorch/pytorch/blob/1f1de20ba965221713a5736dba356539fd6a7217/c10/cuda/CUDAStream.cpp#L278-L288

Here is another example of the same issue, though this one is more subtle:

```
# torchrun --nproc_per_node=1 010.ddp.py

import os
import torch
from torch.nn.parallel import DistributedDataParallel as DDP

def test():
torch.distributed.init_process_group(backend='nccl')
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
torch.distributed.barrier()
model = torch.nn.Linear(128, 128).cuda()
x = torch.randn((32, 128), device="cuda")
stream = torch.cuda.Stream()

side_stream = False
if side_stream: # success
with torch.cuda.stream(stream):
model = DDP(model)
else: # failure
model = DDP(model)

# warmup
for _ in range(15):
model.zero_grad()
y = model(x)
y.sum().backward()

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, capture_error_mode="global"):
model.zero_grad()
y = model(x)
y.sum().backward()

if __name__ == "__main__":
test()
```

So what is the fix?

My hunch is that it is very, very simple. We just need to have special case logic right here: https://github.com/pytorch/pytorch/blob/99e28ffab3b301980b8517e58a772b70d7ac539e/torch/csrc/autograd/engine.cpp#L1069

If the parent stream is the null stream, then simply create a new stream, and then do work in that stream instead. The reality is that, if a user has done work on the null stream before doing stream capture, then that work in the null stream must be ordered before any launches of that cuda graph anyway because you can only add new work to the end of a cuda stream, and never rearrange work on a cuda stream. So I think we are basically guaranteed that

Unfortunately, that is easier said than done because autograd is part of libtorch_cpu.so, which does not depend upon libc10_cuda.so or libtorch_cuda.so. So we need to expose two new virtual function in the c10::Stream class: `bool is_null_stream() const` and `bool is_capturing_to_graph() const`.

I don't think the second function is contentious, since XPU and mindspore both have a concept of "stream capture to a cuda graph": https://github.com/pytorch/pytorch/issues/158827#issuecomment-3281377157 Thus, it is not like only cuda and hip will depend upon this function.

The first function is a bit more annoying, though. I'm not sure if any other implementation has a concept of a null stream, though we can always just return false in those cases.

I will try this out.

cc @ezyang @albanD @gqchen @nikitaved @soulitzer @Varal7 @xmfan @mcarilli @eellison @penguinwu @BoyuanFeng

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.