Batchnorm layers are running in FP32 in an autocast region, so they shouldn’t overflow internally.
Just to make sure I’m understanding your use case: “regular fp16 training” means you are using torch.cuda.amp or are you calling manually half() on any modules?