Repository: NVIDIA/TransformerEngine

Issue:
Export transformer_engine.pytorch.LayerNorm to ONNX gives all zero output
After exporting `transformer_engine.pytorch.LayerNorm` to ONNX, it always give tensors with correct shape but filled with 0.

Example code:
```
import torch
import onnxruntime
from transformer_engine import pytorch as te

model = te.LayerNorm(1000).cuda().eval()
x_sample = torch.randn(3000, 1000)
with torch.inference_mode():
  torch_out = model(x_sample.cuda())    # the result is correct

with torch.inference_mode():
  with te.onnx_export(True):
    torch.onnx.export(model, x_sample.cuda(), "layer_norm.onnx", dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}, opset_version=17, input_names=["input"], output_names=["output"])

ort_session = onnxruntime.InferenceSession("layer_norm.onnx", providers=["CPUExecutionProvider"])
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x_sample)}
ort_output = ort_session.run(None, ort_inputs)[0]
# ort_output is all zero
```

Other pytorch modules like `transformer_engine.pytorch.LayerNormLinear` gives correct results.


OS: RHEL 7
Python: 3.10.11
TransformerEngine: 0.9
Pytorch: 2.0.1+cu118
GPU: 4090


Discussion:
@nzmora-nvidia Is this the problem that was fixed in #293? Does not seem so, since here the error is in FP32.
This is a different problem: the FP32 LN module is not export-ready (meaning it uses non-traceable CPP/Cuda extension code directly).
With the fix, when grad is disabled we use a "leaner" version of LN which doesn't return mu and rsigma, and is thus exportable to ONNX.

Resolving patch:
diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py
--- a/transformer_engine/pytorch/module/layernorm.py
+++ b/transformer_engine/pytorch/module/layernorm.py
@@ -11,7 +11,9 @@
 from torch.nn import init
 
 import transformer_engine_extensions as tex
-
+from ..cpp_extensions import (
+    layernorm_fwd_inf,
+ )
 
 __all__ = ["LayerNorm"]
 
@@ -29,6 +31,7 @@ def forward(
         fwd_ln_sm_margin: int,
         bwd_ln_sm_margin: int,
         zero_centered_gamma: bool,
+        is_grad_enabled: bool,
     ) -> torch.Tensor:
         # Make sure input dimensions are compatible
         in_features = ln_weight.numel()
@@ -36,13 +39,16 @@ def forward(
         assert inp.shape[-1] == in_features, "LayerNorm not possible"
         inputmat = inp.view((-1, in_features))
 
-        ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,
-                                               ln_bias, eps, fwd_ln_sm_margin,
-                                               zero_centered_gamma)
-        ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)
-        ctx.inp_shape = inp.shape
-        ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
-        ctx.zero_centered_gamma = zero_centered_gamma
+        if is_grad_enabled:
+            ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,
+                ln_bias, eps, fwd_ln_sm_margin, zero_centered_gamma)
+            ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)
+            ctx.inp_shape = inp.shape
+            ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
+            ctx.zero_centered_gamma = zero_centered_gamma
+        else:
+            ln_out, mu, rsigma = layernorm_fwd_inf(inputmat, ln_weight,
+                ln_bias, eps, zero_centered_gamma), None, None
         return ln_out.view_as(inp)
 
     @staticmethod
@@ -56,7 +62,7 @@ def backward(
             d_ln_out, inputmat, mu, rsigma, ln_weight,
             ctx.bwd_ln_sm_margin, ctx.zero_centered_gamma
         )
-        return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None
+        return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None, None
 
 
 class LayerNorm(torch.nn.Module):
@@ -162,12 +168,22 @@ def forward(self, inp: torch.Tensor) -> torch.Tensor:
         if hasattr(self, "layer_norm_bias"):
             setattr(self, "bias", self.layer_norm_bias)
 
-        return _LayerNorm.apply(
+        if torch.is_grad_enabled():
+            fwd_fn = _LayerNorm.apply
+            args = []
+        else:
+            fwd_fn = _LayerNorm.forward
+            args = [None]
+
+        args += (
             inp,
             self.weight,
             self.bias,
             self.eps,
             self.fwd_ln_sm_margin,
             self.bwd_ln_sm_margin,
-            self.zero_centered_gamma
+            self.zero_centered_gamma,
+            torch.is_grad_enabled()
         )
+
+        return fwd_fn(*args)


Test patch:
diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py
--- a/tests/pytorch/test_onnx_export.py
+++ b/tests/pytorch/test_onnx_export.py
@@ -620,20 +620,13 @@ def test_export_layernorm(
     class Test_Layernorm(nn.Module):
         def __init__(self) -> None:
             super().__init__()
-            normalized_shape = torch.Size(inp.shape[1:])
-            self.weight = torch.randn(*normalized_shape, device="cuda",
-                dtype=torch.float if fake_bf16_io else precision)
-            self.bias = torch.zeros(*normalized_shape, device="cuda",
-                dtype=torch.float if fake_bf16_io else precision)
-            self.eps = 1e-6 # An arbitrary small value
+            eps = 1e-6 # An arbitrary small value
+            dtype = torch.float if fake_bf16_io else precision
+            self.ln = te.LayerNorm(inp_shape[1], eps, params_dtype=dtype,
+                zero_centered_gamma=False).eval().cuda()
 
         def forward(self, inp):
-            ret = texcpp.layernorm_fwd_inf(
-                inp,
-                self.weight,
-                self.bias,
-                self.eps,
-                zero_centered_gamma)
+            ret = self.ln(inp)
             return ret
 
     class TestFP8_Layernorm(nn.Module):