enable mixed-precision for layernorm

[yxlllc]

Summary

Introduce MixedPrecisionLayerNorm (a drop-in nn.LayerNorm subclass) and use it for the LayerNorms in the LYNXNet / LYNXNet2 diffusion backbones and the ConvNeXt aux decoder, so activations stay in fp16/bf16 under AMP autocast.

Motivation

Under PyTorch AMP autocast (this project trains with pl_trainer_precision: '16-mixed' by default), layer_norm is in autocast's fp32 cast-policy list: every nn.LayerNorm outputs fp32 even with fp16 activations.

The fp32 outputs then flow through ops that autocast does not cast (Transpose, GELU/GLU, residual adds, dropout) until the next conv/linear, so all activations saved for backward in that span are stored in fp32 — doubling a large fraction of activation memory for no precision benefit.

How it works

MixedPrecisionLayerNorm.forward disables autocast locally, casts weight/bias to the activation dtype, and calls F.layer_norm directly. Mean/variance are still accumulated in fp32 inside the CUDA kernel; the only extra rounding is a one-time fp16 quantization of the normalized output and of the affine parameters.

weight/bias must be cast rather than passed through as fp32, because the CUDA layer_norm kernel rejects mixed input/weight dtypes (expected scalar type Half but found Float); upcasting the activation to fp32 instead would discard the memory saving entirely.

Non-AMP fp32 execution is unchanged (bit-identical to nn.LayerNorm), checkpoints stay fully compatible, and ONNX/JIT-exported graphs are identical to before.