Model training with automatic mixed precision is not learning

miken: <p>Using mix precision, the loss flattens out after the first few iterations. The model trains fine when mix precision is not used.</p>
<p>Here is an example of how it’s implemented. I am using a ctc loss function</p>
<pre><code class="lang-python">for i, _data in enumerate(train_loader):
    spectrograms, labels, input_lengths, label_lengths = _data 
    spectrograms, labels = spectrograms.cuda(), labels.cuda()
    with autocast():
        output = model(spectrograms)  # (batch, time, n_class)
        output = F.log_softmax(output, dim=2)

        loss = ctc_loss(output, labels, input_lengths, label_lengths)
        loss = loss / args.grad_acc_steps
        scaler.scale(loss).backward()

    if i % args.grad_acc_steps == 0:
        scaler.unscale_(optimizer)  # unscale to clip gradient
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
</code></pre>

mcarilli: <p>Try moving scaler.scale(loss).backward() out of the autocast context</p>
<pre><code class="lang-auto">    with autocast():
        output = model(spectrograms)  # (batch, time, n_class)
        output = F.log_softmax(output, dim=2)

        loss = ctc_loss(output, labels, input_lengths, label_lengths)
        loss = loss / args.grad_acc_steps
    scaler.scale(loss).backward()
</code></pre>

mcarilli: <p>Also, what does “a few iterations” mean here?  It’s expected that <code>scaler.step(optimizer)</code> may skip the first few steps due to inf/nan gradients as the scale value calibrates, so the loss would not decrease for those iterations.</p>

miken: <p>Thanks, <a class="mention" href="/u/mcarilli">@mcarilli</a>. A few iterations in this context is about 1000 iterations.</p>
<p>And you’re right <code>scaler.scale(loss).backward()</code> should be outside the autocast context. The actual fix was due to how PyTorch did dynamic scaling. For my specific use-case, I had to set the <code>growth_interval</code> parameter to something smaller like 10 iterations. The issue was the scale value was not high enough causing a flat loss. Once the model started learning I’ve reset the <code>growth_interval</code> to the default value 2000 iterations.</p>
<p>my updated code is as follows.</p>
<pre><code class="lang-python">
scaler = GradScaler(growth_interval=10)
for i, _data in enumerate(train_loader):
    spectrograms, labels, input_lengths, label_lengths = _data 
    spectrograms, labels = spectrograms.cuda(), labels.cuda()
    with autocast():
        output = model(spectrograms)  # (batch, time, n_class)
        output = F.log_softmax(output, dim=2)

        loss = ctc_loss(output, labels, input_lengths, label_lengths)
        loss = loss / args.grad_acc_steps
    scaler.scale(loss).backward()

    if i % args.grad_acc_steps == 0:
        scaler.unscale_(optimizer)  # unscale to clip gradient
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
   if scaler.get_growth_interval() != 2000 and total_iter &gt; 1000:
        scaler.set_growth_interval(2000)
</code></pre>
<p>This solution works for me and hopefully it can help someone who runs into the same issue.</p>

mcarilli: <p>Glad it’s working, and an interesting discovery.  However, I don’t see the issue as solved yet.  I think we can make it work better for your model immediately, and also help prevent this issue for future users.</p>
<p>In our experience <code>GradScaler</code>'s default constructor values rarely need to be tuned.  Yours is the first case I’m aware of with the native API, and we tried it with 40ish models spanning many applications before merging.  The intention was to supply default values (and a dynamic scale-finding heuristic) that are effective for the vast majority of networks, so GradScaler’s args don’t become additional “hyperparameters.”</p>
<p>The default <code>init_scale</code> is intended to be larger than the network initially needs.  The large initial value causes inf/nan gradients for the first few iterations but quickly calibrates down to a successful value (because it’s reduced by <code>backoff_factor</code> each time).  After that, the large <code>growth_interval</code> means few iterations should be skipped, and the effect on performance is negligible.</p>
<p>In your case, it appears you’re in the opposite situation:  the default <code>init_scale</code> is smaller than you initially need. <code>growth_interval=10</code> is one way to increase the scale more quickly than it otherwise would, but once the value calibrates/stabilizes, roughly 1 out of 10 iterations will be skipped (a 10% training slowdown).  You work around this by resetting <code>growth_interval</code> later, which is smart, but also inconvenient and not obvious.  If all you need is a higher initial value, I’d construct <code>GradScaler(init_scale=&lt;bigger value&gt;)</code> instead of playing with <code>growth_interval</code> in multiple places.</p>
<p>Per above paragraphs, the best practice is to supply an <code>init_scale</code> that’s larger than your network needs, so the scale quickly calibrates down, then stabilizes.  To do that, we need to figure out the value it calibrates to.  <strong>Can you rerun your existing code (with <code>growth_interval=10</code>) and print <code>scaler.get_scale()</code> just after <code>scaler.update()</code> for the first few dozen steps to get a sense for the scale value it finds, and post the results here?</strong></p>
<p>For you, the best <code>init_scale</code> would then be the next-greatest power of two* above the value it finds, and you can then ignore <code>growth_interval</code>.  For me, the value it finds justifies a PR to increase the default <code>init_scale</code>, reducing the likelihood of this issue in the future.  A larger initial scale value doesn’t do much harm for any network (worst case, it causes a few more iterations at the beginning to be skipped).</p>
<p>(*Powers of two are best for for <code>init_scale</code>, <code>growth_factor</code>, and <code>backoff_factor</code> because multiplication/division by powers of two is a bitwise accurate operation on non-denormal IEEE floats.)</p>

miken: <p><a class="mention" href="/u/mcarilli">@mcarilli</a> thanks for the writeup.</p>
<blockquote>
<p><strong>Can you rerun your existing code (with  <code>growth_interval=10</code> ) and print  <code>scaler.get_scale()</code>  just after  <code>scaler.update()</code>  for the first few dozen steps to get a sense for the scale value it finds, and post the results here?</strong></p>
</blockquote>
<p>So the results for <code>scaler.get_scale()</code> is as follows</p>
<p><strong>iter 1 - 24:</strong> goes from <code>65536.0</code> to <code>0.00390625</code><br>
<strong>iter 24 - 256:</strong> goes from <code> 0.00390625</code> to <code>32768.0</code><br>
<strong>iter 256 - 1000:</strong> bounces around from <code>1024</code> - <code>16384.0</code> but stabilizes at around <code>1024</code> near 1000.<br>
<strong>iter 1000 - 50000:</strong> bounces around from <code>256</code> - <code>2048.0</code> but from eye balling it, the <code>1024</code> seems to be what occurs most frequently.</p>
<p>It seems like my network needs a range of scale values in the beginning to start stabilizing, then bounces around once the model starts learning.</p>

mcarilli: <p>That’s wild.  I’ve never seen behavior like that before.  It doesn’t just decrease to a stable value, it goes way down, then bounces back up…The <code>growth_interval</code> manipulation may be the best approach, so it aggressively tries many new values in the beginning.</p>
<p>Does the full run converge to roughly the same accuracy as FP32?</p>

miken: <p>Yup the model is converging to similar accuracy as FP32! Thanks for the help.</p>

ptrblck: <p>Which learning rate scheduler are you using?<br>
Could you remove the scheduler and check, if you are still seeing this shaky behavior?</p>

Philipp_Singer: <p>Aren’t you missing <code>optimizer.zero_grad()</code> in the loop?</p>

miken: <p>yes i am in the example code, but i have in my actual code. Thanks for pointing it out!</p>

miken: <p>I’m using the one cycle scheduler! Currently moved on from the issue but this could also be an issue so I want to make a statement for others to be aware.</p>