Automatic Mixed Precision (AMP) doesn't run BatchNorm in FP32

[jramapuram]

🐛 Describe the bug

BatchNorm should be kept in FP32 when using mixed precision for numerical stability. This works fine when it is the first layer, eg:

import torch
from torch import nn

net = nn.Sequential(nn.BatchNorm1d(4)).cuda()

o = torch.randn(2, 4).cuda()

# auto-cast should automatically cast stuff
with torch.cuda.amp.autocast():
    for layer in net:
        o = layer(o)
        print(o.dtype)

Result:

torch.float32

But does not when it follows other layers:

import torch
from torch import nn

net = nn.Sequential(nn.Linear(4, 4), nn.BatchNorm1d(4), nn.Linear(4, 5)).cuda()
o = torch.randn(2, 4).cuda()

# auto-cast should automatically cast stuff
with torch.cuda.amp.autocast():
    for layer in net:
        o = layer(o)
        print(o.dtype)

Result:

torch.float16
torch.float16
torch.float16

Expected result:

torch.float16
torch.float32  # BN output
torch.float16

Versions

[pip3] torch==1.12.1

cc @mcarilli @ptrblck @leslie-fang-intel @jgong5

[ngimel]

The behavior is according to the documentation, https://pytorch.org/docs/stable/amp.html#cuda-ops-that-can-autocast-to-float16, batch norm doesn't upcast its inputs to fp32.

[jramapuram]

@ngimel : batch_norm is not listed in the ops that autocast to FP16 or FP32. It should be similar to layer_norm and group_norm and kept in FP32. This is not expected behavior and causes numerous numerical instabilities at scale.

NB: the running_mean and std should at least be kept in FP32, but they don't seem to be when enumerating them inline in the above logic.

[ngimel]

Inputs to batch norm are fp16 in the cases you are talking about, so upcasting them to fp32 won't help accuracy. All intermediate computations in batch norm are done in fp32. If batchnorm outputs are still overflowing, that means that the bn parameter values are unusually large and the training is already unstable, and the next layer that would convert results to fp16 would error out.
Same logic applies to layer_norm and group_norm, and honestly they should be removed from the cast to fp32 list, they are there for legacy reasons.

[jramapuram]

Thanks for the info @ngimel. The problem is w.r.t stability (not accuracy).

Shouldn't the running std and mean be kept in FP32? Division + FP16 seems like a bad recipe. Do you happen to know where in the C++ code I can grep exactly what AMP is doing w.r.t BN?

[ngimel]

Autocast-related code is in https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/autocast_mode.cpp, for batch norm nothing is done, and it falls through to implementation with whatever inputs it received. In amp mode, if cudnn is used iirc it'll produce rstd and mean in fp32 even for low precision inputs.

[jramapuram]

Ok great, as long as cudnn is handling the running mean and std in FP32 👍.

Appreciate your time @ngimel !