Summary

LayerNorm.forward() performs normalization in float32 for numerical stability and then casts the output back to the original input dtype. To avoid dtype mismatch errors when a model has been converted with model.half(), LayerNorm parameters should be explicitly cast to the computation dtype used by the normalization operation.

Motivation

When models are converted to FP16 using:

model.half()

the LayerNorm weights and biases are also converted to float16.

If LayerNorm internally performs computation using:

x.float()

while leaving parameters in FP16, some PyTorch versions may raise dtype mismatch errors because the input tensor and normalization parameters no longer share the same dtype.

Other custom modules in the codebase already handle mixed-precision execution by ensuring operands participate in computation using compatible dtypes. Applying the same pattern to LayerNorm improves consistency and robustness.

Proposed Implementation

Use torch.nn.functional.layer_norm() and explicitly cast the LayerNorm parameters to the computation dtype:

class LayerNorm(nn.LayerNorm):
    def forward(self, x: Tensor) -> Tensor:
        return F.layer_norm(
            x.float(),
            self.normalized_shape,
            self.weight.float() if self.weight is not None else None,
            self.bias.float() if self.bias is not None else None,
            self.eps,
        ).type(x.dtype)

Benefits

• Prevents FP16/FP32 parameter-input dtype mismatches.
• Maintains numerical stability by performing normalization in FP32.
• Preserves existing output dtype behavior.
• Aligns LayerNorm behavior with other mixed-precision-aware modules in the codebase.
• Improves compatibility across PyTorch versions and execution environments.

Compatibility

This change is backward compatible because it preserves the existing API and output dtype while only affecting the internal computation dtype handling.

Expected Outcome

LayerNorm remains numerically stable under mixed-precision execution and avoids failures that can occur when model parameters have been converted to FP16.