We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
The version v1.2 has compatibility issues with PyTorch 2.0.1, causing errors during execution.
Hello,

I encountered the following error while using version v1.2 of Transformer Engine: 

`no_torch_dynamo = lambda recursive=True: lambda f: torch._dynamo.disable(f, recursive=recursive). The error message is TypeError: disable() got an unexpected keyword argument 'recursive'.
`

My environment setup is as follows: CUDA 11.8, PyTorch 2.0.1, and Python 3.10. I am experiencing this issue specifically in the context of using MegatronLM at commit fab0bd6 for Large Language Model (LLM) training. The error occurs right at the start of the training.

Possible cause of the issue:
The `disable` function in PyTorch 2.0.1 is defined as `def disable(f)`, and the recursive parameter was introduced only in later versions starting from PyTorch 2.1.0.

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

{"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@@ -22,7 +22,12 @@\n no_torch_dynamo = lambda recursive=True: lambda func: func\n if torch.__version__ >= \"2\":\n     import torch._dynamo\n-    no_torch_dynamo = lambda recursive=True: lambda f: torch._dynamo.disable(f, recursive=recursive)\n+    if torch.__version__ >= \"2.1\":\n+        no_torch_dynamo = lambda recursive=True: lambda f: \\\n+                                                    torch._dynamo.disable(f, recursive=recursive)\n+    else:\n+        # no \"recursive\" option in pyTorch 2.0 - it acts as if recursive was True\n+        no_torch_dynamo = lambda recursive=True: torch._dynamo.disable\n \n \n def set_jit_fusion_options() -> None:\n", "issue": "The version v1.2 has compatibility issues with PyTorch 2.0.1, causing errors during execution.\nHello,\r\n\r\nI encountered the following error while using version v1.2 of Transformer Engine: \r\n\r\n`no_torch_dynamo = lambda recursive=True: lambda f: torch._dynamo.disable(f, recursive=recursive). The error message is TypeError: disable() got an unexpected keyword argument 'recursive'.\r\n`\r\n\r\nMy environment setup is as follows: CUDA 11.8, PyTorch 2.0.1, and Python 3.10. I am experiencing this issue specifically in the context of using MegatronLM at commit fab0bd6 for Large Language Model (LLM) training. The error occurs right at the start of the training.\r\n\r\nPossible cause of the issue:\r\nThe `disable` function in PyTorch 2.0.1 is defined as `def disable(f)`, and the recursive parameter was introduced only in later versions starting from PyTorch 2.1.0.\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# 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"}], "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    if torch.__version__ >= \"2.1\":\n        no_torch_dynamo = lambda recursive=True: lambda f: \\\n                                                    torch._dynamo.disable(f, recursive=recursive)\n    else:\n        # no \"recursive\" option in pyTorch 2.0 - it acts as if recursive was True\n        no_torch_dynamo = lambda recursive=True: torch._dynamo.disable\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"}]}