Skip to content

GitLab

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

Update jit.py

parent a0e873f8
Pipeline #2223 passed with stage
Show 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
......
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