Why is closure not supported in GradScaler ?

ar90n: <p>In the step function of GradScaler, if a closure is given as a member of kwargs, RuntimeError occurs by the following codes.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/04fe4382ec8e6031d353530d6b8446ac9c9f507e/torch/cuda/amp/grad_scaler.py#L312-L313">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/04fe4382ec8e6031d353530d6b8446ac9c9f507e/torch/cuda/amp/grad_scaler.py#L312-L313" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/04fe4382ec8e6031d353530d6b8446ac9c9f507e/torch/cuda/amp/grad_scaler.py#L312-L313" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/04fe4382ec8e6031d353530d6b8446ac9c9f507e/torch/cuda/amp/grad_scaler.py#L312-L313</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="312" style="counter-reset: li-counter 311 ;">
          <li>if "closure" in kwargs:</li>
          <li>    raise RuntimeError("Closure use is not currently supported if GradScaler is enabled.")</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The error message says “not currently supported”. Are there any plans to support this feature? Or Would you tell me reasons why closure isn’t supported? If possible, I want to try to make a patch to solve this issue.</p>
<p>My motivation is to support SAM with native amp easily.</p>

mcarilli: <p>Making closures work with dynamic gradient scaling (specifically, the fact that dynamic gradient scaling occasionally skips optimizer.step() if any grads were inf/nan) is tricky, and we haven’t heard any use cases that absolutely needed it (LBFGS is the only one I’m aware of and no one’s asked for that).</p>
<p>Can you implement SAM without a closure like this?</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/davda54/sam#usage">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/davda54/sam#usage" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/5fe804143aeb8c614e1fa42c7bb81a4807cadf008707fc36cc2c001105f683a2/davda54/sam" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/davda54/sam#usage" target="_blank" rel="noopener nofollow ugc">GitHub - davda54/sam: SAM: Sharpness-Aware Minimization (PyTorch)</a></h3>

  <p>SAM: Sharpness-Aware Minimization (PyTorch). Contribute to davda54/sam development by creating an account on GitHub.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


ar90n: <p>Thank you for your reply!<br>
As you pointed, I use davda54’s SAM implementation without closure.<br>
And it works well.</p>
<p>But I use pytorch-lightning. In pytorch-lightning environment, I have to write some extra codes to calculate gradients two times. My motivation is to solve this problem. I want to use SAM as other optimizers.</p>
<p>I have understood that closure support is tricky. I decided to use SAM without closure. Thank!!</p>

Alex_Li: <p>Hi, I have the same problem you did. How did you code fp16 support without closure successfully using the linked github repo? It seems to me that we have to call <code>scaler.step(optimizer)</code>, and this in turn will call <code>optimizer.step()</code> with no closure parameter. To use the library, we need to call <code>first_step()</code> and <code>second_step()</code>, so I don’t see how it is possible to use this implementation and float16 support.</p>

ar90n: <p>Hi, I don’t remember clearly but I’m sure I used codes from pytorch-lightning’s example as follows in README.md.</p>
<pre><code class="lang-auto">def training_step(self, batch, batch_idx):
    optimizer = self.optimizers()

    # first forward-backward pass
    loss_1 = self.compute_loss(batch)
    self.manual_backward(loss_1, optimizer)
    optimizer.first_step(zero_grad=True)

    # second forward-backward pass
    loss_2 = self.compute_loss(batch)
    self.manual_backward(loss_2, optimizer)
    optimizer.second_step(zero_grad=True)

    return loss_1
</code></pre>

An_Tran_Nhut: <p>Hi Alex, I ran into the same problem recently, and figured out a way that uses vanilla torch. You still run scaler.scale(loss).backward(), but instead of calling scaler.step(optimizer), you first unscale the optimizer, check if the grads are non-finite, then call optimizer.first_step() and scaler.update(). You then repeat for the second step. Here’s an example:</p>
<pre><code class="lang-auto">def are_grads_finite(params):
    grads = [p.grad for p in params if p.grad is not None]
    vec = torch.nn.utils.parameters_to_vector(grads)
    return bool(torch.isfinite(vec).all())

for input, label in dataset:
    # First pass
    with torch.amp.autocast():
        output = model(input)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    if not are_grads_finite(model.parameters()):
        optimizer.zero_grad()
        scaler.update()
        continue # skip batch if grads are non-finite, like in scaler.step()
    optimizer.first_step(zero_grad=True)
    scaler.update()

    # Second pass
    with torch.amp.autocast():
        output = model(input)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    if not are_grads_finite(model.parameters()):
        optimizer.zero_grad()
        scaler.update()
        continue
    optimizer.second_step(zero_grad=True)
    scaler.update()

</code></pre>
<p>Basically you will need to manually do what scaler.step(optimizer) does, since you can’t tell it to run optmizer.first_step() or second_step(). I tested this training some classification model, and they converged without issue.</p>