Automatic Mixed Precision with two optimisers that step unevenly

ClaartjeBarkhof: <p>Hi there,</p>
<p>I have a question regarding the amp scaler working with multiple optimisers that step unevenly. The patterns looks as follows:</p>
<pre><code class="lang-auto">    
def do_train_step(wrapper, model_optimiser, contraint_optimiser, inputs):
    
    # wrapper contains two nn.Modules: model and constraint, 
    # that both have their own optimiser: model_optimiser, constraint_optimiser
    # and both interact with the inputs to form the loss

    with torch.set_grad_enabled(True):
        with autocast(enabled=True):
            loss = wrapper(inputs)

    # Gradients accumulate in here
    scaler.scale(loss).backward()

    # The model optimsier should only update once every &lt; accumulate_n_batches_grad &gt;
    if (global_step + 1) % accumulate_n_batches_grad == 0:

        # GRADIENT CLIPPING
        if gradient_clipping:
            # Unscales the loss of optimizer's assigned params in-place
            scaler.unscale_(model_optimiser)

            # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
            torch.nn.utils.clip_grad_norm_(wrapper.model.parameters(), 1.0)

        # OPTIMISER STEP (first unscaled)
        scaler.step(model_optimiser)

        # AMP SCALER UPDATE
        scaler.update()

        # Zero the gradients only here
        model_optimiser.zero_grad()

        # LR Scheduler
        lr_scheduler.step()

    # The constraint optimiser should update every train step
    scaler.step(constraint_optimiser)
    constraint_optimiser.zero_grad()

    return loss
</code></pre>
<p>But, this alternating pattern seems to break the amp scaling. I get the following error:</p>
<pre><code class="lang-auto">  File "/home/cbarkhof/.local/lib/python3.6/site-packages/torch/cuda/amp/grad_scaler.py", line 302, in step
    raise RuntimeError("step() has already been called since the last update().")
RuntimeError: step() has already been called since the last update().
</code></pre>
<p>which is about the <code>scaler.step(constraint_optimiser)</code>.</p>
<p>If there is anyone that spots what might be going wrong here that would be very helpful.</p>
<p>Thanks in advance!</p>
<p>Cheers, Claartje</p>

ptrblck: <p>I think you could either use two different scalers for each optimizer or alternatively use <code>scaler.step()</code> on the <code>constraint_optimizer</code>, if the inner block wasn’t executed.<br>
After the inner code was already executed, the gradients should be already unscaled, which are used in the first optimizer.<br>
However, I’m currently unsure about your use case and don’t know, if both optimizers are using a different subset of all parameters, are partially overlapping, or reuse the same parameters.<br>
Could you explain the use case a bit more?</p>
<p>Also, I think you have a typo in the posted code snippet, as you are calling <code>scaler.step(loss)</code>, while an optimizer would be needed.</p>

ClaartjeBarkhof: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, thanks for your answer. You are right about the typo, I corrected it for clarity.</p>
<p>A bit more context:</p>
<p>The wrapper wraps two <code>nn.Module</code>s, the constraint and the model. They do not have any shared parameters. They do both interact with (parts of) the total loss. The loss composition looks something like:</p>
<pre><code class="lang-auto">loss_1, loss_2 = model(input)
loss_2 = constraint(loss_2)
total_loss = loss_1 + loss_2
</code></pre>
<p>The model optimises with gradients from the <code>total_loss</code>, while the constraint optimises only part of the loss (<code>loss_2</code>). I hope this makes it clearer.</p>
<p>On your suggested solutions:</p>
<p><strong>1. Using two scalers.</strong><br>
Would the following be valid?</p>
<pre><code class="lang-auto">model_scaled_loss = model_scaler.scale(loss)
constraint_scaled_loss = constraint_scaler.scale(loss)
loss.backward()

if (global_step + 1) % accumulate_n_batches_grad == 0:
        # OPTIMISER STEP (first unscaled)
        model_scaler.step(model_optimiser)

        # AMP SCALER UPDATE
        model_scaler.update()

        # Zero the gradients only here
        model_optimiser.zero_grad()

constraint_scaler.step(constraint_optimiser)
constraint_optimiser.zero_grad()
</code></pre>
<p><strong>2. use <code>scaler.step()</code> on the <code>constraint_optimizer</code> , if the inner block wasn’t executed.</strong><br>
I am not sure what you mean by this. Are you saying only to execute <code>scaler.step(constraint_optimiser)</code> if <code>scaler.step(model_optimiser)</code> is not called (inner block execution)? I would think that both optimisers need explicit unscaling of the gradients, or don’t they? If both optimisers do not share parameters, would <code>scaler.unscale_(model_optimiser)</code> affect <code>scaler.unscale_(constraint_optimiser)</code> in any way then?</p>
<p>Hopefully this is clear enough, otherwise I can provide more information.</p>
<p>Cheers,</p>
<p>Claartje</p>

Zaigham_Randhawa: <p>Hey, any updates on this issue? I am also in a similar situation with two optimisers and mixed precision training. Getting the same error. Not sure why i am getting that error as well.  <a class="mention" href="/u/ptrblck">@ptrblck</a>  <a class="mention" href="/u/claartjebarkhof">@ClaartjeBarkhof</a></p>

ptrblck: <p>Did you try to use my suggestion from my previous post using different optimizers and scalers?</p>