Run backward on a different CUDA stream than forward

karttikeya_mangalam: <p>Hi,</p>
<p>I was wondering if it is possible to run forward on a block on one stream and backward on the same block on a different stream (but still on the same GPU)?</p>
<p>It seems like it used to possible before this PR: <a href="https://github.com/pytorch/pytorch/pull/8354" class="inline-onebox" rel="noopener nofollow ugc">Updates autograd engine to respect streams set in forward by mruberry · Pull Request #8354 · pytorch/pytorch · GitHub</a> I am wondering if there’s a way to support this use-case now?</p>

ptrblck: <p>Could you check if simply wrapping the <code>backward()</code> call into a separate stream would use it?<br>
You should be able to use e.g. Nsight Systems to check the streams.</p>

karttikeya_mangalam: <p>No, it still happens on the stream forward happens on. (check with nsys)</p>

ptrblck: <p>It seems you are right given <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-streams">this explanation from the docs</a>:</p>
<blockquote>
<p>Each backward CUDA op runs on the same stream that was used for its corresponding forward op. If your forward pass runs independent ops in parallel on different streams, this helps the backward pass exploit that same parallelism.</p>
</blockquote>
<p>What’s your use case that you would like to avoid using the same streams from the forward?</p>

shadowshadow: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="156238">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>What’s your use case that you would like to avoid using the same streams from the forward?</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> Hello, I have the same requirement: I want an operator’s forward and backward to run on different streams. I tried to achieve this with an autograd.Function, using <code>torch.cuda.set_stream</code> in both the forward and backward methods to switch streams, but it doesn’t work completely—PyTorch implicitly switches streams in some cases.</p>
<p>for example</p>
<pre data-code-wrap="python"><code class="lang-python">class ForwardStream(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_tensor: torch.Tensor, stream: torch.cuda.Stream) -&gt; torch.Tensor:
        stream.wait_stream(torch.cuda.current_stream())
        torch.cuda.set_stream(stream)
        return input_tensor

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -&gt; torch.Tensor:
        return grad_output

class BackwardStream(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_tensor: torch.Tensor, stream: torch.cuda.Stream) -&gt; torch.Tensor:
        ctx.stream = stream
        return input_tensor

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -&gt; torch.Tensor:
        stream = ctx.stream
        stream.wait_stream(torch.cuda.current_stream())
        torch.cuda.set_stream(stream)
        return grad_output

class TestModel(nn.Module):
    
    def __init__(self, input_size=5120, hidden_size=27648, output_size=5120):
        super().__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.layer2 = nn.Linear(hidden_size, hidden_size)
        self.layer3 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
        self.stream = create_external_stream() # a cuda stream
        self.default_stream = torch.cuda.current_stream()
        
        
        
    def forward(self, x):
        x = self.layer1(x)
        x = BackwardStream(x, self.default_stream) 
        x = self.relu(x) # relu backward is supposed to launch on stream, but actually on default stream?
        x = ForwardStream(x, self.stream)
        x = self.layer2(x)
        x = ForwardStream(x, self.default_stream)
        x = BackwardStream(x, self.stream)
        x = self.layer3(x)
        return x
  
</code></pre>
<p>In this example, I try to add autograd function to dispatch ops’<br>
forward and backward to different streams. relu backward is supposed to launch on stream, but actually on default stream.</p>
<p>Can you tell how to launch op’s forward and backward on different streams</p>