check te version

54740897 · dongcl · 2862a32a · 54740897 · 54740897
Commit 54740897 authored May 07, 2025 by dongcl
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

.gitignore .gitignore +1 -0

dcu_megatron/adaptor/megatron_adaptor.py dcu_megatron/adaptor/megatron_adaptor.py +1 -1

No files found.
--- a/.gitignore
+++ b/.gitignore
+__pycache__
--- a/dcu_megatron/adaptor/megatron_adaptor.py
+++ b/dcu_megatron/adaptor/megatron_adaptor.py
@@ -140,7 +140,7 @@ class CoreAdaptation(MegatronAdaptationABC):
                                        TEDotProductAttentionPatch.__init__)

        if int(os.getenv("GROUPED_GEMM_BatchLinear", '0')):
-            TEGroupedLinear.__bases__ = (te.pytorch.BatchedLinear if is_te_min_version("2.3.0") else te.pytorch.BatchLinear,)
+            TEGroupedLinear.__bases__ = (te.pytorch.BatchedLinear if is_te_min_version("2.3.0.dev0") else te.pytorch.BatchLinear,)

    def patch_tensor_parallel(self):
        from ..core.tensor_parallel.cross_entropy import VocabParallelCrossEntropy