diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index c176985254..d09d2871e9 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -142,6 +142,34 @@ def run( # 2nd return is the activation to save for wgrad: MXFP8 (fp8 path) or bf16 (16-bit path). return query, x_saved + @classmethod + def backward_linear( + cls, + grad_output, + x_saved, + w_q, + act_dtype, + wgrad_store, + fuse_wgrad_accumulation, + tp_group, + sequence_parallel, + **kwargs, + ): + """Linear backward for the fused Q up-proj — delegates to :func:`~transformer_engine.pytorch.module.linear.backward_linear`.""" + from ..module.linear import backward_linear as _bwd + + return _bwd( + grad_output, + x_saved, + w_q, + act_dtype, + wgrad_store, + fuse_wgrad_accumulation, + tp_group, + sequence_parallel, + **kwargs, + ) + @classmethod def wrap_mxfp8( cls, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 56622db5e6..1950096fb2 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -786,6 +786,92 @@ def _linear_setup_ctx( return (saved_inputmat, wt_save, saved_weight, saved_bias) +def backward_linear( + grad_output: torch.Tensor, + x_saved, + w_q, + act_dtype: torch.dtype, + wgrad_store, + fuse_wgrad_accumulation: bool, + tp_group, + sequence_parallel: bool, + *, + use_bias: bool = False, + requires_dgrad: bool = True, + requires_wgrad: bool = True, + parallel_mode: str = "column", + backward_input_needs_gather: bool = False, +) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """Linear backward for fused operations that bypass TE's autograd chain. + + Wraps :func:`_linear_backward` with a simplified interface for callers + (e.g. Megatron's fused MLA Q up-proj) that run their own forward kernel + and need to delegate the projection backward to TE. + + Args: + grad_output: upstream gradient (e.g. post-RoPE-backward) ``[tokens, out_features]``. + x_saved: activation saved from the forward (``MXFP8Tensor`` or bf16). + w_q: weight (``MXFP8Tensor`` for FP8 path, bf16 tensor otherwise). + act_dtype: output dtype for the dgrad tensor. + wgrad_store: optional deferred weight-grad store. + fuse_wgrad_accumulation: accumulate wgrad directly into ``w_q.main_grad``. + tp_group: tensor-parallel process group (or ``None``). + sequence_parallel: whether sequence parallelism is active. + use_bias: compute a bias gradient (default ``False``). + requires_dgrad: compute dgrad (default ``True``). + requires_wgrad: compute wgrad (default ``True``). + parallel_mode: cuBLAS parallel mode (default ``"column"``). + backward_input_needs_gather: all-gather ``x_saved`` before the wgrad + GEMM (default ``False`` — assumes fused forward pre-gathers). + + Returns: + ``(dgrad, wgrad, grad_bias)`` — ``wgrad`` is a typed dummy when + ``fuse_wgrad_accumulation=True``; ``grad_bias`` is ``None`` when + ``use_bias=False``. + """ + tp_size = get_distributed_world_size(tp_group) if tp_group is not None else 1 + fp8 = isinstance(w_q, QuantizedTensor) + + grad_output_quantizer = None + if fp8: + grad_output_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ) + grad_output_quantizer.optimize_for_gemm = True + + bwd_args = LinearBwdArgs( + grad_output=grad_output, + inputmat=x_saved, + weight_fp8=w_q, + saved_weight=w_q, + bias=None, + grad_output_quantizer=grad_output_quantizer, + use_bias=use_bias, + requires_dgrad=requires_dgrad, + requires_wgrad=requires_wgrad, + inp_shape=x_saved.shape, + activation_dtype=act_dtype, + fp8=fp8, + dgrad_use_split_accumulator=_2X_ACC_DGRAD, + wgrad_use_split_accumulator=_2X_ACC_WGRAD, + is_weight_param_quantized=fp8, + parallel_mode=parallel_mode, + tp_group=tp_group, + tp_size=tp_size, + tensor_parallel=tp_size > 1, + sequence_parallel=sequence_parallel, + backward_input_needs_gather=backward_input_needs_gather, + is_fsdp2=False, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + wgrad_store=wgrad_store, + origin_weight_ref=weakref.ref(w_q) if fuse_wgrad_accumulation else None, + main_grad_func=(lambda: w_q.main_grad) if fuse_wgrad_accumulation else None, + ) + + wgrad, dgrad, grad_bias = _linear_backward(bwd_args) + return dgrad, wgrad, grad_bias + + def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: """Backward implementation for the linear layer.