ADLR/megatron-lm!3340 - Fix the TE assertion for release

Co-authored-by: Gao Deng <gdeng@cw-dfw-cs-001-vscode-01.cm.cluster>

Author: 2 people

megatron/core/extensions/transformer_engine.py

@@ -1355,13 +1355,13 @@ def fused_apply_rotary_pos_emb(
             warnings.warn(
                 "transpose_output_memory is not supported by TE's fused RoPE and will be ignored."
             )
-        if is_te_min_version("2.3.0.dev0"):
+        if is_te_min_version("2.2.0.dev0"):
             return apply_rotary_pos_emb(
                 t, freqs, tensor_format="sbhd", interleaved=interleaved, fused=True
             )
         else:
             if interleaved:
-                raise ValueError("Only TE >= 2.3.0.dev0 supports interleaved fused RoPE.")
+                raise ValueError("Only TE >= 2.2.0.dev0 supports interleaved fused RoPE.")
             return apply_rotary_pos_emb(t, freqs, tensor_format="sbhd", fused=True)
 
     def fused_apply_rotary_pos_emb_thd(

megatron/core/models/common/embeddings/rope_utils.py

@@ -211,7 +211,7 @@ def apply_rotary_pos_emb(
                     except ImportError:
                         raise ImportError(
                             "TE interleaved fused RoPE is not available."
-                            "Please install TE >= 2.3.0.dev0."
+                            "Please install TE >= 2.2.0.dev0."
                         )
                 else:
                     assert (

megatron/core/transformer/transformer_config.py

@@ -910,10 +910,10 @@ def __post_init__(self):
 
         if self.apply_rope_fusion:
             if self.rotary_interleaved:
-                if not is_te_min_version("2.3.0.dev0"):
+                if not is_te_min_version("2.2.0.dev0"):
                     raise ValueError(
                         "rotary_interleaved does not work with apply_rope_fusion for "
-                        "TE < 2.3.0.dev0. Please install TE >= 2.3.0.dev0"
+                        "TE < 2.2.0.dev0. Please install TE >= 2.2.0.dev0"
                     )
 
             from megatron.core.models.common.embeddings.rope_utils import (