[DDP + amp + gradient accumulation] Calling optimizer.step() leads to NaN loss

qsj287068067: <p>I want to add gradient accumulation feature to my DDP + amp training program for constant batch size when training large models. The model is wrapped by <code>torch.nn.parallel.DistributedDataParallel</code>.</p>
<p>Refer to</p>
<ol>
<li>2nd case of <a href="https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/20">albanD’s reply</a>;</li>
<li>pytorch document of <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation" rel="noopener nofollow ugc">amp working with gradient accumulation</a>,</li>
</ol>
<p>I implemented my code like</p>
<pre><code class="lang-python">optimizer = torch.optim.AdamW(model.parameters(), args.lr,weight_decay=args.weight_decay)
for i, input in enumerate(train_loader):
    with torch.cuda.amp.autocast(True):
        loss = model(input) / iters_to_accumulate   # Normalize our loss (if averaged). See https://gist.github.com/thomwolf/ac7a7da6b1888c2eeac8ac8b9b05d3d3#file-gradient_accumulation-py-L5.
    scaler.scale(loss).backward()
    if (i+1) % iters_to_accumulate == 0:                # Wait for several backward steps
        optimizer.step()                            # Now we can do an optimizer step
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()                           # Reset gradients tensors only if we have done a step. See https://gist.github.com/thomwolf/ac7a7da6b1888c2eeac8ac8b9b05d3d3?permalink_comment_id=2921188#gistcomment-2921188. And https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation.
</code></pre>
<p>The losses turn out to be <code>NaN</code> since second iteration (when i==1).<br>
But When I comment out line <code>optimizer.step()</code>, everything works like a charm.<br>
Later, I found that in both<br>
<a href="https://discuss.pytorch.org/u/BramVanroy">BramVanroy’s thread</a>, and<br>
pytorch document of <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation" rel="noopener nofollow ugc">amp working with gradient accumulation</a><br>
<code>optimizer.step()</code> is not called intentionally.</p>
<p>My question is, should <code>optimizer.step()</code> be called in my case?<br>
If yes, how to solve the <code>NaN</code> loss problem.<br>
If not, why? and would it influence the performance of the model?</p>

ptrblck: <aside class="quote no-group" data-username="qsj287068067" data-post="1" data-topic="162624">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/q/9fc348/48.png" class="avatar"> qsj287068067:</div>
<blockquote>
<p>My question is, should <code>optimizer.step()</code> be called in my case?</p>
</blockquote>
</aside>
<p>No, it shouldn’t as given in the <code>amp</code> tutorial. The <a href="https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.step"><code>GradScaler.step</code></a> function calls <code>optimizer.step()</code> internally after unscaling the gradients and making sure no Infs/NaNs were found.</p>

qsj287068067: <p>Thank you.</p>
<p>For those who may encounter the same issue, <a href="https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.step" rel="noopener nofollow ugc">Automatic Mixed Precision package - torch.amp#torch.cuda.amp.GradScaler.step</a> says</p>
<blockquote>
<ol start="2">
<li>If no inf/NaN gradients are found, invokes <code>optimizer.step()</code> using the unscaled gradients. Otherwise, <code>optimizer.step()</code> is skipped to avoid corrupting the params.</li>
</ol>
</blockquote>
<p>It should not be called manually.</p>