Merge branch 'ko3n1g/chore/cherry-pick-3596' into 'core_r0.13.0'

Cherry-pick `!3596 - fix: Guard TE import`

See merge request ADLR/megatron-lm!3666

Author: ko3n1g

megatron/core/models/T5/t5_spec.py

@@ -16,6 +16,8 @@
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelLinear,
         TEDotProductAttention,

megatron/core/models/bert/bert_layer_specs.py

@@ -12,6 +12,8 @@
 from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEDotProductAttention,
         TELayerNormColumnParallelLinear,

megatron/core/models/gpt/gpt_layer_specs.py

@@ -34,6 +34,8 @@
 )
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
@@ -42,6 +44,8 @@
     HAVE_TE = False
 
 try:
+    import nvidia_kitchen  # pylint: disable=unused-import
+
     from megatron.core.extensions.kitchen import KitchenSpecProvider
 
     HAVE_KITCHEN = True

megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py

@@ -29,6 +29,8 @@
 from megatron.core.utils import is_te_min_version
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEDotProductAttention,
         TELayerNormColumnParallelLinear,

megatron/core/models/gpt/moe_module_specs.py

@@ -9,6 +9,8 @@
 from megatron.core.transformer.spec_utils import ModuleSpec
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
 
     HAVE_TE = True

megatron/core/models/retro/decoder_spec.py

@@ -40,6 +40,8 @@
     HAVE_APEX = False
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelLinear,
         TEDotProductAttention,

megatron/core/models/retro/encoder_spec.py

@@ -21,6 +21,8 @@
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelLinear,
         TEDotProductAttention,

megatron/core/transformer/cuda_graphs.py

@@ -23,6 +23,7 @@
 from megatron.core.utils import is_te_min_version
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
     from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, fp8_autocast
     from transformer_engine.pytorch.graph import restore_fp8_tensors, save_fp8_tensors
     from transformer_engine.pytorch.graph import set_capture_end as te_set_capture_end

megatron/core/transformer/heterogeneous/linear_replacements.py

@@ -16,6 +16,8 @@
 from megatron.core.utils import divide
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
+
     from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear
 
     HAVE_TE = True
@@ -67,6 +69,7 @@ def __init__(self, config: TransformerConfig, tp_comm_buffer_name: str, *args, *
             )
 
         def forward(self, x, **kwargs):
+            """Forward of TELayerNormColumnParallelLinearGathered"""
             out, bias = super().forward(x)
             assert bias is None, "bias should be None since we set skip_bias_add=False"
 
@@ -100,6 +103,7 @@ def forward(
         runtime_gather_output: bool | None = None,
         **kwargs,
     ):
+        """Forward of ColumnParallelLinearGathered"""
         out, bias = super().forward(input_, weight, runtime_gather_output)
         assert bias is None, "bias should be None since we set skip_bias_add=False"
 

megatron/core/transformer/moe/experts.py

@@ -40,6 +40,7 @@
 )
 
 try:
+    import transformer_engine as te  # pylint: disable=unused-import
 
     from megatron.core.extensions.transformer_engine import Fp8Padding, Fp8Unpadding