CUDA Graph Error with Residual Connections in `torch.compile` (RuntimeError: accessing tensor output of CUDAGraphs)

lixt: <p>Hey,<br>
I’m encountering a persistent <code>RuntimeError: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run</code> when using <code>torch.compile</code> (specifically with <code>mode="max-autotune"</code>) on my transformer model. This error appears to be related to residual connections within my <code>Block</code> module, even after implementing common solutions.</p>
<p><strong>Environment:</strong></p>
<ul>
<li>PyTorch version: <code>2.7.1+cu126</code></li>
<li>CUDA version: <code>12.9</code></li>
<li>GPU: <code>NVIDIA GeForce GTX 1650 Mobile / Max-Q</code></li>
<li>Operating System: <code>Linux Mint 22.1 x86_64</code></li>
<li>Python version: <code>3.12</code></li>
<li><code>bitsandbytes</code> is used for 4-bit linear layers.</li>
</ul>
<p><strong>Model Architecture Snippets (simplified for relevance):</strong></p>
<pre data-code-wrap="python"><code class="lang-python">class Block(nn.Module):
    def __init__(self, d_model: int, n_heads: int, ffn_hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.attention_norm = nn.RMSNorm(d_model)
        self.self_attn = MultiHeadAttention(n_embd, n_heads, block_size, dropout)
        self.ffn_norm = nn.RMSNorm(d_model)
        self.ffn = FeedForward(d_model, ffn_hidden_dim, dropout)
        self.residual_dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        # First residual connection
        x = x + self.residual_dropout(self.self_attn(self.attention_norm(x)))
        # Second residual connection
        x = x + self.residual_dropout(self.ffn(self.ffn_norm(x)))
        return x

class LanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        ffn_hidden_dim = int(4 * n_embd * (2/3))
        self.blocks = nn.ModuleList([
            Block(d_model=n_embd, n_heads=n_head, ffn_hidden_dim=ffn_hidden_dim, dropout=dropout)
            for _ in range(n_layer)
        ])
        self.ln_f = nn.RMSNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.token_embedding_table.weight = self.lm_head.weight

    def forward(self, idx, targets=None, smoothing=0.0):
        x = self.token_embedding_table(idx)
        for block in self.blocks:
            x = block(x) # This loop is likely where the issue originates within the graph
        x = self.ln_f(x)
        logits = self.lm_head(x)
        # ... loss calculation ...
        return logits, loss
</code></pre>
<p><strong>Training Loop Snippet:</strong></p>
<pre data-code-wrap="python"><code class="lang-python"># ... (inside main function) ...
model = LanguageModel().to(device)
model = torch.compile(model, mode="max-autotune") # Compilation happens here

    optimizer.zero_grad(set_to_none=True)

while iter_num &lt; max_iters:
    try:
        xb, yb = next(train_iter)
    except StopIteration:
        train_iter = iter(train_loader)
        xb, yb = next(train_iter)

    xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)

    torch.compiler.cudagraph_mark_step_begin()

    with autocast(device_type=device, dtype=pt_dtype):
        _, loss = model(xb, yb, smoothing=label_smoothing)
        loss = loss / gradient_accumulation_steps

    scaler.scale(loss).backward()
</code></pre>
<p><strong>Problem Description:</strong></p>
<p>I’m encountering</p>
<pre><code class="lang-auto">RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. Stack trace: File "/home/lixt/Desktop/ㅤ/projects/AI/CHAT/train.py", line 246, in forward
    norm_x = self.attention_norm(x)
  File "/home/lixt/.local/lib/python3.12/site-packages/torch/nn/functional.py", line 2929, in rms_norm
    return torch.rms_norm(input, normalized_shape, weight, eps). To prevent overwriting, clone the tensor outside of torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.
</code></pre>
<p>every time I use <code>mode='max-autotune'</code> or <code>reduce-overhead</code>.</p>
<p>The stack trace points to RMSNorm inside the Block.</p>
<p><strong>What I’ve tried (and did not fix it):</strong></p>
<ol>
<li>
<p><strong>Refactoring Block.forward</strong> : Changed h = x + … to x = x + … to ensure explicit data flow and avoid potential in-place issues within the compiled graph.</p>
</li>
<li>
<p><strong>torch.compiler.cudagraph_mark_step_begin()</strong> : Placed this at the beginning of each training step before the model call, which is standard practice for CUDA graphs and torch.compile.</p>
</li>
<li>
<p><strong>Adding .clone()</strong> : I initially tried adding .clone() on the x input to the block or before the residual addition, but this did not resolve the issue and often resulted in other graph-related errors or performance degradation.</p>
</li>
</ol>
<p>Given that the error persists even after restructuring the Block’s forward pass to avoid explicit intermediate variables and using <code>cudagraph_mark_step_begin</code>, I’m looking for help to understand why torch.compile is still perceiving an overwrite issue within the graph, especially when the stack trace points to the input of a normalized layer and how I could fix it. Thank you.</p>