[Megatron-LM] feat: add grouped gemm fp4 support, skip cache trans weight and santize code

examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh

@@ -421,7 +421,6 @@ for v in DOCKER_IMAGE NVTE_FUSED_ATTN NVTE_FUSED_ATTN_CK NVTE_FUSED_ATTN_AOTRITO
          PRIMUS_STACK_GROUPED_WEIGHT_TRITON PRIMUS_ROPE_TRITON \
          PRIMUS_SINKHORN_TRITON PRIMUS_HC_TRITON PRIMUS_INDEXER_TRITON \
          PRIMUS_INDEXER_TRITON_FULL PRIMUS_V4_ROUTER_TRITON \
-         PRIMUS_TURBO_FUSE_GROUPED_WGRAD PRIMUS_TURBO_FUSE_WGRAD_DEBUG \
          PRIMUS_MUON_BATCHED_NS PRIMUS_COMPRESS_ROPE_CACHE PRIMUS_COMPRESS_POOL_TRITON; do
     ENV_ARGS+=("--env" "$v")
 done

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

@@ -3,9 +3,7 @@
 #
 # See LICENSE for license information.
 ###############################################################################
-import contextlib
 import gc
-import os
 from contextlib import contextmanager, nullcontext
 from typing import Callable, Iterable, List, Optional, Tuple, Union
 
@@ -159,25 +157,6 @@ def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tenso
     return _dummy_wgrads[key].detach()
 
 
-class _MainGradShim:
-    """Per-expert handle for primus_turbo's ``fused_grouped_wgrad`` over a
-    *consolidated* ``[E, N, K]`` grouped-expert weight (OPT-1).
-
-    ``PrimusTurboGroupedLinear`` keeps one consolidated weight with a single
-    ``main_grad`` block, but ``fused_grouped_wgrad`` / ``_expert_main_grad_view``
-    expect a list of per-expert handles, each exposing a 2-D ``main_grad`` and a
-    ``grad_added_to_main_grad`` flag. These shims point at the contiguous 2-D
-    slices ``main_grad[i]`` of the consolidated block, so the grouped GEMM
-    backward accumulates each expert's wgrad straight into the right slice.
-    """
-
-    __slots__ = ("main_grad", "grad_added_to_main_grad")
-
-    def __init__(self, main_grad_slice: torch.Tensor) -> None:
-        self.main_grad = main_grad_slice
-        self.grad_added_to_main_grad = False
-
-
 def _bridge_weight_grad(
     x: torch.Tensor, weight: torch.nn.Parameter, weight_buffer: PrimusTurboQuantizedTensorPair
 ):
@@ -207,33 +186,8 @@ def backward(ctx, grad_x, grad_quantized_weight, grad_quantized_weight_trans):
                 weight, "grad_added_to_main_grad"
             ), "weight.grad_added_to_main_grad don't have grad_added_to_main_grad attribute."
 
-            # NOTE: Set weight.grad_added_to_main_grad to True to avoid adding the
-            # quantized weight gradient to main_grad twice.
-            if grad_quantized_weight is None:
-                # OPT-1 fused path: the grouped GEMM backward already accumulated the
-                # expert wgrad straight into main_grad (under fused_grouped_wgrad) and
-                # returned grad_b=None, so there is nothing to add here -- just flag it.
-                weight.grad_added_to_main_grad = True
-            else:
-                # `is_gfx1250` only exists in newer primus_turbo builds (it gates a
-                # gfx1250-specific elementwise-add workaround). Older / feature-branch
-                # primus_turbo that predate it (e.g. the flydsl sparse-MLA attention branch)
-                # don't define it; treat a missing symbol as False so those builds still work
-                # on non-gfx1250 archs (gfx942 / gfx950) instead of raising ImportError.
-                try:
-                    from primus_turbo.pytorch.core.utils import is_gfx1250
-
-                    _use_triton_inplace_add = is_gfx1250()
-                except ImportError:
-                    _use_triton_inplace_add = False
-
-                if _use_triton_inplace_add:
-                    # NOTE: The bandwith of torch's elementwise add kernel has issue. Use triton to temporary workaround for gfx1250.
-                    _triton_inplace_add_(weight.main_grad, grad_quantized_weight)
-                else:
-                    weight.main_grad.add_(grad_quantized_weight)
-
-                weight.grad_added_to_main_grad = True
+            weight.main_grad.add_(grad_quantized_weight)
+            weight.grad_added_to_main_grad = True
 
             return grad_x, _get_dummy_wgrad(list(weight.shape), weight.dtype), None, None
 
