We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `transformer_engine/pytorch/module/layernorm.py`
Content:
```
1 # Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2 #
3 # See LICENSE for license information.
4 
5 """LayerNorm API"""
6 import os
7 from typing import Union, Tuple, Any, Mapping, Optional
8 
9 import torch
10 from torch.nn.parameter import Parameter
11 from torch.nn import init
12 
13 import transformer_engine_extensions as tex
14 
15 
16 __all__ = ["LayerNorm"]
17 
18 
19 class _LayerNorm(torch.autograd.Function):
20     """functional LayerNorm"""
21 
22     @staticmethod
23     def forward(
24         ctx,
25         inp: torch.Tensor,
26         ln_weight: torch.Tensor,
27         ln_bias: torch.Tensor,
28         eps: float,
29         fwd_ln_sm_margin: int,
30         bwd_ln_sm_margin: int,
31         zero_centered_gamma: bool,
32     ) -> torch.Tensor:
33         # Make sure input dimensions are compatible
34         in_features = ln_weight.numel()
35         assert inp.is_cuda, "TransformerEngine needs CUDA."
36         assert inp.shape[-1] == in_features, "LayerNorm not possible"
37         inputmat = inp.view((-1, in_features))
38 
39         ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,
40                                                ln_bias, eps, fwd_ln_sm_margin,
41                                                zero_centered_gamma)
42         ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)
43         ctx.inp_shape = inp.shape
44         ctx.bwd_ln_sm_margin = bwd_ln_sm_margin
45         ctx.zero_centered_gamma = zero_centered_gamma
46         return ln_out.view_as(inp)
47 
48     @staticmethod
49     def backward(
50         ctx, grad_output: torch.Tensor
51     ) -> Tuple[Union[torch.Tensor, None], ...]:
52         inputmat, ln_weight, mu, rsigma = ctx.saved_tensors
53         grad_output = grad_output.contiguous()
54         d_ln_out = grad_output.view(inputmat.shape)
55         dxmat, dgamma, dbeta = tex.layernorm_bwd(
56             d_ln_out, inputmat, mu, rsigma, ln_weight,
57             ctx.bwd_ln_sm_margin, ctx.zero_centered_gamma
58         )
59         return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None
60 
61 
62 class LayerNorm(torch.nn.Module):
63     r"""
64     Applies Layer Normalization over a mini-batch of inputs as described in
65     the paper `Layer Normalization <https://arxiv.org/abs/1607.06450>`__
66 
67     .. math::
68         y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma + \beta
69 
70     :math:`\gamma` and :math:`\beta` are learnable affine transform parameters of
71     size :attr:`hidden_size`
72 
73     Parameters
74     ----------
75     hidden_size : int
76                 size of each input sample.
77     eps : float, default = 1e-5
78         a value added to the denominator of layer normalization for numerical stability.
79     sequence_parallel : bool, default = `False`
80                         if set to `True`, uses sequence parallelism.
81     params_dtype : torch.dtype, default = `torch.get_default_dtype()`
82                     it controls the type used to allocate the initial parameters. Useful when
83                     the model is trained with lower precision and the original FP32 parameters
84                     would not fit in GPU memory.
85     zero_centered_gamma : bool, default = 'False'
86                          if set to 'True', gamma parameter in LayerNorm is initialized to 0 and
87                          the LayerNorm formula changes to
88 
89                          .. math::
90                             y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} *
91                             (1 + \gamma) + \beta
92     """
93 
94     def __init__(
95         self,
96         hidden_size: int,
97         eps: float = 1e-5,
98         sequence_parallel: bool = False,
99         params_dtype: Optional[torch.dtype] = None,
100         zero_centered_gamma: bool = False,
101     ) -> None:
102         super().__init__()
103         params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
104         self.eps = eps
105         self.zero_centered_gamma = zero_centered_gamma
106         self.weight = Parameter(
107             torch.empty(
108                 hidden_size,
109                 device=torch.cuda.current_device(),
110                 dtype=params_dtype,
111             )
112         )
113         self.bias = Parameter(
114             torch.empty(
115                 hidden_size,
116                 device=torch.cuda.current_device(),
117                 dtype=params_dtype,
118             )
119         )
120         setattr(self.weight, "sequence_parallel", sequence_parallel)
121         setattr(self.bias, "sequence_parallel", sequence_parallel)
122         self.reset_layer_norm_parameters()
123 
124         # These many SMs are subtracted from the total SM count when calling forward
125         # and backward LayerNorm C APIs. These envvars can be used to prevent the LN
126         # kernels from using all SMs in the device. This is useful for cases such as
127         # communication overlap with LN.
128         self.fwd_ln_sm_margin = int(os.getenv("NVTE_FWD_LAYERNORM_SM_MARGIN", "0"))
129         self.bwd_ln_sm_margin = int(os.getenv("NVTE_BWD_LAYERNORM_SM_MARGIN", "0"))
130 
131     def load_state_dict(
132         self,
133         state_dict: Mapping[str, Any],
134         strict: bool = True,
135     ) -> None:
136         """Override PyTorch loader to maintain backward compatibility
137         with previous version of LayerNorm parameter names.
138         """
139         if "layer_norm_weight" in state_dict:
140             state_dict["weight"] = state_dict["layer_norm_weight"]
141             del state_dict["layer_norm_weight"]
142         if "layer_norm_bias" in state_dict:
143             state_dict["bias"] = state_dict["layer_norm_bias"]
144             del state_dict["layer_norm_bias"]
145 
146         super().load_state_dict(state_dict, strict)
147 
148     def reset_layer_norm_parameters(self) -> None:
149         """Init LN params"""
150         if not self.zero_centered_gamma:
151             init.ones_(self.weight)
152         else:
153             init.zeros_(self.weight)
154         init.zeros_(self.bias)
155 
156 
157     def forward(self, inp: torch.Tensor) -> torch.Tensor:
158         """LayerNorm FWD"""
159         # Maintain backward compatibility.
160         if hasattr(self, "layer_norm_weight"):
161             setattr(self, "weight", self.layer_norm_weight)
162         if hasattr(self, "layer_norm_bias"):
163             setattr(self, "bias", self.layer_norm_bias)
164 
165         return _LayerNorm.apply(
166             inp,
167             self.weight,
168             self.bias,
169             self.eps,
170             self.fwd_ln_sm_margin,
171             self.bwd_ln_sm_margin,
172             self.zero_centered_gamma
173         )
174
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
         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 @@
             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 @@
         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)

{"golden_diff": "diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py\n--- a/transformer_engine/pytorch/module/layernorm.py\n+++ b/transformer_engine/pytorch/module/layernorm.py\n@@ -11,7 +11,9 @@\n from torch.nn import init\n \n import transformer_engine_extensions as tex\n-\n+from ..cpp_extensions import (\n+    layernorm_fwd_inf,\n+ )\n \n __all__ = [\"LayerNorm\"]\n \n@@ -29,6 +31,7 @@\n         fwd_ln_sm_margin: int,\n         bwd_ln_sm_margin: int,\n         zero_centered_gamma: bool,\n+        is_grad_enabled: bool,\n     ) -> torch.Tensor:\n         # Make sure input dimensions are compatible\n         in_features = ln_weight.numel()\n@@ -36,13 +39,16 @@\n         assert inp.shape[-1] == in_features, \"LayerNorm not possible\"\n         inputmat = inp.view((-1, in_features))\n \n-        ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,\n-                                               ln_bias, eps, fwd_ln_sm_margin,\n-                                               zero_centered_gamma)\n-        ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)\n-        ctx.inp_shape = inp.shape\n-        ctx.bwd_ln_sm_margin = bwd_ln_sm_margin\n-        ctx.zero_centered_gamma = zero_centered_gamma\n+        if is_grad_enabled:\n+            ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,\n+                ln_bias, eps, fwd_ln_sm_margin, zero_centered_gamma)\n+            ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)\n+            ctx.inp_shape = inp.shape\n+            ctx.bwd_ln_sm_margin = bwd_ln_sm_margin\n+            ctx.zero_centered_gamma = zero_centered_gamma\n+        else:\n+            ln_out, mu, rsigma = layernorm_fwd_inf(inputmat, ln_weight,\n+                ln_bias, eps, zero_centered_gamma), None, None\n         return ln_out.view_as(inp)\n \n     @staticmethod\n@@ -56,7 +62,7 @@\n             d_ln_out, inputmat, mu, rsigma, ln_weight,\n             ctx.bwd_ln_sm_margin, ctx.zero_centered_gamma\n         )\n-        return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None\n+        return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None, None\n \n \n class LayerNorm(torch.nn.Module):\n@@ -162,12 +168,22 @@\n         if hasattr(self, \"layer_norm_bias\"):\n             setattr(self, \"bias\", self.layer_norm_bias)\n \n-        return _LayerNorm.apply(\n+        if torch.is_grad_enabled():\n+            fwd_fn = _LayerNorm.apply\n+            args = []\n+        else:\n+            fwd_fn = _LayerNorm.forward\n+            args = [None]\n+\n+        args += (\n             inp,\n             self.weight,\n             self.bias,\n             self.eps,\n             self.fwd_ln_sm_margin,\n             self.bwd_ln_sm_margin,\n-            self.zero_centered_gamma\n+            self.zero_centered_gamma,\n+            torch.is_grad_enabled()\n         )\n+\n+        return fwd_fn(*args)\n", "issue": "Export transformer_engine.pytorch.LayerNorm to ONNX gives all zero output\nAfter exporting `transformer_engine.pytorch.LayerNorm` to ONNX, it always give tensors with correct shape but filled with 0.\r\n\r\nExample code:\r\n```\r\nimport torch\r\nimport onnxruntime\r\nfrom transformer_engine import pytorch as te\r\n\r\nmodel = te.LayerNorm(1000).cuda().eval()\r\nx_sample = torch.randn(3000, 1000)\r\nwith torch.inference_mode():\r\n  torch_out = model(x_sample.cuda())    # the result is correct\r\n\r\nwith torch.inference_mode():\r\n  with te.onnx_export(True):\r\n    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\"])\r\n\r\nort_session = onnxruntime.InferenceSession(\"layer_norm.onnx\", providers=[\"CPUExecutionProvider\"])\r\nort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x_sample)}\r\nort_output = ort_session.run(None, ort_inputs)[0]\r\n# ort_output is all zero\r\n```\r\n\r\nOther pytorch modules like `transformer_engine.pytorch.LayerNormLinear` gives correct results.\r\n\r\n\r\nOS: RHEL 7\r\nPython: 3.10.11\r\nTransformerEngine: 0.9\r\nPytorch: 2.0.1+cu118\r\nGPU: 4090\n", "before_files": [{"content": "# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# See LICENSE for license information.\n\n\"\"\"LayerNorm API\"\"\"\nimport os\nfrom typing import Union, Tuple, Any, Mapping, Optional\n\nimport torch\nfrom torch.nn.parameter import Parameter\nfrom torch.nn import init\n\nimport transformer_engine_extensions as tex\n\n\n__all__ = [\"LayerNorm\"]\n\n\nclass _LayerNorm(torch.autograd.Function):\n    \"\"\"functional LayerNorm\"\"\"\n\n    @staticmethod\n    def forward(\n        ctx,\n        inp: torch.Tensor,\n        ln_weight: torch.Tensor,\n        ln_bias: torch.Tensor,\n        eps: float,\n        fwd_ln_sm_margin: int,\n        bwd_ln_sm_margin: int,\n        zero_centered_gamma: bool,\n    ) -> torch.Tensor:\n        # Make sure input dimensions are compatible\n        in_features = ln_weight.numel()\n        assert inp.is_cuda, \"TransformerEngine needs CUDA.\"\n        assert inp.shape[-1] == in_features, \"LayerNorm not possible\"\n        inputmat = inp.view((-1, in_features))\n\n        ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,\n                                               ln_bias, eps, fwd_ln_sm_margin,\n                                               zero_centered_gamma)\n        ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)\n        ctx.inp_shape = inp.shape\n        ctx.bwd_ln_sm_margin = bwd_ln_sm_margin\n        ctx.zero_centered_gamma = zero_centered_gamma\n        return ln_out.view_as(inp)\n\n    @staticmethod\n    def backward(\n        ctx, grad_output: torch.Tensor\n    ) -> Tuple[Union[torch.Tensor, None], ...]:\n        inputmat, ln_weight, mu, rsigma = ctx.saved_tensors\n        grad_output = grad_output.contiguous()\n        d_ln_out = grad_output.view(inputmat.shape)\n        dxmat, dgamma, dbeta = tex.layernorm_bwd(\n            d_ln_out, inputmat, mu, rsigma, ln_weight,\n            ctx.bwd_ln_sm_margin, ctx.zero_centered_gamma\n        )\n        return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None\n\n\nclass LayerNorm(torch.nn.Module):\n    r\"\"\"\n    Applies Layer Normalization over a mini-batch of inputs as described in\n    the paper `Layer Normalization <https://arxiv.org/abs/1607.06450>`__\n\n    .. math::\n        y = \\frac{x - \\mathrm{E}[x]}{ \\sqrt{\\mathrm{Var}[x] + \\varepsilon}} * \\gamma + \\beta\n\n    :math:`\\gamma` and :math:`\\beta` are learnable affine transform parameters of\n    size :attr:`hidden_size`\n\n    Parameters\n    ----------\n    hidden_size : int\n                size of each input sample.\n    eps : float, default = 1e-5\n        a value added to the denominator of layer normalization for numerical stability.\n    sequence_parallel : bool, default = `False`\n                        if set to `True`, uses sequence parallelism.\n    params_dtype : torch.dtype, default = `torch.get_default_dtype()`\n                    it controls the type used to allocate the initial parameters. Useful when\n                    the model is trained with lower precision and the original FP32 parameters\n                    would not fit in GPU memory.\n    zero_centered_gamma : bool, default = 'False'\n                         if set to 'True', gamma parameter in LayerNorm is initialized to 0 and\n                         the LayerNorm formula changes to\n\n                         .. math::\n                            y = \\frac{x - \\mathrm{E}[x]}{ \\sqrt{\\mathrm{Var}[x] + \\varepsilon}} *\n                            (1 + \\gamma) + \\beta\n    \"\"\"\n\n    def __init__(\n        self,\n        hidden_size: int,\n        eps: float = 1e-5,\n        sequence_parallel: bool = False,\n        params_dtype: Optional[torch.dtype] = None,\n        zero_centered_gamma: bool = False,\n    ) -> None:\n        super().__init__()\n        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype\n        self.eps = eps\n        self.zero_centered_gamma = zero_centered_gamma\n        self.weight = Parameter(\n            torch.empty(\n                hidden_size,\n                device=torch.cuda.current_device(),\n                dtype=params_dtype,\n            )\n        )\n        self.bias = Parameter(\n            torch.empty(\n                hidden_size,\n                device=torch.cuda.current_device(),\n                dtype=params_dtype,\n            )\n        )\n        setattr(self.weight, \"sequence_parallel\", sequence_parallel)\n        setattr(self.bias, \"sequence_parallel\", sequence_parallel)\n        self.reset_layer_norm_parameters()\n\n        # These many SMs are subtracted from the total SM count when calling forward\n        # and backward LayerNorm C APIs. These envvars can be used to prevent the LN\n        # kernels from using all SMs in the device. This is useful for cases such as\n        # communication overlap with LN.\n        self.fwd_ln_sm_margin = int(os.getenv(\"NVTE_FWD_LAYERNORM_SM_MARGIN\", \"0\"))\n        self.bwd_ln_sm_margin = int(os.getenv(\"NVTE_BWD_LAYERNORM_SM_MARGIN\", \"0\"))\n\n    def load_state_dict(\n        self,\n        state_dict: Mapping[str, Any],\n        strict: bool = True,\n    ) -> None:\n        \"\"\"Override PyTorch loader to maintain backward compatibility\n        with previous version of LayerNorm parameter names.\n        \"\"\"\n        if \"layer_norm_weight\" in state_dict:\n            state_dict[\"weight\"] = state_dict[\"layer_norm_weight\"]\n            del state_dict[\"layer_norm_weight\"]\n        if \"layer_norm_bias\" in state_dict:\n            state_dict[\"bias\"] = state_dict[\"layer_norm_bias\"]\n            del state_dict[\"layer_norm_bias\"]\n\n        super().load_state_dict(state_dict, strict)\n\n    def reset_layer_norm_parameters(self) -> None:\n        \"\"\"Init LN params\"\"\"\n        if not self.zero_centered_gamma:\n            init.ones_(self.weight)\n        else:\n            init.zeros_(self.weight)\n        init.zeros_(self.bias)\n\n\n    def forward(self, inp: torch.Tensor) -> torch.Tensor:\n        \"\"\"LayerNorm FWD\"\"\"\n        # Maintain backward compatibility.\n        if hasattr(self, \"layer_norm_weight\"):\n            setattr(self, \"weight\", self.layer_norm_weight)\n        if hasattr(self, \"layer_norm_bias\"):\n            setattr(self, \"bias\", self.layer_norm_bias)\n\n        return _LayerNorm.apply(\n            inp,\n            self.weight,\n            self.bias,\n            self.eps,\n            self.fwd_ln_sm_margin,\n            self.bwd_ln_sm_margin,\n            self.zero_centered_gamma\n        )\n", "path": "transformer_engine/pytorch/module/layernorm.py"}], "after_files": [{"content": "# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# See LICENSE for license information.\n\n\"\"\"LayerNorm API\"\"\"\nimport os\nfrom typing import Union, Tuple, Any, Mapping, Optional\n\nimport torch\nfrom torch.nn.parameter import Parameter\nfrom torch.nn import init\n\nimport transformer_engine_extensions as tex\nfrom ..cpp_extensions import (\n    layernorm_fwd_inf,\n )\n\n__all__ = [\"LayerNorm\"]\n\n\nclass _LayerNorm(torch.autograd.Function):\n    \"\"\"functional LayerNorm\"\"\"\n\n    @staticmethod\n    def forward(\n        ctx,\n        inp: torch.Tensor,\n        ln_weight: torch.Tensor,\n        ln_bias: torch.Tensor,\n        eps: float,\n        fwd_ln_sm_margin: int,\n        bwd_ln_sm_margin: int,\n        zero_centered_gamma: bool,\n        is_grad_enabled: bool,\n    ) -> torch.Tensor:\n        # Make sure input dimensions are compatible\n        in_features = ln_weight.numel()\n        assert inp.is_cuda, \"TransformerEngine needs CUDA.\"\n        assert inp.shape[-1] == in_features, \"LayerNorm not possible\"\n        inputmat = inp.view((-1, in_features))\n\n        if is_grad_enabled:\n            ln_out, mu, rsigma = tex.layernorm_fwd(inputmat, ln_weight,\n                ln_bias, eps, fwd_ln_sm_margin, zero_centered_gamma)\n            ctx.save_for_backward(inputmat, ln_weight, mu, rsigma)\n            ctx.inp_shape = inp.shape\n            ctx.bwd_ln_sm_margin = bwd_ln_sm_margin\n            ctx.zero_centered_gamma = zero_centered_gamma\n        else:\n            ln_out, mu, rsigma = layernorm_fwd_inf(inputmat, ln_weight,\n                ln_bias, eps, zero_centered_gamma), None, None\n        return ln_out.view_as(inp)\n\n    @staticmethod\n    def backward(\n        ctx, grad_output: torch.Tensor\n    ) -> Tuple[Union[torch.Tensor, None], ...]:\n        inputmat, ln_weight, mu, rsigma = ctx.saved_tensors\n        grad_output = grad_output.contiguous()\n        d_ln_out = grad_output.view(inputmat.shape)\n        dxmat, dgamma, dbeta = tex.layernorm_bwd(\n            d_ln_out, inputmat, mu, rsigma, ln_weight,\n            ctx.bwd_ln_sm_margin, ctx.zero_centered_gamma\n        )\n        return dxmat.view(ctx.inp_shape), dgamma, dbeta, None, None, None, None, None\n\n\nclass LayerNorm(torch.nn.Module):\n    r\"\"\"\n    Applies Layer Normalization over a mini-batch of inputs as described in\n    the paper `Layer Normalization <https://arxiv.org/abs/1607.06450>`__\n\n    .. math::\n        y = \\frac{x - \\mathrm{E}[x]}{ \\sqrt{\\mathrm{Var}[x] + \\varepsilon}} * \\gamma + \\beta\n\n    :math:`\\gamma` and :math:`\\beta` are learnable affine transform parameters of\n    size :attr:`hidden_size`\n\n    Parameters\n    ----------\n    hidden_size : int\n                size of each input sample.\n    eps : float, default = 1e-5\n        a value added to the denominator of layer normalization for numerical stability.\n    sequence_parallel : bool, default = `False`\n                        if set to `True`, uses sequence parallelism.\n    params_dtype : torch.dtype, default = `torch.get_default_dtype()`\n                    it controls the type used to allocate the initial parameters. Useful when\n                    the model is trained with lower precision and the original FP32 parameters\n                    would not fit in GPU memory.\n    zero_centered_gamma : bool, default = 'False'\n                         if set to 'True', gamma parameter in LayerNorm is initialized to 0 and\n                         the LayerNorm formula changes to\n\n                         .. math::\n                            y = \\frac{x - \\mathrm{E}[x]}{ \\sqrt{\\mathrm{Var}[x] + \\varepsilon}} *\n                            (1 + \\gamma) + \\beta\n    \"\"\"\n\n    def __init__(\n        self,\n        hidden_size: int,\n        eps: float = 1e-5,\n        sequence_parallel: bool = False,\n        params_dtype: Optional[torch.dtype] = None,\n        zero_centered_gamma: bool = False,\n    ) -> None:\n        super().__init__()\n        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype\n        self.eps = eps\n        self.zero_centered_gamma = zero_centered_gamma\n        self.weight = Parameter(\n            torch.empty(\n                hidden_size,\n                device=torch.cuda.current_device(),\n                dtype=params_dtype,\n            )\n        )\n        self.bias = Parameter(\n            torch.empty(\n                hidden_size,\n                device=torch.cuda.current_device(),\n                dtype=params_dtype,\n            )\n        )\n        setattr(self.weight, \"sequence_parallel\", sequence_parallel)\n        setattr(self.bias, \"sequence_parallel\", sequence_parallel)\n        self.reset_layer_norm_parameters()\n\n        # These many SMs are subtracted from the total SM count when calling forward\n        # and backward LayerNorm C APIs. These envvars can be used to prevent the LN\n        # kernels from using all SMs in the device. This is useful for cases such as\n        # communication overlap with LN.\n        self.fwd_ln_sm_margin = int(os.getenv(\"NVTE_FWD_LAYERNORM_SM_MARGIN\", \"0\"))\n        self.bwd_ln_sm_margin = int(os.getenv(\"NVTE_BWD_LAYERNORM_SM_MARGIN\", \"0\"))\n\n    def load_state_dict(\n        self,\n        state_dict: Mapping[str, Any],\n        strict: bool = True,\n    ) -> None:\n        \"\"\"Override PyTorch loader to maintain backward compatibility\n        with previous version of LayerNorm parameter names.\n        \"\"\"\n        if \"layer_norm_weight\" in state_dict:\n            state_dict[\"weight\"] = state_dict[\"layer_norm_weight\"]\n            del state_dict[\"layer_norm_weight\"]\n        if \"layer_norm_bias\" in state_dict:\n            state_dict[\"bias\"] = state_dict[\"layer_norm_bias\"]\n            del state_dict[\"layer_norm_bias\"]\n\n        super().load_state_dict(state_dict, strict)\n\n    def reset_layer_norm_parameters(self) -> None:\n        \"\"\"Init LN params\"\"\"\n        if not self.zero_centered_gamma:\n            init.ones_(self.weight)\n        else:\n            init.zeros_(self.weight)\n        init.zeros_(self.bias)\n\n\n    def forward(self, inp: torch.Tensor) -> torch.Tensor:\n        \"\"\"LayerNorm FWD\"\"\"\n        # Maintain backward compatibility.\n        if hasattr(self, \"layer_norm_weight\"):\n            setattr(self, \"weight\", self.layer_norm_weight)\n        if hasattr(self, \"layer_norm_bias\"):\n            setattr(self, \"bias\", self.layer_norm_bias)\n\n        if torch.is_grad_enabled():\n            fwd_fn = _LayerNorm.apply\n            args = []\n        else:\n            fwd_fn = _LayerNorm.forward\n            args = [None]\n\n        args += (\n            inp,\n            self.weight,\n            self.bias,\n            self.eps,\n            self.fwd_ln_sm_margin,\n            self.bwd_ln_sm_margin,\n            self.zero_centered_gamma,\n            torch.is_grad_enabled()\n        )\n\n        return fwd_fn(*args)\n", "path": "transformer_engine/pytorch/module/layernorm.py"}]}