feat: megatron support turbo fp8 grouped gemm

Author: xiaobochen-amd

.github/workflows/ci.yaml

@@ -10,7 +10,7 @@ on:
   pull_request:
 
 env:
-  PRIMUS_TURBO_COMMIT: 3acc0fe3271b7e41f0646382311f36bdffca0554 # feat(permute): permute op support to compute tokens_per_expert (#140)
+  PRIMUS_TURBO_COMMIT: 0385cdd615cb4eff28a1cbbf3583fccf95d11fe9 # chore: refactor grouped gemm blockwise python code (#142)
 
 jobs:
   code-lint:

examples/megatron/configs/MI300X/deepseek_v2_lite-pretrain.yaml

@@ -84,6 +84,9 @@ modules:
       use_turbo_attention: true
       use_turbo_grouped_mlp: true
 
+      # fp8: e4m3
+      # fp8_recipe: blockwise  # tensorwise, blockwise
+
       # deepep
       use_turbo_deepep: false
 

primus/backends/megatron/core/extensions/primus_turbo.py

@@ -7,7 +7,6 @@
 from contextlib import contextmanager
 from typing import Callable, List, Optional, Tuple
 
-import grouped_gemm
 import primus_turbo.pytorch as pt
 import torch
 import torch.nn.functional as F
@@ -731,22 +730,15 @@ def __init__(
             pg_collection,
         )
         args = get_args()
-        grouped_gemm_backend = args.grouped_gemm_backend
-        self.grouped_gemm_backend = grouped_gemm_backend
 
         if args.patch_zero_bubble and args.enable_zero_bubble:
             from .zbpp_gemm import grouped_gemm_with_weight_gradient_store
 
             self.grouped_gemm = functools.partial(
-                grouped_gemm_with_weight_gradient_store, gg_backend=grouped_gemm_backend
+                grouped_gemm_with_weight_gradient_store, gg_backend="turbo-gg"
             )
         else:
-            if grouped_gemm_backend == "turbo-gg":
-                self.grouped_gemm = pt.ops.grouped_gemm
-            elif grouped_gemm_backend == "lagacy-gg":
-                self.grouped_gemm = grouped_gemm.ops.gmm
-            else:
-                raise NotImplementedError(f"Grouped gemm backend {grouped_gemm_backend} not implemented")
+            self.grouped_gemm = pt.ops.grouped_gemm
 
         if args.use_turbo_fused_act_with_probs:
             assert self.config.gated_linear_unit, "turbo_fused_act_with_probs only support with GLU."
