apex amp not working in 1.2.0

[ruotianluo]

🐛 Bug

Got RuntimeError: Invoked 'with amp.scale_loss, but internal Amp state has not been initialized. model, optimizer = amp.initialize(model, optimizer, opt_level=...) must be called before with amp.scale_loss.`

When training with apex amp

No problem in 1.1.8.

Also got amp.load_state_dict when resuming from checkpoints trained with earlier version.

Please reproduce using the BoringModel

https://colab.research.google.com/drive/1OKQgh1gDUp78HMXX-mmrkxktVC2qV5oE?usp=sharing

To Reproduce

Expected behavior

Environment

Note: Bugs with code are solved faster ! Colab Notebook should be made public !

• IDE: Please, use our python bug_report_model.py template.

• Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[awaelchli]

import torch
from pytorch_lightning import LightningModule
from torch.utils.data import DataLoader, Dataset
import pytorch_lightning as pl


class RandomDataset(Dataset):
    def __init__(self, size, num_samples):
        self.len = num_samples
        self.data = torch.randn(num_samples, size)

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

    def __len__(self):
        return self.len


num_samples = 10000

train = RandomDataset(32, num_samples)
train = DataLoader(train, batch_size=32)

val = RandomDataset(32, num_samples)
val = DataLoader(val, batch_size=32)

test = RandomDataset(32, num_samples)
test = DataLoader(test, batch_size=32)


class BoringModel(LightningModule):

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

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

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def training_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"loss": loss}

    def training_step_end(self, training_step_outputs):
        return training_step_outputs

    def training_epoch_end(self, outputs) -> None:
        torch.stack([x["loss"] for x in outputs]).mean()

    def validation_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def validation_epoch_end(self, outputs) -> None:
        torch.stack([x['x'] for x in outputs]).mean()

    def test_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        self.log('fake_test_acc', loss)
        return {"y": loss}

    def test_epoch_end(self, outputs) -> None:
        torch.stack([x["y"] for x in outputs]).mean()

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]


if __name__ == "__main__":
    model = BoringModel()
    trainer = pl.Trainer(
        max_epochs=1,
        progress_bar_refresh_rate=20,
        amp_backend='apex',  # change
        precision=16,
        amp_level='O2',
    )
    trainer.fit(model, train, val)
    trainer.test(test_dataloaders=test)

[awaelchli]

@ruotianluo apex is only supported on the gpu. You need to specify gpus to run this training.

[ruotianluo]

Sorry, I forgot to add that in the colab. (I did meet similar error locally, any way.)

[Jongchan]

Is this problem resolved?
I have run the code @awaelchli attached above(I have added gpus=-1 option to pl.Trainer.),
yet it produces the same error.

RuntimeError: Invoked 'with amp.scale_loss`, but internal Amp state has not been initialized.  model, optimizer = amp.initialize(model, optimizer, opt_level=...) must be called before `with amp.scale_loss`.

I am using PyTorch==1.6.0 and pytorch-lightning==1.2.5, pytorch-lightning-bolts==0.3.0 with 8 Titan Xp GPUs.

UPDATE-20210330
In version 1.1.6
There is no problem of using apex, because amp.initialize is properly called.
However, there is a warning shown in the command line LightningOptimizer doesn't support Apex, but the program runs without errors. The training speed is 10~20% slower than our native PyTorch implementation (with DDP, APEX, amp_level=O1). (checked again, it is almost identical to native PyTorch apex speed)
https://github.com/PyTorchLightning/pytorch-lightning/blob/c462b274ee193a40ab0172931e0c5102a7be51bc/pytorch_lightning/plugins/apex.py#L34-L37

In version 1.2.5
When ApexMixedPrecisionPlugin.connect is called, the model is in device cpu, and amp.initialize is never called.
https://github.com/PyTorchLightning/pytorch-lightning/blob/cc40fa306e8cec8822579246c1fd6ca7cef0edc4/pytorch_lightning/plugins/precision/apex_amp.py#L39-L48