Export transformer_engine.pytorch.LayerNorm to ONNX gives all zero output

[sgsdxzy]

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

[ptrendx]

@nzmora-nvidia Is this the problem that was fixed in #293? Does not seem so, since here the error is in FP32.

[nzmora-nvidia]

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.