Repository: NVIDIA/TransformerEngine

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 


Discussion:
What is the PyTorch version being used in these profiles? @yaox12 
@ksivaman I'm using the NGC PyTorch 23.10 container, so the PyTorch version is `2.1.0a0+32f93b1`.
The performance of `torch.compile` has improved drastically when fusing kernels with dropout op and using 23.12 container, you should see the performance uplifts. Perhaps we can decide which fuser to use based on the torch version and the corresponding performance?
I believe 23.11 with PyTorch version `2.2.0a0+6a974be` is when nvFuser support is fully deprecated in PyTorch JIT, causing kernels decorated with `torch.jit.script` to be unfused.
So one method I can think of is to add a `jit_fuser_2.2` or a similar decorator for certain fusions including bias-dropout-add that switches from `torch.jit.script` to `torch.compile` when PyTorch version >= 2.2 (I believe the current `jit_fuser` does the switch with PyTorch version >= 2.0).

Resolving patch:
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 @@ def _bias_dropout_add(x, bias, residual, prob):
     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 @@ def bias_dropout_add_fused_train(
             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:


Test patch: