Split transformer engine benchmark suite (#2677)

Author: tengqm

benchmark/performance_utils.py

@@ -580,6 +580,67 @@ def get_input_iter(self, cur_dtype) -> Generator:
                 yield inp,
 
 
+class TexGluBenchmark(Benchmark):
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+    # Triton grid_y is capped at 65535, BLOCK_SIZE_H=64 -> last dim <= 8388480.
+    MAX_LAST_DIM = 2 * 64 * 65535
+
+    def set_more_shapes(self):
+        # Last dim must be even for GLU operations to split
+        special_shapes_2d = [[1024, 2**i] for i in range(1, 20, 4)]
+        sp_shapes_3d = [[64, 64, 2**i] for i in range(1, 15, 4)]
+
+        return special_shapes_2d + sp_shapes_3d
+
+    def init_user_config(self):
+        super().init_user_config()
+        supported = []
+        for shape in self.shapes:
+            last_dim = shape[-1]
+            if last_dim % 2 != 0:
+                continue
+            if last_dim > self.MAX_LAST_DIM:
+                continue
+            supported.append(shape)
+        if not supported:
+            pytest.skip(
+                "No geglu shapes satisfy the constraints of FlagGems implementation."
+            )
+        self.shapes = supported
+
+
+class TexGluForwardBenchmark(TexGluBenchmark):
+    def get_input_iter(self, dtype):
+        for shape in self.shapes:
+            x = generate_tensor_input(shape, dtype, self.device)
+            # TE GLU APIs typically accept (input, quantizer).
+            yield (x, None)
+
+    def get_tflops(self, op, *args, **kwargs):
+        # args[0] is the input tensor x
+        shape = list(args[0].shape)
+        return torch.tensor(shape).prod().item()
+
+
+class TexGluBackwardBenchmark(TexGluBenchmark):
+    def get_input_iter(self, dtype):
+        for shape in self.shapes:
+            inp = generate_tensor_input(shape, dtype, self.device)
+
+            out_shape = list(shape)
+            out_shape[-1] = out_shape[-1] // 2
+
+            grad_out = torch.randn(out_shape, dtype=dtype, device=self.device)
+
+            yield grad_out, inp, None
+
+    def get_tflops(self, op, *args, **kwargs):
+        # args[1] is the original input tensor 'inp'
+        inp_shape = list(args[1].shape)
+        # Proxy FLOPs estimate: forward + backward cost roughly approximated
+        return torch.tensor(inp_shape).prod().item() * 2
+
+
 def generate_tensor_input(shape, dtype, device):
     if dtype in FLOAT_DTYPES:
         return torch.randn(shape, dtype=dtype, device=device)

benchmark/test_dgeglu.py

@@ -0,0 +1,36 @@
+import pytest
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import TexGluBackwardBenchmark
+
+# Note: Importing transformer_engine (especially in some versions like py 3.10) may automatically
+# configure the Root Logger (adding handlers). This may cause subsequent `logging.basicConfig`
+# calls (used by FlagGems benchmark) to be ignored/no-op, leading to missing result log files.
+# See: https://github.com/NVIDIA/TransformerEngine/issues/1065
+try:
+    from transformer_engine.pytorch import cpp_extensions as tex
+
+    TE_OP = getattr(tex, "dgeglu")
+    TE_AVAILABLE = True
+    GEMS_OP = getattr(flag_gems, "dgeglu")
+except ImportError:
+    TE_AVAILABLE = False
+    TE_OP = None
+    GEMS_OP = None
+
+
+@pytest.mark.dgeglu
+@pytest.mark.skipif(not TE_AVAILABLE, reason="TransformerEngine not installed")
+@pytest.mark.skipif(TE_OP is None, reason="'dgeglu' not found in TransformerEngine")
+@pytest.mark.skipif(GEMS_OP is None, reason="'dgeglu' not found in FlagGems")
+def test_dgeglu():
+    bench = TexGluBackwardBenchmark(
+        op_name="dgeglu",
+        torch_op=TE_OP,
+        gems_op=GEMS_OP,
+        dtypes=FLOAT_DTYPES,
+        # TODO(Qiming): Is this flag correct?
+        is_backward=False,
+    )
+    bench.run()

benchmark/test_dreglu.py

@@ -0,0 +1,36 @@
+import pytest
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import TexGluBackwardBenchmark
+
+# Note: Importing transformer_engine (especially in some versions like py 3.10) may automatically
+# configure the Root Logger (adding handlers). This may cause subsequent `logging.basicConfig`
+# calls (used by FlagGems benchmark) to be ignored/no-op, leading to missing result log files.
+# See: https://github.com/NVIDIA/TransformerEngine/issues/1065
+try:
+    from transformer_engine.pytorch import cpp_extensions as tex
+
+    TE_OP = getattr(tex, "dreglu")
+    TE_AVAILABLE = True
+    GEMS_OP = getattr(flag_gems, "dreglu")
+except ImportError:
+    TE_AVAILABLE = False
+    TE_OP = None
+    GEMS_OP = None
+
+
+@pytest.mark.dreglu
+@pytest.mark.skipif(not TE_AVAILABLE, reason="TransformerEngine not installed")
+@pytest.mark.skipif(TE_OP is None, reason="'dreglu' not found in TransformerEngine")
+@pytest.mark.skipif(GEMS_OP is None, reason="'dreglu' not found in FlagGems")
+def test_dreglu():
+    bench = TexGluBackwardBenchmark(
+        op_name="dreglu",
+        torch_op=TE_OP,
+        gems_op=GEMS_OP,
+        dtypes=FLOAT_DTYPES,
+        # TODO(Qiming): Is this flag correct?
+        is_backward=False,
+    )
+    bench.run()

benchmark/test_dswiglu.py

@@ -0,0 +1,36 @@
+import pytest
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import TexGluBackwardBenchmark
+
+# Note: Importing transformer_engine (especially in some versions like py 3.10) may automatically
+# configure the Root Logger (adding handlers). This may cause subsequent `logging.basicConfig`
+# calls (used by FlagGems benchmark) to be ignored/no-op, leading to missing result log files.
+# See: https://github.com/NVIDIA/TransformerEngine/issues/1065
+try:
+    from transformer_engine.pytorch import cpp_extensions as tex
+
+    TE_OP = getattr(tex, "dswiglu")
+    TE_AVAILABLE = True
+    GEMS_OP = getattr(flag_gems, "dswiglu")
+except ImportError:
+    TE_AVAILABLE = False
+    TE_OP = None
+    GEMS_OP = None
+
+
+@pytest.mark.dswiglu
+@pytest.mark.skipif(not TE_AVAILABLE, reason="TransformerEngine not installed")
+@pytest.mark.skipif(TE_OP is None, reason="'dswiglu' not found in TransformerEngine")
+@pytest.mark.skipif(GEMS_OP is None, reason="'dswiglu' not found in FlagGems")
+def test_swigglu():
+    bench = TexGluBackwardBenchmark(
+        op_name="dswiglu",
+        torch_op=TE_OP,
+        gems_op=GEMS_OP,
+        dtypes=FLOAT_DTYPES,
+        # TODO(Qiming): Is this flag correct?
+        is_backward=False,
+    )
+    bench.run()

benchmark/test_geglu.py

@@ -0,0 +1,34 @@
+import pytest
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import TexGluForwardBenchmark
+
+# Note: Importing transformer_engine (especially in some versions like py 3.10) may automatically
+# configure the Root Logger (adding handlers). This may cause subsequent `logging.basicConfig`
+# calls (used by FlagGems benchmark) to be ignored/no-op, leading to missing result log files.
+# See: https://github.com/NVIDIA/TransformerEngine/issues/1065
+try:
+    from transformer_engine.pytorch import cpp_extensions as tex
+
+    TE_OP = getattr(tex, "geglu")
+    TE_AVAILABLE = True
+    GEMS_OP = getattr(flag_gems, "geglu")
+except ImportError:
+    TE_AVAILABLE = False
+    TE_OP = None
+    GEMS_OP = None
+
+
+@pytest.mark.geglu
+@pytest.mark.skipif(not TE_AVAILABLE, reason="TransformerEngine not installed")
+@pytest.mark.skipif(TE_OP is None, reason="'geglu' not found in TransformerEngine")
+@pytest.mark.skipif(GEMS_OP is None, reason="'geglu' not found in FlagGems")
+def test_geglu():
+    bench = TexGluForwardBenchmark(
+        op_name="geglu",
+        torch_op=TE_OP,
+        gems_op=GEMS_OP,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_reglu.py

@@ -0,0 +1,34 @@
+import pytest
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import TexGluForwardBenchmark
+
+# Note: Importing transformer_engine (especially in some versions like py 3.10) may automatically
+# configure the Root Logger (adding handlers). This may cause subsequent `logging.basicConfig`
+# calls (used by FlagGems benchmark) to be ignored/no-op, leading to missing result log files.
+# See: https://github.com/NVIDIA/TransformerEngine/issues/1065
+try:
+    from transformer_engine.pytorch import cpp_extensions as tex
+
+    TE_OP = getattr(tex, "reglu")
+    TE_AVAILABLE = True
+    GEMS_OP = getattr(flag_gems, "reglu")
+except ImportError:
+    TE_AVAILABLE = False
+    TE_OP = None
+    GEMS_OP = None
+
+
+@pytest.mark.reglu
+@pytest.mark.skipif(not TE_AVAILABLE, reason="TransformerEngine not installed")
+@pytest.mark.skipif(TE_OP is None, reason="'reglu' not found in TransformerEngine")
+@pytest.mark.skipif(GEMS_OP is None, reason="'reglu' not found in FlagGems")
+def test_reglu():
+    bench = TexGluForwardBenchmark(
+        op_name="reglu",
+        torch_op=TE_OP,
+        gems_op=GEMS_OP,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_swiglu.py

@@ -0,0 +1,34 @@
+import pytest
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import TexGluForwardBenchmark
+
+# Note: Importing transformer_engine (especially in some versions like py 3.10) may automatically
+# configure the Root Logger (adding handlers). This may cause subsequent `logging.basicConfig`
+# calls (used by FlagGems benchmark) to be ignored/no-op, leading to missing result log files.
+# See: https://github.com/NVIDIA/TransformerEngine/issues/1065
+try:
+    from transformer_engine.pytorch import cpp_extensions as tex
+
+    TE_OP = getattr(tex, "swiglu")
+    TE_AVAILABLE = True
+    GEMS_OP = getattr(flag_gems, "swiglu")
+except ImportError:
+    TE_AVAILABLE = False
+    TE_OP = None
+    GEMS_OP = None
+
+
+@pytest.mark.swiglu
+@pytest.mark.skipif(not TE_AVAILABLE, reason="TransformerEngine not installed")
+@pytest.mark.skipif(TE_OP is None, reason="'swilu' not found in TransformerEngine")
+@pytest.mark.skipif(GEMS_OP is None, reason="'swiglu' not found in FlagGems")
+def test_swiglu():
+    bench = TexGluForwardBenchmark(
+        op_name="swiglu",
+        torch_op=TE_OP,
+        gems_op=GEMS_OP,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()