accept mxfp8 grad_output

Signed-off-by: Xin Yao <[email protected]>

Author: yaox12

tests/pytorch/test_numerics.py

@@ -4,7 +4,6 @@
 
 import math
 import os
-from pickletools import genops
 from typing import Dict, List, Tuple, Optional
 import pytest
 import random
@@ -108,8 +107,6 @@
     "swiglu": te.ops.SwiGLU,
 }
 
-gated_act_ops = (te.ops.GEGLU, te.ops.QGEGLU, te.ops.ReGLU, te.ops.SReGLU, te.ops.SwiGLU)
-
 all_normalizations = ["LayerNorm", "RMSNorm"]
 
 mask_types = ["causal", "no_mask"]
@@ -1729,20 +1726,34 @@ def _test_grouped_linear_accuracy(
     fp8,
     fuse_wgrad_accumulation,
     delay_wgrad_compute=False,
-    activation_func=te.ops.Identity(),
+    activation_func=None,
+    fp8_activation=False,
 ):
     reset_rng_states()
     if fp8:
         FP8GlobalStateManager.reset()
 
-    if isinstance(activation_func, gated_act_ops) or (
-        isinstance(activation_func, te.ops.Sequential)
-        and isinstance(activation_func[0], gated_act_ops)
-    ):
+    if activation_func in ("geglu", "qgeglu", "reglu", "sreglu", "swiglu"):
         hidden_size = 2 * config.hidden_size
     else:
         hidden_size = config.hidden_size
 
