Repository: pytorch/TensorRT

Issue:
❓ [Question] Running LayerNorm in fp16
## ❓ Question

<!-- Your question -->

## What you have already tried

I am trying to convert a transformer model to TRT in fp16 (fp32 works fine 🙂). It includes bunch of LayerNorms, all of them have explicit casting of inputs to fp32, i.e:
``` python
class LayerNormFP32(nn.LayerNorm):
    def forward(self, x):
        return super().forward(x.float()).type(x.dtype)
``` 
I am getting warnings about precisions of the layers:
```
WARNING: [Torch-TensorRT TorchScript Conversion Context] - Detected layernorm nodes in FP16: %126 : Tensor = aten::layer_norm(%input.9, %127, %self.decoder.layers.0.attn_ln.weight.1, %370, %129, %130), scope: __module.decoder/__module.decoder.layers.0/__module.decoder.layers.0.attn_ln
...
WARNING: [Torch-TensorRT TorchScript Conversion Context] - Running layernorm after self-attention in FP16 may cause overflow. Exporting the model to the latest available ONNX opset (later than opset 17) to use the INormalizationLayer, or forcing layernorm layers to run in FP32 precision can help with preserving accuracy.
WARNING: [Torch-TensorRT TorchScript Conversion Context] - TensorRT encountered issues when converting weights between types and that could affect accuracy.
WARNING: [Torch-TensorRT TorchScript Conversion Context] - If this is not the desired behavior, please modify the weights or retrain with regularization to adjust the magnitude of the weights.
WARNING: [Torch-TensorRT TorchScript Conversion Context] - Check verbose logs for the list of affected weights.
WARNING: [Torch-TensorRT TorchScript Conversion Context] - - 2 weights are affected by this issue: Detected FP32 infinity values and converted them to corresponding FP16 infinity.
WARNING: [Torch-TensorRT TorchScript Conversion Context] - - 27 weights are affected by this issue: Detected subnormal FP16 values.
WARNING: [Torch-TensorRT TorchScript Conversion Context] - - 3 weights are affected by this issue: Detected values less than smallest positive FP16 subnormal value and converted them to the FP16 minimum subnormalized value.
```
I checked dtype of the mentioned weights in the trace that I pass to `torch_tensorrt.compile` and they are correctly in fp32, even though the warnings state the opposite.

The warning suggets two solutions (use INormalizationLayer or force FP32 precisions) but I have no idea ho to achieve it.
This might be a related: https://github.com/pytorch/TensorRT/pull/2509 (or https://github.com/NVIDIA/TensorRT/issues/3101)

Any ideas how to resolve or debug this issue?

## Environment

- Python 3.11.8
- torch 2.2.1
- torch_tensorrt 2.2.0
- a100




Discussion:
Here is a minimal reproducible example:
``` python
import torch
import torch.nn as nn


class LayerNormFP32(nn.LayerNorm):
    def forward(self, x):
        return super().forward(x.float()).type(x.dtype)


class Model(nn.Module):
    def __init__(self, hidden_dim: int = 1024):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.ln = LayerNormFP32(hidden_dim, bias=False)

    def forward(self, x: torch.Tensor):
        return self.ln(x)

    def to_jit_trace(
        self,
        device: str = "cpu",
        dtype: torch.dtype = torch.float, 
        batch_size: int = 2,
    ) -> torch.jit.ScriptModule:
        
        dummy_inputs = torch.randn((batch_size, self.hidden_dim), dtype=dtype, device=device)
        
        self.to(device)
        self.eval()
    
        with torch.no_grad():
            outputs1 = self(*dummy_inputs)
            trace = torch.jit.trace(self, dummy_inputs, check_trace=False)
            outputs2 = trace(*dummy_inputs)
        assert torch.allclose(outputs1, outputs2)
        
        return trace, dummy_inputs

    def to_tensorrt(
        self,
        batch_size,
        precisions: set[torch.dtype] = {
            torch.float,
            torch.half
        },
    ):
        import torch_tensorrt

        dtype = torch.float
        if torch.half in precisions:
            dtype = torch.half

        with torch.cuda.amp.autocast(enabled=True):
            trace, dummy_inputs = self.to_jit_trace("cuda", dtype, batch_size=batch_size)
        
        trt = torch_tensorrt.compile(
            trace,
            input_signature=(torch_tensorrt.Input(shape=dummy_inputs.shape, dtype=dummy_inputs.dtype),),
            enabled_precisions=precisions,
            require_full_compilation=True,
            truncate_long_and_double=True,
        )
            
        return trt
```

