Computing gradients with Apex

kl_divergence: <p>I want to compute gradients in the standard manner:</p>
<pre><code class="lang-auto">grad = torch.autograd.grad(loss, model.parameters()
</code></pre>
<p>I want to use AMP. To get the scaled loss, apex api suggests:</p>
<pre><code class="lang-auto">with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
</code></pre>
<p>But I don’t want to call loss.backward() directly. I compute the gradients in the manner I showed above<br>
I tried this:</p>
<pre><code class="lang-auto">with amp.scale_loss(loss, optimizer) as scaled_loss:
    grad = torch.autograd.grad(scaled_lossloss, model.parameters()
</code></pre>
<p>This makes the loss nan. How to compute gradients with apex ?</p>

ptrblck: <p>We recommend to use the native automatic mixed-precision training using a master build or the nightly binaries.<br>
For your use case, you might have forgotten to unscale the gradients or might have manually updated the parameters with invalid gradients.</p>
<p>Have a look at the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty">gradient penalty example</a> for an example use case of <code>autograd.grad</code>.</p>

kl_divergence: <p>Thanks for your reply. I did what you said. Before I can test your suggestion, now I’m facing some other issue.<br>
Without fp16, I’m able to train with batch size 128.<br>
After <code>amp.initialize()</code> is called, I am getting an OOM error with the same batch size which is not ideal since RAM usage should decrease. A similar issue has been raised <a href="https://github.com/NVIDIA/apex/issues/439" rel="nofollow noopener">here</a>.</p>

ptrblck: <p>The linked issue initialized models in a loop, which is not supported in apex/amp.<br>
Could you post a reproducible code snippet using native amp, so that we could have a look?</p>

kl_divergence: <p>Regarding first issue, as per your suggestion, I modified my code as:</p>
<pre><code class="lang-auto">if self.args.fp16:
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_grad_params = torch.autograd.grad(self.scaler.scale(scaled_loss), params, allow_unused=True)
</code></pre>
<p>Before I can multiply params by <code>inv_scale</code>, it throws an error because <code>scaled_grad_params</code> is <code>None</code>.<br>
I assume that <code>amp.scale_loss</code> is not supposed to be used like this. I searched in the documentation, this method is used with <code>loss.backward()</code> which I’m not calling. Can you please suggest what am I missing ?</p>

ptrblck: <p>Is your code working without amp, i.e. are you getting the sum of gradients from <code>autograd.grad</code>?</p>

kl_divergence: <p>Yes, my code works fine without AMP.<br>
I compute <code>grads</code>:</p>
<pre><code class="lang-auto">grads = torch.autograd.grad(loss, model.parameters())
</code></pre>
<p>Then update my model’s parameters:</p>
<pre><code class="lang-auto">for param in model.parameters():  
    param.data = ....  # (these params come from grads after some operation)
</code></pre>

ptrblck: <p>I don’t know, why the output is apparently empty when you are using amp, so we would need to debug it.<br>
Could you post a minimal, executable code snippet so that we could have a look?</p>