+    if activation_func is None:
+        input_act_func = te.ops.Identity()
+        output_act_func = te.ops.Identity()
+    elif fp8_activation:
+        input_act_func = te.ops.Sequential(
+            act_ops[activation_func](),
+            te.ops.Quantize(forward=True, backward=False),  # Output QuantizedTensor in forward
+        )
+        output_act_func = te.ops.Sequential(
+            te.ops.Quantize(forward=False, backward=True),  # Output QuantizedTensor in backward
+            act_ops[activation_func](),
+        )
+    else:
+        input_act_func = act_ops[activation_func]()
+        output_act_func = act_ops[activation_func]()
+
     inp_hidden_states = torch.randn(
         (config.max_seqlen_q, bs, hidden_size),
         dtype=dtype,
@@ -1769,11 +1780,11 @@ def _test_grouped_linear_accuracy(
     with fp8_autocast(enabled=fp8, fp8_recipe=recipe):
         if isinstance(block, GroupedLinear):
             m_splits = m_splits * bs
-            out = block(activation_func(inp_hidden_states), m_splits.tolist())
+            out = output_act_func(block(input_act_func(inp_hidden_states), m_splits.tolist()))
         else:
             out = torch.cat(
                 [
-                    block[i](activation_func(inp))
+                    output_act_func(block[i](input_act_func(inp)))
                     for i, inp in enumerate(torch.split(inp_hidden_states, m_splits.tolist()))
                 ]
             )
@@ -2063,12 +2074,6 @@ def test_grouped_linear_fp8_input(
             weight_i_copy = getattr(grouped_linear_fp8_input, f"weight{i}")
             weight_i_copy.main_grad = weight_i.main_grad.clone()
 
-    bf16_activation = act_ops[activation]()
-    fp8_activation = te.ops.Sequential(
-        bf16_activation,
-        te.ops.Quantize(forward=True, backward=False),  # Output QuantizedTensor in forward
-    )
-
     outputs_ref = _test_grouped_linear_accuracy(
         grouped_linear_bf16_input,
         num_gemms,
@@ -2078,7 +2083,8 @@ def test_grouped_linear_fp8_input(
         recipe,
         fp8=True,
         fuse_wgrad_accumulation=True,
-        activation_func=bf16_activation,
+        activation_func=activation,
+        fp8_activation=False,
     )
     outputs = _test_grouped_linear_accuracy(
         grouped_linear_fp8_input,
@@ -2089,7 +2095,8 @@ def test_grouped_linear_fp8_input(
         recipe,
         fp8=True,
         fuse_wgrad_accumulation=True,
-        activation_func=fp8_activation,
+        activation_func=activation,
+        fp8_activation=True,
     )
     # Shoule be bit-wise match
     for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)):

transformer_engine/common/util/cast_gated_kernels.cuh

@@ -44,8 +44,6 @@ constexpr size_t BUFFER_STAGES_NUM = BUFFER_DIM_Y / THREADS_PER_CHUNK_Y;  //  8
 constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y;                 //   4 = 128 / 32
 static_assert(ITERATIONS >= 1);
 
-__device__ inline float sigmoidf(const float x) { return __frcp_rn(1.0f + __expf(-x)); }
-
 template <bool IS_DGATED, typename ParamOP, float (*ActOP)(float, const ParamOP &),
           float (*DActOP)(float, const ParamOP &), typename IType, typename OType>
 __global__ void __launch_bounds__(THREADS_PER_CHUNK)

transformer_engine/common/util/math.h

@@ -9,6 +9,8 @@
 
 namespace transformer_engine {
 
+__device__ inline float sigmoidf(const float x) { return __frcp_rn(1.0f + __expf(-x)); }
+
 struct Empty {};
 
 template <typename OType, typename IType>
@@ -28,7 +30,7 @@ __device__ inline OType dgelu(const IType val, const Empty&) {
 template <typename OType, typename IType>
 __device__ inline OType sigmoid(const IType val, const Empty&) {
   const float cval = val;
-  return 1.f / (1.f + expf(-cval));
+  return sigmoidf(cval);
 }
 
 template <typename OType, typename IType>

transformer_engine/common/util/vectorized_pointwise.h

@@ -11,6 +11,7 @@
 
 #include "../common.h"
 #include "../utils.cuh"
+#include "math.h"
 
 namespace transformer_engine {
 
@@ -531,8 +532,18 @@ __launch_bounds__(unary_kernel_threads) __global__
       const ComputeType gelu_in = static_cast<ComputeType>(input_loader0.separate()[i]);
       const ComputeType gate_in = static_cast<ComputeType>(input_loader1.separate()[i]);
 
-      ComputeType after_dgelu = Dactivation(gelu_in, p) * grad_val * gate_in;
-      ComputeType after_dgate = grad_val * Activation(gelu_in, p);
+      ComputeType act_in, dact_in;
+      if constexpr ((Activation == &silu<fp32, fp32>) && (Dactivation == &dsilu<fp32, fp32>)) {
+        const float s = sigmoidf(gelu_in);
+        dact_in = gelu_in * s * (1 - s) + s;
+        act_in = gelu_in * s;
+      } else {
+        dact_in = Dactivation(gelu_in, p);
+        act_in = Activation(gelu_in, p);
+      }
+
+      ComputeType after_dgelu = dact_in * grad_val * gate_in;
+      ComputeType after_dgate = grad_val * act_in;
 
       if (requires_amax) {
         __builtin_assume(max >= 0);

transformer_engine/pytorch/module/grouped_linear.py

@@ -282,10 +282,16 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
                     weights[i] = w
 
             # Preprocess grad output
-            grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1])
+            grad_output_view = grad_output
             grad_output = [None] * ctx.num_gemms
             grad_biases = [None] * ctx.num_gemms
-            if ctx.fp8:
+            if isinstance(grad_output_view, QuantizedTensorBase):
+                assert not ctx.use_bias, "Bias is not supported for quantized grad output"
+                grad_output = tex.split_quantized_tensor(grad_output_view, ctx.m_splits)
+            elif ctx.fp8:
+                grad_output_view = grad_output_view.contiguous().view(
+                    -1, grad_output_view.shape[-1]
+                )
                 if ctx.use_bias:
                     grad_output_mats = torch.split(grad_output_view, ctx.m_splits)
                     recipe = ctx.fp8_recipe
@@ -313,6 +319,9 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None],
                         ctx.grad_output_quantizers,
                     )
             else:
+                grad_output_view = grad_output_view.contiguous().view(
+                    -1, grad_output_view.shape[-1]
+                )
                 # Only split grad output. Grad bias is fused with
                 # wgrad GEMM.
                 grad_output = torch.split(