Merge branch 'mcore_te_api_fixes' into 'main'

Use TE provided APIs See merge request ADLR/megatron-lm!690
NVIDIA · Jul 24, 2023 · 65da5be · 65da5be
2 parents 5f03f6d + 35b2860
commit 65da5be
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/megatron/core/transformer/custom_layers/transformer_engine.py b/megatron/core/transformer/custom_layers/transformer_engine.py
@@ -9,7 +9,7 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 
 
-class TELayerNorm(te.pytorch.module.LayerNorm):
+class TELayerNorm(te.pytorch.LayerNorm):
     """
     Wrapper for the Transformer-Engine's `LayerNorm`.
     """
@@ -20,7 +20,7 @@ def __init__(
         super().__init__(hidden_size=hidden_size, eps=eps, sequence_parallel=sequence_parallel)
 
 
-class TELinear(te.pytorch.module.Linear):
+class TELinear(te.pytorch.Linear):
     """
     Wrapper for the Transformer-Engine's `Linear` layer.
 
@@ -111,7 +111,7 @@ def __init__(self, input_size: int, output_size: int, config: TransformerConfig,
         )
 
 
-class TECoreAttention(te.pytorch.transformer.DotProductAttention):
+class TECoreAttention(te.pytorch.DotProductAttention):
     """
     Wrapper for the Transformer-Engine's `DotProductAttention` layer that also
     has "flash attention" enabled.

diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -1516,7 +1516,7 @@ def custom_forward(*args, **kwargs):
             l = 0
             while l < self.num_layers:
                 if self.transformer_impl == 'transformer_engine':
-                    hidden_states = transformer_engine.pytorch.distributed.checkpoint(
+                    hidden_states = transformer_engine.pytorch.checkpoint(
                         custom(l, l + self.recompute_num_layers),
                         self.distribute_saved_activations,
                         tensor_parallel.get_cuda_rng_tracker,
@@ -1540,7 +1540,7 @@ def custom_forward(*args, **kwargs):
             for l in range(self.num_layers):
                 if l < self.recompute_num_layers:
                     if self.transformer_impl == 'transformer_engine':
-                        hidden_states = transformer_engine.pytorch.distributed.checkpoint(
+                        hidden_states = transformer_engine.pytorch.checkpoint(
                             custom(l, l + 1),
                             self.distribute_saved_activations,
                             tensor_parallel.get_cuda_rng_tracker,