LayerNorm / BatchNorm fp16 behavior is different in Pytorch Native and Deepspeed

🐛 Bug

Currently, LayerNorm and BatchNorm behave differently when using fp16 in pytorch and deepspeed. Currently, in native pytorch, LayerNorm and BatchNorm retain fp32 weights, but in deepspeed it is fp16 weights.

In some cases, fp32 for certain layers are critical

To Reproduce

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer

import fire
import socket
import datetime

from pytorch_lightning.utilities.rank_zero import _get_rank
from pytorch_lightning.loggers import WandbLogger


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

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

    def __len__(self):
        return len(self.data)


class BoringModel(LightningModule):
    def __init__(self, ds):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.norm_layer = torch.nn.LayerNorm(2)
        self.batch_norm_layer = torch.nn.BatchNorm1d(2)
        self.ds = ds

    def train_dataloader(self):
        train_ds = self.ds
        self.train_dl = DataLoader(train_ds, batch_size=2, num_workers=10)
        return self.train_dl

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

    def training_step(self, batch, batch_idx):
        print(f'dtypes batch={batch.dtype}, norm_layer={self.norm_layer.dtype}')
        loss = self(batch).sum()
        batch_sum = batch.sum()
        _B = len(batch)
        self.log(
            "train_loss",
            loss,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            batch_size=_B,
        )
        self.log(
            "batch_sum",
            batch_sum,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            logger=True,
        )
        return {"loss": loss}

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

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

def main():

    hostname = socket.gethostname()
    ngpus = torch.cuda.device_count()

    num_nodes = int(os.environ.get("SL_NUM_NODES", 1))
    wsize = ngpus * num_nodes
    grank = _get_rank()

    train_ds = RandomDataset(32, 32000)

    model = BoringModel(ds=train_ds)
    trainer = Trainer(
        devices=ngpus,
        num_nodes=num_nodes,
        accelerator="gpu",
        strategy="deepspeed_stage_2", # change it to ddp
        # strategy="ddp",
        limit_train_batches=1000,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=3,
        log_every_n_steps=1,
        reload_dataloaders_every_n_epochs=1,
        precision=16,
    )
    train_data = model.train_dataloader()
    trainer.fit(model, train_dataloaders=train_data)

if __name__ == "__main__":
    fire.Fire(main)

Expected behavior

Ideally, there should be some flag which the user can set to get consistent behavior.

Environment

• PyTorch Lightning Version: 1.6.4
• PyTorch Version: 1.10.1
• Python version: 3.9.12
• OS: Linux
• CUDA/cuDNN version: 11.3
• GPU models and configuration: 8x A100
• How you installed PyTorch: conda

cc @awaelchli @rohitgr7 @akihironitta