LR scheduler called before optimizer step when using DDP sharded plugin

[timothybrooks]

I am using a learning rate scheduler called every step:

    def configure_optimizers(self):
        optimizer = Adam(...)
        scheduler = LambdaLR(optimizer, ...)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {"scheduler": scheduler, "interval": "step"},
        }

When using DDPPlugin, there is no issue, but when using DDPShardedPlugin, I get the following warning:

lib/python3.8/site-packages/torch/optim/lr_scheduler.py:129: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
  warnings.warn("Detected call of `lr_scheduler.step()` before `optimizer.step()`. "

pytorch 1.8.1
lighting 1.2.9

cc @tchaton @rohitgr7 @akihironitta @SeanNaren @awaelchli

[SeanNaren]

Thanks @timothybrooks for the issue! Here is a reproduce:

import os

import torch
from pytorch_lightning import LightningModule, Trainer
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import Dataset, DataLoader


class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):

    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)

        scheduler = StepLR(optimizer, step_size=1)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {"scheduler": scheduler, "interval": "step"},
        }


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        gpus=1,
        plugins='ddp_sharded',
        weights_summary=None,
    )
    trainer.fit(model, train_dataloader=train_data, val_dataloaders=val_data)
    trainer.test(model, test_dataloaders=test_data)


if __name__ == '__main__':
    run()

The issue stems from us replacing the optimizer with a Sharded Optimizer, but not updating the scheduler. However after updating the reference in the scheduler we run into the issue that counters have not been setup in the optimizer as the scheduler is created before the optimizer.

Schedulers do not maintain their initialization parameters however, making it harder to parse here. What might work for now is to just copy some of the metadata that is required for this check, and update the optimizer that the scheduler is pointing to.

[SeanNaren]

Just a followup here! This not super high priority, but just a matter of the LR scheduler not being instantiated with the sharded optimizer, which throws this warning even though the sharded optimizer is called correctly.

We'll need to still address the above to fix this, however this is primarily for suppressing the warning.