We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Performance regression of `bias_dropout_add_fused_train` after switching to `torch.compile`
There is a performance regression of `bias_dropout_add_fused_train()` after switching to `torch.compile` for jit fusion in this [PR](https://github.com/NVIDIA/TransformerEngine/pull/589).

Here are my profiling results. As we can see, from either host or device perspective, there's a big regression.

- torch.compile
![image](https://github.com/NVIDIA/TransformerEngine/assets/3831900/7bf647b5-71d6-4b52-a264-c6d48633bc78)

- nvFuser
![image](https://github.com/NVIDIA/TransformerEngine/assets/3831900/c3654308-f914-490e-a52d-2676308e0071)

My profiling scripts:
```python
import torch
import transformer_engine as te

def profile_bias_dropout_add_fused_train(
     hidden_size: int, dtype: torch.dtype, seq_length: int, micro_batch_size: int
) -> None:
    # warm up
    te.pytorch.jit.warmup_jit_bias_dropout_add(hidden_size, dtype, seq_length, micro_batch_size)

    inp = torch.rand(
        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda"
    )
    residual = torch.rand(
        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda"
    )
    bias = torch.rand((hidden_size), dtype=dtype, device="cuda")
    dropout_rate = 0.1

    torch.cuda.profiler.start()
    for _ in range(10):
        torch.cuda.nvtx.range_push("fused bda")
        output = te.pytorch.jit.bias_dropout_add_fused_train(inp, bias, residual, dropout_rate)
        torch.cuda.nvtx.range_pop()
    torch.cuda.profiler.stop()

te.pytorch.jit.set_jit_fusion_options()
profile_bias_dropout_add_fused_train(1024, torch.float32, 1024, 2)
```

cc @ksivaman @minitu 

--- 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/jit.py`
Content:
```
1 # Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2 #
3 # See LICENSE for license information.
4 
5 """NVFuser functions and JIT utilities"""
6 import os
7 from typing import Callable, Optional, Tuple
8 
9 import torch
10 
11 jit_fuser = torch.jit.script
12 if torch.__version__ >= "2" and bool(int(os.getenv("NVTE_TORCH_COMPILE", "1"))):
13     jit_fuser = torch.compile
14 
15 # Decorator to disable Torch Dynamo
16 # See: https://github.com/NVIDIA/TransformerEngine/issues/308
17 no_torch_dynamo = lambda recursive=True: lambda func: func
18 if torch.__version__ >= "2":
19     import torch._dynamo
20     no_torch_dynamo = lambda recursive=True: lambda f: torch._dynamo.disable(f, recursive=recursive)
21 
22 
23 def set_jit_fusion_options() -> None:
24     """Set PyTorch JIT layer fusion options."""
25     # flags required to enable jit fusion kernels
26     TORCH_MAJOR = int(torch.__version__.split(".")[0])
27     TORCH_MINOR = int(torch.__version__.split(".")[1])
28     if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10):
29         # nvfuser
30         torch._C._jit_set_profiling_executor(True)
31         torch._C._jit_set_profiling_mode(True)
32         torch._C._jit_override_can_fuse_on_cpu(False)
33         torch._C._jit_override_can_fuse_on_gpu(False)
34         torch._C._jit_set_texpr_fuser_enabled(False)
35         torch._C._jit_set_nvfuser_enabled(True)
36         torch._C._debug_set_autodiff_subgraph_inlining(False)
37     else:
38         # legacy pytorch fuser
39         torch._C._jit_set_profiling_mode(False)
40         torch._C._jit_set_profiling_executor(False)
41         torch._C._jit_override_can_fuse_on_cpu(True)
42         torch._C._jit_override_can_fuse_on_gpu(True)
43 
44 
45 @jit_fuser
46 def bias_gelu_fused_(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
47     """Bias-GeLU fused"""
48     x = inp + bias
49     return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
50 
51 
52 @jit_fuser
53 def gelu_fused_(inp: torch.Tensor) -> torch.Tensor:
54     """
55     GeLU fused, this is copy of bias_gelu_fused cause jit fusion doesn't allow conditioning.
56     """
57     x = inp
58     return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
59 
60 
61 # gradient of tanh approximation of gelu
62 # gradient of actual gelu is:
63 # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
64 @jit_fuser
65 def bgrad_dgelu_fused_(
66     grad_output: torch.Tensor, inp: torch.Tensor, bias: torch.Tensor
67 ) -> Tuple[torch.Tensor, torch.Tensor]:
68     """Bgrad-Dgelu fused"""
69     x = inp + bias
70     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
71     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
72     ff = 0.5 * x * (
73         (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)
74     ) + 0.5 * (1 + tanh_out)
75     dgelu = ff * grad_output
76     bgrad = dgelu.sum(dim=0)
77     return bgrad, dgelu
78 
79 
80 @jit_fuser
81 def dgelu_fused_(
82     grad_output: torch.Tensor, inp: torch.Tensor
83 ) -> torch.Tensor:
84     """
85     Dgelu fused, this is copy of bgrad_dgelu_fused_ cause jit fusion doesn't allow conditioning.
86     """
87     x = inp
88     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
89     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
90     ff = 0.5 * x * (
91         (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)
92     ) + 0.5 * (1 + tanh_out)
93     dgelu = ff * grad_output
94     return dgelu
95 
96 
97 def bias_gelu_fused(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
98     """Disable native AMP for bias_gelu_fused_"""
99     with torch.cuda.amp.autocast(enabled=False):
100         if bias.numel() != 0:
101             return bias_gelu_fused_(inp, bias)
102         return gelu_fused_(inp)
103 
104 
105 def bgrad_dgelu_fused(
106     grad_output: torch.Tensor, inp: torch.Tensor, bias: torch.Tensor
107 ) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
108     """Disable native AMP for `bgrad_dgelu_fused_`"""
109     with torch.cuda.amp.autocast(enabled=False):
110         if bias.numel() != 0:
111             return bgrad_dgelu_fused_(grad_output, inp, bias)
112         return None, dgelu_fused_(grad_output, inp)
113 
114 
115 def bias_dropout_add(
116     x: torch.Tensor,
117     bias: torch.Tensor,
118     residual: torch.Tensor,
119     prob: float,
120     training: bool,
121 ) -> torch.Tensor:
122     """dropout(inp + bias) + residual"""
123     out = torch.nn.functional.dropout(x + bias, p=prob, training=training)
124     out = residual + out
125     return out
126 
127 
128 def get_bias_dropout_add(training: bool) -> Callable:
129     """bias_dropout_add based on training or not"""
130 
131     def _bias_dropout_add(x, bias, residual, prob):
132         return bias_dropout_add(x, bias, residual, prob, training)
133 
134     return _bias_dropout_add
135 
136 
137 @jit_fuser
138 def bias_dropout_add_fused_train_(
139     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
140 ) -> torch.Tensor:
141     """Jit fused bias_dropout_add for training"""
142     return bias_dropout_add(x, bias, residual, prob, True)
143 
144 
145 def bias_dropout_add_fused_train(
146     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
147 ) -> torch.Tensor:
148     """Disable native AMP and enable grad for BDA"""
149     with torch.enable_grad():
150         with torch.cuda.amp.autocast(enabled=False):
151             return bias_dropout_add_fused_train_(x, bias, residual, prob)
152 
153 
154 @jit_fuser
155 def bias_dropout_add_fused_inference_(
156     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
157 ) -> torch.Tensor:
158     """Jit fused bias_dropout_add for inference"""
159     return bias_dropout_add(x, bias, residual, prob, False)
160 
161 
162 def bias_dropout_add_fused_inference(
163     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
164 ) -> torch.Tensor:
165     """Disable native AMP for BDA"""
166     with torch.cuda.amp.autocast(enabled=False):
167         return bias_dropout_add_fused_inference_(x, bias, residual, prob)
168 
169 
170 def warmup_jit_bias_dropout_add(
171     hidden_size: int, dtype: torch.dtype, seq_length: int, micro_batch_size: int
172 ) -> None:
173     """Compile BDA JIT function before the main training steps"""
174 
175     # Save cuda RNG state to ensure warmup does not affect reproducibility.
176     rng_state = torch.cuda.get_rng_state()
177 
178     inp = torch.rand(
179         (seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda"
180     )
181     residual = torch.rand(
182         (seq_length, micro_batch_size, hidden_size), dtype=dtype, device="cuda"
183     )
184     bias = torch.rand((hidden_size), dtype=dtype, device="cuda")
185     dropout_rate = 0.1
186     # Warmup JIT fusions with the input grad_enable state of both forward
187     # prop and recomputation
188     for input_grad, bias_grad, residual_grad in zip(
189         [False, True], [True, True], [True, True]
190     ):
191         inp.requires_grad = input_grad
192         bias.requires_grad = bias_grad
193         residual.requires_grad = residual_grad
194         for _ in range(5):
195             output = bias_dropout_add_fused_train(inp, bias, residual, dropout_rate)
196     del bias, inp, residual, output
197 
198     torch.cuda.empty_cache()
199     torch.cuda.set_rng_state(rng_state)
200 
201 
202 def warmup_jit_bias_dropout_add_all_dtypes(
203     hidden_size: int, seq_length: int, micro_batch_size: int
204 ) -> None:
205     """Call `warmup_jit_bias_dropout_add` for all training dtypes"""
206     for dtype in [torch.float32, torch.bfloat16, torch.float16]:
207         warmup_jit_bias_dropout_add(hidden_size, dtype, seq_length, micro_batch_size)
208 
209 
210 def warmup_jit_bias_gelu(
211     ffn_hidden_size_per_partition: int,
212     dtype: torch.dtype,
213     seq_length: int,
214     micro_batch_size: int,
215 ) -> None:
216     """Compile bias-gelu JIT function before the main training steps"""
217 
218     # Save cuda RNG state to ensure warmup does not affect reproducibility.
219     rng_state = torch.cuda.get_rng_state()
220 
221     bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device="cuda")
222     inp = torch.rand(
223         (seq_length * micro_batch_size, ffn_hidden_size_per_partition),
224         dtype=dtype,
225         device="cuda",
226     )
227     # Warmup JIT fusions with the input grad_enable state of both forward
228     # prop and recomputation
229     for bias_grad, input_grad in zip([True, True], [False, True]):
230         bias.requires_grad, inp.requires_grad = bias_grad, input_grad
231         for _ in range(5):
232             _ = bias_gelu_fused_(inp, bias)
233             _ = gelu_fused_(inp)
234     del bias, inp
235 
236     torch.cuda.empty_cache()
237     torch.cuda.set_rng_state(rng_state)
238 
239 
240 def warmup_jit_bias_gelu_all_dtypes(
241     ffn_hidden_size: int, seq_length: int, micro_batch_size: int
242 ) -> None:
243     """Call `warmup_jit_bias_gelu` for all training dtypes"""
244     for dtype in [torch.float32, torch.bfloat16, torch.float16]:
245         warmup_jit_bias_gelu(ffn_hidden_size, dtype, seq_length, micro_batch_size)
246
```
--- 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/jit.py b/transformer_engine/pytorch/jit.py
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -12,6 +12,11 @@
 if torch.__version__ >= "2" and bool(int(os.getenv("NVTE_TORCH_COMPILE", "1"))):
     jit_fuser = torch.compile
 
+# See: https://github.com/NVIDIA/TransformerEngine/issues/597
+dropout_fuser = torch.jit.script
+if torch.__version__ >= "2.2" and bool(int(os.getenv("NVTE_TORCH_COMPILE", "1"))):
+    dropout_fuser = torch.compile
+
 # Decorator to disable Torch Dynamo
 # See: https://github.com/NVIDIA/TransformerEngine/issues/308
 no_torch_dynamo = lambda recursive=True: lambda func: func
@@ -134,7 +139,7 @@
     return _bias_dropout_add
 
 
-@jit_fuser
+@dropout_fuser
 def bias_dropout_add_fused_train_(
     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
 ) -> torch.Tensor:
@@ -151,7 +156,7 @@
             return bias_dropout_add_fused_train_(x, bias, residual, prob)
 
 
-@jit_fuser
+@dropout_fuser
 def bias_dropout_add_fused_inference_(
     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
 ) -> torch.Tensor:

{"golden_diff": "diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py\n--- a/transformer_engine/pytorch/jit.py\n+++ b/transformer_engine/pytorch/jit.py\n@@ -12,6 +12,11 @@\n if torch.__version__ >= \"2\" and bool(int(os.getenv(\"NVTE_TORCH_COMPILE\", \"1\"))):\n     jit_fuser = torch.compile\n \n+# See: https://github.com/NVIDIA/TransformerEngine/issues/597\n+dropout_fuser = torch.jit.script\n+if torch.__version__ >= \"2.2\" and bool(int(os.getenv(\"NVTE_TORCH_COMPILE\", \"1\"))):\n+    dropout_fuser = torch.compile\n+\n # Decorator to disable Torch Dynamo\n # See: https://github.com/NVIDIA/TransformerEngine/issues/308\n no_torch_dynamo = lambda recursive=True: lambda func: func\n@@ -134,7 +139,7 @@\n     return _bias_dropout_add\n \n \n-@jit_fuser\n+@dropout_fuser\n def bias_dropout_add_fused_train_(\n     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n ) -> torch.Tensor:\n@@ -151,7 +156,7 @@\n             return bias_dropout_add_fused_train_(x, bias, residual, prob)\n \n \n-@jit_fuser\n+@dropout_fuser\n def bias_dropout_add_fused_inference_(\n     x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n ) -> torch.Tensor:\n", "issue": "Performance regression of `bias_dropout_add_fused_train` after switching to `torch.compile`\nThere is a performance regression of `bias_dropout_add_fused_train()` after switching to `torch.compile` for jit fusion in this [PR](https://github.com/NVIDIA/TransformerEngine/pull/589).\r\n\r\nHere are my profiling results. As we can see, from either host or device perspective, there's a big regression.\r\n\r\n- torch.compile\r\n![image](https://github.com/NVIDIA/TransformerEngine/assets/3831900/7bf647b5-71d6-4b52-a264-c6d48633bc78)\r\n\r\n- nvFuser\r\n![image](https://github.com/NVIDIA/TransformerEngine/assets/3831900/c3654308-f914-490e-a52d-2676308e0071)\r\n\r\nMy profiling scripts:\r\n```python\r\nimport torch\r\nimport transformer_engine as te\r\n\r\ndef profile_bias_dropout_add_fused_train(\r\n     hidden_size: int, dtype: torch.dtype, seq_length: int, micro_batch_size: int\r\n) -> None:\r\n    # warm up\r\n    te.pytorch.jit.warmup_jit_bias_dropout_add(hidden_size, dtype, seq_length, micro_batch_size)\r\n\r\n    inp = torch.rand(\r\n        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device=\"cuda\"\r\n    )\r\n    residual = torch.rand(\r\n        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device=\"cuda\"\r\n    )\r\n    bias = torch.rand((hidden_size), dtype=dtype, device=\"cuda\")\r\n    dropout_rate = 0.1\r\n\r\n    torch.cuda.profiler.start()\r\n    for _ in range(10):\r\n        torch.cuda.nvtx.range_push(\"fused bda\")\r\n        output = te.pytorch.jit.bias_dropout_add_fused_train(inp, bias, residual, dropout_rate)\r\n        torch.cuda.nvtx.range_pop()\r\n    torch.cuda.profiler.stop()\r\n\r\nte.pytorch.jit.set_jit_fusion_options()\r\nprofile_bias_dropout_add_fused_train(1024, torch.float32, 1024, 2)\r\n```\r\n\r\ncc @ksivaman @minitu \n", "before_files": [{"content": "# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# See LICENSE for license information.\n\n\"\"\"NVFuser functions and JIT utilities\"\"\"\nimport os\nfrom typing import Callable, Optional, Tuple\n\nimport torch\n\njit_fuser = torch.jit.script\nif torch.__version__ >= \"2\" and bool(int(os.getenv(\"NVTE_TORCH_COMPILE\", \"1\"))):\n    jit_fuser = torch.compile\n\n# Decorator to disable Torch Dynamo\n# See: https://github.com/NVIDIA/TransformerEngine/issues/308\nno_torch_dynamo = lambda recursive=True: lambda func: func\nif torch.__version__ >= \"2\":\n    import torch._dynamo\n    no_torch_dynamo = lambda recursive=True: lambda f: torch._dynamo.disable(f, recursive=recursive)\n\n\ndef set_jit_fusion_options() -> None:\n    \"\"\"Set PyTorch JIT layer fusion options.\"\"\"\n    # flags required to enable jit fusion kernels\n    TORCH_MAJOR = int(torch.__version__.split(\".\")[0])\n    TORCH_MINOR = int(torch.__version__.split(\".\")[1])\n    if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10):\n        # nvfuser\n        torch._C._jit_set_profiling_executor(True)\n        torch._C._jit_set_profiling_mode(True)\n        torch._C._jit_override_can_fuse_on_cpu(False)\n        torch._C._jit_override_can_fuse_on_gpu(False)\n        torch._C._jit_set_texpr_fuser_enabled(False)\n        torch._C._jit_set_nvfuser_enabled(True)\n        torch._C._debug_set_autodiff_subgraph_inlining(False)\n    else:\n        # legacy pytorch fuser\n        torch._C._jit_set_profiling_mode(False)\n        torch._C._jit_set_profiling_executor(False)\n        torch._C._jit_override_can_fuse_on_cpu(True)\n        torch._C._jit_override_can_fuse_on_gpu(True)\n\n\n@jit_fuser\ndef bias_gelu_fused_(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:\n    \"\"\"Bias-GeLU fused\"\"\"\n    x = inp + bias\n    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))\n\n\n@jit_fuser\ndef gelu_fused_(inp: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n    GeLU fused, this is copy of bias_gelu_fused cause jit fusion doesn't allow conditioning.\n    \"\"\"\n    x = inp\n    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))\n\n\n# gradient of tanh approximation of gelu\n# gradient of actual gelu is:\n# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)\n@jit_fuser\ndef bgrad_dgelu_fused_(\n    grad_output: torch.Tensor, inp: torch.Tensor, bias: torch.Tensor\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Bgrad-Dgelu fused\"\"\"\n    x = inp + bias\n    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))\n    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243\n    ff = 0.5 * x * (\n        (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)\n    ) + 0.5 * (1 + tanh_out)\n    dgelu = ff * grad_output\n    bgrad = dgelu.sum(dim=0)\n    return bgrad, dgelu\n\n\n@jit_fuser\ndef dgelu_fused_(\n    grad_output: torch.Tensor, inp: torch.Tensor\n) -> torch.Tensor:\n    \"\"\"\n    Dgelu fused, this is copy of bgrad_dgelu_fused_ cause jit fusion doesn't allow conditioning.\n    \"\"\"\n    x = inp\n    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))\n    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243\n    ff = 0.5 * x * (\n        (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)\n    ) + 0.5 * (1 + tanh_out)\n    dgelu = ff * grad_output\n    return dgelu\n\n\ndef bias_gelu_fused(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:\n    \"\"\"Disable native AMP for bias_gelu_fused_\"\"\"\n    with torch.cuda.amp.autocast(enabled=False):\n        if bias.numel() != 0:\n            return bias_gelu_fused_(inp, bias)\n        return gelu_fused_(inp)\n\n\ndef bgrad_dgelu_fused(\n    grad_output: torch.Tensor, inp: torch.Tensor, bias: torch.Tensor\n) -> Tuple[Optional[torch.Tensor], torch.Tensor]:\n    \"\"\"Disable native AMP for `bgrad_dgelu_fused_`\"\"\"\n    with torch.cuda.amp.autocast(enabled=False):\n        if bias.numel() != 0:\n            return bgrad_dgelu_fused_(grad_output, inp, bias)\n        return None, dgelu_fused_(grad_output, inp)\n\n\ndef bias_dropout_add(\n    x: torch.Tensor,\n    bias: torch.Tensor,\n    residual: torch.Tensor,\n    prob: float,\n    training: bool,\n) -> torch.Tensor:\n    \"\"\"dropout(inp + bias) + residual\"\"\"\n    out = torch.nn.functional.dropout(x + bias, p=prob, training=training)\n    out = residual + out\n    return out\n\n\ndef get_bias_dropout_add(training: bool) -> Callable:\n    \"\"\"bias_dropout_add based on training or not\"\"\"\n\n    def _bias_dropout_add(x, bias, residual, prob):\n        return bias_dropout_add(x, bias, residual, prob, training)\n\n    return _bias_dropout_add\n\n\n@jit_fuser\ndef bias_dropout_add_fused_train_(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Jit fused bias_dropout_add for training\"\"\"\n    return bias_dropout_add(x, bias, residual, prob, True)\n\n\ndef bias_dropout_add_fused_train(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Disable native AMP and enable grad for BDA\"\"\"\n    with torch.enable_grad():\n        with torch.cuda.amp.autocast(enabled=False):\n            return bias_dropout_add_fused_train_(x, bias, residual, prob)\n\n\n@jit_fuser\ndef bias_dropout_add_fused_inference_(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Jit fused bias_dropout_add for inference\"\"\"\n    return bias_dropout_add(x, bias, residual, prob, False)\n\n\ndef bias_dropout_add_fused_inference(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Disable native AMP for BDA\"\"\"\n    with torch.cuda.amp.autocast(enabled=False):\n        return bias_dropout_add_fused_inference_(x, bias, residual, prob)\n\n\ndef warmup_jit_bias_dropout_add(\n    hidden_size: int, dtype: torch.dtype, seq_length: int, micro_batch_size: int\n) -> None:\n    \"\"\"Compile BDA JIT function before the main training steps\"\"\"\n\n    # Save cuda RNG state to ensure warmup does not affect reproducibility.\n    rng_state = torch.cuda.get_rng_state()\n\n    inp = torch.rand(\n        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device=\"cuda\"\n    )\n    residual = torch.rand(\n        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device=\"cuda\"\n    )\n    bias = torch.rand((hidden_size), dtype=dtype, device=\"cuda\")\n    dropout_rate = 0.1\n    # Warmup JIT fusions with the input grad_enable state of both forward\n    # prop and recomputation\n    for input_grad, bias_grad, residual_grad in zip(\n        [False, True], [True, True], [True, True]\n    ):\n        inp.requires_grad = input_grad\n        bias.requires_grad = bias_grad\n        residual.requires_grad = residual_grad\n        for _ in range(5):\n            output = bias_dropout_add_fused_train(inp, bias, residual, dropout_rate)\n    del bias, inp, residual, output\n\n    torch.cuda.empty_cache()\n    torch.cuda.set_rng_state(rng_state)\n\n\ndef warmup_jit_bias_dropout_add_all_dtypes(\n    hidden_size: int, seq_length: int, micro_batch_size: int\n) -> None:\n    \"\"\"Call `warmup_jit_bias_dropout_add` for all training dtypes\"\"\"\n    for dtype in [torch.float32, torch.bfloat16, torch.float16]:\n        warmup_jit_bias_dropout_add(hidden_size, dtype, seq_length, micro_batch_size)\n\n\ndef warmup_jit_bias_gelu(\n    ffn_hidden_size_per_partition: int,\n    dtype: torch.dtype,\n    seq_length: int,\n    micro_batch_size: int,\n) -> None:\n    \"\"\"Compile bias-gelu JIT function before the main training steps\"\"\"\n\n    # Save cuda RNG state to ensure warmup does not affect reproducibility.\n    rng_state = torch.cuda.get_rng_state()\n\n    bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device=\"cuda\")\n    inp = torch.rand(\n        (seq_length * micro_batch_size, ffn_hidden_size_per_partition),\n        dtype=dtype,\n        device=\"cuda\",\n    )\n    # Warmup JIT fusions with the input grad_enable state of both forward\n    # prop and recomputation\n    for bias_grad, input_grad in zip([True, True], [False, True]):\n        bias.requires_grad, inp.requires_grad = bias_grad, input_grad\n        for _ in range(5):\n            _ = bias_gelu_fused_(inp, bias)\n            _ = gelu_fused_(inp)\n    del bias, inp\n\n    torch.cuda.empty_cache()\n    torch.cuda.set_rng_state(rng_state)\n\n\ndef warmup_jit_bias_gelu_all_dtypes(\n    ffn_hidden_size: int, seq_length: int, micro_batch_size: int\n) -> None:\n    \"\"\"Call `warmup_jit_bias_gelu` for all training dtypes\"\"\"\n    for dtype in [torch.float32, torch.bfloat16, torch.float16]:\n        warmup_jit_bias_gelu(ffn_hidden_size, dtype, seq_length, micro_batch_size)\n", "path": "transformer_engine/pytorch/jit.py"}], "after_files": [{"content": "# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# See LICENSE for license information.\n\n\"\"\"NVFuser functions and JIT utilities\"\"\"\nimport os\nfrom typing import Callable, Optional, Tuple\n\nimport torch\n\njit_fuser = torch.jit.script\nif torch.__version__ >= \"2\" and bool(int(os.getenv(\"NVTE_TORCH_COMPILE\", \"1\"))):\n    jit_fuser = torch.compile\n\n# See: https://github.com/NVIDIA/TransformerEngine/issues/597\ndropout_fuser = torch.jit.script\nif torch.__version__ >= \"2.2\" and bool(int(os.getenv(\"NVTE_TORCH_COMPILE\", \"1\"))):\n    dropout_fuser = torch.compile\n\n# Decorator to disable Torch Dynamo\n# See: https://github.com/NVIDIA/TransformerEngine/issues/308\nno_torch_dynamo = lambda recursive=True: lambda func: func\nif torch.__version__ >= \"2\":\n    import torch._dynamo\n    no_torch_dynamo = lambda recursive=True: lambda f: torch._dynamo.disable(f, recursive=recursive)\n\n\ndef set_jit_fusion_options() -> None:\n    \"\"\"Set PyTorch JIT layer fusion options.\"\"\"\n    # flags required to enable jit fusion kernels\n    TORCH_MAJOR = int(torch.__version__.split(\".\")[0])\n    TORCH_MINOR = int(torch.__version__.split(\".\")[1])\n    if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10):\n        # nvfuser\n        torch._C._jit_set_profiling_executor(True)\n        torch._C._jit_set_profiling_mode(True)\n        torch._C._jit_override_can_fuse_on_cpu(False)\n        torch._C._jit_override_can_fuse_on_gpu(False)\n        torch._C._jit_set_texpr_fuser_enabled(False)\n        torch._C._jit_set_nvfuser_enabled(True)\n        torch._C._debug_set_autodiff_subgraph_inlining(False)\n    else:\n        # legacy pytorch fuser\n        torch._C._jit_set_profiling_mode(False)\n        torch._C._jit_set_profiling_executor(False)\n        torch._C._jit_override_can_fuse_on_cpu(True)\n        torch._C._jit_override_can_fuse_on_gpu(True)\n\n\n@jit_fuser\ndef bias_gelu_fused_(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:\n    \"\"\"Bias-GeLU fused\"\"\"\n    x = inp + bias\n    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))\n\n\n@jit_fuser\ndef gelu_fused_(inp: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n    GeLU fused, this is copy of bias_gelu_fused cause jit fusion doesn't allow conditioning.\n    \"\"\"\n    x = inp\n    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))\n\n\n# gradient of tanh approximation of gelu\n# gradient of actual gelu is:\n# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)\n@jit_fuser\ndef bgrad_dgelu_fused_(\n    grad_output: torch.Tensor, inp: torch.Tensor, bias: torch.Tensor\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Bgrad-Dgelu fused\"\"\"\n    x = inp + bias\n    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))\n    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243\n    ff = 0.5 * x * (\n        (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)\n    ) + 0.5 * (1 + tanh_out)\n    dgelu = ff * grad_output\n    bgrad = dgelu.sum(dim=0)\n    return bgrad, dgelu\n\n\n@jit_fuser\ndef dgelu_fused_(\n    grad_output: torch.Tensor, inp: torch.Tensor\n) -> torch.Tensor:\n    \"\"\"\n    Dgelu fused, this is copy of bgrad_dgelu_fused_ cause jit fusion doesn't allow conditioning.\n    \"\"\"\n    x = inp\n    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))\n    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243\n    ff = 0.5 * x * (\n        (1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)\n    ) + 0.5 * (1 + tanh_out)\n    dgelu = ff * grad_output\n    return dgelu\n\n\ndef bias_gelu_fused(inp: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:\n    \"\"\"Disable native AMP for bias_gelu_fused_\"\"\"\n    with torch.cuda.amp.autocast(enabled=False):\n        if bias.numel() != 0:\n            return bias_gelu_fused_(inp, bias)\n        return gelu_fused_(inp)\n\n\ndef bgrad_dgelu_fused(\n    grad_output: torch.Tensor, inp: torch.Tensor, bias: torch.Tensor\n) -> Tuple[Optional[torch.Tensor], torch.Tensor]:\n    \"\"\"Disable native AMP for `bgrad_dgelu_fused_`\"\"\"\n    with torch.cuda.amp.autocast(enabled=False):\n        if bias.numel() != 0:\n            return bgrad_dgelu_fused_(grad_output, inp, bias)\n        return None, dgelu_fused_(grad_output, inp)\n\n\ndef bias_dropout_add(\n    x: torch.Tensor,\n    bias: torch.Tensor,\n    residual: torch.Tensor,\n    prob: float,\n    training: bool,\n) -> torch.Tensor:\n    \"\"\"dropout(inp + bias) + residual\"\"\"\n    out = torch.nn.functional.dropout(x + bias, p=prob, training=training)\n    out = residual + out\n    return out\n\n\ndef get_bias_dropout_add(training: bool) -> Callable:\n    \"\"\"bias_dropout_add based on training or not\"\"\"\n\n    def _bias_dropout_add(x, bias, residual, prob):\n        return bias_dropout_add(x, bias, residual, prob, training)\n\n    return _bias_dropout_add\n\n\n@dropout_fuser\ndef bias_dropout_add_fused_train_(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Jit fused bias_dropout_add for training\"\"\"\n    return bias_dropout_add(x, bias, residual, prob, True)\n\n\ndef bias_dropout_add_fused_train(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Disable native AMP and enable grad for BDA\"\"\"\n    with torch.enable_grad():\n        with torch.cuda.amp.autocast(enabled=False):\n            return bias_dropout_add_fused_train_(x, bias, residual, prob)\n\n\n@dropout_fuser\ndef bias_dropout_add_fused_inference_(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Jit fused bias_dropout_add for inference\"\"\"\n    return bias_dropout_add(x, bias, residual, prob, False)\n\n\ndef bias_dropout_add_fused_inference(\n    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float\n) -> torch.Tensor:\n    \"\"\"Disable native AMP for BDA\"\"\"\n    with torch.cuda.amp.autocast(enabled=False):\n        return bias_dropout_add_fused_inference_(x, bias, residual, prob)\n\n\ndef warmup_jit_bias_dropout_add(\n    hidden_size: int, dtype: torch.dtype, seq_length: int, micro_batch_size: int\n) -> None:\n    \"\"\"Compile BDA JIT function before the main training steps\"\"\"\n\n    # Save cuda RNG state to ensure warmup does not affect reproducibility.\n    rng_state = torch.cuda.get_rng_state()\n\n    inp = torch.rand(\n        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device=\"cuda\"\n    )\n    residual = torch.rand(\n        (seq_length, micro_batch_size, hidden_size), dtype=dtype, device=\"cuda\"\n    )\n    bias = torch.rand((hidden_size), dtype=dtype, device=\"cuda\")\n    dropout_rate = 0.1\n    # Warmup JIT fusions with the input grad_enable state of both forward\n    # prop and recomputation\n    for input_grad, bias_grad, residual_grad in zip(\n        [False, True], [True, True], [True, True]\n    ):\n        inp.requires_grad = input_grad\n        bias.requires_grad = bias_grad\n        residual.requires_grad = residual_grad\n        for _ in range(5):\n            output = bias_dropout_add_fused_train(inp, bias, residual, dropout_rate)\n    del bias, inp, residual, output\n\n    torch.cuda.empty_cache()\n    torch.cuda.set_rng_state(rng_state)\n\n\ndef warmup_jit_bias_dropout_add_all_dtypes(\n    hidden_size: int, seq_length: int, micro_batch_size: int\n) -> None:\n    \"\"\"Call `warmup_jit_bias_dropout_add` for all training dtypes\"\"\"\n    for dtype in [torch.float32, torch.bfloat16, torch.float16]:\n        warmup_jit_bias_dropout_add(hidden_size, dtype, seq_length, micro_batch_size)\n\n\ndef warmup_jit_bias_gelu(\n    ffn_hidden_size_per_partition: int,\n    dtype: torch.dtype,\n    seq_length: int,\n    micro_batch_size: int,\n) -> None:\n    \"\"\"Compile bias-gelu JIT function before the main training steps\"\"\"\n\n    # Save cuda RNG state to ensure warmup does not affect reproducibility.\n    rng_state = torch.cuda.get_rng_state()\n\n    bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device=\"cuda\")\n    inp = torch.rand(\n        (seq_length * micro_batch_size, ffn_hidden_size_per_partition),\n        dtype=dtype,\n        device=\"cuda\",\n    )\n    # Warmup JIT fusions with the input grad_enable state of both forward\n    # prop and recomputation\n    for bias_grad, input_grad in zip([True, True], [False, True]):\n        bias.requires_grad, inp.requires_grad = bias_grad, input_grad\n        for _ in range(5):\n            _ = bias_gelu_fused_(inp, bias)\n            _ = gelu_fused_(inp)\n    del bias, inp\n\n    torch.cuda.empty_cache()\n    torch.cuda.set_rng_state(rng_state)\n\n\ndef warmup_jit_bias_gelu_all_dtypes(\n    ffn_hidden_size: int, seq_length: int, micro_batch_size: int\n) -> None:\n    \"\"\"Call `warmup_jit_bias_gelu` for all training dtypes\"\"\"\n    for dtype in [torch.float32, torch.bfloat16, torch.float16]:\n        warmup_jit_bias_gelu(ffn_hidden_size, dtype, seq_length, micro_batch_size)\n", "path": "transformer_engine/pytorch/jit.py"}]}