Mixed precision not working with v 1.2

[harpone]

🐛 Bug

After updating to 1.2 from 1.1.1, automatic mixed precision stopped working. Everything's float32 and getting CUDA OOM when I shouldn't get it (with float16 tensors). Worked fine on 1.1.1.

Here's my Trainer args (maybe there's a conflicting combo of args or something):

Trainer(logger=logger,
        callbacks=[checkpoint_callback, lr_monitor],
        default_root_dir=None,
        gradient_clip_val=args.gradient_clip_val,
        gpus=args.gpus,
        auto_select_gpus=False,
        log_gpu_memory=None,
        progress_bar_refresh_rate=1,
        check_val_every_n_epoch=args.check_val_every_n_epoch,
        overfit_batches=0.,
        fast_dev_run=False,
        accumulate_grad_batches=1,
        max_epochs=args.max_epochs,
        limit_train_batches=vars(args).get('limit_train_batches', 1.),
        val_check_interval=args.val_check_interval,
        limit_val_batches=args.limit_val_batches,
        accelerator='ddp',
        sync_batchnorm=True,
        precision=args.precision,
        weights_summary='top',
        weights_save_path=None,
        num_sanity_val_steps=args.num_sanity_val_steps,
        resume_from_checkpoint=None,
        benchmark=False,
        deterministic=False,
        reload_dataloaders_every_epoch=True,
        terminate_on_nan=False,
        prepare_data_per_node=True,
        amp_backend='native',
        profiler=args.profiler)

Environment

• GCP VM with V100 GPU(s)
• NVIDIA-SMI 450.51.06 Driver Version: 450.51.06 CUDA Version: 11.0

Don't really have the time to go deeper than this, but just rolled back to 1.1.1 and everything's fine.

[awaelchli]

Minimal example to reproduce:

from argparse import ArgumentParser
import torch
from torch.nn import functional as F

import pytorch_lightning as pl
from pl_examples.basic_examples.mnist_datamodule import MNISTDataModule


class LitClassifier(pl.LightningModule):

    def __init__(self, hidden_dim=128, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        print(x.dtype)
        y_hat = self(x)
        print(y_hat.dtype)
        loss = F.cross_entropy(y_hat, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument('--hidden_dim', type=int, default=128)
        parser.add_argument('--learning_rate', type=float, default=0.0001)
        return parser


def cli_main():
    pl.seed_everything(1234)
    parser = ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
    parser = LitClassifier.add_model_specific_args(parser)
    parser = MNISTDataModule.add_argparse_args(parser)
    args = parser.parse_args()
    
    dm = MNISTDataModule.from_argparse_args(args)
    model = LitClassifier(args.hidden_dim, args.learning_rate)
    
    trainer = pl.Trainer.from_argparse_args(args, precision=16, gpus=1, max_steps=1)
    print(trainer.training_type_plugin)
    print(trainer.precision_plugin.precision)
    
    trainer.fit(model, datamodule=dm)


if __name__ == '__main__':
    cli_main()

[awaelchli]

@harpone thanks for reporting, should be fixed on master :)