modify TEGroupedLinear base

dcu_megatron/adaptor/megatron_adaptor.py

@@ -5,6 +5,8 @@ import types
 import argparse
 import torch
 
+from megatron.core.utils import is_te_min_version
+
 
 class MegatronAdaptation:
     """
@@ -132,12 +134,13 @@ class CoreAdaptation(MegatronAdaptationABC):
         from ..core.extensions.transformer_engine import TEDotProductAttentionPatch
         from megatron.core.extensions.transformer_engine import TEGroupedLinear
 
-        # kv channels, te_min_version 1.10.0 -> 1.9.0
-        MegatronAdaptation.register('megatron.core.extensions.transformer_engine.TEDotProductAttention.__init__',
-                                    TEDotProductAttentionPatch.__init__)
+        if not is_te_min_version("1.10.0"):
+            # kv channels, te_min_version 1.10.0 -> 1.9.0
+            MegatronAdaptation.register('megatron.core.extensions.transformer_engine.TEDotProductAttention.__init__',
+                                        TEDotProductAttentionPatch.__init__)
 
         if int(os.getenv("GROUPED_GEMM_BatchLinear", '0')):
-            TEGroupedLinear.__bases__ = (te.pytorch.BatchLinear,)
+            TEGroupedLinear.__bases__ = (te.pytorch.BatchedLinear if is_te_min_version("2.3.0") else te.pytorch.BatchLinear,)
 
     def patch_tensor_parallel(self):
         from ..core.tensor_parallel.cross_entropy import VocabParallelCrossEntropy