GradScaler.unscale_, autograd.grad and second differentiation

username: morg
cooked: <p>Hi!</p>
<p>I have come across a Mixed Precision issue while testing the AdaHessian optimizer (<a href="https://arxiv.org/abs/2006.00719" rel="nofollow noopener">https://arxiv.org/abs/2006.00719</a>), a second-order optimizer, which I think is arising out of <code>Gradscaler.unscale_</code></p>
<h3>This <a href="https://colab.research.google.com/drive/1UKvlAAvchsTpVKvQtc6AsdHrGMtSDQsj?usp=sharing" rel="nofollow noopener">colab working example</a> shows the results of using <code>unscale_</code> or just manually unscaling</h3>
<h2>AdaHessian</h2>
<p>A core element AdaHessian is calculating <code>hvs</code> as below after scaling up the loss and performing the backward pass. The gradients need to be unscaled before being passed to <code>autograd.grad</code>.</p>
<pre><code class="lang-auto">hvs = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, retain_graph=False)
</code></pre>
<h2>Using GradScaler.unscale_ produces <code>nan</code>
</h2>
<p>However despite manually verifying (printing min/mean of p.grad) that the gradients being passed to <code>autograd.grad</code> have been scaled down via <code>scaler.unscale_(optimizer)</code> it seems like <strong><code>autograd.grad</code> doesn’t “respond”(?) to this unscaling</strong> and subsequently produces <code>nan</code> or very high values for <code>hvs</code></p>
<h2>Manual unscaling works</h2>
<p>However if, after gathering the gradients from the optimizer, I <strong>manually</strong> unscale them using the same scaling factor used in <code>GradScaler</code>, the calculation of <code>hvs</code> works as expected</p>
<h2>Pseudo code</h2>
<p>(see colab link for full version)</p>
<pre><code class="lang-auto"># AUTOCAST
with autocast(enabled=fp16_enabled):   
    output = net(input)
    loss = criterion(output, target)
      
# SCALE UP
scaler.scale(loss).backward(create_graph=True)

# SCALE DOWN
if not manual_unscale:
    scaler.unscale_(optimizer)     # UNSCALE GRADS WITH GRADSCALER   

# GATHER GRADSH AND PARAMS
gradsH = []
params = optimizer.param_groups[0]['params']
for p in optimizer.param_groups[0]['params']:
    gradsH.append(0. if p.grad is None else p.grad + 0.)

# MANUALLY UNSCALE (if manual_unscale==True)
if manual_unscale:    
    gradscaler_scale = scaler.get_scale()
    for g in gradsH: 
        g.div_(gradscaler_scale/scale_reduction)

# CALCULATE hvs
hvs = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, retain_graph=False)
</code></pre>
<h2>Question</h2>
<p>Does using <code>GradScaler.unscale_</code> to unscale the gradients mean that this unscaling operation is not "recorded* " in the graph? Whereas gathering the gradients from the optimizer and manually unscaling them is “recorded”?</p>
<p>*(sorry not sure what the correct terminology is here)</p>
<p>Is this behaviour expected and is there a workaround to unscaling that involves using <code>GradScaler</code> or should I just stick with manually unscaling the gradients i have gathered?</p>
<h2>Comment</h2>
<p>(I know that if I manually unscale then that unscaling will not be recorded by <code>GradScaler</code> and the optimizer gradients will be subsequently unscaled in <code>GradScaler.step</code>, however this should be ok I think as I will have manually unscaled a new list of gradients and the gradients in the optimizer will still be untouched by this manual scaling)</p>
<p>Thanks for reading this far, looking forward to hearing what anyone thinks!</p>

username: ptrblck
cooked: <p>How did you define <code>scale_reduction</code> as this seems to be the main difference between both approaches, no?</p>

username: morg
cooked: <p>Oh sorry meant to remove that, it is just set manually for now to be able to test different down-scaling values, <code>scale_reduction=1</code> for now</p>

username: mcarilli
cooked: <p><code>scaler.unscale_(optimizer)</code> unscales the <code>.grad</code> attributes of all params owned by <code>optimizer</code>, after those <code>.grads</code> have been fully accumulated for those parameters this iteration and are about to be applied.  If you intend to accumulate more gradients into <code>.grad</code>s later in the iteration, <code>scaler.unscale_</code> is premature.</p>
<p>Also, the unscale+inf/nan check kernel used by <code>scaler.unscale_</code> is not autograd-aware.</p>
<p>In use cases I’ve seen, creating out-of-place gradients via <code>torch.autograd.grad</code> with <code>create_graph=True</code> is typically a setup for a double-backward that will accumulate gradients into the <code>param.grad</code> attributes.  For these cases, manually unscaling the out-of-place grads before using them to set up the double-backward is expected:  see the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty" rel="nofollow noopener">gradient penalty example</a>.</p>
<p>In your snippet the control flow looks strange to me:</p>
<ol>
<li>first you run <code>backward(create_graph=True)</code> that accumulates gradients into <code>param.grad</code>s.</li>
<li>then you collect <code>gradsH.append(0. if p.grad is None else p.grad + 0.)</code>, presumably the p.grad+0. ensures gradsH are separate tensors from <code>param.grad</code>s (but if so why not also use an out-of-place <code>torch.autograd.grad</code> to create gradsH?)</li>
<li>Then you use those <code>.grad</code>s as <code>torch.autograd.grad</code> graph roots to create <code>hvs</code> out of place.</li>
</ol>
<p>Your optimizer is vanilla SGD and you don’t appear to <em>use</em> <code>hvs</code>, so it’s hard for me to recommend a way to create it.</p>
<p>FYI <code>GradScaler</code> also has an <code>enabled</code> argument, ie</p>
<pre><code class="lang-auto">      if fp16_enabled: 
        scaler.scale(loss).backward(create_graph=True)
      else: 
        loss.backward(create_graph=True)
