- Updated transformer_engine requirments

- Import megablocks only if num_experts > 1

megatron/model/transformer.py

@@ -26,8 +26,6 @@
 from pkg_resources import packaging
 from importlib.metadata import version
 
-from megatron.model.moe import ParallelDroplessMoE
-
 from .norms import get_norm
 from megatron import mpu
 from megatron.model.fused_softmax import FusedScaleMaskSoftmax
@@ -947,6 +945,9 @@ def __init__(
             else 1
         )
 
+        if self.num_experts > 1:
+            from megatron.model.moe import ParallelDroplessMoE
+
         if self.gpt_j_residual:
             # GPT-J style layers allow us to defer the reduction of results across TP ranks until the end of the two sublayers.
             # the reduction we use is a simple allreduce for pure Tensor Parallel,

requirements/requirements-transformerengine.txt

@@ -1 +1,2 @@
-transformer-engine[pytorch]
+transformer-engine[pytorch]==1.12
+flash-attn==2.6.3