Update transformer.py

megatron/legacy/model/transformer.py

@@ -165,7 +165,7 @@ class ParallelMLP(MegatronModule):
             is_expert=is_expert,
         )
 
-    @torch.compile(mode="max-autotune-no-cudagraphs")
+    # @torch.compile(mode="max-autotune-no-cudagraphs")
     def forward(self, hidden_states):
 
         # [s, b, 4hp]
@@ -1213,7 +1213,7 @@ class ParallelTransformerLayer(MegatronModule):
 
         return retriever_output, norm_input, norm_output
 
-    @torch.compile(mode="max-autotune-no-cudagraphs")
+    # @torch.compile(mode="max-autotune-no-cudagraphs")
     def forward(self, hidden_states, attention_mask,
                 encoder_output=None, enc_dec_attn_mask=None,
                 retriever_input=None,