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

Nobody has claimed this yet.

## 🐛 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

Start with the provided reproduction, especially Trainer configured with strategy="deepspeed_stage_2" and precision=16, and compare it with the commented DDP strategy. Investigate how LayerNorm and BatchNorm parameter dtypes differ between the two paths; done means users can select consistent fp32 behavior for these layers.

Written by the indexing model from the issue text.