Torch compile: optimizer.step Generates Excessive Warning Messages

username: hassonofer
cooked: <p>I’m experiencing an issue with excessive warning messages when using <code>torch.compile()</code> with optimizer step function.</p>
<p>While compiling the optimizer step provides a significant speed improvement during training (particularly with AdamW), I’m encountering challenges when using it with a learning rate scheduler.</p>
<p>Following the approach outlined in this <a href="https://dev-discuss.pytorch.org/t/compiling-the-optimizer-with-pt2/1669/6" rel="noopener nofollow ugc">discussion</a> and the official <a href="https://pytorch.org/tutorials/recipes/compiling_optimizer_lr_scheduler.html" rel="noopener nofollow ugc">PyTorch tutorial</a>, I’ve wrapped the learning rate with a tensor to prevent recompilations. However, this results in thousands of warning messages being generated, such as:</p>
<pre data-code-wrap="python"><code class="lang-python">('Grad tensors ["L['self'].param_groups[394]['params'][0].grad"] will be copied during cudagraphs execution. If using cudagraphs and the grad tensor addresses will be the same across runs, use torch._dynamo.decorators.mark_static_address to elide this copy.',)
</code></pre>
<p>These warnings appear even when running the simple linear network example from the official tutorial.</p>
<p>Is there a recommended way to:</p>
<ol>
<li>Prevent these warnings from occurring?</li>
<li>Suppress these specific warning messages?</li>
<li>Handle this situation differently?</li>
</ol>
<p>Any guidance would be greatly appreciated.</p>

username: janeyx99
cooked: <p>This looks CUDA graphs related, cc <a class="mention" href="/u/mlazos">@mlazos</a> and <a class="mention" href="/u/elias_ellison">@Elias_Ellison</a>.</p>

username: mlazos
cooked: <p><a class="mention" href="/u/hassonofer">@hassonofer</a> can you share your code?</p>
<p>This warning can be suppressed by calling torch._dynamo.decorators.mark_static_address on each of the grads of the parameters. This is a warning because when running the optimizer standalone with cudagraphs, you can have very bad perf without doing this due to host-device copies of the gradients for every parameter, which can be quite large.</p>

username: hassonofer
cooked: <p>Thanks for your response.</p>
<p>I don’t have a minimal code example to share. It happens on my <a href="https://gitlab.com/birder/birder/" rel="noopener nofollow ugc">Birder project</a>.<br>
I’ll try to illustrate.<br>
I’m, training a swin transformer v2 (small). Using <code>torch.optim.AdamW</code> and a <code>torch.optim.lr_scheduler.CosineAnnealingLR</code> scheduler.</p>
<p>The model is being compiled without any special arguments:<br>
<a href="https://gitlab.com/birder/birder/-/blob/39693d818a0c7444843c4061ed35013a5bb631ee/birder/scripts/train.py#L199" rel="noopener nofollow ugc"><code>net = torch.compile(net)</code></a></p>
<p>While the optimizer step function set to compile without fullgraph (as documented):<br>
<a href="https://gitlab.com/birder/birder/-/blob/39693d818a0c7444843c4061ed35013a5bb631ee/birder/scripts/train.py#L225" rel="noopener nofollow ugc"><code>optimizer.step = torch.compile(optimizer.step, fullgraph=False)</code></a></p>
<p>Note: the current <code>main</code> branch does not convert the learning rate to torch.tensor, but the logging happens either way. The optimizer defined <a href="https://gitlab.com/birder/birder/-/blob/39693d818a0c7444843c4061ed35013a5bb631ee/birder/common/training_utils.py#L337" rel="noopener nofollow ugc">here</a>.</p>
<p>I think that the only thing which is not “standard” in this training code, is the way the parameter groups are being built. It’s a bit more complex in order to support features like layer-decay, custom weight decay per layer type, etc.<br>
The function defined <a href="https://gitlab.com/birder/birder/-/blob/39693d818a0c7444843c4061ed35013a5bb631ee/birder/common/training_utils.py#L179" rel="noopener nofollow ugc">here</a> while a specific key can be seen <a href="https://gitlab.com/birder/birder/-/blob/39693d818a0c7444843c4061ed35013a5bb631ee/birder/common/training_utils.py#L271" rel="noopener nofollow ugc">here</a>.</p>
<p>I assume going through all of this is too much of a hassle <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20">, sorry about that…</p>
<p>I’ll try to create a straight forward code to reproduce it tomorrow (it’s kinda late here now).</p>
<p>Thanks a lot.</p>

username: mlazos
cooked: <p>Alternatively you can just call torch._dynamo.decorators.mark_static_address on each of the .grad attributes of the parameters. This is the recommended way to remove the warning.</p>

username: hassonofer
cooked: <p>Does this mean, something like this ?</p>
<pre data-code-wrap="python"><code class="lang-python">for p in model.parameters()
    mark_static_address(p)
</code></pre>
<p>Is it safe ? generally model parameters have static address ?</p>
<p>Thanks again</p>

username: mlazos
cooked: <p>You should mark <code>p.grad</code> as static, you can verify that this is safe and the grads have the same address by printing <code>p.grad.data_ptr()</code> on each iter. This is generally safe across iterations.</p>

username: cora-codes
cooked: <p><a class="mention" href="/u/mlazos">@mlazos</a> this solution is quite awkward. Let’s imagine we want to compile our step function. We cannot mark our parameter’s gradients as static beforehand (they don’t exist) and we cannot place <code>mark_static_address</code> into our step function since it cannot be traced.</p>

username: mlazos
cooked: <p>I actually submitted this <a href="https://github.com/pytorch/pytorch/pull/146374" rel="noopener nofollow ugc">PR</a> to make it less spammy. It will now only show up if you set TORCH_LOGS=“perf_hints”</p>

username: cora-codes
cooked: <p>That’s great thank you!</p>
<p>However, I’d also love a solution to the issue I’ve pointed out - it seems pertinent that users have a way to get the best performance possible in an ergonomic/sane fashion.</p>

username: Tiru_B
cooked: <p>where exactly this copy host-device happens on grads, what i have seen is all other tensors have static address  for each iter except this grads,  am building similar thing fro custom hardware, based on adderss of the tensors i want avoid mem copies, but i was not able to fine where this grads are copied new tensors that will go as input to optimizer,    i,e backward graph outputs are updated to param.grad address, i also verified params.grad is address coming to optimizer for only first time , for second itr this addess diffr at optimizer input</p>

username: Lianghong_Huang
cooked: <p>When I run torchbench.py in training mode, some model can generate 3 compiled module: forward, backward and inference. But the others only generate forward and backward. I notice that for those which generate 3 compiled module, also have this warning message. I am confused about the relationship between them, and why they produce module file with “AOT ID: [‘1_inference’]” during training mode?</p>