Fixed layernorm when weight and bias is None in Stable Diffusion 3 (#2936)

Author: cehongwang

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -27,7 +27,7 @@
 def args_bounds_check(
     args: Tuple[Argument, ...], i: int, replacement: Optional[Any] = None
 ) -> Any:
-    return args[i] if len(args) > i else replacement
+    return args[i] if len(args) > i and args[i] is not None else replacement
 
 
 def get_ir(target: Target) -> SourceIR:
@@ -156,8 +156,8 @@ def aten_ops_layer_norm(
         name,
         input=args[0],
         normalized_shape=args[1],
-        weight=args_bounds_check(args, 2),
-        bias=args_bounds_check(args, 3),
+        weight=args_bounds_check(args, 2, 1.0),
+        bias=args_bounds_check(args, 3, 0.0),
         eps=args_bounds_check(args, 4, 1e-05),
         cudnn_enable=args_bounds_check(args, 5, True),
         return_mean_rstd=(target == torch.ops.aten.native_layer_norm.default),

tests/py/dynamo/conversion/test_layer_norm_aten.py

@@ -119,6 +119,24 @@ def forward(self, x):
             input_specs,
         )
 
+    @parameterized.expand([((5, 3, 2, 4), [2, 4])])
+    def test_layer_norm_without_Scaling(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,
+                    None,
+                    None,
+                    eps,
+                )[0]
+
+        inputs = [torch.randn(input_shape)]
+        self.run_test(
+            LayerNorm(),
+            inputs,
+        )
+
 
 if __name__ == "__main__":
     run_tests()