</code></pre>
<p>is equivalent to</p>
<pre><code class="lang-auto">  scaler = GradScaler(init_scale=init_scale)
  ...
  scaler.scale(loss).backward(create_graph=True)
</code></pre>

username: morg
cooked: <p>Thank you for all this, it really helps alot!</p>
<aside class="quote no-group" data-username="mcarilli" data-post="4" data-topic="95953">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mcarilli/48/28517_2.png" class="avatar"> mcarilli:</div>
<blockquote>
<p>manually unscaling the out-of-place grads before using them to set up the double-backward is expected</p>
</blockquote>
</aside>
<p>Good to know, thanks</p>
<aside class="quote no-group" data-username="mcarilli" data-post="4" data-topic="95953">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mcarilli/48/28517_2.png" class="avatar"> mcarilli:</div>
<blockquote>
<p>why not also use an out-of-place <code>torch.autograd.grad</code> to create gradsH</p>
</blockquote>
</aside>
<p>I tested this today and couldn’t get it to work for some reason, collecting <code>gradsH</code> and <code>params</code> from the optimizer seemed to be the only thing that worked (i.e. it trained with no errors but results were poorer).</p>
<aside class="quote no-group" data-username="mcarilli" data-post="4" data-topic="95953">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mcarilli/48/28517_2.png" class="avatar"> mcarilli:</div>
<blockquote>
<p>FYI <code>GradScaler</code> also has an <code>enabled</code> argument</p>
</blockquote>
</aside>
<p>Missed that, thanks</p>
<aside class="quote no-group" data-username="mcarilli" data-post="4" data-topic="95953">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mcarilli/48/28517_2.png" class="avatar"> mcarilli:</div>
<blockquote>
<p>Your optimizer is vanilla SGD and you don’t appear to <em>use</em> <code>hvs</code> , so it’s hard for me to recommend a way to create it.</p>
</blockquote>
</aside>
<p>Oh yep, that was just some bare bones code to be able to run a single step. <code>hvs</code> is used to calculate the Hutchinson trace, which is subsequently used in a modification of AdamW</p>
<pre><code class="lang-auto">v = [torch.randint_like(p, high=2, device=device) for p in params]

for v_i in v: v_i[v_i == 0] = -1
        
hvs = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, create_graph=False, retain_graph=False)

hutchinson_trace = []
for hv, vi in zip(hvs, v):
    param_size = hv.size()
    if len(param_size) &lt;= 1:  
        tmp_output = torch.abs(hv * vi)
        hutchinson_trace.append(tmp_output)
</code></pre>
<p><code>hutchinson_trace</code> subsequently replaces the <code>gradient</code> in the denominator for the calculation of the AdamW step</p>
<h2><a name="batchnorm-had-to-be-set-to-fp32-1" class="anchor" href="#batchnorm-had-to-be-set-to-fp32-1"></a>BatchNorm had to be set to FP32</h2>
<p>After implementing manual scaling, I got native AMP to work with my <code>ResNet18</code>, however I had to modify the input to my BatchNorm to explicitly be FP32 for it to work with larger input sizes (it worked for toy-size problems without modifying the input to BatchNorm).</p>
<p><strong>Have you heard of others with this BN issue?</strong><br>
I see in the <code>autocast</code> docs its not explicitly defined whether BN should be FP16 or FP32. Maybe it just wasn’t working because of my double backward case…</p>
<h2><a name="memory-usage-autocast-debugging-2" class="anchor" href="#memory-usage-autocast-debugging-2"></a>Memory usage - autocast debugging</h2>
<p>Despite successful training, when checking memory use via ‘nvidia-smi’ or ‘nvtop’ it doesn’t seem like the GPU memory usage has decreased. Its about the same as without native AMP and in fact after about 30 steps GPU memory usage <strong>increases</strong> by approx 800mb. I have identified that the increase happens after the second backward below is called for the 30th time.</p>
<pre><code class="lang-auto">hvs = torch.autograd.grad(gradsH, params, grad_outputs=v, only_inputs=True, retain_graph=False)
</code></pre>
<p>Note that 800mb also appears to be the increase in memory when the <code>hvs</code> is calculated in the very first step for the very first time, so maybe a copy is being made?</p>
<p><strong>What might the best way to check what precision different operations are done in under <code>autocast</code>?</strong><br>
Checking dtype of the inputs/outputs of each layer?</p>
<p><strong>Is there anything that you are aware of that would cause <code>torch.autograd.grad</code> memory usage to increase?</strong><br>
I checked the size of the tensors being input as below, but the inputs stay the exact same size from start to finish (as expected). The size of the output also remains the same.</p>
<p><code>gsize += g.element_size() * g.nelement()</code></p>
<p>Thanks again for your detailed reply above!</p>