- Disable torch.compile explicitly when using TE

EleutherAI · Feb 1, 2025 · b51938d · b51938d
1 parent 2118a14
commit b51938d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/megatron/model/transformer_engine.py b/megatron/model/transformer_engine.py
@@ -41,6 +41,10 @@
 from megatron.model.positional_embeddings import RotaryEmbedding
 from megatron import mpu
 
+# https://github.com/NVIDIA/TransformerEngine/issues/405
+import os
+os.environ['NVTE_TORCH_COMPILE'] = str(0)
+
 try:
     import transformer_engine as te
 except ImportError:
@@ -50,6 +54,7 @@
     )
 
 
+
 class TERMSNorm(torch.nn.Module):
     def __init__(self, dim, eps=1e-8, **kwargs):
         """