A problem about cuda graph when input and output are the same

hello_e: <pre data-code-wrap="python"><code class="lang-python">     def _forward(self, hidden_states_a, hidden_states_b, temb):
        xxxx
        return hidden_states_a, hidden_states_b

     def forward(self, hidden_states_a, hidden_states_b, temb):
        if self.cuda_graph is None:
            hidden_states_a_bk, hidden_states_b_bk, temb_bk = hidden_states_a.clone(), hidden_states_b.clone(), temb.clone()
            self.cuda_graph = torch.cuda.CUDAGraph()
            #input static placeholder
            self.graph_input['hidden_states_a'] = hidden_states_a_bk
            self.graph_input['hidden_states_b'] = hidden_states_b_bk
            self.graph_input['temb'] = temb_bk
            
            #capture the graph
            with torch.cuda.graph(self.cuda_graph):
                hidden_states_a_bk, hidden_states_b_bk = self._forward(hidden_states_a_bk, hidden_states_b_bk, temb_bk)
            torch.cuda.synchronize()

            # output placeholder
            self.graph_output['hidden_states_a'] = hidden_states_a_bk
            self.graph_output['hidden_states_b'] = hidden_states_b_bk

            return self.execute_model(hidden_states_a, hidden_states_b, temb)
        else:
            return self.execute_model(hidden_states_a, hidden_states_b, temb)
    def execute_model(self, hidden_states_a, hidden_states_b, temb):
        self.graph_input['hidden_states_a'].copy_(hidden_states_a)
        self.graph_input['hidden_states_b'].copy_(hidden_states_b)
        self.graph_input['temb'].copy_(temb)
        self.cuda_graph.replay()
        return self.graph_output['hidden_states_a'].clone(), self.graph_output['hidden_states_b'].clone()
</code></pre>
<p>Here is my code. In the original forward function, the inputs are <code>hidden_states_a</code>, <code>hidden_states_b</code>, and <code>temb</code>, and the outputs are <code>hidden_states_a</code>, <code>hidden_states_b</code>. That is to say, the first two inputs are also the outputs. I renamed the original forward function to <code>_forward</code>. In the new <code>forward</code> function, when it runs for the first time, <code>self.cuda_graph</code> is None. I clone all three input tensors, then put them into placeHolder. After that, I capture the CUDA graph through the <code>self._forward</code> function and put the outputs into placeHolder. In the <code>execute_model</code> function, I simply replay the CUDA graph. Are there any bugs in this approach? After writing it this way, the code runs faster, but the results are incorrect.</p>

ptrblck: <p>From the <a href="https://pytorch.org/docs/stable/notes/cuda.html#graph-memory-management">docs</a>:</p>
<blockquote>
<p>A captured graph acts on the same virtual addresses every time it replays. If PyTorch frees the memory, a later replay can hit an illegal memory access. If PyTorch reassigns the memory to new tensors, the replay can corrupt the values seen by those tensors. Therefore, the virtual addresses used by the graph must be reserved for the graph across replays.</p>
</blockquote>
<p>Did you check if <code>self.graph_output['hidden_states_X</code>]` changed after the replay?</p>

hello_e: <p>Hello,</p>
<p>I printed the value and address of <code>hidden_states_a_bk</code> before and after the graph capture, and then printed the value of the first tensor returned after executing the model, they are in the comment of the following code. In my <code>_forward</code> function, there is a scenario where <code>hidden_states_a = f(hidden_states_a)</code> occurs, and <em>there are some custom kernel and inplace kernel</em>. I’m wondering if this could be the cause of the bug. Thank you for your response! Below is my code, which is identical to the code described in the issue except for the added <code>print</code> statements.</p>
<pre data-code-wrap="python"><code class="lang-python">    def forward(self, hidden_states_a, hidden_states_b, temb):
        if self.cuda_graph is None:
            hidden_states_a_bk, hidden_states_b_bk, temb_bk = hidden_states_a.clone(), hidden_states_b.clone(), temb.clone()
            self.cuda_graph = torch.cuda.CUDAGraph()
            #input static placeholder
            self.graph_input['hidden_states_a'] = hidden_states_a_bk
            self.graph_input['hidden_states_b'] = hidden_states_b_bk
            self.graph_input['temb'] = temb_bk
            #before graph capture hidden_states_a_bk.float().abs().sum():3843304.5
            print(f"before graph capture hidden_states_a_bk.float().abs().sum():{hidden_states_a_bk.float().abs().sum()}")
            #before graph capture hidden_states_a_bk.data_ptr():139705777651712
            print(f"before graph capture hidden_states_a_bk.data_ptr():{hidden_states_a_bk.data_ptr()}")
            
            #capture the graph
            with torch.cuda.graph(self.cuda_graph):
                hidden_states_a_bk, hidden_states_b_bk = self._forward(hidden_states_a_bk, hidden_states_b_bk, temb_bk)
            torch.cuda.synchronize()
            
            #after graph capture hidden_states_a_bk.float().abs().sum():0.0
            print(f"after graph capture hidden_states_a_bk.float().abs().sum():{hidden_states_a_bk.float().abs().sum()}")
            #after graph capture hidden_states_a_bk.data_ptr():139705943326720
            print(f"after graph capture hidden_states_a_bk.data_ptr():{hidden_states_a_bk.data_ptr()}")
            
            # output placeholder
            self.graph_output['hidden_states_a'] = hidden_states_a_bk
            self.graph_output['hidden_states_b'] = hidden_states_b_bk

            res = self.execute_model(hidden_states_a, hidden_states_b, temb)

            #after replay res[0].float().abs().sum():3859856.75
            print(f"after replay res[0].float().abs().sum():{res[0].float().abs().sum()}")
            return res
        else:
            return self.execute_model(hidden_states_a, hidden_states_b, temb)
</code></pre>

hello_e: <p>yes I check, in the first execution, before replay it is zero, and after it is some incorrect value…</p>
<p>and I found that if I delete the custom kernel in the <code>_forward</code> function, everything just goes normally, is it just because torch cuda graph do not support custom kernel?</p>

hello_e: <p>finally i know where the problem is, in my custom kernel i did not pass the current cuda stream when launching the cuda kernel, after passing the current cuda stream, everything is normal…</p>

ptrblck: <p>Great debugging and thanks for the update!</p>