and fp32 gives the same outputs, fp16 does not (while producing the warnings):
``` python
model = Model()
batch_size = 1

trt_16 = model.to_tensorrt(batch_size=batch_size, precisions={torch.float, torch.half})
with torch.cuda.amp.autocast(enabled=True):
    trace_fp16, dummy_inputs_16 = model.to_jit_trace("cuda", torch.half, batch_size=batch_size)

trt_32 = model.to_tensorrt(batch_size=batch_size, precisions={torch.float})
trace_fp32, dummy_inputs_32 = model.to_jit_trace("cuda", torch.float, batch_size=batch_size)

with torch.no_grad():

    # False
    # tensor(0.0020, device='cuda:0', dtype=torch.float16)
    print(torch.allclose(trace_fp16(dummy_inputs_16), trt_16(dummy_inputs_16)))
    print((trace_fp16(dummy_inputs_16) - trt_16(dummy_inputs_16)).abs().max())

    # True
    # tensor(2.9802e-08, device='cuda:0')
    print(torch.allclose(trace_fp32(dummy_inputs_32), trt_32(dummy_inputs_32)))
    print((trace_fp32(dummy_inputs_32) - trt_32(dummy_inputs_32)).abs().max())
```

Resolving patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py
--- a/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py
@@ -9,6 +9,7 @@
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     cast_trt_tensor,
+    get_axes_for_reduce_op,
     get_positive_dim,
     get_trt_tensor,
     to_numpy,
@@ -105,102 +106,30 @@ def layer_norm(
     cudnn_enable: bool,
     return_mean_rstd: bool,
 ) -> Union[TRTTensor, Tuple[TRTTensor, torch.Tensor, torch.Tensor]]:
-    if weight is None:
-        weight = to_numpy(1.0)
-
-    if bias is None:
-        bias = to_numpy(0.0)
-
-    shape = weight.shape
-    gamma = to_numpy(weight).reshape(shape)
-    beta = to_numpy(bias).reshape(shape)
-
-    dims = list(range(len(input.shape) - len(shape), len(input.shape)))
-
-    # E[x]
-    mean_expected_trt = impl.reduce.mean(
-        ctx, target, source_ir, f"{name}_mean_expected", input, dims, True
-    )
-
-    # X-E[x]
-    sub_trt = impl.elementwise.sub(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_sub",
-        input,
-        mean_expected_trt,
-    )
-
-    # Variance = mean(pow(x_sub_mean, 2))
-    pow_trt = get_trt_tensor(ctx, 2, f"{name}_power", np.float32)
-    pow_var = impl.elementwise.pow(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_pow_var",
-        sub_trt,
-        pow_trt,
-    )
-    mean_trt = impl.reduce.mean(
-        ctx, target, source_ir, f"{name}_mean", pow_var, dims, True
-    )
-
-    # sqrt((var + eps))
-    eps_trt = get_trt_tensor(ctx, eps, f"{name}_eps", np.float32)
-    add_trt = impl.elementwise.add(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_add",
-        mean_trt,
-        eps_trt,
-    )
-    sqrt_trt = impl.unary.sqrt(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_sqrt",
-        add_trt,
-    )
-
-    # (X - E[X]) / sqrt((var + eps))
-    div_trt = impl.elementwise.div(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_div",
-        sub_trt,
-        sqrt_trt,
-    )
-
-    gamma_trt = get_trt_tensor(ctx, weight, f"{name}_gamma")
-    beta_trt = get_trt_tensor(ctx, bias, f"{name}_beta")
-
-    # y * gamma + beta
-    scaled_y = impl.elementwise.mul(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_mul_gamma",
-        div_trt,
-        gamma_trt,
-    )
+    dims = list(range(len(input.shape) - len(normalized_shape), len(input.shape)))
+    axes = get_axes_for_reduce_op(dims)
+
+    weight = get_trt_tensor(ctx, weight, f"{name}_weight")
+    bias = get_trt_tensor(ctx, bias, f"{name}_bias")
+    if tuple(input.shape) != tuple(weight.shape):
+        weight = impl.slice.expand(
+            ctx, target, source_ir, f"{name}_expand_weight", weight, input.shape
+        )
+    if tuple(input.shape) != tuple(bias.shape):
+        bias = impl.slice.expand(
+            ctx, target, source_ir, f"{name}_expand_bias", bias, input.shape
+        )
 
