fix(tp-overlap): adapt transformer_engine 2.4 for Megatron backend (#259)

Co-authored-by: Xiaoming-AMD <[email protected]>

Author: zhenhuang12

primus/backends/transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -16,11 +16,20 @@
 
 import primus.backends.transformer_engine.transformer_engine_torch as ptex
 
-if is_te_min_version("2.0"):
-    from transformer_engine.pytorch.cpp_extensions.gemm import (
-        reset_swizzled_inputs,
-        swizzle_inputs,
+if is_te_min_version("2.3", check_equality=False):
+    from transformer_engine.debug.pytorch.debug_quantization import DebugQuantizer
+    from transformer_engine.pytorch.tensor._internal.float8_blockwise_tensor_base import (
+        Float8BlockwiseQTensorBase,
     )
+
+if is_te_min_version("2.0"):
+
+    # TE version >= 2.0 and <= 2.3
+    if not is_te_min_version("2.3", check_equality=False):
+        from transformer_engine.pytorch.cpp_extensions.gemm import (
+            reset_swizzled_inputs,
+            swizzle_inputs,
+        )
     from transformer_engine.pytorch.tensor.quantized_tensor import Quantizer
 
     def general_gemm(
@@ -86,6 +95,19 @@ def general_gemm(
             if not out.is_contiguous():
                 raise ValueError("Output tensor is not contiguous.")
 
+        # TE version > 2.3
+        if is_te_min_version("2.3", check_equality=False):
+            debug_quantizer = None
+            if isinstance(quantization_params, DebugQuantizer):
+                debug_quantizer = quantization_params
+                quantization_params = quantization_params.parent_quantizer
+                A = A.get_tensor(not transa)
+                B = B.get_tensor(transb)
+            if isinstance(A, Float8BlockwiseQTensorBase) or isinstance(B, Float8BlockwiseQTensorBase):
+                # There is not use_split_accumulator == False
+                # implementation for Float8BlockwiseQTensorBase GEMM
+                use_split_accumulator = True
+
         # Use bfloat16 as default bias_dtype
         bias_dtype = torch.bfloat16 if bias is None else bias.dtype
 
@@ -114,9 +136,18 @@ def general_gemm(
             "bulk_overlap": bulk_overlap,
         }
 
-        original_scale_inverses = swizzle_inputs(A, B, layout)
+        # TE version >= 2.0 and <= 2.3
+        if not is_te_min_version("2.3", check_equality=False):
+            original_scale_inverses = swizzle_inputs(A, B, layout)
+
         out, bias_grad, gelu_input, extra_output = ptex.generic_gemm(*args, **kwargs)
-        reset_swizzled_inputs(A, B, original_scale_inverses)
+
+        # TE version >= 2.0 and <= 2.3
+        if not is_te_min_version("2.3", check_equality=False):
+            reset_swizzled_inputs(A, B, original_scale_inverses)
+        elif debug_quantizer is not None:
+            # TE version >= 2.4
+            out = debug_quantizer.process_gemm_output(out)
 
         return out, bias_grad, gelu_input, extra_output
 

primus/backends/transformer_engine/transformer_engine_torch/comm_overlap.py

@@ -124,7 +124,9 @@ def is_p2p_overlap(self) -> bool: ...
         def is_fp8_ubuf(self) -> bool:
             return self.buf_dtype.itemsize == 1
 
-        def copy_into_buffer(self, input: torch.Tensor, quantizer: Quantizer, local_chunk: bool = False):
+        def copy_into_buffer(
+            self, input: torch.Tensor, quantizer: Quantizer = None, local_chunk: bool = False
+        ):
             """copy input to local buffer
 
             Args:

tests/run_unit_tests.py

@@ -13,17 +13,15 @@
 
 UNIT_TEST_PASS = True
 
+EXCLUDE_UNIT_TESTS = []
+
 
 def get_all_unit_tests():
-    global DISTRIBUTED_UNIT_TESTS
+    global DISTRIBUTED_UNIT_TESTS, EXCLUDE_UNIT_TESTS
 
     cur_dir = "./tests"
     unit_tests = {}
 
-    EXCLUDE_UNIT_TESTS = [
-        "unit_tests/megatron/cco/test_tp_overlap.py",
-    ]
-
     for root, dirs, files in os.walk(cur_dir):
         for file_name in files:
             if not file_name.endswith(".py") or not file_name.startswith("test_"):