bug fix: restore_optimizers correctly handles non-mapping values in optimizer.state.values()

[circlecrystal]

What does this PR do?

• This PR fixed the bug reported in Bug report: This line causes error when using MADGRAD optimizer & resume training #11741
• To reproduce this bug (and this is exactly what the test code below does): 1) use MADGRAD optmizer 2) use multiple GPUs for training, then resume training
• The main root cause for this bug to exist is because the original restore_optimizers() method cannot correctly handles non-mapping values in optmizer.state.values().

The test code is provided below:

import os

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.core.lightning import LightningModule

from madgrad import MADGRAD


class LitMNIST(LightningModule):
    def __init__(self):
        super().__init__()

        # mnist images are (1, 28, 28) (channels, height, width)
        self.layer_1 = nn.Linear(28 * 28, 128)
        self.layer_2 = nn.Linear(128, 256)
        self.layer_3 = nn.Linear(256, 10)

    def forward(self, x):
        batch_size, channels, height, width = x.size()

        # (b, 1, 28, 28) -> (b, 1*28*28)
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        x = F.relu(x)
        x = self.layer_3(x)

        x = F.log_softmax(x, dim=1)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return loss 

    def configure_optimizers(self):
        return MADGRAD(self.parameters(), lr=1e-6)

    def train_dataloader(self):
        # transforms
        # prepare transforms standard to MNIST
        transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        # data
        mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
        return DataLoader(mnist_train, batch_size=64)


model = LitMNIST()
checkpoint_callback = ModelCheckpoint(save_last=True, save_top_k=1, monitor="loss", mode="min", filename="best")
trainer1 = Trainer(gpus=-1, strategy="dp", max_epochs=1, callbacks=[checkpoint_callback])
trainer1.fit(model)
trainer2 = Trainer(gpus=-1, strategy="dp", max_epochs=2, callbacks=[checkpoint_callback])
trainer2.fit(model, ckpt_path=checkpoint_callback.last_model_path)

(Notice that the madgrad.py being used in the test code can be freely downloaded from the official repo: https://github.com/facebookresearch/madgrad/blob/main/madgrad/madgrad.py)

Without this bug fix, the test code will fail and report error message as:

Fixes #11741

Does your PR introduce any breaking changes? If yes, please list them.

No.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[ananthsub]

Please update the changelog to document this fix

[awaelchli]

We didn't add a test for this? We could have easily mocked the behavior of e.g. MADGRAD optimizer. Can we follow up on this to avoid regressions?