`optimizer.step()` before `lr_scheduler.step()` error using GradScaler

gil_fernandes: <p>Even though I think my code calls the <code>optimizer.step</code> via <code>Gradscaler</code> function before the <code>lr_scheduler.step()</code> function I am still getting this warning:</p>
<blockquote>
<p>/opt/anaconda3/envs/huggingface/lib/python3.7/site-packages/torch/optim/lr_scheduler.py:123: UserWarning: Detected call of <code>lr_scheduler.step()</code> before <code>optimizer.step()</code>. In PyTorch 1.1.0 and later, you should call them in the opposite order: <code>optimizer.step()</code> before <code>lr_scheduler.step()</code>.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule.</p>
</blockquote>
<p>As you can see in my training code <code>scaler.step(optimizer)</code> gets called before <code>scheduler.step()</code>, but I am still getting this warning. Any ideas what might be wrong?</p>
<p>Here is my training code.</p>
<pre><code class="lang-auto">    steps = len(train_dl) * epochs
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=lr, steps_per_epoch=len(train_dl), epochs=epochs)
    avg_train_losses = []
    avg_val_losses = []
    avg_val_scores = []
    lr = []
    best_avg_val_score = -1000
    scaler = torch.cuda.amp.GradScaler() # mixed precision support
    for epoch in tqdm(range(epochs), total=epochs):
        model.train()
        total_train_loss = 0.0
        for i, (x, y, image_tensor) in enumerate(train_dl):
            x, y, image_tensor = move_to_dev(x, y, image_tensor)
            model.zero_grad()
            with torch.cuda.amp.autocast():
                output = model(x, image_tensor)
                loss = criterion(y, output)
            total_train_loss += loss.item()
            
            # Backward Pass and Optimization
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            lr.append(get_lr(optimizer))
</code></pre>

ptrblck: <p>If the first iteration creates NaN gradients (e.g. due to a high scaling factor and thus gradient overflow), the <code>optimizer.step()</code> will be skipped and you might get this warning.</p>
<p>You could check the scaling factor via <code>scaler.get_scale()</code> and skip the learning rate scheduler, if it was decreased. I think it might be useful to add a utility function (or return value in <code>scaler.step()</code>) to indicate, if the current <code>optimizer.step()</code> was skipped.</p>

gil_fernandes: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> many thanks for your answer, really appreciate you took time for this. I am using the <code>torch.optim.lr_scheduler.OneCycleLR</code> scheduler (<code>torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, steps_per_epoch=len(train_dl), epochs=epochs)</code>) and at the beginning of the training the learning rate might be very small. The actual learning rate curve looks like this:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/d/9d8b3e45a182032e721a04f94a806cbed62ba8a9.png" alt="image" data-base62-sha1="mtHhZkkbO8lhG4duZZ0aYN0dUud" width="682" height="411"></p>
<p>So could it be that because the learning rate is very small at the first iteration, that this is causing this <code>NaN</code> gradients?</p>

ptrblck: <p>A small learning rate should reduce the probability of overflows in the gradients, which could create NaNs. However, if the scale factor is still too high, even a low learning rate might cause it.<br>
Could you check, if the scale factor was reduced in the first iteration(s)?</p>

gil_fernandes: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Just checked the documentation of the <code>GradScaler</code> class and found this:</p>
<blockquote>
<p>The scale factor often causes infs/NaNs to appear in gradients for the first few iterations as its value calibrates.  <code>scaler.step</code>  will skip the underlying  <code>optimizer.step()</code>  for these iterations. After that, step skipping should occur rarely (once every few hundred or thousand iterations).</p>
</blockquote>
<p>Could this be the cause for such warnings?</p>
<p>And another question: do you get the scale factor using <code>scaler.get_scale()</code> where <code>scaler</code> is an instance of <code>torch.cuda.amp.GradScaler</code>?</p>

ptrblck: <aside class="quote no-group" data-username="gil_fernandes" data-post="5" data-topic="92930">
<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/gil_fernandes/48/27322_2.png" class="avatar"> gil_fernandes:</div>
<blockquote>
<p>Could this be the cause for such warnings?</p>
</blockquote>
</aside>
<p>Yes, exactly. That’s what I was trying to say here:</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="92930">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If the first iteration creates NaN gradients (e.g. due to a high scaling factor and thus gradient overflow), the <code>optimizer.step()</code> will be skipped and you might get this warning.</p>
</blockquote>
</aside>
<p>Sorry, if I couldn’t describe it clearly. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="gil_fernandes" data-post="5" data-topic="92930">
<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/gil_fernandes/48/27322_2.png" class="avatar"> gil_fernandes:</div>
<blockquote>
<p>do you get the scale factor using <code>scaler.get_scale()</code> where <code>scaler</code> is an instance of <code>torch.cuda.amp.GradScaler</code> ?</p>
</blockquote>
</aside>
<p>Yes, that’s what I had in mind.</p>

gil_fernandes: <p>I have tried to capture the scale factor on each batch (100 epochs, batch size: 64, 24 batches per epoch) and then plot it:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8a8f8b1e497cb29599a9799098d1ba46ad307f54.png" alt="image" data-base62-sha1="jLLqJVTehMXWsglxRZcldBFn1KQ" width="686" height="417"></p>
<p>Right now I am not getting the warning, but still it shows up some now and then. Anyway I think that we know what is causing this issue. <a class="mention" href="/u/ptrblck">@ptrblck</a> Many thanks for your pointers.</p>

nofreewill: <p>Thank you all very much for the help. It wasn’t that clear but I could figure it out. The way I did it:</p>
<pre><code>        ...

        scale = scaler.get_scale()
        scaler.update()
        skip_lr_sched = (scale != scaler.get_scale())
        ...
        
        if not skip_lr_sched:
            lr_sched.step()
        ...</code></pre>

JamesDickens: <p>So this solution checks that the scale is NONE rather than just really large? Does this work, since I am encountering the same thing.</p>

rdutta: <p>I know this is an old post, but just wanted to get this out there in case it helps someone.</p>
<p>I would suggest checking:<br>
<code>skip_lr_sched = (scale &gt; scaler.get_scale())</code><br>
instead of<br>
<code>skip_lr_sched = (scale != scaler.get_scale())</code></p>
<p>because according to the <a href="https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.update" rel="noopener nofollow ugc">docs</a>, <code>scaler.update()</code> decreases the <code>scale_factor</code> when <code>optimizer.step()</code> is skipped, as well as increases the <code>scale_factor</code> when <code>optimizer.step()</code> has not been skipped for <code>growth_interval</code> consecutive iterations.</p>
<p>Simply checking <code>scale != scaler.get_scale()</code> will return <code>False</code> even when the <code>scale_factor</code> is increased (and <code>optimizer.step</code> has NOT been skipped), which we don’t want.</p>