@@ -802,13 +794,18 @@ def forward(
                 w1 = self.weight1.view(self.num_local_experts, self.config.hidden_size, -1)
                 w2 = self.weight2.view(self.num_local_experts, -1, self.config.hidden_size)
 
-            if self.grouped_gemm_backend == "turbo-gg":
-                tokens_per_expert = tokens_per_expert.cuda()
+            tokens_per_expert = tokens_per_expert.to(w1.device)
             assert w1.is_contiguous(), "w1 must be contiguous"
             assert w2.is_contiguous(), "w2 must be contiguous"
-            fc1_output = self.grouped_gemm(
-                permuted_local_hidden_states, w1, tokens_per_expert, trans_b=False, **(gemm_kargs[0])
-            )
+            if PrimusTurboFP8GlobalStateManager.is_turbo_fp8_enabled():
+                quant_config = PrimusTurboFP8GlobalStateManager.get_turbo_fp8_quant_config()
+                fc1_output = pt.ops.grouped_gemm_fp8(
+                    permuted_local_hidden_states, w1, tokens_per_expert, trans_b=False, config=quant_config
+                )
+            else:
+                fc1_output = self.grouped_gemm(
+                    permuted_local_hidden_states, w1, tokens_per_expert, trans_b=False, **(gemm_kargs[0])
+                )
             if self.activation_recompute:
                 if args.use_turbo_fused_act_with_probs:
                     intermediate_parallel = self.activation_checkpoint.checkpoint(
@@ -821,9 +818,15 @@ def forward(
                     intermediate_parallel = self.activation_checkpoint.checkpoint(
                         self.activation_func_with_probs, fc1_output, permuted_probs.unsqueeze(-1)
                     )
-                fc2_output = self.grouped_gemm(
-                    intermediate_parallel, w2, tokens_per_expert, trans_b=False, **(gemm_kargs[1])
-                )
+                if PrimusTurboFP8GlobalStateManager.is_turbo_fp8_enabled():
+                    quant_config = PrimusTurboFP8GlobalStateManager.get_turbo_fp8_quant_config()
+                    fc2_output = pt.ops.grouped_gemm_fp8(
+                        intermediate_parallel, w2, tokens_per_expert, trans_b=False, config=quant_config
+                    )
+                else:
+                    fc2_output = self.grouped_gemm(
+                        intermediate_parallel, w2, tokens_per_expert, trans_b=False, **(gemm_kargs[1])
+                    )
                 self.activation_checkpoint.discard_output_and_register_recompute(fc2_output)
             else:
                 if args.use_turbo_fused_act_with_probs:
@@ -834,9 +837,15 @@ def forward(
                     intermediate_parallel = self.activation_func_with_probs(
                         fc1_output, permuted_probs.unsqueeze(-1)
                     )
-                fc2_output = self.grouped_gemm(
-                    intermediate_parallel, w2, tokens_per_expert, trans_b=False, **(gemm_kargs[1])
-                )
+                if PrimusTurboFP8GlobalStateManager.is_turbo_fp8_enabled():
+                    quant_config = PrimusTurboFP8GlobalStateManager.get_turbo_fp8_quant_config()
+                    fc2_output = pt.ops.grouped_gemm_fp8(
+                        intermediate_parallel, w2, tokens_per_expert, trans_b=False, config=quant_config
+                    )
+                else:
+                    fc2_output = self.grouped_gemm(
+                        intermediate_parallel, w2, tokens_per_expert, trans_b=False, **(gemm_kargs[1])
+                    )
         else:
             # No token is allocated for local experts.
             assert torch.count_nonzero(tokens_per_expert) == 0
@@ -925,9 +934,7 @@ def __init__(
             deepep_num_use_cu=args.turbo_deepep_num_cu,
             deepep_num_worst_tokens=num_worst_tokens,
             deepep_use_cuda_num_tokens_per_expert=(
-                args.use_turbo_grouped_mlp
-                and args.moe_use_legacy_grouped_gemm
-                and args.grouped_gemm_backend == "turbo-gg"
+                args.use_turbo_grouped_mlp and args.moe_use_legacy_grouped_gemm
             ),
             deepep_async_finish=True,
             deepep_allocate_on_comm_stream=True,

primus/backends/megatron/core/pipeline_parallel/zerobubble/README.md

@@ -9,16 +9,8 @@ Zero bubbles is a state-of-art technique aiming to reduce the bubble time and me
 
 - Zero bubble patch the gemm OP and the grouped gemm OP for splitting the backward of the inputs and weights, support TE & Primus-turbo backend.
 
-- We suggest to use primus-turbo gemm & grouped gemm to patch the original TE implementation, the following flags is needed to turn on.
-```
-enable_primus_turbo: true
-use_turbo_parallel_linear: true
-use_turbo_grouped_mlp: true
-```
-    - If it is for MoE model, you can specify group gemm backend by `grouped_gemm_backend: "turbo-gg" # turbo-gg, lagacy-gg`.
-
 - Some other flags need to be specified
-```
+
 overlap_grad_reduce: false
 overlap_param_gather: false
 no_persist_layer_norm: true

primus/configs/models/megatron/llama4_17B128E.yaml

@@ -26,4 +26,4 @@ tokenizer_model: meta-llama/Llama-4-Maverick-17B-128E
 expert_model_parallel_size: 8
 expert_tensor_parallel_size: null # int
 moe_permute_fusion: true
-moe_shared_expert_overlap: true
+moe_shared_expert_overlap: true

primus/configs/models/megatron/llama4_17B16E.yaml

@@ -30,4 +30,4 @@ tokenizer_model: meta-llama/Llama-4-Scout-17B-16E
 expert_model_parallel_size: 8
 expert_tensor_parallel_size: null # int
 moe_permute_fusion: true
-moe_shared_expert_overlap: true
+moe_shared_expert_overlap: true

primus/configs/models/megatron/llama4_base.yaml

@@ -12,9 +12,7 @@ moe_router_topk: 1
 # moe_router_pre_softmax needs to be set to be true for moe_router_topk=1
 moe_router_pre_softmax: true
 moe_router_load_balancing_type: aux_loss
-moe_aux_loss_coeff: 0.001 
+moe_aux_loss_coeff: 0.001
 moe_grouped_gemm: true
 moe_use_legacy_grouped_gemm: false
 moe_token_dispatcher_type: alltoall
-
-

primus/configs/modules/megatron/primus_turbo.yaml

@@ -19,8 +19,5 @@ turbo_deepep_use_comm_stream: false
 # sync-free moe
 turbo_sync_free_moe_stage: 0
 
-# group-gemm
-grouped_gemm_backend: "turbo-gg" # turbo-gg, lagacy-gg
-
 # use turbo fused activation_with_probs to optmize redundant computation
 use_turbo_fused_act_with_probs: false

tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py

@@ -34,7 +34,6 @@
 def create_args():
     """Setup dummy args."""
     args = SimpleNamespace()
-    args.grouped_gemm_backend = "turbo-gg"
     args.turbo_sync_free_moe_stage = 0
     args.sequence_parallel = False
     args.seq_length = 4096