Skip to content

GitLab

Commit 35cb3cc1 authored by silencealiang's avatar silencealiang
Browse files

Update jit.py

parent a0e873f8
Pipeline #2223 passed with stage
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
megatron/core/jit.py
View file @ 35cb3cc1
...@@ -7,7 +7,7 @@ from megatron.core.utils import is_torch_min_version ...@@ -7,7 +7,7 @@ from megatron.core.utils import is_torch_min_version
jit_fuser = torch.jit.script jit_fuser = torch.jit.script
# nvFuser is deprecated in PyTorch JIT starting from 2.2 # nvFuser is deprecated in PyTorch JIT starting from 2.2
if is_torch_min_version("2.2.0a0"): 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 # Decorator to disable Torch Dynamo
# See: https://github.com/NVIDIA/TransformerEngine/issues/308 # See: https://github.com/NVIDIA/TransformerEngine/issues/308
...@@ -21,4 +21,4 @@ if torch.__version__ >= "2": ...@@ -21,4 +21,4 @@ if torch.__version__ >= "2":
) )
else: else:
# no "recursive" option in pyTorch 2.0 - it acts as if recursive was True # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True
no_torch_dynamo = lambda recursive=True: torch._dynamo.disable no_torch_dynamo = lambda recursive=True: torch._dynamo.disable
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment