Update jit.py

megatron/core/jit.py

@@ -7,7 +7,7 @@ from megatron.core.utils import is_torch_min_version
 jit_fuser = torch.jit.script
 # nvFuser is deprecated in PyTorch JIT starting from 2.2
 if is_torch_min_version("2.2.0a0"):
-    jit_fuser = torch.compile
+    jit_fuser = torch.compile(mode='max-autotune-no-cudagraphs')
 
 # Decorator to disable Torch Dynamo
 # See: https://github.com/NVIDIA/TransformerEngine/issues/308
@@ -21,4 +21,4 @@ if torch.__version__ >= "2":
         )
     else:
         # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True
-        no_torch_dynamo = lambda recursive=True: torch._dynamo.disable
\ No newline at end of file
+        no_torch_dynamo = lambda recursive=True: torch._dynamo.disable