SAM with lr scheduler results in MisconfigurationException

[maxmatical]

This is what my configure_optimizers() looks like, trying to use the SAM optimizer with pytorch lightning

def configure_optimizers(self):
    base_optimizer = AdamW
    optimizer = SAM(self.pytorch_model.parameters(), base_optimizer=base_optimizer, lr=self.lr)

    scheduler = OneCycleLR(
      optimizer.base_optimizer,
      max_lr=self.lr,
      pct_start=0.3,
      total_steps=self.total_steps

    )
return [optimizer], [scheduler]

but when i try to fit the model, i get the following error

---------------------------------------------------------------------------

MisconfigurationException                 Traceback (most recent call last)

<ipython-input-62-c883f9dc1920> in <module>()
----> 1 trainer.fit(toxic_comment_model, data_module)

9 frames

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/optimizers.py in _validate_scheduler_optimizer(optimizers, lr_schedulers)
    231     if any(sch["scheduler"].optimizer not in optimizers for sch in lr_schedulers):
    232         raise MisconfigurationException(
--> 233             "Some schedulers are attached with an optimizer that wasn't returned from `configure_optimizers`."
    234         )
    235 

MisconfigurationException: Some schedulers are attached with an optimizer that wasn't returned from `configure_optimizers`.

has anyone been successful in using SAM with lr schedulers in pytorch lightning?

[maxmatical]

i modified configure_optimizers to return [optimizer, optimizer.base_optimizer], [scheduler], which now passes the validation checks, and the training step looks like

def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]

    opt, _ = self.optimizers()
    out = self(input_ids, attention_mask)

    # manual optimization with SAM
    loss = self.criterion(out, labels.to(dtype=torch.float32))
    acc = self.accuracy(out, labels)
    self.manual_backward(loss)
    opt.first_step(zero_grad=True)
    # 2nd pass
    out2 = self(input_ids, attention_mask)
    loss2 = self.criterion(out2, labels.to(dtype=torch.float32))
    self.manual_backward(loss2)
    opt.second_step(zero_grad=True)

    # lr schedule
    lr_sch = self.lr_schedulers()
    lr_sch.step()
 
    log_values = {"train_loss": loss, "train_acc": acc}
    self.log_dict(log_values, sync_dist=True, prog_bar=True, on_step=True, on_epoch=True)
    return {"loss": loss, "accuracy": acc, "predictions": out, "labels": labels}

that seems to work, albeit a bit hacky

[rohitgr7]

what if you set your scheduler like:

scheduler = OneCycleLR(
      optimizer,
      max_lr=self.lr,
      pct_start=0.3,
      total_steps=self.total_steps

    )

it won't work?

[rohitgr7]

yep it says:

 LR scheduling should be either applied to the base optimizer or you should use SAM with a single step call (with a closure):

lightning provides a closure so it should work. Can you try?

[rohitgr7]

also, I think all that schedulers needs from an optimizer is param_groups so and it has them
https://github.com/davda54/sam/blob/295977586de6a6e38f4730adb2ae496fd75f94ec/sam.py#L12

probably should work.

[maxmatical]

i actually tried using a closure, so the training step looks like this

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]

    opt = self.optimizers()
    out = self(input_ids, attention_mask)


    loss = self.criterion(out, labels.to(dtype=torch.float32))
    acc = self.accuracy(out, labels)
    def closure():
      out = self(input_ids, attention_mask)
      loss = self.criterion(out, labels.to(dtype=torch.float32))
      self.manual_backward(loss)
      return loss
    
    self.manual_backward(loss)
    opt.step(closure)
    opt.zero_grad()

    lr_sch = self.lr_schedulers()
    lr_sch.step()
    log_values = {"train_loss": loss, "train_acc": acc}
    self.log_dict(log_values, prog_bar=True, on_step=True, on_epoch=True)
    return {"loss": loss, "accuracy": acc, "predictions": out, "labels": labels}

but for some reason the opt.step(closure) step triggers assertion error that closure is None

<ipython-input-18-c4cd28162513> in step(self, closure)
     36     @torch.no_grad()
     37     def step(self, closure=None):
---> 38         assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
     39         closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass
     40 

AssertionError: Sharpness Aware Minimization requires closure, but it was not provided

[maxmatical]

i did try manually calling the 2 forward and backward passes, using the LR scheduler on the SAM optimizer directly instead of the base optimizer, and the learning rate schedule does appear to be behaving as intended, so perhaps it's not an issue after all

[rohitgr7]

opt.step(closure)

the closure isn't None here. Plus I think you don't need a manual optimization here and don't need to create a closure by yourself because lightning internally uses a closure to do forward and backward pass, you can just rely on automatic optimization.

I think

def training_ste(...):
    loss = self.criterion(out, labels.to(dtype=torch.float32))
    return {'loss': loss, ...}

should work just fine.