@@ -1070,7 +1024,8 @@ def forward_internal(
                         weight,
                         float8_e4m3,
                         quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                        or quant_config.current_scaling(),
                     )
 
                 x, quantized_weight = _bridge_weight_grad(
@@ -1241,7 +1196,8 @@ def forward_internal(
                         weight,
                         float8_e4m3,
                         quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                        or quant_config.current_scaling(),
                     )
 
                 x, quantized_weight = _bridge_weight_grad(
@@ -1405,7 +1361,8 @@ def forward_internal(
                         weight,
                         float8_e4m3,
                         quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                        or quant_config.current_scaling(),
                     )
 
                 x, quantized_weight = _bridge_weight_grad(
@@ -1582,7 +1539,8 @@ def forward_internal(self, x, is_first_microbatch: bool = False):
                         weight,
                         float8_e4m3,
                         quant_config,
-                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                        disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                        or quant_config.current_scaling(),
                     )
 
                 inp, quantized_weight = _bridge_weight_grad(
@@ -1814,7 +1772,8 @@ def forward_internal(
                     weights,
                     float8_e4m3,
                     quant_config,
-                    disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
+                    disable_parameter_transpose_cache=self.disable_parameter_transpose_cache
+                    or quant_config.current_scaling(),
                 )
 
             x, quantized_weights = _bridge_weight_grad(
@@ -1825,67 +1784,43 @@ def forward_internal(
                 ),
             )
 
-            # OPT-1 (opt-in, single-GPU): accumulate the expert wgrad straight into
-            # main_grad in the grouped GEMM backward (beta=1 ACCUMULATE) instead of
-            # GEMM-wgrad -> _WeightGradBridge.main_grad.add_(). Per-expert shims point
-            # at the consolidated [E,N,K] main_grad slices; the grouped backward
-            # accumulates into them and returns grad_b=None, then the bridge backward
-            # just flags grad_added. ONLY safe with no gradient all-reduce /
-            # reduce-scatter (TP=1 / DP=1 / EP=1), since grad_b=None skips the reduce.
-            # Needs a turbo wheel carrying fused_grouped_wgrad (else falls back).
-            _wgrad_ctx = contextlib.nullcontext()
-            _fwg_dbg = (
-                os.environ.get("PRIMUS_TURBO_FUSE_WGRAD_DEBUG") == "1"
-                and getattr(type(self), "_fwg_logn", 0) < 10
+            out = primus_turbo_torch.ops.grouped_gemm_fp8(
+                x,
+                quantized_weights,
+                m_splits,
+                trans_b=True,
+                config=quant_config.data(),
             )
-            _fwg_flag = os.environ.get("PRIMUS_TURBO_FUSE_GROUPED_WGRAD", "0") == "1"
-            _mg = getattr(weights, "main_grad", None)
-            if _fwg_flag and _mg is not None and _mg.dim() == 3:
-                try:
-                    from primus_turbo.pytorch.ops.grouped_gemm_fp8 import (
-                        _expert_main_grad_view,
-                        fused_grouped_wgrad,
-                    )
+        elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
+            quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config()
+            assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4."
 
-                    _shims = [_MainGradShim(_mg[i]) for i in range(_mg.shape[0])]
-                    if _fwg_dbg:
-                        import sys
-
-                        _v = _expert_main_grad_view(_shims)
-                        print(
-                            f"[OPT-1] gate PASS shape={tuple(_mg.shape)} contig={_mg.is_contiguous()} "
-                            f"stride={_mg.stride()} view={'OK' if _v is not None else 'REJECTED'}",
-                            file=sys.stderr,
-                            flush=True,
-                        )
-                        type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1
-                    _wgrad_ctx = fused_grouped_wgrad(_shims)
-                except ImportError as _e:
-                    if _fwg_dbg:
-                        import sys
-
-                        print(f"[OPT-1] ImportError: {_e}", file=sys.stderr, flush=True)
-                        type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1
-            elif _fwg_dbg:
-                import sys
-
-                print(
-                    f"[OPT-1] gate FAIL flag={_fwg_flag} main_grad={'None' if _mg is None else f'dim={_mg.dim()}'}",
-                    file=sys.stderr,
-                    flush=True,
+            if is_first_microbatch:
+                (
+                    self.quantized_weight_buffer,
+                    self.quantized_weight_t_buffer,
+                ) = _maybe_create_quantized_weight_buffers(
+                    weights,
+                    float4_e2m1fn_x2,
+                    quant_config,
+                    disable_parameter_transpose_cache=self.disable_parameter_transpose_cache,
                 )
-                type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1
 
-            with _wgrad_ctx:
-                out = primus_turbo_torch.ops.grouped_gemm_fp8(
-                    x,
-                    quantized_weights,
-                    m_splits,
-                    trans_b=True,
-                    config=quant_config.data(),
-                )
-        elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled():
-            assert False, "FP4 is not supported in PrimusTurboGroupedLinear"
+            x, quantized_weights = _bridge_weight_grad(
+                x,
+                weights,
+                PrimusTurboQuantizedTensorPair(
+                    data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer
+                ),
+            )
+
+            out = primus_turbo_torch.ops.grouped_gemm_fp4(
+                x,
+                quantized_weights,
+                m_splits,
+                trans_b=True,
+                config=quant_config.data(),
+            )
         else:
             out = primus_turbo_torch.ops.grouped_gemm(x, weights, m_splits, trans_b=True)