-    output = impl.elementwise.add(
-        ctx,
-        target,
-        source_ir,
-        f"{name}_add_beta",
-        scaled_y,
-        beta_trt,
-    )
+    layer_norm = ctx.net.add_normalization(input, weight, bias, axes)
+    layer_norm.epsilon = eps
+    layer_norm.compute_precision = input.dtype
+    set_layer_name(layer_norm, target, f"{name}_layer_norm", source_ir)
 
     if return_mean_rstd:
         # return fake mean and rstd for now
-        return output, None, None
+        return layer_norm.get_output(0), None, None
 
-    return output
+    return layer_norm.get_output(0)
 
 
 def native_group_norm(


Test patch:
diff --git a/tests/py/dynamo/conversion/test_layer_norm_aten.py b/tests/py/dynamo/conversion/test_layer_norm_aten.py
--- a/tests/py/dynamo/conversion/test_layer_norm_aten.py
+++ b/tests/py/dynamo/conversion/test_layer_norm_aten.py
@@ -1,4 +1,5 @@
 import torch
+from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 
@@ -6,19 +7,31 @@
 
 
 class TestLayerNormConverter(DispatchTestCase):
-    def test_layer_norm(self):
+    @parameterized.expand(
+        [
+            (
+                (5, 3, 2, 4),
+                [
+                    4,
+                ],
+            ),
+            ((5, 3, 2, 4), [2, 4]),
+            ((5, 3, 2, 4), [3, 2, 4]),
+            ((5, 3, 2, 4), [5, 3, 2, 4]),
+        ]
+    )
+    def test_layer_norm(self, input_shape, normalized_shape, eps=1e-05):
         class LayerNorm(torch.nn.Module):
             def forward(self, x):
                 return torch.ops.aten.layer_norm.default(
                     x,
-                    torch.tensor([3, 224, 224]),
-                    torch.ones((3, 224, 224)),
-                    torch.zeros((3, 224, 224)),
-                    1e-05,
-                    True,
+                    normalized_shape,
+                    torch.randn(normalized_shape),
+                    torch.randn(normalized_shape),
+                    eps,
                 )
 
-        inputs = [torch.randn(1, 3, 224, 224)]
+        inputs = [torch.randn(input_shape)]
         self.run_test(
             LayerNorm(),
             inputs,
@@ -26,7 +39,37 @@ def forward(self, x):
 
 
 class TestNativeLayerNormConverter(DispatchTestCase):
-    def test_layer_norm(self):
+    @parameterized.expand(
+        [
+            (
+                (5, 3, 2, 4),
+                [
+                    4,
+                ],
+            ),
+            ((5, 3, 2, 4), [2, 4]),
+            ((5, 3, 2, 4), [3, 2, 4]),
+            ((5, 3, 2, 4), [5, 3, 2, 4]),
+        ]
+    )
+    def test_layer_norm(self, input_shape, normalized_shape, eps=1e-05):
+        class LayerNorm(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.native_layer_norm.default(
+                    x,
+                    normalized_shape,
+                    torch.randn(normalized_shape),
+                    torch.randn(normalized_shape),
+                    eps,
+                )[0]
+
+        inputs = [torch.randn(input_shape)]
+        self.run_test(
+            LayerNorm(),
+            inputs,
+        )
+
+    def test_layernorm_with_dynamic_shape(self):
         class LayerNorm(torch.nn.Module):
             def forward(self, x):
                 return torch.ops.aten.native_layer_norm.default(
@@ -37,10 +80,17 @@ def forward(self, x):
                     1e-05,
                 )[0]
 
-        inputs = [torch.randn(1, 3, 224, 224)]
-        self.run_test(
+        input_specs = [
+            Input(
+                shape=(-1, 3, 224, 224),
+                dtype=torch.float32,
+                shape_ranges=[((1, 3, 224, 224), (5, 3, 224, 224), (10, 3, 224, 224))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
             LayerNorm(),
-            inputs,
+            input_specs,
         )