From 5959b64d3b0dc5e62677b5a8264c1c07e4935f68 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Thu, 23 Apr 2026 18:21:19 +0000 Subject: [PATCH 01/18] starting effort Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 758 +++++++++++++++++- 1 file changed, 741 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index fe5997a71e..af243eedc1 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -14,7 +14,7 @@ import torch import transformer_engine_torch as tex -from ...cpp_extensions import general_grouped_gemm +from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore from ...module.base import ( @@ -41,6 +41,24 @@ ) +# Quantizer types that support graph-safe `tex.group_quantize` / +# `tex.bgrad_group_quantize` and are eligible for the GroupedTensor + cublas +# `general_grouped_gemm_for_grouped_tensor` flow. All other recipes fall back +# to the legacy `tex.split_quantize` + `general_grouped_gemm` flow that uses +# CPU-side `m_splits` and is not CUDA-graph capturable. +# Quantizer types eligible for the graph-safe grouped-tensor flow +# (``tex.group_quantize`` + ``general_grouped_gemm_for_grouped_tensor``). +# +# NVFP4 is intentionally NOT included yet: +# * ``tex.group_quantize`` only implements the RHT (Random Hadamard Transform) +# code path for NVFP4 today (non-RHT raises in cast.cpp), and +# * grouped-swizzle support for NVFP4 in the cublas grouped GEMM is not +# implemented yet. +# When both land, add ``NVFP4Quantizer`` here (and update +# ``_is_grouped_quantize_supported`` accordingly). +_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS = (MXFP8Quantizer,) + + class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -696,6 +714,23 @@ def op_backward(self, *args, **kwargs): "It overrides `fuser_backward` instead of `op_backward`." ) + # ------------------------------------------------------------------ + # Path-selection helpers for fuser_forward/backward + # ------------------------------------------------------------------ + + @staticmethod + def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer]]) -> bool: + """Whether all input quantizers support the graph-safe grouped-tensor flow. + + See ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS`` for the gating rationale. + Currently this is MXFP8 only; every other recipe (bf16, fp8 delayed / + current scaling, fp8 block scaling, NVFP4, ...) falls back to the + legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. + """ + return all( + isinstance(q, _GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS) for q in input_quantizers + ) + def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -707,7 +742,6 @@ def fuser_forward( basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: num_groups = self.num_groups - has_bias = self.has_bias weight_param = self.weight if self.single_grouped_weight else self.weight0 device = weight_param.device @@ -739,17 +773,92 @@ def fuser_forward( else: dtype = weight_param.dtype - # Extract split sizes from extra input + # Extract split sizes from extra input. Keep on GPU for graph safety. split_sizes = basic_op_extra_inputs[0][0] - split_sizes_int = [int(s) for s in split_sizes.tolist()] - if len(split_sizes_int) != num_groups: - raise ValueError(f"Expected {num_groups} splits, but got {len(split_sizes_int)}.") + if int(split_sizes.numel()) != num_groups: + raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") + if split_sizes.dtype != torch.int64: + split_sizes = split_sizes.to(dtype=torch.int64) + if split_sizes.device != device: + split_sizes = split_sizes.to(device=device) # Extract scales tensor for bias scaling scales = None if self._scale_bias: scales = basic_op_extra_inputs[0][1] + # Dispatch: graph-safe GroupedTensor flow whenever it can be used -- + # * for the unquantized (bf16/fp16) compute path we just wrap the + # existing high-precision data as a ``GroupedTensor`` (no quantize + # call). FP32 is excluded because the cublasLt grouped GEMM only + # accepts FP8 / BF16 / FP16 inputs, and + # * for quantized compute only when the quantizer supports the + # graph-safe ``tex.group_quantize`` kernel (see + # ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS``; currently MXFP8). + # All remaining cases (fp32 unquantized, fp8 delayed / current scaling, + # fp8 block scaling, NVFP4) fall back to the legacy + # ``split_quantize`` + ``general_grouped_gemm`` flow. + use_grouped_tensor_path = ( + with_quantized_compute and self._is_grouped_quantize_supported(input_quantizers) + ) or (not with_quantized_compute and dtype in (torch.bfloat16, torch.float16)) + if use_grouped_tensor_path: + return self._fuser_forward_grouped_tensor( + ctx=ctx, + input_=input_, + split_sizes=split_sizes, + scales=scales, + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + grad_output_quantizers=grad_output_quantizers, + dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + ) + return self._fuser_forward_split_quantize( + ctx=ctx, + input_=input_, + split_sizes=split_sizes, + scales=scales, + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + grad_output_quantizers=grad_output_quantizers, + dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + ) + + # ================================================================== + # Legacy `tex.split_quantize` + `general_grouped_gemm` flow. + # Used for bf16/fp16 inputs and FP8 recipes other than MXFP8/NVFP4. + # ``m_splits`` is needed on CPU here, so this flow is NOT cuda-graphable. + # ================================================================== + + def _fuser_forward_split_quantize( + self, + *, + ctx: OperationContext, + input_: torch.Tensor, + split_sizes: torch.Tensor, + scales: Optional[torch.Tensor], + with_quantized_compute: bool, + input_quantizers: list[Optional[Quantizer]], + weight_quantizers: list[Optional[Quantizer]], + grad_output_quantizers: list[Optional[Quantizer]], + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + device: torch.device, + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + num_groups = self.num_groups + has_bias = self.has_bias + + # Need CPU split sizes for split_quantize / general_grouped_gemm. + split_sizes_int = [int(s) for s in split_sizes.tolist()] + # Extract params if self.single_grouped_weight: weights = self.weight.quantized_tensors @@ -759,16 +868,7 @@ def fuser_forward( weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] bs = None if has_bias: - if self.single_grouped_bias: - bias_parts = self.bias.quantized_tensors - if bias_parts is None: - bias_parts = self.bias.split_into_quantized_tensors() - bs = [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts] - else: - bs = [ - maybe_dequantize(getattr(self, f"bias{idx}"), dtype) - for idx in range(num_groups) - ] + bs = self._get_bias_tensors(dtype) # Convert weight dtype if needed ws = [] @@ -843,6 +943,304 @@ def fuser_forward( saved.extend(xs) saved.extend(ws) ctx.save_for_backward(*saved) + ctx.use_grouped_tensor_path = False + ctx.with_quantized_compute = with_quantized_compute + ctx.input_quantizers = input_quantizers + ctx.weight_quantizers = weight_quantizers + ctx.grad_output_quantizers = grad_output_quantizers + ctx.grad_input_quantizers = None + ctx.dtype = dtype + ctx.input_requires_grad = input_requires_grad + ctx.weight_requires_grad = weight_requires_grad + + return out, [()] + + # ================================================================== + # Graph-safe `tex.group_quantize` + `general_grouped_gemm_for_grouped_tensor` + # flow. Used only for MXFP8 / NVFP4 recipes, where ``m_splits`` stays on + # GPU and the flow can be captured into a CUDA graph. + # ``general_grouped_gemm_for_grouped_tensor`` handles bias / scaled-bias + # addition inside the GEMM. + # ================================================================== + + def _build_grouped_bias_for_gemm( + self, + dtype: torch.dtype, + device: torch.device, + ) -> Optional[torch.Tensor]: + """Build a uniform GroupedTensor of per-group biases for the cublas + grouped GEMM. + + Each group expects a (1, out_features) bias vector. Returns ``None`` + when no additive bias is configured. + """ + if not self.has_bias: + return None + num_groups = self.num_groups + + if self.single_grouped_bias: + # Already a contiguous (num_groups * out_features) buffer. + bias_data = self.bias.rowwise_data + if bias_data.dtype != dtype: + bias_data = bias_data.to(dtype=dtype) + else: + bias_list = [ + maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) + ] + bias_data = torch.stack(bias_list, dim=0).contiguous() + + return GroupedTensor( + shape=(num_groups, self.out_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(1, self.out_features)] * num_groups, + quantizer=None, + data=bias_data.reshape(-1), + ) + + def _quantize_weights_for_grouped_gemm( + self, + weight_quantizers: list[Optional[Quantizer]], + *, + columnwise_usage: bool, + ) -> "GroupedTensor | list[torch.Tensor]": + """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. + + Mirrors the ``forward_grouped_mlp.py`` pattern. Four distinct cases: + + ``single_grouped_weight=True``: + * **quantized_model_init**: ``self.weight`` is already a quantized + ``GroupedTensor`` (``self.weight.quantizer is not None``). We do + NOT re-quantize; we just refresh its ``quantizer`` reference so + usage flags propagate, and return it as-is. Dispatches the GEMM + to ``no_discrete`` mode. + * **normal init**: ``self.weight`` is a high-precision + ``GroupedTensor`` (``quantizer is None``). We call + ``tex.group_quantize`` once on its packed ``rowwise_data`` and + return the resulting quantized ``GroupedTensor``. Dispatches the + GEMM to ``no_discrete`` mode. + + ``single_grouped_weight=False``: + * **quantized_model_init**: each ``weight{idx}`` is already a + quantized tensor (e.g. ``MXFP8Tensor``); use as-is. + * **normal init**: each ``weight{idx}`` is a high-precision + ``torch.Tensor``; quantize per-group via ``quantizer(weight)``. + + In both subcases we return a Python list, which dispatches the GEMM + to ``discrete_in`` mode. + """ + num_groups = self.num_groups + + if self.single_grouped_weight: + grouped_weight = self.weight + if not isinstance(grouped_weight, GroupedTensor): + raise RuntimeError( + "GroupedLinear with single_grouped_weight=True expected the weight " + f"parameter to be a GroupedTensor, got {type(grouped_weight).__name__}." + ) + weight_quantizer = weight_quantizers[0] + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + + if grouped_weight.quantizer is not None: + # quantized_model_init: weight is already quantized. + grouped_weight.quantizer = weight_quantizer + return grouped_weight + + # Normal init: high-precision GroupedTensor. + if grouped_weight.rowwise_data is None: + raise RuntimeError( + "GroupedLinear: grouped weight has no rowwise_data to quantize." + ) + return tex.group_quantize( + grouped_weight.rowwise_data.view(grouped_weight.logical_shape), + weight_quantizer, + num_groups, + None, + ) + + # single_grouped_weight=False: per-group parameters. + out: list[torch.Tensor] = [] + for idx, quantizer in enumerate(weight_quantizers): + w = getattr(self, f"weight{idx}") + if is_quantized_tensor(w): + # quantized_model_init: weight{idx} is already quantized. + out.append(w) + else: + # Normal init: quantize the high-precision parameter. + quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + out.append(quantizer(w)) + return out + + def _fuser_forward_grouped_tensor( + self, + *, + ctx: OperationContext, + input_: torch.Tensor, + split_sizes: torch.Tensor, + scales: Optional[torch.Tensor], + with_quantized_compute: bool, + input_quantizers: list[Optional[Quantizer]], + weight_quantizers: list[Optional[Quantizer]], + grad_output_quantizers: list[Optional[Quantizer]], + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + device: torch.device, + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + """Graph-safe GroupedTensor forward path. + + Handles both: + * **Quantized compute** (currently MXFP8): input/weights/grad_output + are quantized via ``tex.group_quantize`` and the GEMM consumes the + packed scale/data buffers. + * **Unquantized compute** (bf16/fp16/fp32): no quantize calls -- we + simply wrap the contiguous high-precision data as a + ``GroupedTensor`` so the cublas grouped GEMM path is still used, + making the operation CUDA-graphable end-to-end. + """ + num_groups = self.num_groups + has_bias = self.has_bias + + # Cumulative-token offsets on GPU (graph-safe). Computed once and + # scaled by the relevant last dim to obtain `tensor_offsets` for any + # GroupedTensor with that last dim. Reused in backward. + base_offsets = tex.splits_to_offsets(split_sizes, 1) + + # Flatten to 2D so the first dim is the total token count, matching + # the convention used by forward_grouped_mlp.py. + original_shape = list(input_.size()) + x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) + total_tokens = x.size(0) + + # Build the input GroupedTensor. + if with_quantized_compute: + # Quantize input as a single GroupedTensor (one quantizer for all + # groups; OK for MXFP8 because its config is data-dependent). + input_quantizer = input_quantizers[0] + input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + input_quantizer.optimize_for_gemm = True + grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes) + else: + # No quantize: wrap the contiguous high-precision buffer. + grouped_x = GroupedTensor( + shape=(total_tokens, self.in_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=x.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_offsets * self.in_features, + ) + + # Build the weight GroupedTensor / list. In the unquantized path we + # consume the parameters directly: ``self.weight`` (a high-precision + # ``GroupedTensor``) for ``single_grouped_weight=True`` -- dispatches + # the GEMM to ``no_discrete`` mode -- or the per-group ``weight{idx}`` + # tensors otherwise -- ``discrete_in`` mode. + if with_quantized_compute: + grouped_weights = self._quantize_weights_for_grouped_gemm( + weight_quantizers, + columnwise_usage=input_requires_grad, + ) + elif self.single_grouped_weight: + # ``self.weight`` may itself be a quantized GroupedTensor when the + # caller used ``quantized_model_init``. Dequantize back to the + # compute dtype so cublas grouped GEMM sees matching A/B dtypes. + grouped_weights = maybe_dequantize(self.weight, dtype) + else: + grouped_weights = [ + maybe_dequantize(getattr(self, f"weight{idx}"), dtype) + for idx in range(num_groups) + ] + + # Allocate output buffer and wrap as a GroupedTensor view (no copy). + out_shape = original_shape[:-1] + [self.out_features] + out = torch.empty(out_shape, dtype=dtype, device=device) + grouped_out = GroupedTensor( + shape=(total_tokens, self.out_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=out.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_offsets * self.out_features, + ) + + # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias + # uses ``bias=``; scaled bias also passes per-token ``bias_scale=``. + grouped_bias = None + bias_scale: Optional[torch.Tensor] = None + if has_bias: + grouped_bias = self._build_grouped_bias_for_gemm(dtype, device) + if self._scale_bias: + bias_scale = scales.reshape(-1) + if bias_scale.dtype != torch.float32: + bias_scale = bias_scale.to(dtype=torch.float32) + + # Forward grouped GEMM (TN layout: out[i] = x[i] @ w[i]^T) + general_grouped_gemm_for_grouped_tensor( + grouped_weights, + grouped_x, + grouped_out, + layout="TN", + use_split_accumulator=_2X_ACC_FPROP, + bias=grouped_bias, + bias_scale=bias_scale, + ) + + # Prepare weight tensor(s) for the dgrad pass. + # + # NOTE: We deliberately do NOT call ``update_usage(rowwise_usage=False, ...)`` + # on quantized weights here. In the ``quantized_model_init`` case, + # ``grouped_weights`` IS the user's ``self.weight`` parameter (or its + # per-group ``weight{idx}`` parameters), and dropping its rowwise data + # would corrupt the parameter for subsequent forward passes. + # + # For the quantized compute path, the columnwise representation needed + # by dgrad has already been populated up-front via + # ``quantizer.set_usage(rowwise=True, columnwise=input_requires_grad)``. + # This mirrors the behavior in ``forward_grouped_mlp.py``. + weight_is_grouped = isinstance(grouped_weights, GroupedTensor) + if not input_requires_grad: + grouped_weights = None if weight_is_grouped else [None] * num_groups + + # ``grouped_x`` already has all the data we need to save for backward + # (rowwise for the unquantized path; columnwise too for the quantized + # path because we configured the input quantizer with + # ``columnwise=weight_requires_grad`` above). + if not weight_requires_grad: + grouped_x = None + + # Save state for backward pass. Following the fused grouped MLP + # pattern, we save the GroupedTensor's component buffers (rather than + # the wrapper) and rebuild it in backward. ``base_offsets`` is saved + # so backward can reuse it without recomputing via ``splits_to_offsets``. + if ctx.requires_grad: + saved: list[Optional[torch.Tensor]] = [split_sizes, base_offsets] + if self._scale_bias: + saved.append(scales) + # For the wgrad input we save (data, scale_inv). + # * Quantized path saves columnwise data + scale. + # * Unquantized path saves the raw rowwise data and a None scale. + if grouped_x is not None: + if with_quantized_compute: + saved.extend( + [ + grouped_x.columnwise_data, + grouped_x.columnwise_scale_inv, + ] + ) + else: + saved.extend([grouped_x.rowwise_data, None]) + else: + saved.extend([None, None]) + if weight_is_grouped: + saved.append(grouped_weights) + else: + saved.extend(grouped_weights) + ctx.save_for_backward(*saved) + ctx.use_grouped_tensor_path = True + ctx.weight_is_grouped = weight_is_grouped ctx.with_quantized_compute = with_quantized_compute ctx.input_quantizers = input_quantizers ctx.weight_quantizers = weight_quantizers @@ -864,6 +1262,28 @@ def fuser_backward( torch.Tensor, Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + # Dispatch to the path used in forward (saved as ``ctx.use_grouped_tensor_path``). + if getattr(ctx, "use_grouped_tensor_path", False): + return self._fuser_backward_grouped_tensor( + ctx=ctx, + grad_output=grad_output, + ) + return self._fuser_backward_split_quantize( + ctx=ctx, + grad_output=grad_output, + ) + + def _fuser_backward_split_quantize( + self, + *, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], ]: num_groups = self.num_groups has_bias = self.has_bias @@ -871,7 +1291,6 @@ def fuser_backward( device = weight_param.device # Saved tensors from forward pass - ctx = basic_op_ctxs[0] saved_tensors = ctx.saved_tensors split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] scales = None @@ -1084,3 +1503,308 @@ def fuser_backward( grad_extra = (None, grad_scales) if self._scale_bias else (None,) return grad_input, [grad_params], [grad_extra] + + # ================================================================== + # Graph-safe backward: counterpart of `_fuser_forward_grouped_tensor`. + # Uses ``tex.bgrad_group_quantize`` / ``tex.group_quantize`` on the grad + # output and ``general_grouped_gemm_for_grouped_tensor`` for both dgrad + # and wgrad. ``m_splits`` stays on GPU throughout. + # ================================================================== + + def _fuser_backward_grouped_tensor( + self, + *, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + num_groups = self.num_groups + has_bias = self.has_bias + weight_param = self.weight if self.single_grouped_weight else self.weight0 + device = weight_param.device + dtype = ctx.dtype + + with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) + + # Saved tensors from forward pass (see _fuser_forward_grouped_tensor). + # ``base_offsets`` is the cumulative-token offset tensor (computed once + # in forward via ``tex.splits_to_offsets``); we scale it by the + # appropriate last dim to get any GroupedTensor's ``tensor_offsets``. + # ``ws`` is either a single ``GroupedTensor`` (single_grouped_weight) + # or a list of per-group tensors; it is passed straight to + # ``general_grouped_gemm_for_grouped_tensor`` in dgrad. + # ``x_data`` / ``x_scale`` carry the saved input for the wgrad pass: + # * Quantized path: columnwise data + scale. + # * Unquantized path: raw rowwise data + ``None`` scale. + saved_tensors = ctx.saved_tensors + split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] + base_offsets, saved_tensors = saved_tensors[0], saved_tensors[1:] + scales = None + if self._scale_bias: + scales, saved_tensors = saved_tensors[0], saved_tensors[1:] + x_data, saved_tensors = saved_tensors[0], saved_tensors[1:] + x_scale, saved_tensors = saved_tensors[0], saved_tensors[1:] + if getattr(ctx, "weight_is_grouped", False): + ws, saved_tensors = saved_tensors[0], saved_tensors[1:] + else: + ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] + + # Flatten grad_output to 2D (total_tokens, out_features), matching the + # forward convention. The grouped GEMMs and dbias kernels all operate + # on this 2D view. + dy_2d = maybe_dequantize(grad_output, dtype).reshape(-1, self.out_features) + total_tokens = dy_2d.size(0) + + # Rebuild the grouped input tensor used by wgrad. + grouped_x = None + if ctx.weight_requires_grad and x_data is not None: + if with_quantized_compute: + grouped_x = GroupedTensor( + shape=(total_tokens, self.in_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=ctx.input_quantizers[0], + columnwise_data=x_data, + columnwise_scale_inv=x_scale, + first_dims=split_sizes, + tensor_offsets=base_offsets * self.in_features, + with_gemm_swizzled_scales=True, + ) + else: + grouped_x = GroupedTensor( + shape=(total_tokens, self.in_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=x_data, + first_dims=split_sizes, + tensor_offsets=base_offsets * self.in_features, + ) + + # Build the grad_output GroupedTensor. + # * Quantized: ``tex.group_quantize`` (or ``bgrad_group_quantize`` + # when bias gradient can be fused for MXFP8). + # * Unquantized: just wrap the contiguous bf16/fp16/fp32 buffer. + grad_biases: list[Optional[torch.Tensor]] = [None] * num_groups + grad_scales = None + dbias_packed = None + + if with_quantized_compute: + grad_output_quantizer = ctx.grad_output_quantizers[0] + grad_output_quantizer.set_usage( + rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad + ) + grad_output_quantizer.optimize_for_gemm = True + + use_fused_bgrad = ( + has_bias + and not self._scale_bias + and isinstance(grad_output_quantizer, MXFP8Quantizer) + ) + if use_fused_bgrad: + grouped_dy, dbias_packed = tex.bgrad_group_quantize( + dy_2d, grad_output_quantizer, num_groups, split_sizes + ) + else: + grouped_dy = tex.group_quantize( + dy_2d, grad_output_quantizer, num_groups, split_sizes + ) + else: + grouped_dy = GroupedTensor( + shape=(total_tokens, self.out_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=dy_2d.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_offsets * self.out_features, + ) + + # Bias / scale-bias gradients (graph-safe; offsets are GPU tensors). + if has_bias: + if self._scale_bias: + bias_packed = torch.stack(self._get_bias_tensors(dtype)) + scales_f32 = scales.to(dtype=torch.float32) + dbias_packed, grad_scales = compute_grouped_dbias_dscales( + dy_2d, + scales_f32, + bias_packed, + offsets=base_offsets, + ) + elif dbias_packed is None: + # NVFP4 (or any path without a fused bgrad kernel) falls back + # to the standalone Triton dbias. + dbias_packed = compute_grouped_dbias(dy_2d, base_offsets, num_groups) + grad_biases = [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)] + + # ---- dgrad GEMM ---------------------------------------------------- + # grad_input mirrors grad_output's leading dims (batch / seq) and uses + # in_features for the trailing dim. + grad_input = None + if ctx.input_requires_grad: + grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] + grad_input = torch.empty(grad_input_shape, dtype=dtype, device=device) + grouped_grad_input = GroupedTensor( + shape=(total_tokens, self.in_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=grad_input.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_offsets * self.in_features, + ) + general_grouped_gemm_for_grouped_tensor( + ws, + grouped_dy, + grouped_grad_input, + layout="NN", + use_split_accumulator=_2X_ACC_DGRAD, + ) + + # ---- wgrad GEMM ---------------------------------------------------- + accumulate_into_main_grad = self._accumulate_into_main_grad + weight_shape = (self.out_features, self.in_features) + wgrad_output: Any = None + grouped_wgrad: Optional[GroupedTensor] = None + w_list: list[Optional[torch.Tensor]] = ( + [None] if self.single_grouped_weight else [None] * num_groups + ) + + if ctx.weight_requires_grad: + if self.single_grouped_weight: + if accumulate_into_main_grad: + if hasattr(weight_param, "__fsdp_param__"): + weight_param.main_grad = weight_param.get_main_grad() + main_grad = weight_param.main_grad + grouped_shape = (num_groups, *weight_shape) + if main_grad.shape != grouped_shape: + if main_grad.numel() != math.prod(grouped_shape): + raise RuntimeError( + "GroupedLinear expected grouped weight main_grad to have " + f"shape {grouped_shape} or matching numel, " + f"but got shape {tuple(main_grad.shape)}" + ) + main_grad = main_grad.view(grouped_shape) + accumulate_into_main_grad = not getattr( + weight_param, "overwrite_main_grad", False + ) + if accumulate_into_main_grad: + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=num_groups, + tensor_shape=weight_shape, + rowwise_data=main_grad, + dtype=main_grad.dtype, + ) + if grouped_wgrad is None: + grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_groups, + shapes=[weight_shape] * num_groups, + quantizer=None, + device=device, + dtype=dtype, + ) + wgrad_output = grouped_wgrad + else: + if accumulate_into_main_grad: + for idx in range(num_groups): + wp = getattr(self, f"weight{idx}") + if hasattr(wp, "__fsdp_param__"): + wp.main_grad = wp.get_main_grad() + w_list[idx] = wp.main_grad + accumulate_into_main_grad = not getattr( + self.weight0, "overwrite_main_grad", False + ) + else: + for idx in range(num_groups): + w_list[idx] = torch.empty(weight_shape, dtype=dtype, device=device) + wgrad_output = w_list + else: + accumulate_into_main_grad = False + + delay_wgrad = ( + ctx.weight_requires_grad + and self.wgrad_store is not None + and self.wgrad_store.delay_wgrad_compute() + ) + if ctx.weight_requires_grad: + wgrad_gemm = functools.partial( + general_grouped_gemm_for_grouped_tensor, + layout="NT", + accumulate=accumulate_into_main_grad, + use_split_accumulator=_2X_ACC_WGRAD, + ) + if delay_wgrad: + self.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], wgrad_gemm) + else: + wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + + # ---- Assemble grad_params (mirrors split-quantize path layout) ---- + grad_weights: list[Optional[torch.Tensor]] = [None] * num_groups + if ctx.weight_requires_grad and not delay_wgrad: + if self.single_grouped_weight: + if grouped_wgrad is not None and not accumulate_into_main_grad: + packed = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) + grad_weights = [packed[idx] for idx in range(num_groups)] + else: + grad_weights = list(w_list) + + if accumulate_into_main_grad: + if self.single_grouped_weight: + if hasattr(weight_param, "grad_added_to_main_grad"): + weight_param.grad_added_to_main_grad = True + grad_weight_dummy = get_dummy_wgrad( + list(weight_param.size()), + weight_param.dtype, + zero=getattr(weight_param, "zero_out_wgrad", False), + ) + else: + grad_weight_dummy = None + if has_bias: + if self.single_grouped_bias: + final_bias_grads = torch.stack(grad_biases, dim=0).to(dtype) + grad_params = [grad_weight_dummy, final_bias_grads] + else: + grad_params = grad_biases + [grad_weight_dummy] + else: + grad_params = [grad_weight_dummy] + grad_extra = (None, grad_scales) if self._scale_bias else (None,) + return grad_input, [grad_params], [grad_extra] + + grad_weights = [None] * num_groups + for group_idx in range(num_groups): + wp = getattr(self, f"weight{group_idx}") + if hasattr(wp, "grad_added_to_main_grad"): + wp.grad_added_to_main_grad = True + grad_weights[group_idx] = get_dummy_wgrad( + list(wp.size()), + wp.dtype, + zero=getattr(wp, "zero_out_wgrad", False), + ) + + if self.single_grouped_weight: + grad_weight = None + if ctx.weight_requires_grad and not delay_wgrad: + grad_weight = torch.stack(grad_weights, dim=0) + final_weight_grads = [grad_weight] + else: + if delay_wgrad and ctx.weight_requires_grad and not accumulate_into_main_grad: + final_weight_grads = [None] * num_groups + else: + final_weight_grads = grad_weights + + if not has_bias: + grad_params = list(final_weight_grads) + elif self.single_grouped_bias: + final_bias_grads = torch.stack(grad_biases, dim=0).to(dtype) + grad_params = list(final_weight_grads) + [final_bias_grads] + else: + if self.single_grouped_weight: + grad_params = list(grad_biases) + list(final_weight_grads) + else: + grad_params = list(final_weight_grads) + list(grad_biases) + + grad_extra = (None, grad_scales) if self._scale_bias else (None,) + return grad_input, [grad_params], [grad_extra] From a5739428c6aa423934f50ad3c8bc5b84e07aabdf Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 24 Apr 2026 03:26:15 +0000 Subject: [PATCH 02/18] all tests seem to be working Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 71 ++++- .../pytorch/ops/basic/grouped_linear.py | 256 ++++++++---------- 2 files changed, 165 insertions(+), 162 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 0f40e92183..c541026c5e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2015,6 +2015,8 @@ def test_dropout( @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("weight_requires_grad", (False, True)) @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) + @pytest.mark.parametrize("single_grouped_bias", (False, True)) def test_grouped_linear( self, *, @@ -2030,6 +2032,8 @@ def test_grouped_linear( input_requires_grad: bool, weight_requires_grad: bool, delay_wgrad_compute: bool, + single_grouped_weight: bool, + single_grouped_bias: bool, ) -> None: """Grouped GEMM""" @@ -2053,6 +2057,18 @@ def test_grouped_linear( if quantization is not None and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") + if single_grouped_bias and not bias: + pytest.skip("single_grouped_bias requires bias=True") + if ( + single_grouped_weight + and quantized_weight + and quantization in ("fp8_delayed_scaling", "fp8_current_scaling") + ): + pytest.skip( + "single_grouped_weight does not support FP8 delayed/current scaling " + "with quantized_model_init" + ) + # Random data x_ref, x_test = make_reference_and_test_tensors( in_shape, @@ -2111,12 +2127,26 @@ def test_grouped_linear( device=device, dtype=dtype, delay_wgrad_compute=delay_wgrad_compute, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, ) with torch.no_grad(): + if single_grouped_weight: + op_weights = op.weight.quantized_tensors + if op_weights is None: + op_weights = op.weight.split_into_quantized_tensors() + if single_grouped_bias: + op_bias_parts = op.bias.split_into_quantized_tensors() for group_idx in range(group_size): - getattr(op, f"weight{group_idx}").copy_(ws_test[group_idx]) + if single_grouped_weight: + op_weights[group_idx].copy_(ws_test[group_idx]) + else: + getattr(op, f"weight{group_idx}").copy_(ws_test[group_idx]) if bias: - getattr(op, f"bias{group_idx}").copy_(bs_test[group_idx]) + if single_grouped_bias: + op_bias_parts[group_idx].reshape(-1).copy_(bs_test[group_idx]) + else: + getattr(op, f"bias{group_idx}").copy_(bs_test[group_idx]) del ws_test, bs_test for param in op.parameters(): param.requires_grad_(requires_grad=weight_requires_grad) @@ -2144,20 +2174,37 @@ def test_grouped_linear( torch.testing.assert_close(dx_test, x_ref.grad, **tols) else: assert x_test.grad is None - for group_idx in range(group_size): - w_test = getattr(op, f"weight{group_idx}") + if single_grouped_weight: if weight_requires_grad: - dw_test = w_test.grad.to(dtype=torch.float64, device="cpu") - torch.testing.assert_close(dw_test, ws_ref[group_idx].grad, **tols) + dw_test_all = op.weight.grad.to(dtype=torch.float64, device="cpu") + w_ref_grad = torch.stack([w.grad for w in ws_ref], dim=0) + torch.testing.assert_close(dw_test_all, w_ref_grad, **tols) else: - assert w_test.grad is None - if bias: - b_test = getattr(op, f"bias{group_idx}") + assert op.weight.grad is None + else: + for group_idx in range(group_size): + w_test = getattr(op, f"weight{group_idx}") + if weight_requires_grad: + dw_test = w_test.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(dw_test, ws_ref[group_idx].grad, **tols) + else: + assert w_test.grad is None + if bias: + if single_grouped_bias: if weight_requires_grad: - db_test = b_test.grad.to(dtype=torch.float64, device="cpu") - torch.testing.assert_close(db_test, bs_ref[group_idx].grad, **tols) + db_test_all = op.bias.grad.to(dtype=torch.float64, device="cpu") + b_ref_grad = torch.stack([b.grad for b in bs_ref], dim=0) + torch.testing.assert_close(db_test_all, b_ref_grad, **tols) else: - assert b_test.grad is None + assert op.bias.grad is None + else: + for group_idx in range(group_size): + b_test = getattr(op, f"bias{group_idx}") + if weight_requires_grad: + db_test = b_test.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(db_test, bs_ref[group_idx].grad, **tols) + else: + assert b_test.grad is None @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index af243eedc1..d41ff5de0a 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -41,24 +41,6 @@ ) -# Quantizer types that support graph-safe `tex.group_quantize` / -# `tex.bgrad_group_quantize` and are eligible for the GroupedTensor + cublas -# `general_grouped_gemm_for_grouped_tensor` flow. All other recipes fall back -# to the legacy `tex.split_quantize` + `general_grouped_gemm` flow that uses -# CPU-side `m_splits` and is not CUDA-graph capturable. -# Quantizer types eligible for the graph-safe grouped-tensor flow -# (``tex.group_quantize`` + ``general_grouped_gemm_for_grouped_tensor``). -# -# NVFP4 is intentionally NOT included yet: -# * ``tex.group_quantize`` only implements the RHT (Random Hadamard Transform) -# code path for NVFP4 today (non-RHT raises in cast.cpp), and -# * grouped-swizzle support for NVFP4 in the cublas grouped GEMM is not -# implemented yet. -# When both land, add ``NVFP4Quantizer`` here (and update -# ``_is_grouped_quantize_supported`` accordingly). -_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS = (MXFP8Quantizer,) - - class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -728,7 +710,7 @@ def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. """ return all( - isinstance(q, _GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS) for q in input_quantizers + isinstance(q, MXFP8Quantizer) for q in input_quantizers ) def fuser_forward( @@ -754,7 +736,7 @@ def fuser_forward( # Check which grads are required ctx = basic_op_ctxs[0] input_requires_grad = ctx.requires_grad - weight_requires_grad = ctx.requires_grad and weight_param.requires_grad + weight_requires_grad = weight_param.requires_grad # Quantizers input_quantizers = [None] * num_groups @@ -1001,13 +983,12 @@ def _build_grouped_bias_for_gemm( def _quantize_weights_for_grouped_gemm( self, weight_quantizers: list[Optional[Quantizer]], - *, columnwise_usage: bool, + with_quantized_compute: bool, + dtype: torch.dtype, ) -> "GroupedTensor | list[torch.Tensor]": """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. - Mirrors the ``forward_grouped_mlp.py`` pattern. Four distinct cases: - ``single_grouped_weight=True``: * **quantized_model_init**: ``self.weight`` is already a quantized ``GroupedTensor`` (``self.weight.quantizer is not None``). We do @@ -1033,24 +1014,46 @@ def _quantize_weights_for_grouped_gemm( if self.single_grouped_weight: grouped_weight = self.weight - if not isinstance(grouped_weight, GroupedTensor): - raise RuntimeError( - "GroupedLinear with single_grouped_weight=True expected the weight " - f"parameter to be a GroupedTensor, got {type(grouped_weight).__name__}." + is_weight_quantized = grouped_weight.quantizer is not None + if is_weight_quantized and with_quantized_compute: + return grouped_weight + if is_weight_quantized and not with_quantized_compute: + # Quantized parameter (``quantized_model_init``) but the + # forward is running in unquantized compute. Dequantize each + # member back to high precision and re-wrap as a + # high-precision ``GroupedTensor`` so the grouped GEMM gets + # matching A/B element types. + weight_parts = grouped_weight.quantized_tensors + if weight_parts is None: + weight_parts = grouped_weight.split_into_quantized_tensors() + dequantized = [maybe_dequantize(w, dtype) for w in weight_parts] + weight_data = torch.stack(dequantized, dim=0).contiguous() + return GroupedTensor( + shape=(num_groups * self.out_features, self.in_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(self.out_features, self.in_features)] * num_groups, + quantizer=None, + data=weight_data.reshape(-1), + ) + if not with_quantized_compute: + # High-precision parameter and unquantized compute: just hand + # the existing high-precision ``GroupedTensor`` to the GEMM. + # Cast its rowwise buffer to ``dtype`` only if it differs (no + # copy in the common case). + if grouped_weight.rowwise_data.dtype == dtype: + return grouped_weight + weight_data = grouped_weight.rowwise_data.to(dtype=dtype) + return GroupedTensor( + shape=(num_groups * self.out_features, self.in_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(self.out_features, self.in_features)] * num_groups, + quantizer=None, + data=weight_data.reshape(-1), ) weight_quantizer = weight_quantizers[0] weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - - if grouped_weight.quantizer is not None: - # quantized_model_init: weight is already quantized. - grouped_weight.quantizer = weight_quantizer - return grouped_weight - - # Normal init: high-precision GroupedTensor. - if grouped_weight.rowwise_data is None: - raise RuntimeError( - "GroupedLinear: grouped weight has no rowwise_data to quantize." - ) return tex.group_quantize( grouped_weight.rowwise_data.view(grouped_weight.logical_shape), weight_quantizer, @@ -1062,13 +1065,12 @@ def _quantize_weights_for_grouped_gemm( out: list[torch.Tensor] = [] for idx, quantizer in enumerate(weight_quantizers): w = getattr(self, f"weight{idx}") - if is_quantized_tensor(w): - # quantized_model_init: weight{idx} is already quantized. - out.append(w) - else: - # Normal init: quantize the high-precision parameter. + if not with_quantized_compute: + w = maybe_dequantize(w, dtype) + elif with_quantized_compute and not is_quantized_tensor(w): quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - out.append(quantizer(w)) + w = quantizer(w) + out.append(w) return out def _fuser_forward_grouped_tensor( @@ -1137,21 +1139,12 @@ def _fuser_forward_grouped_tensor( # ``GroupedTensor``) for ``single_grouped_weight=True`` -- dispatches # the GEMM to ``no_discrete`` mode -- or the per-group ``weight{idx}`` # tensors otherwise -- ``discrete_in`` mode. - if with_quantized_compute: - grouped_weights = self._quantize_weights_for_grouped_gemm( - weight_quantizers, - columnwise_usage=input_requires_grad, - ) - elif self.single_grouped_weight: - # ``self.weight`` may itself be a quantized GroupedTensor when the - # caller used ``quantized_model_init``. Dequantize back to the - # compute dtype so cublas grouped GEMM sees matching A/B dtypes. - grouped_weights = maybe_dequantize(self.weight, dtype) - else: - grouped_weights = [ - maybe_dequantize(getattr(self, f"weight{idx}"), dtype) - for idx in range(num_groups) - ] + grouped_weights = self._quantize_weights_for_grouped_gemm( + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) # Allocate output buffer and wrap as a GroupedTensor view (no copy). out_shape = original_shape[:-1] + [self.out_features] @@ -1200,9 +1193,8 @@ def _fuser_forward_grouped_tensor( # by dgrad has already been populated up-front via # ``quantizer.set_usage(rowwise=True, columnwise=input_requires_grad)``. # This mirrors the behavior in ``forward_grouped_mlp.py``. - weight_is_grouped = isinstance(grouped_weights, GroupedTensor) if not input_requires_grad: - grouped_weights = None if weight_is_grouped else [None] * num_groups + grouped_weights = None if self.single_grouped_weight else [None] * num_groups # ``grouped_x`` already has all the data we need to save for backward # (rowwise for the unquantized path; columnwise too for the quantized @@ -1234,13 +1226,12 @@ def _fuser_forward_grouped_tensor( saved.extend([grouped_x.rowwise_data, None]) else: saved.extend([None, None]) - if weight_is_grouped: + if self.single_grouped_weight: saved.append(grouped_weights) else: saved.extend(grouped_weights) ctx.save_for_backward(*saved) ctx.use_grouped_tensor_path = True - ctx.weight_is_grouped = weight_is_grouped ctx.with_quantized_compute = with_quantized_compute ctx.input_quantizers = input_quantizers ctx.weight_quantizers = weight_quantizers @@ -1547,15 +1538,14 @@ def _fuser_backward_grouped_tensor( scales, saved_tensors = saved_tensors[0], saved_tensors[1:] x_data, saved_tensors = saved_tensors[0], saved_tensors[1:] x_scale, saved_tensors = saved_tensors[0], saved_tensors[1:] - if getattr(ctx, "weight_is_grouped", False): + if self.single_grouped_weight: ws, saved_tensors = saved_tensors[0], saved_tensors[1:] else: ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] - # Flatten grad_output to 2D (total_tokens, out_features), matching the - # forward convention. The grouped GEMMs and dbias kernels all operate - # on this 2D view. - dy_2d = maybe_dequantize(grad_output, dtype).reshape(-1, self.out_features) + # Flatten grad_output to 2D (total_tokens, out_features) + # to figure out total tokens and use it to build the grouped input tensor. + dy_2d = grad_output.reshape(-1, self.out_features) total_tokens = dy_2d.size(0) # Rebuild the grouped input tensor used by wgrad. @@ -1585,13 +1575,8 @@ def _fuser_backward_grouped_tensor( ) # Build the grad_output GroupedTensor. - # * Quantized: ``tex.group_quantize`` (or ``bgrad_group_quantize`` - # when bias gradient can be fused for MXFP8). - # * Unquantized: just wrap the contiguous bf16/fp16/fp32 buffer. - grad_biases: list[Optional[torch.Tensor]] = [None] * num_groups - grad_scales = None + # Optionally get dbias is fusion available with bgrad_group_quantize dbias_packed = None - if with_quantized_compute: grad_output_quantizer = ctx.grad_output_quantizers[0] grad_output_quantizer.set_usage( @@ -1599,12 +1584,7 @@ def _fuser_backward_grouped_tensor( ) grad_output_quantizer.optimize_for_gemm = True - use_fused_bgrad = ( - has_bias - and not self._scale_bias - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ) - if use_fused_bgrad: + if has_bias and not self._scale_bias: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, num_groups, split_sizes ) @@ -1613,6 +1593,8 @@ def _fuser_backward_grouped_tensor( dy_2d, grad_output_quantizer, num_groups, split_sizes ) else: + dy_2d = maybe_dequantize(dy_2d, dtype) + # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm grouped_dy = GroupedTensor( shape=(total_tokens, self.out_features), dtype=dtype, @@ -1623,7 +1605,9 @@ def _fuser_backward_grouped_tensor( tensor_offsets=base_offsets * self.out_features, ) - # Bias / scale-bias gradients (graph-safe; offsets are GPU tensors). + # Bias Grads compute if not already computed in bgrad_group_quantize + final_bias_grads: Optional[torch.Tensor] = None + grad_scales: Optional[torch.Tensor] = None if has_bias: if self._scale_bias: bias_packed = torch.stack(self._get_bias_tensors(dtype)) @@ -1635,14 +1619,15 @@ def _fuser_backward_grouped_tensor( offsets=base_offsets, ) elif dbias_packed is None: - # NVFP4 (or any path without a fused bgrad kernel) falls back - # to the standalone Triton dbias. + # BF16/FP16 path dbias_packed = compute_grouped_dbias(dy_2d, base_offsets, num_groups) - grad_biases = [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)] + + if self.single_grouped_bias: + final_bias_grads = [dbias_packed.to(dtype=dtype)] + else: + final_bias_grads = [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)] # ---- dgrad GEMM ---------------------------------------------------- - # grad_input mirrors grad_output's leading dims (batch / seq) and uses - # in_features for the trailing dim. grad_input = None if ctx.input_requires_grad: grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] @@ -1664,15 +1649,17 @@ def _fuser_backward_grouped_tensor( use_split_accumulator=_2X_ACC_DGRAD, ) - # ---- wgrad GEMM ---------------------------------------------------- + # params init for wgrad GEMM accumulate_into_main_grad = self._accumulate_into_main_grad weight_shape = (self.out_features, self.in_features) wgrad_output: Any = None grouped_wgrad: Optional[GroupedTensor] = None - w_list: list[Optional[torch.Tensor]] = ( + final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups ) + # Get the right wgrad buffers for grouped gemm. + # Can we GroupedTensor of list of tensors based on single_grouped_weight. if ctx.weight_requires_grad: if self.single_grouped_weight: if accumulate_into_main_grad: @@ -1680,14 +1667,12 @@ def _fuser_backward_grouped_tensor( weight_param.main_grad = weight_param.get_main_grad() main_grad = weight_param.main_grad grouped_shape = (num_groups, *weight_shape) - if main_grad.shape != grouped_shape: - if main_grad.numel() != math.prod(grouped_shape): - raise RuntimeError( - "GroupedLinear expected grouped weight main_grad to have " - f"shape {grouped_shape} or matching numel, " - f"but got shape {tuple(main_grad.shape)}" - ) - main_grad = main_grad.view(grouped_shape) + if main_grad.numel() != math.prod(grouped_shape): + raise RuntimeError( + "GroupedLinear expected grouped weight main_grad to have " + f"shape {grouped_shape} or matching numel, " + f"but got shape {tuple(main_grad.shape)}" + ) accumulate_into_main_grad = not getattr( weight_param, "overwrite_main_grad", False ) @@ -1695,7 +1680,7 @@ def _fuser_backward_grouped_tensor( grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( num_tensors=num_groups, tensor_shape=weight_shape, - rowwise_data=main_grad, + rowwise_data=main_grad.view(-1), dtype=main_grad.dtype, ) if grouped_wgrad is None: @@ -1706,6 +1691,7 @@ def _fuser_backward_grouped_tensor( device=device, dtype=dtype, ) + final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: if accumulate_into_main_grad: @@ -1713,17 +1699,18 @@ def _fuser_backward_grouped_tensor( wp = getattr(self, f"weight{idx}") if hasattr(wp, "__fsdp_param__"): wp.main_grad = wp.get_main_grad() - w_list[idx] = wp.main_grad + final_weight_grads[idx] = wp.main_grad accumulate_into_main_grad = not getattr( self.weight0, "overwrite_main_grad", False ) else: for idx in range(num_groups): - w_list[idx] = torch.empty(weight_shape, dtype=dtype, device=device) - wgrad_output = w_list + final_weight_grads[idx] = torch.empty(weight_shape, dtype=dtype, device=device) + wgrad_output = final_weight_grads else: accumulate_into_main_grad = False - + + # wgrad GEMM delay_wgrad = ( ctx.weight_requires_grad and self.wgrad_store is not None @@ -1741,70 +1728,39 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) - # ---- Assemble grad_params (mirrors split-quantize path layout) ---- - grad_weights: list[Optional[torch.Tensor]] = [None] * num_groups - if ctx.weight_requires_grad and not delay_wgrad: - if self.single_grouped_weight: - if grouped_wgrad is not None and not accumulate_into_main_grad: - packed = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - grad_weights = [packed[idx] for idx in range(num_groups)] - else: - grad_weights = list(w_list) - + # Dummy out the weight grad params if accumulating into main grad is true. if accumulate_into_main_grad: if self.single_grouped_weight: + final_weight_grads[0] = None if hasattr(weight_param, "grad_added_to_main_grad"): weight_param.grad_added_to_main_grad = True - grad_weight_dummy = get_dummy_wgrad( + final_weight_grads[0] = get_dummy_wgrad( list(weight_param.size()), weight_param.dtype, zero=getattr(weight_param, "zero_out_wgrad", False), ) - else: - grad_weight_dummy = None - if has_bias: - if self.single_grouped_bias: - final_bias_grads = torch.stack(grad_biases, dim=0).to(dtype) - grad_params = [grad_weight_dummy, final_bias_grads] - else: - grad_params = grad_biases + [grad_weight_dummy] - else: - grad_params = [grad_weight_dummy] - grad_extra = (None, grad_scales) if self._scale_bias else (None,) - return grad_input, [grad_params], [grad_extra] - - grad_weights = [None] * num_groups - for group_idx in range(num_groups): - wp = getattr(self, f"weight{group_idx}") - if hasattr(wp, "grad_added_to_main_grad"): - wp.grad_added_to_main_grad = True - grad_weights[group_idx] = get_dummy_wgrad( - list(wp.size()), - wp.dtype, - zero=getattr(wp, "zero_out_wgrad", False), - ) - - if self.single_grouped_weight: - grad_weight = None - if ctx.weight_requires_grad and not delay_wgrad: - grad_weight = torch.stack(grad_weights, dim=0) - final_weight_grads = [grad_weight] - else: - if delay_wgrad and ctx.weight_requires_grad and not accumulate_into_main_grad: - final_weight_grads = [None] * num_groups else: - final_weight_grads = grad_weights + final_weight_grads = [None] * num_groups + for group_idx in range(num_groups): + wp = getattr(self, f"weight{group_idx}") + if hasattr(wp, "grad_added_to_main_grad"): + wp.grad_added_to_main_grad = True + final_weight_grads[group_idx] = get_dummy_wgrad( + list(wp.size()), + wp.dtype, + zero=getattr(wp, "zero_out_wgrad", False), + ) + # Assemble grad params in parameter registration order and return. if not has_bias: - grad_params = list(final_weight_grads) + grad_params = final_weight_grads elif self.single_grouped_bias: - final_bias_grads = torch.stack(grad_biases, dim=0).to(dtype) - grad_params = list(final_weight_grads) + [final_bias_grads] + grad_params = final_weight_grads + final_bias_grads else: if self.single_grouped_weight: - grad_params = list(grad_biases) + list(final_weight_grads) + grad_params = final_bias_grads + final_weight_grads else: - grad_params = list(final_weight_grads) + list(grad_biases) + grad_params = final_weight_grads + final_bias_grads grad_extra = (None, grad_scales) if self._scale_bias else (None,) return grad_input, [grad_params], [grad_extra] From 2082ce6e14177e147995757854ee1c032f241953 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 24 Apr 2026 18:58:19 +0000 Subject: [PATCH 03/18] cuda graph test + clean ups Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 183 ++++++ .../pytorch/ops/basic/grouped_linear.py | 548 ++++++++---------- .../pytorch/ops/fused/backward_grouped_mlp.py | 52 +- 3 files changed, 474 insertions(+), 309 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index c541026c5e..f741157081 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2206,6 +2206,189 @@ def test_grouped_linear( else: assert b_test.grad is None + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) + @pytest.mark.parametrize( + "quantization", + [None] + (["mxfp8"] if mxfp8_available else []), + ) + @pytest.mark.parametrize("quantized_weight", (False, True)) + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) + @pytest.mark.parametrize("single_grouped_bias", (False, True)) + @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) + def test_grouped_linear_cuda_graph_safe( + self, + *, + dtype: torch.dtype, + quantization: Optional[str], + quantized_weight: bool, + bias: bool, + single_grouped_weight: bool, + single_grouped_bias: bool, + accumulate_into_main_grad: bool, + device: torch.device = "cuda", + group_size: int = 4, + in_features: int = 128, + out_features: int = 128, + split_alignment: int = 128, + token_padding: int = 256, + ) -> None: + """GroupedLinear forward+backward should be CUDA graph capturable. + + Exercises the grouped-tensor / cublas-grouped-gemm path which uses + GPU-resident split offsets and is the only flow safe to capture. + """ + + # Skip invalid configurations + if quantization is None and quantized_weight: + pytest.skip("quantized_weight requires a quantization recipe") + if single_grouped_bias and not bias: + pytest.skip("single_grouped_bias requires bias=True") + + # Split sizes (statically pinned for graph capture) + split_sizes = [split_alignment * (i + 1) for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + # Pad input tokens to validate the sync-free flow + in_shape = (split_sizes.sum().item() + token_padding, in_features) + out_shape = (in_shape[0], out_features) + + recipe = make_recipe(quantization) + with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): + op = te_ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=bias, + device=device, + dtype=dtype, + accumulate_into_main_grad=accumulate_into_main_grad, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + ) + + def _weight_params() -> list[torch.nn.Parameter]: + if single_grouped_weight: + return [op.weight] + return [getattr(op, f"weight{i}") for i in range(group_size)] + + def _bias_params() -> list[torch.nn.Parameter]: + if not bias: + return [] + if single_grouped_bias: + return [op.bias] + return [getattr(op, f"bias{i}") for i in range(group_size)] + + def _init_main_grads(value: float = 0.0) -> None: + if not accumulate_into_main_grad: + return + with torch.no_grad(): + for w in _weight_params(): + if getattr(w, "main_grad", None) is None: + w.main_grad = torch.empty( + w.size(), device=device, dtype=torch.float32 + ) + w.main_grad.fill_(value) + + def _collect_main_grads() -> list[torch.Tensor]: + return [w.main_grad.detach().clone() for w in _weight_params()] + + def _zero_param_grads() -> None: + for param in op.parameters(): + if param.grad is None: + param.grad = torch.zeros_like(param) + else: + param.grad.zero_() + + static_split_sizes = split_sizes.clone() + + def train_step( + x: torch.Tensor, + dy: torch.Tensor, + out_buf: torch.Tensor, + *, + use_graphed: bool, + ) -> torch.Tensor: + with te.autocast(enabled=quantization is not None, recipe=recipe): + out = ( + graphed_module(x, static_split_sizes) + if use_graphed + else op(x, static_split_sizes) + ) + out.backward(dy) + out_buf.copy_(out) + return out_buf + + _init_main_grads(0.0) + + static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True) + static_dy = torch.randn(out_shape, device=device, dtype=dtype) + static_out_buf = torch.empty(out_shape, device=device, dtype=dtype) + + graphed_module = te.make_graphed_callables( + op, + (static_x, static_split_sizes), + num_warmup_iters=3, + enabled=quantization is not None, + recipe=recipe, + ) + + # Replace static buffers with fresh data (graph captures must replay + # against new inputs without re-recording). + fresh_x = torch.randn_like(static_x) + fresh_dy = torch.randn_like(static_dy) + with torch.no_grad(): + static_x.copy_(fresh_x) + static_dy.copy_(fresh_dy) + + # Reset grads & main_grads so the captured iteration starts fresh. + _zero_param_grads() + _init_main_grads(0.5) + if static_x.grad is not None: + static_x.grad.zero_() + + # Replay the graph + graph_out = ( + train_step(static_x, static_dy, static_out_buf, use_graphed=True) + .detach() + .clone() + ) + torch.cuda.synchronize() + graph_dx = static_x.grad.detach().clone() + if accumulate_into_main_grad: + graph_main_grads = _collect_main_grads() + graph_param_grads: list[torch.Tensor] = [] + else: + graph_main_grads = [] + graph_param_grads = [ + param.grad.detach().clone() for param in op.parameters() + ] + + # Reference: same op invoked eagerly with the same fresh inputs and + # the same starting grad/main_grad state. + _zero_param_grads() + _init_main_grads(0.5) + static_x.grad.zero_() + + expected_x = fresh_x.detach().clone().requires_grad_(True) + expected_dy = fresh_dy.detach().clone() + with te.autocast(enabled=quantization is not None, recipe=recipe): + expected_out = op(expected_x, static_split_sizes) + expected_out.backward(expected_dy) + + tols = dtype_tols(dtype) + if quantization is not None: + tols = quantization_tols(quantization) + + assert_close(graph_out, expected_out, **tols) + assert_close(graph_dx, expected_x.grad, **tols) + if accumulate_into_main_grad: + for g, w in zip(graph_main_grads, _weight_params()): + assert_close(g, w.main_grad, **tols) + else: + for g, param in zip(graph_param_grads, op.parameters()): + assert_close(g, param.grad, **tols) + @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("scales_requires_grad", (False, True)) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index d41ff5de0a..21f88ce40e 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -696,10 +696,6 @@ def op_backward(self, *args, **kwargs): "It overrides `fuser_backward` instead of `op_backward`." ) - # ------------------------------------------------------------------ - # Path-selection helpers for fuser_forward/backward - # ------------------------------------------------------------------ - @staticmethod def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer]]) -> bool: """Whether all input quantizers support the graph-safe grouped-tensor flow. @@ -712,6 +708,149 @@ def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer return all( isinstance(q, MXFP8Quantizer) for q in input_quantizers ) + + def _get_grouped_weight_for_gemm( + self, + weight_param: GroupedTensor, + weight_quantizers: list[Optional[Quantizer]], + columnwise_usage: bool, + with_quantized_compute: bool, + dtype: torch.dtype, + ) -> "GroupedTensor | list[torch.Tensor]": + """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. + Supports MXFP8/BF16/FP16 compute paths. + """ + num_groups = self.num_groups + is_weight_quantized = weight_param.quantizer is not None + if is_weight_quantized and with_quantized_compute: + # GGEMM can use it as it is + return weight_param + if is_weight_quantized and not with_quantized_compute: + # This use-case isnt optimized yet. Involves a per-group + # dequantize loop and a torch.stack copy. + weight_parts = weight_param.quantized_tensors + if weight_parts is None: + weight_parts = weight_param.split_into_quantized_tensors() + dequantized = [maybe_dequantize(w, dtype) for w in weight_parts] + weight_data = torch.stack(dequantized, dim=0).contiguous() + return GroupedTensor( + shape=(num_groups * self.out_features, self.in_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(self.out_features, self.in_features)] * num_groups, + quantizer=None, + data=weight_data.reshape(-1), + ) + if not with_quantized_compute: + # Make sure that weight param is the correct dtype, + # otherwise cast it to the correct dtype. + if weight_param.rowwise_data.dtype == dtype: + return weight_param + weight_data = weight_param.rowwise_data.to(dtype=dtype) + return GroupedTensor( + shape=(num_groups * self.out_features, self.in_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(self.out_features, self.in_features)] * num_groups, + quantizer=None, + data=weight_data.reshape(-1), + ) + # Quantized compute path, use the fused group quantize kernel. + weight_quantizer = weight_quantizers[0] + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + return tex.group_quantize( + weight_param.rowwise_data.view(weight_param.logical_shape), + weight_quantizer, + num_groups, + None, + ) + + + def _get_discrete_weights_for_gemm(self, + weight_params: Optional[GroupedTensor] | list[torch.Tensor], + weight_quantizers: list[Optional[Quantizer]], + columnwise_usage: bool, + with_quantized_compute: bool, + dtype: torch.dtype, + ) -> list[torch.Tensor]: + """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. + Returns a Python list, which dispatches the GEMM to ``discrete_in`` mode. + """ + out: list[torch.Tensor] = [] + for w, quantizer in zip(weight_params, weight_quantizers): + if not with_quantized_compute: + w = maybe_dequantize(w, dtype) + elif with_quantized_compute and not is_quantized_tensor(w): + quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + w = quantizer(w) + out.append(w) + return out + + def _dummy_main_grad_wgrads(self) -> list[Optional[torch.Tensor]]: + """Return a NEW list of per-output dummy weight gradients for the + ``accumulate_into_main_grad`` (Megatron-LM wgrad fusion) path. + Length: 1 for ``single_grouped_weight=True``, ``num_groups`` + otherwise. + """ + if self.single_grouped_weight: + weight_param = self.weight + if hasattr(weight_param, "grad_added_to_main_grad"): + weight_param.grad_added_to_main_grad = True + return [ + get_dummy_wgrad( + list(weight_param.size()), + weight_param.dtype, + zero=getattr(weight_param, "zero_out_wgrad", False), + ) + ] + return [None] + + out: list[Optional[torch.Tensor]] = [None] * self.num_groups + for group_idx in range(self.num_groups): + wp = getattr(self, f"weight{group_idx}") + if hasattr(wp, "grad_added_to_main_grad"): + wp.grad_added_to_main_grad = True + out[group_idx] = get_dummy_wgrad( + list(wp.size()), + wp.dtype, + zero=getattr(wp, "zero_out_wgrad", False), + ) + return out + + def _get_grouped_bias_for_gemm( + self, + dtype: torch.dtype, + device: torch.device, + ) -> Optional[torch.Tensor]: + """Build a uniform GroupedTensor of per-group biases for the cublas + grouped GEMM. + + Each group expects a (1, out_features) bias vector. Returns ``None`` + when no additive bias is configured. + """ + if not self.has_bias: + return None + num_groups = self.num_groups + + if self.single_grouped_bias: + # Already a contiguous (num_groups * out_features) buffer. + bias_data = self.bias.rowwise_data + if bias_data.dtype != dtype: + bias_data = bias_data.to(dtype=dtype) + else: + bias_list = [ + maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) + ] + bias_data = torch.stack(bias_list, dim=0).contiguous() + + return GroupedTensor( + shape=(num_groups, self.out_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(1, self.out_features)] * num_groups, + quantizer=None, + data=bias_data.reshape(-1), + ) def fuser_forward( self, @@ -815,10 +954,8 @@ def fuser_forward( # ================================================================== # Legacy `tex.split_quantize` + `general_grouped_gemm` flow. - # Used for bf16/fp16 inputs and FP8 recipes other than MXFP8/NVFP4. # ``m_splits`` is needed on CPU here, so this flow is NOT cuda-graphable. # ================================================================== - def _fuser_forward_split_quantize( self, *, @@ -852,15 +989,13 @@ def _fuser_forward_split_quantize( if has_bias: bs = self._get_bias_tensors(dtype) - # Convert weight dtype if needed - ws = [] - for w, quantizer in zip(weights, weight_quantizers): - if not with_quantized_compute: - w = maybe_dequantize(w, dtype) - elif with_quantized_compute and not is_quantized_tensor(w): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - w = quantizer(w) - ws.append(w) + ws = self._get_discrete_weights_for_gemm( + weights, + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) # Split input tensor and convert dtypes if needed x = maybe_dequantize(input_, dtype) @@ -893,7 +1028,6 @@ def _fuser_forward_split_quantize( ) # Add bias * scales when scale_bias is enabled - # TODO(vthumbe): Need to use GroupedBiasAdd kernel here. # Would be done as part of larger refactor for GroupedLinear + GroupedTensor # integration. if self._scale_bias and has_bias: @@ -937,141 +1071,6 @@ def _fuser_forward_split_quantize( return out, [()] - # ================================================================== - # Graph-safe `tex.group_quantize` + `general_grouped_gemm_for_grouped_tensor` - # flow. Used only for MXFP8 / NVFP4 recipes, where ``m_splits`` stays on - # GPU and the flow can be captured into a CUDA graph. - # ``general_grouped_gemm_for_grouped_tensor`` handles bias / scaled-bias - # addition inside the GEMM. - # ================================================================== - - def _build_grouped_bias_for_gemm( - self, - dtype: torch.dtype, - device: torch.device, - ) -> Optional[torch.Tensor]: - """Build a uniform GroupedTensor of per-group biases for the cublas - grouped GEMM. - - Each group expects a (1, out_features) bias vector. Returns ``None`` - when no additive bias is configured. - """ - if not self.has_bias: - return None - num_groups = self.num_groups - - if self.single_grouped_bias: - # Already a contiguous (num_groups * out_features) buffer. - bias_data = self.bias.rowwise_data - if bias_data.dtype != dtype: - bias_data = bias_data.to(dtype=dtype) - else: - bias_list = [ - maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) - ] - bias_data = torch.stack(bias_list, dim=0).contiguous() - - return GroupedTensor( - shape=(num_groups, self.out_features), - dtype=dtype, - num_tensors=num_groups, - shapes=[(1, self.out_features)] * num_groups, - quantizer=None, - data=bias_data.reshape(-1), - ) - - def _quantize_weights_for_grouped_gemm( - self, - weight_quantizers: list[Optional[Quantizer]], - columnwise_usage: bool, - with_quantized_compute: bool, - dtype: torch.dtype, - ) -> "GroupedTensor | list[torch.Tensor]": - """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. - - ``single_grouped_weight=True``: - * **quantized_model_init**: ``self.weight`` is already a quantized - ``GroupedTensor`` (``self.weight.quantizer is not None``). We do - NOT re-quantize; we just refresh its ``quantizer`` reference so - usage flags propagate, and return it as-is. Dispatches the GEMM - to ``no_discrete`` mode. - * **normal init**: ``self.weight`` is a high-precision - ``GroupedTensor`` (``quantizer is None``). We call - ``tex.group_quantize`` once on its packed ``rowwise_data`` and - return the resulting quantized ``GroupedTensor``. Dispatches the - GEMM to ``no_discrete`` mode. - - ``single_grouped_weight=False``: - * **quantized_model_init**: each ``weight{idx}`` is already a - quantized tensor (e.g. ``MXFP8Tensor``); use as-is. - * **normal init**: each ``weight{idx}`` is a high-precision - ``torch.Tensor``; quantize per-group via ``quantizer(weight)``. - - In both subcases we return a Python list, which dispatches the GEMM - to ``discrete_in`` mode. - """ - num_groups = self.num_groups - - if self.single_grouped_weight: - grouped_weight = self.weight - is_weight_quantized = grouped_weight.quantizer is not None - if is_weight_quantized and with_quantized_compute: - return grouped_weight - if is_weight_quantized and not with_quantized_compute: - # Quantized parameter (``quantized_model_init``) but the - # forward is running in unquantized compute. Dequantize each - # member back to high precision and re-wrap as a - # high-precision ``GroupedTensor`` so the grouped GEMM gets - # matching A/B element types. - weight_parts = grouped_weight.quantized_tensors - if weight_parts is None: - weight_parts = grouped_weight.split_into_quantized_tensors() - dequantized = [maybe_dequantize(w, dtype) for w in weight_parts] - weight_data = torch.stack(dequantized, dim=0).contiguous() - return GroupedTensor( - shape=(num_groups * self.out_features, self.in_features), - dtype=dtype, - num_tensors=num_groups, - shapes=[(self.out_features, self.in_features)] * num_groups, - quantizer=None, - data=weight_data.reshape(-1), - ) - if not with_quantized_compute: - # High-precision parameter and unquantized compute: just hand - # the existing high-precision ``GroupedTensor`` to the GEMM. - # Cast its rowwise buffer to ``dtype`` only if it differs (no - # copy in the common case). - if grouped_weight.rowwise_data.dtype == dtype: - return grouped_weight - weight_data = grouped_weight.rowwise_data.to(dtype=dtype) - return GroupedTensor( - shape=(num_groups * self.out_features, self.in_features), - dtype=dtype, - num_tensors=num_groups, - shapes=[(self.out_features, self.in_features)] * num_groups, - quantizer=None, - data=weight_data.reshape(-1), - ) - weight_quantizer = weight_quantizers[0] - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - return tex.group_quantize( - grouped_weight.rowwise_data.view(grouped_weight.logical_shape), - weight_quantizer, - num_groups, - None, - ) - - # single_grouped_weight=False: per-group parameters. - out: list[torch.Tensor] = [] - for idx, quantizer in enumerate(weight_quantizers): - w = getattr(self, f"weight{idx}") - if not with_quantized_compute: - w = maybe_dequantize(w, dtype) - elif with_quantized_compute and not is_quantized_tensor(w): - quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - w = quantizer(w) - out.append(w) - return out def _fuser_forward_grouped_tensor( self, @@ -1090,26 +1089,13 @@ def _fuser_forward_grouped_tensor( device: torch.device, ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: """Graph-safe GroupedTensor forward path. - - Handles both: - * **Quantized compute** (currently MXFP8): input/weights/grad_output - are quantized via ``tex.group_quantize`` and the GEMM consumes the - packed scale/data buffers. - * **Unquantized compute** (bf16/fp16/fp32): no quantize calls -- we - simply wrap the contiguous high-precision data as a - ``GroupedTensor`` so the cublas grouped GEMM path is still used, - making the operation CUDA-graphable end-to-end. """ num_groups = self.num_groups has_bias = self.has_bias - # Cumulative-token offsets on GPU (graph-safe). Computed once and - # scaled by the relevant last dim to obtain `tensor_offsets` for any - # GroupedTensor with that last dim. Reused in backward. base_offsets = tex.splits_to_offsets(split_sizes, 1) - # Flatten to 2D so the first dim is the total token count, matching - # the convention used by forward_grouped_mlp.py. + # Flatten to 2D so the first dim is the total token count. original_shape = list(input_.size()) x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) total_tokens = x.size(0) @@ -1139,12 +1125,22 @@ def _fuser_forward_grouped_tensor( # ``GroupedTensor``) for ``single_grouped_weight=True`` -- dispatches # the GEMM to ``no_discrete`` mode -- or the per-group ``weight{idx}`` # tensors otherwise -- ``discrete_in`` mode. - grouped_weights = self._quantize_weights_for_grouped_gemm( - weight_quantizers, - columnwise_usage=input_requires_grad, - with_quantized_compute=with_quantized_compute, - dtype=dtype, - ) + if self.single_grouped_weight: + grouped_weights = self._get_grouped_weight_for_gemm( + self.weight, + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) + else: + grouped_weights = self._get_discrete_weights_for_gemm( + [getattr(self, f"weight{idx}") for idx in range(num_groups)], + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) # Allocate output buffer and wrap as a GroupedTensor view (no copy). out_shape = original_shape[:-1] + [self.out_features] @@ -1164,7 +1160,8 @@ def _fuser_forward_grouped_tensor( grouped_bias = None bias_scale: Optional[torch.Tensor] = None if has_bias: - grouped_bias = self._build_grouped_bias_for_gemm(dtype, device) + # Bias always needs to be passed as a GroupedTensor for the grouped GEMM. + grouped_bias = self._get_grouped_bias_for_gemm(dtype, device) if self._scale_bias: bias_scale = scales.reshape(-1) if bias_scale.dtype != torch.float32: @@ -1181,18 +1178,6 @@ def _fuser_forward_grouped_tensor( bias_scale=bias_scale, ) - # Prepare weight tensor(s) for the dgrad pass. - # - # NOTE: We deliberately do NOT call ``update_usage(rowwise_usage=False, ...)`` - # on quantized weights here. In the ``quantized_model_init`` case, - # ``grouped_weights`` IS the user's ``self.weight`` parameter (or its - # per-group ``weight{idx}`` parameters), and dropping its rowwise data - # would corrupt the parameter for subsequent forward passes. - # - # For the quantized compute path, the columnwise representation needed - # by dgrad has already been populated up-front via - # ``quantizer.set_usage(rowwise=True, columnwise=input_requires_grad)``. - # This mirrors the behavior in ``forward_grouped_mlp.py``. if not input_requires_grad: grouped_weights = None if self.single_grouped_weight else [None] * num_groups @@ -1323,27 +1308,44 @@ def _fuser_backward_split_quantize( dbias_packed = compute_grouped_dbias(dy_2d, offsets, num_groups) grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] - # Initialize grad weight buffers - accumulate_into_main_grad = self._accumulate_into_main_grad + # Initialize grad weight buffers and the autograd-return + # ``final_weight_grads`` upfront (mirrors the grouped-tensor + # flow). For ``single_grouped_weight=True`` we pre-allocate a + # stacked ``[num_groups, out, in]`` buffer and feed per-group + # views of it to the GEMM, so the buffer itself is the autograd + # return -- no post-GEMM ``torch.stack`` copy needed. + # ``request_main_grad_fusion`` records the user-facing opt-in to + # Megatron-LM main-grad fusion; ``accumulate_into_main_grad`` is the + # local GEMM ``accumulate`` flag (downgraded to ``False`` when + # ``weight.overwrite_main_grad`` is set, e.g. Megatron-FSDP). The + # post-GEMM bookkeeping (dummy ``.grad`` + ``grad_added_to_main_grad``) + # always fires when fusion was requested -- see comment at the + # post-GEMM dummy block below. + request_main_grad_fusion = ( + ctx.weight_requires_grad and self._accumulate_into_main_grad + ) + accumulate_into_main_grad = request_main_grad_fusion grad_weights = [None] * num_groups + final_weight_grads: list[Optional[torch.Tensor]] = ( + [None] if self.single_grouped_weight else [None] * num_groups + ) if ctx.weight_requires_grad: - if accumulate_into_main_grad: - # Megatron-LM wgrad fusion - # Note: Get grad tensors from params so we can - # accumulate directly into it. - if self.single_grouped_weight: + weight_shape = (self.out_features, self.in_features) + grouped_shape = (num_groups, *weight_shape) + if self.single_grouped_weight: + if accumulate_into_main_grad: + # Megatron-LM wgrad fusion: GEMM accumulates into the + # parameter's ``main_grad`` directly. if hasattr(weight_param, "__fsdp_param__"): weight_param.main_grad = weight_param.get_main_grad() main_grad = weight_param.main_grad if isinstance(main_grad, GroupedTensor): + # Legacy path: no contiguous backing; per-group + # quantized tensors as discrete GEMM outputs. grad_weights = main_grad.quantized_tensors if grad_weights is None: grad_weights = main_grad.split_into_quantized_tensors() else: - # main_grad may be [num_groups, out, in] or a flat buffer. - # Canonicalize to grouped layout before slicing per-group views. - weight_shape = (self.out_features, self.in_features) - grouped_shape = (num_groups, *weight_shape) if main_grad.shape != grouped_shape: if main_grad.numel() != math.prod(grouped_shape): raise RuntimeError( @@ -1352,11 +1354,18 @@ def _fuser_backward_split_quantize( f"but got shape {tuple(main_grad.shape)}" ) main_grad = main_grad.reshape(grouped_shape) + final_weight_grads[0] = main_grad grad_weights = [main_grad[idx] for idx in range(num_groups)] accumulate_into_main_grad = not getattr( weight_param, "overwrite_main_grad", False ) else: + final_weight_grads[0] = torch.empty( + grouped_shape, dtype=ctx.dtype, device=device + ) + grad_weights = [final_weight_grads[0][idx] for idx in range(num_groups)] + else: + if accumulate_into_main_grad: for group_idx in range(num_groups): weight_param = getattr(self, f"weight{group_idx}") if hasattr(weight_param, "__fsdp_param__"): @@ -1365,16 +1374,12 @@ def _fuser_backward_split_quantize( accumulate_into_main_grad = not getattr( self.weight0, "overwrite_main_grad", False ) - else: - weight_shape = (self.out_features, self.in_features) - for group_idx in range(num_groups): - grad_weights[group_idx] = torch.empty( - weight_shape, - dtype=ctx.dtype, - device=device, - ) - else: - accumulate_into_main_grad = False + else: + for group_idx in range(num_groups): + grad_weights[group_idx] = torch.empty( + weight_shape, dtype=ctx.dtype, device=device + ) + final_weight_grads = list(grad_weights) # Perform dgrad GEMMs grad_input = None @@ -1432,54 +1437,16 @@ def _fuser_backward_split_quantize( if not delay_wgrad: clear_tensor_data(*xs) - # Megatron-LM wgrad fusion - # Note: Return dummy tensor for grad weight if needed. - if accumulate_into_main_grad: - grad_weights = [None] * num_groups - if self.single_grouped_weight: - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - else: - grad_weight = None - # Be mindful of param registration order. - if has_bias: - if self.single_grouped_bias: - final_bias_grads = torch.stack(grad_biases, dim=0).to(ctx.dtype) - grad_params = [grad_weight, final_bias_grads] - else: - grad_params = grad_biases + [grad_weight] - else: - grad_params = [grad_weight] - grad_extra = (None, grad_scales) if self._scale_bias else (None,) - return grad_input, [grad_params], [grad_extra] - for group_idx in range(num_groups): - weight_param = getattr(self, f"weight{group_idx}") - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weights[group_idx] = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - - if self.single_grouped_weight: - grad_weight = None - if ctx.weight_requires_grad: - if delay_wgrad: - grad_weight = None - else: - grad_weight = torch.stack(grad_weights, dim=0) - final_weight_grads = [grad_weight] - else: - if delay_wgrad and ctx.weight_requires_grad and not accumulate_into_main_grad: - final_weight_grads = [None] * num_groups - else: - final_weight_grads = grad_weights + # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, + # signal that ``main_grad`` already carries the wgrad and replace + # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into + # ``main_grad`` again. Matches the convention in ``module/linear.py`` + # and the grouped-tensor backward path. We REBIND + # ``final_weight_grads`` rather than mutate it -- ``grad_weights`` + # (captured by ``wgrad_store.put`` for delayed wgrad) still holds + # the ``main_grad`` views the GEMM writes into. + if request_main_grad_fusion: + final_weight_grads = self._dummy_main_grad_wgrads() if not has_bias: grad_params = list(final_weight_grads) @@ -1497,11 +1464,7 @@ def _fuser_backward_split_quantize( # ================================================================== # Graph-safe backward: counterpart of `_fuser_forward_grouped_tensor`. - # Uses ``tex.bgrad_group_quantize`` / ``tex.group_quantize`` on the grad - # output and ``general_grouped_gemm_for_grouped_tensor`` for both dgrad - # and wgrad. ``m_splits`` stays on GPU throughout. # ================================================================== - def _fuser_backward_grouped_tensor( self, *, @@ -1621,7 +1584,6 @@ def _fuser_backward_grouped_tensor( elif dbias_packed is None: # BF16/FP16 path dbias_packed = compute_grouped_dbias(dy_2d, base_offsets, num_groups) - if self.single_grouped_bias: final_bias_grads = [dbias_packed.to(dtype=dtype)] else: @@ -1649,8 +1611,21 @@ def _fuser_backward_grouped_tensor( use_split_accumulator=_2X_ACC_DGRAD, ) - # params init for wgrad GEMM - accumulate_into_main_grad = self._accumulate_into_main_grad + # params init for wgrad GEMM + # ``request_main_grad_fusion`` records the user-facing opt-in to + # Megatron-LM main-grad fusion. ``accumulate_into_main_grad`` is the + # local GEMM ``accumulate`` flag, which gets downgraded to ``False`` + # when ``weight.overwrite_main_grad`` is set (e.g. Megatron-FSDP). + # The two must stay separated: the GEMM flag controls overwrite vs. + # accumulate, but the post-GEMM bookkeeping (dummy ``.grad`` + + # ``grad_added_to_main_grad=True``) must always fire when fusion was + # requested -- otherwise FSDP's post-backward hook would also touch + # ``main_grad`` and double-count (or, with ``delay_wgrad``, copy + # uninitialized data before the deferred GEMM ever runs). + request_main_grad_fusion = ( + ctx.weight_requires_grad and self._accumulate_into_main_grad + ) + accumulate_into_main_grad = request_main_grad_fusion weight_shape = (self.out_features, self.in_features) wgrad_output: Any = None grouped_wgrad: Optional[GroupedTensor] = None @@ -1659,10 +1634,14 @@ def _fuser_backward_grouped_tensor( ) # Get the right wgrad buffers for grouped gemm. - # Can we GroupedTensor of list of tensors based on single_grouped_weight. + # Can be a GroupedTensor or list of tensors based on single_grouped_weight. if ctx.weight_requires_grad: if self.single_grouped_weight: if accumulate_into_main_grad: + # Main-grad fusion: GEMM writes directly into ``main_grad``. + # ``overwrite_main_grad`` only flips the GEMM's + # ``accumulate`` flag (overwrite vs. accumulate); it does + # not change the output buffer. if hasattr(weight_param, "__fsdp_param__"): weight_param.main_grad = weight_param.get_main_grad() main_grad = weight_param.main_grad @@ -1673,17 +1652,16 @@ def _fuser_backward_grouped_tensor( f"shape {grouped_shape} or matching numel, " f"but got shape {tuple(main_grad.shape)}" ) + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=num_groups, + tensor_shape=weight_shape, + rowwise_data=main_grad.view(-1), + dtype=main_grad.dtype, + ) accumulate_into_main_grad = not getattr( weight_param, "overwrite_main_grad", False ) - if accumulate_into_main_grad: - grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( - num_tensors=num_groups, - tensor_shape=weight_shape, - rowwise_data=main_grad.view(-1), - dtype=main_grad.dtype, - ) - if grouped_wgrad is None: + else: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_groups, shapes=[weight_shape] * num_groups, @@ -1707,9 +1685,7 @@ def _fuser_backward_grouped_tensor( for idx in range(num_groups): final_weight_grads[idx] = torch.empty(weight_shape, dtype=dtype, device=device) wgrad_output = final_weight_grads - else: - accumulate_into_main_grad = False - + # wgrad GEMM delay_wgrad = ( ctx.weight_requires_grad @@ -1728,28 +1704,12 @@ def _fuser_backward_grouped_tensor( else: wgrad_gemm(grouped_x, grouped_dy, wgrad_output) - # Dummy out the weight grad params if accumulating into main grad is true. - if accumulate_into_main_grad: - if self.single_grouped_weight: - final_weight_grads[0] = None - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - final_weight_grads[0] = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - else: - final_weight_grads = [None] * num_groups - for group_idx in range(num_groups): - wp = getattr(self, f"weight{group_idx}") - if hasattr(wp, "grad_added_to_main_grad"): - wp.grad_added_to_main_grad = True - final_weight_grads[group_idx] = get_dummy_wgrad( - list(wp.size()), - wp.dtype, - zero=getattr(wp, "zero_out_wgrad", False), - ) + # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, + # signal that ``main_grad`` already carries the wgrad and replace + # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into + # ``main_grad`` again. Matches the convention in ``module/linear.py``. + if request_main_grad_fusion: + final_weight_grads = self._dummy_main_grad_wgrads() # Assemble grad params in parameter registration order and return. if not has_bias: diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index aca49e9866..8e7fd566c7 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -143,7 +143,20 @@ def _compute_grad_params( Returns the grad_params list in parameter registration order. """ - # Allocate grad buffers, determine accumulate flag + # Allocate grad buffers, determine accumulate flag. + # ``request_main_grad_fusion`` records the user-facing opt-in to + # Megatron-LM main-grad fusion. ``accumulate_into_main_grad`` is the + # local GEMM ``accumulate`` flag, which gets downgraded to ``False`` when + # ``weight.overwrite_main_grad`` is set (e.g. Megatron-FSDP). The two must + # stay separated: the GEMM flag controls overwrite vs. accumulate, but the + # post-GEMM bookkeeping (dummy ``.grad`` + ``grad_added_to_main_grad=True``) + # must always fire when fusion was requested -- otherwise FSDP's + # post-backward hook would also touch ``main_grad`` and double-count + # (or, with ``delay_wgrad``, copy uninitialized data before the deferred + # GEMM ever runs). + request_main_grad_fusion = ( + ctx.weight_requires_grad and fc_op._accumulate_into_main_grad + ) accumulate_into_main_grad = False grouped_wgrad = None wgrad_output = None @@ -152,6 +165,10 @@ def _compute_grad_params( if ctx.weight_requires_grad: weight_param = fc_op.weight if fc_op._accumulate_into_main_grad: + # Main-grad fusion: GEMM writes directly into ``main_grad``. + # ``overwrite_main_grad`` only flips the GEMM's ``accumulate`` + # flag (overwrite vs. accumulate); it does not change the + # output buffer. if hasattr(weight_param, "__fsdp_param__"): weight_param.main_grad = weight_param.get_main_grad() main_grad = weight_param.main_grad @@ -172,16 +189,16 @@ def _compute_grad_params( f" {tuple(main_grad.shape)} and stride" f" {tuple(main_grad.stride())}" ) from e - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if accumulate_into_main_grad: - grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( - num_tensors=num_groups, - tensor_shape=weight_shape, - rowwise_data=main_grad, - dtype=main_grad.dtype, - ) - - if grouped_wgrad is None: + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=num_groups, + tensor_shape=weight_shape, + rowwise_data=main_grad, + dtype=main_grad.dtype, + ) + accumulate_into_main_grad = not getattr( + weight_param, "overwrite_main_grad", False + ) + else: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_groups, shapes=[weight_shape] * num_groups, @@ -232,12 +249,17 @@ def _compute_grad_params( else: gemm_fn(grouped_x, grouped_dy, wgrad_output) - # Extract results, mark accumulated if needed + # Extract results, mark accumulated if needed. Gate the post-GEMM + # ``grad_added_to_main_grad`` + dummy ``.grad`` on the user-requested + # fusion flag, NOT on the (possibly-downgraded) GEMM accumulate flag. + # This matches ``module/linear.py`` and prevents FSDP's post-backward + # hook from re-touching ``main_grad`` when ``overwrite_main_grad=True`` + # (Megatron-FSDP) or when wgrad is delayed. if fc_op.single_grouped_weight: packed_wgrad = None if not delay_wgrad: packed_wgrad = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - if accumulate_into_main_grad and hasattr(weight_param, "grad_added_to_main_grad"): + if request_main_grad_fusion and hasattr(weight_param, "grad_added_to_main_grad"): weight_param.grad_added_to_main_grad = True packed_wgrad = get_dummy_wgrad( list(weight_param.size()), @@ -246,9 +268,9 @@ def _compute_grad_params( ) w_list = [packed_wgrad] else: - if delay_wgrad or accumulate_into_main_grad: + if delay_wgrad or request_main_grad_fusion: w_list = [None] * num_groups - if accumulate_into_main_grad: + if request_main_grad_fusion: for idx in range(num_groups): wp = getattr(fc_op, f"weight{idx}") if hasattr(wp, "grad_added_to_main_grad"): From b263845b42e8864c7394f64ce7260887fb6d33d3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Apr 2026 20:05:54 +0000 Subject: [PATCH 04/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 12 +++------ .../pytorch/ops/basic/grouped_linear.py | 26 +++++++------------ .../pytorch/ops/fused/backward_grouped_mlp.py | 8 ++---- 3 files changed, 15 insertions(+), 31 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index f741157081..8090321cc3 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2285,9 +2285,7 @@ def _init_main_grads(value: float = 0.0) -> None: with torch.no_grad(): for w in _weight_params(): if getattr(w, "main_grad", None) is None: - w.main_grad = torch.empty( - w.size(), device=device, dtype=torch.float32 - ) + w.main_grad = torch.empty(w.size(), device=device, dtype=torch.float32) w.main_grad.fill_(value) def _collect_main_grads() -> list[torch.Tensor]: @@ -2349,9 +2347,7 @@ def train_step( # Replay the graph graph_out = ( - train_step(static_x, static_dy, static_out_buf, use_graphed=True) - .detach() - .clone() + train_step(static_x, static_dy, static_out_buf, use_graphed=True).detach().clone() ) torch.cuda.synchronize() graph_dx = static_x.grad.detach().clone() @@ -2360,9 +2356,7 @@ def train_step( graph_param_grads: list[torch.Tensor] = [] else: graph_main_grads = [] - graph_param_grads = [ - param.grad.detach().clone() for param in op.parameters() - ] + graph_param_grads = [param.grad.detach().clone() for param in op.parameters()] # Reference: same op invoked eagerly with the same fresh inputs and # the same starting grad/main_grad state. diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 21f88ce40e..5ed4242274 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -705,10 +705,8 @@ def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer current scaling, fp8 block scaling, NVFP4, ...) falls back to the legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. """ - return all( - isinstance(q, MXFP8Quantizer) for q in input_quantizers - ) - + return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) + def _get_grouped_weight_for_gemm( self, weight_param: GroupedTensor, @@ -765,8 +763,8 @@ def _get_grouped_weight_for_gemm( None, ) - - def _get_discrete_weights_for_gemm(self, + def _get_discrete_weights_for_gemm( + self, weight_params: Optional[GroupedTensor] | list[torch.Tensor], weight_quantizers: list[Optional[Quantizer]], columnwise_usage: bool, @@ -1071,7 +1069,6 @@ def _fuser_forward_split_quantize( return out, [()] - def _fuser_forward_grouped_tensor( self, *, @@ -1088,8 +1085,7 @@ def _fuser_forward_grouped_tensor( weight_requires_grad: bool, device: torch.device, ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - """Graph-safe GroupedTensor forward path. - """ + """Graph-safe GroupedTensor forward path.""" num_groups = self.num_groups has_bias = self.has_bias @@ -1321,9 +1317,7 @@ def _fuser_backward_split_quantize( # post-GEMM bookkeeping (dummy ``.grad`` + ``grad_added_to_main_grad``) # always fires when fusion was requested -- see comment at the # post-GEMM dummy block below. - request_main_grad_fusion = ( - ctx.weight_requires_grad and self._accumulate_into_main_grad - ) + request_main_grad_fusion = ctx.weight_requires_grad and self._accumulate_into_main_grad accumulate_into_main_grad = request_main_grad_fusion grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( @@ -1622,9 +1616,7 @@ def _fuser_backward_grouped_tensor( # requested -- otherwise FSDP's post-backward hook would also touch # ``main_grad`` and double-count (or, with ``delay_wgrad``, copy # uninitialized data before the deferred GEMM ever runs). - request_main_grad_fusion = ( - ctx.weight_requires_grad and self._accumulate_into_main_grad - ) + request_main_grad_fusion = ctx.weight_requires_grad and self._accumulate_into_main_grad accumulate_into_main_grad = request_main_grad_fusion weight_shape = (self.out_features, self.in_features) wgrad_output: Any = None @@ -1683,7 +1675,9 @@ def _fuser_backward_grouped_tensor( ) else: for idx in range(num_groups): - final_weight_grads[idx] = torch.empty(weight_shape, dtype=dtype, device=device) + final_weight_grads[idx] = torch.empty( + weight_shape, dtype=dtype, device=device + ) wgrad_output = final_weight_grads # wgrad GEMM diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 8e7fd566c7..744766321e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -154,9 +154,7 @@ def _compute_grad_params( # post-backward hook would also touch ``main_grad`` and double-count # (or, with ``delay_wgrad``, copy uninitialized data before the deferred # GEMM ever runs). - request_main_grad_fusion = ( - ctx.weight_requires_grad and fc_op._accumulate_into_main_grad - ) + request_main_grad_fusion = ctx.weight_requires_grad and fc_op._accumulate_into_main_grad accumulate_into_main_grad = False grouped_wgrad = None wgrad_output = None @@ -195,9 +193,7 @@ def _compute_grad_params( rowwise_data=main_grad, dtype=main_grad.dtype, ) - accumulate_into_main_grad = not getattr( - weight_param, "overwrite_main_grad", False - ) + accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) else: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_groups, From 9ab69060a759b64b038fc96b99100ca63f1f8a6b Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sat, 25 Apr 2026 00:34:31 +0000 Subject: [PATCH 05/18] clean up Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 74 +++++-------------- 1 file changed, 20 insertions(+), 54 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 21f88ce40e..20531965b0 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -701,7 +701,7 @@ def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer """Whether all input quantizers support the graph-safe grouped-tensor flow. See ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS`` for the gating rationale. - Currently this is MXFP8 only; every other recipe (bf16, fp8 delayed / + Currently this is MXFP8 only; every other quantization recipe (fp8 delayed / current scaling, fp8 block scaling, NVFP4, ...) falls back to the legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. """ @@ -716,7 +716,7 @@ def _get_grouped_weight_for_gemm( columnwise_usage: bool, with_quantized_compute: bool, dtype: torch.dtype, - ) -> "GroupedTensor | list[torch.Tensor]": + ) -> GroupedTensor: """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. Supports MXFP8/BF16/FP16 compute paths. """ @@ -765,7 +765,6 @@ def _get_grouped_weight_for_gemm( None, ) - def _get_discrete_weights_for_gemm(self, weight_params: Optional[GroupedTensor] | list[torch.Tensor], weight_quantizers: list[Optional[Quantizer]], @@ -909,16 +908,12 @@ def fuser_forward( scales = basic_op_extra_inputs[0][1] # Dispatch: graph-safe GroupedTensor flow whenever it can be used -- - # * for the unquantized (bf16/fp16) compute path we just wrap the - # existing high-precision data as a ``GroupedTensor`` (no quantize - # call). FP32 is excluded because the cublasLt grouped GEMM only - # accepts FP8 / BF16 / FP16 inputs, and - # * for quantized compute only when the quantizer supports the - # graph-safe ``tex.group_quantize`` kernel (see - # ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS``; currently MXFP8). - # All remaining cases (fp32 unquantized, fp8 delayed / current scaling, - # fp8 block scaling, NVFP4) fall back to the legacy - # ``split_quantize`` + ``general_grouped_gemm`` flow. + # Unquantized (bf16/fp16) compute path: + # * We just wrap the existing high-precision data as a ``GroupedTensor`` (no quantize call). + # * FP32 is excluded because the cublasLt grouped GEMM doesnt support it. + # Quantized compute path: + # * Only when the quantizer supports the graph-safe ``tex.group_quantize`` kernel + # (see ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS``; currently MXFP8). use_grouped_tensor_path = ( with_quantized_compute and self._is_grouped_quantize_supported(input_quantizers) ) or (not with_quantized_compute and dtype in (torch.bfloat16, torch.float16)) @@ -1102,8 +1097,6 @@ def _fuser_forward_grouped_tensor( # Build the input GroupedTensor. if with_quantized_compute: - # Quantize input as a single GroupedTensor (one quantizer for all - # groups; OK for MXFP8 because its config is data-dependent). input_quantizer = input_quantizers[0] input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) input_quantizer.optimize_for_gemm = True @@ -1120,12 +1113,9 @@ def _fuser_forward_grouped_tensor( tensor_offsets=base_offsets * self.in_features, ) - # Build the weight GroupedTensor / list. In the unquantized path we - # consume the parameters directly: ``self.weight`` (a high-precision - # ``GroupedTensor``) for ``single_grouped_weight=True`` -- dispatches - # the GEMM to ``no_discrete`` mode -- or the per-group ``weight{idx}`` - # tensors otherwise -- ``discrete_in`` mode. + # Build the weight GroupedTensor / list. if self.single_grouped_weight: + # GroupedTensor grouped_weights = self._get_grouped_weight_for_gemm( self.weight, weight_quantizers, @@ -1134,6 +1124,7 @@ def _fuser_forward_grouped_tensor( dtype=dtype, ) else: + # Discrete weights grouped_weights = self._get_discrete_weights_for_gemm( [getattr(self, f"weight{idx}") for idx in range(num_groups)], weight_quantizers, @@ -1142,7 +1133,7 @@ def _fuser_forward_grouped_tensor( dtype=dtype, ) - # Allocate output buffer and wrap as a GroupedTensor view (no copy). + # Allocate output buffer and wrap as a GroupedTensor view. out_shape = original_shape[:-1] + [self.out_features] out = torch.empty(out_shape, dtype=dtype, device=device) grouped_out = GroupedTensor( @@ -1167,7 +1158,7 @@ def _fuser_forward_grouped_tensor( if bias_scale.dtype != torch.float32: bias_scale = bias_scale.to(dtype=torch.float32) - # Forward grouped GEMM (TN layout: out[i] = x[i] @ w[i]^T) + # Forward grouped GEMM. general_grouped_gemm_for_grouped_tensor( grouped_weights, grouped_x, @@ -1181,10 +1172,6 @@ def _fuser_forward_grouped_tensor( if not input_requires_grad: grouped_weights = None if self.single_grouped_weight else [None] * num_groups - # ``grouped_x`` already has all the data we need to save for backward - # (rowwise for the unquantized path; columnwise too for the quantized - # path because we configured the input quantizer with - # ``columnwise=weight_requires_grad`` above). if not weight_requires_grad: grouped_x = None @@ -1321,10 +1308,7 @@ def _fuser_backward_split_quantize( # post-GEMM bookkeeping (dummy ``.grad`` + ``grad_added_to_main_grad``) # always fires when fusion was requested -- see comment at the # post-GEMM dummy block below. - request_main_grad_fusion = ( - ctx.weight_requires_grad and self._accumulate_into_main_grad - ) - accumulate_into_main_grad = request_main_grad_fusion + accumulate_into_main_grad = self._accumulate_into_main_grad grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( [None] if self.single_grouped_weight else [None] * num_groups @@ -1440,12 +1424,8 @@ def _fuser_backward_split_quantize( # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into - # ``main_grad`` again. Matches the convention in ``module/linear.py`` - # and the grouped-tensor backward path. We REBIND - # ``final_weight_grads`` rather than mutate it -- ``grad_weights`` - # (captured by ``wgrad_store.put`` for delayed wgrad) still holds - # the ``main_grad`` views the GEMM writes into. - if request_main_grad_fusion: + # ``main_grad`` again. + if ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = self._dummy_main_grad_wgrads() if not has_bias: @@ -1612,20 +1592,7 @@ def _fuser_backward_grouped_tensor( ) # params init for wgrad GEMM - # ``request_main_grad_fusion`` records the user-facing opt-in to - # Megatron-LM main-grad fusion. ``accumulate_into_main_grad`` is the - # local GEMM ``accumulate`` flag, which gets downgraded to ``False`` - # when ``weight.overwrite_main_grad`` is set (e.g. Megatron-FSDP). - # The two must stay separated: the GEMM flag controls overwrite vs. - # accumulate, but the post-GEMM bookkeeping (dummy ``.grad`` + - # ``grad_added_to_main_grad=True``) must always fire when fusion was - # requested -- otherwise FSDP's post-backward hook would also touch - # ``main_grad`` and double-count (or, with ``delay_wgrad``, copy - # uninitialized data before the deferred GEMM ever runs). - request_main_grad_fusion = ( - ctx.weight_requires_grad and self._accumulate_into_main_grad - ) - accumulate_into_main_grad = request_main_grad_fusion + accumulate_into_main_grad = self._accumulate_into_main_grad weight_shape = (self.out_features, self.in_features) wgrad_output: Any = None grouped_wgrad: Optional[GroupedTensor] = None @@ -1640,8 +1607,7 @@ def _fuser_backward_grouped_tensor( if accumulate_into_main_grad: # Main-grad fusion: GEMM writes directly into ``main_grad``. # ``overwrite_main_grad`` only flips the GEMM's - # ``accumulate`` flag (overwrite vs. accumulate); it does - # not change the output buffer. + # ``accumulate`` flag. if hasattr(weight_param, "__fsdp_param__"): weight_param.main_grad = weight_param.get_main_grad() main_grad = weight_param.main_grad @@ -1707,8 +1673,8 @@ def _fuser_backward_grouped_tensor( # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into - # ``main_grad`` again. Matches the convention in ``module/linear.py``. - if request_main_grad_fusion: + # ``main_grad`` again. + if ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = self._dummy_main_grad_wgrads() # Assemble grad params in parameter registration order and return. From f48a5562abafea3795978ede64a146cfd5e87958 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sat, 25 Apr 2026 00:50:05 +0000 Subject: [PATCH 06/18] cleanup Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 14 +------------- .../pytorch/ops/fused/backward_grouped_mlp.py | 18 ++++-------------- 2 files changed, 5 insertions(+), 27 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index a0bbec3067..b0214bb2b0 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1292,19 +1292,7 @@ def _fuser_backward_split_quantize( dbias_packed = compute_grouped_dbias(dy_2d, offsets, num_groups) grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] - # Initialize grad weight buffers and the autograd-return - # ``final_weight_grads`` upfront (mirrors the grouped-tensor - # flow). For ``single_grouped_weight=True`` we pre-allocate a - # stacked ``[num_groups, out, in]`` buffer and feed per-group - # views of it to the GEMM, so the buffer itself is the autograd - # return -- no post-GEMM ``torch.stack`` copy needed. - # ``request_main_grad_fusion`` records the user-facing opt-in to - # Megatron-LM main-grad fusion; ``accumulate_into_main_grad`` is the - # local GEMM ``accumulate`` flag (downgraded to ``False`` when - # ``weight.overwrite_main_grad`` is set, e.g. Megatron-FSDP). The - # post-GEMM bookkeeping (dummy ``.grad`` + ``grad_added_to_main_grad``) - # always fires when fusion was requested -- see comment at the - # post-GEMM dummy block below. + # Initialize grad weight buffers. accumulate_into_main_grad = self._accumulate_into_main_grad grad_weights = [None] * num_groups final_weight_grads: list[Optional[torch.Tensor]] = ( diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 744766321e..da94fd3c4b 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -144,17 +144,6 @@ def _compute_grad_params( """ # Allocate grad buffers, determine accumulate flag. - # ``request_main_grad_fusion`` records the user-facing opt-in to - # Megatron-LM main-grad fusion. ``accumulate_into_main_grad`` is the - # local GEMM ``accumulate`` flag, which gets downgraded to ``False`` when - # ``weight.overwrite_main_grad`` is set (e.g. Megatron-FSDP). The two must - # stay separated: the GEMM flag controls overwrite vs. accumulate, but the - # post-GEMM bookkeeping (dummy ``.grad`` + ``grad_added_to_main_grad=True``) - # must always fire when fusion was requested -- otherwise FSDP's - # post-backward hook would also touch ``main_grad`` and double-count - # (or, with ``delay_wgrad``, copy uninitialized data before the deferred - # GEMM ever runs). - request_main_grad_fusion = ctx.weight_requires_grad and fc_op._accumulate_into_main_grad accumulate_into_main_grad = False grouped_wgrad = None wgrad_output = None @@ -255,7 +244,8 @@ def _compute_grad_params( packed_wgrad = None if not delay_wgrad: packed_wgrad = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - if request_main_grad_fusion and hasattr(weight_param, "grad_added_to_main_grad"): + if ctx.weight_requires_grad and fc_op._accumulate_into_main_grad\ + and hasattr(weight_param, "grad_added_to_main_grad"): weight_param.grad_added_to_main_grad = True packed_wgrad = get_dummy_wgrad( list(weight_param.size()), @@ -264,9 +254,9 @@ def _compute_grad_params( ) w_list = [packed_wgrad] else: - if delay_wgrad or request_main_grad_fusion: + if delay_wgrad or (ctx.weight_requires_grad and fc_op._accumulate_into_main_grad): w_list = [None] * num_groups - if request_main_grad_fusion: + if ctx.weight_requires_grad and fc_op._accumulate_into_main_grad: for idx in range(num_groups): wp = getattr(fc_op, f"weight{idx}") if hasattr(wp, "grad_added_to_main_grad"): From 00a3521855b60aacb8ff38afa043e3e02529674c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 25 Apr 2026 00:52:35 +0000 Subject: [PATCH 07/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/ops/fused/backward_grouped_mlp.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index da94fd3c4b..4c0dd894da 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -244,8 +244,11 @@ def _compute_grad_params( packed_wgrad = None if not delay_wgrad: packed_wgrad = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - if ctx.weight_requires_grad and fc_op._accumulate_into_main_grad\ - and hasattr(weight_param, "grad_added_to_main_grad"): + if ( + ctx.weight_requires_grad + and fc_op._accumulate_into_main_grad + and hasattr(weight_param, "grad_added_to_main_grad") + ): weight_param.grad_added_to_main_grad = True packed_wgrad = get_dummy_wgrad( list(weight_param.size()), From b618d6f2092ae735770c8285d39c48ba45e119fd Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 27 Apr 2026 18:10:06 +0000 Subject: [PATCH 08/18] clean up main grad business Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/_common.py | 94 +++++++++++ .../pytorch/ops/basic/basic_linear.py | 32 ++-- .../pytorch/ops/basic/grouped_linear.py | 149 +++++++----------- .../pytorch/ops/fused/backward_grouped_mlp.py | 93 +++-------- .../pytorch/ops/fused/backward_linear_add.py | 29 ++-- .../ops/fused/backward_linear_scale.py | 29 ++-- .../ops/fused/userbuffers_backward_linear.py | 34 ++-- 7 files changed, 214 insertions(+), 246 deletions(-) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index e21915a5a6..2b52e6f792 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -6,6 +6,7 @@ from __future__ import annotations import functools +import math from importlib.metadata import PackageNotFoundError, version as get_pkg_version from typing import Optional @@ -94,6 +95,99 @@ def get_fp8_meta_from_fp8_tensor(tensor: Float8Tensor) -> tuple[FP8TensorMeta, i return fp8_meta, 0 +def get_main_grad_from_param( + weight_param: torch.nn.Parameter, + *, + op_label: str = "", +) -> torch.Tensor: + """Refresh ``main_grad`` from FSDP (if applicable) and return it. + Used by Megatron-LM-style wgrad fusion paths + (``accumulate_into_main_grad=True``) to obtain the buffer the wgrad GEMM + will write into. + Raises if the parameter does not have a ``main_grad`` attribute or if it + is ``None``. + """ + if hasattr(weight_param, "__fsdp_param__"): + weight_param.main_grad = weight_param.get_main_grad() + if not hasattr(weight_param, "main_grad") or weight_param.main_grad is None: + prefix = f"{op_label} " if op_label else "" + raise RuntimeError( + f"{prefix}operation is configured with accumulate_into_main_grad=True, " + "but weight parameter does not have a valid main_grad attribute" + ) + return weight_param.main_grad + + +def get_accumulate_flag_in_param(weight_param: torch.nn.Parameter) -> bool: + """Return whether the wgrad GEMM should accumulate into ``main_grad``. + + Returns ``False`` (i.e. overwrite) when the parameter has + ``overwrite_main_grad=True`` (used in Megatron-FSDP), and ``True`` + otherwise. + """ + return not getattr(weight_param, "overwrite_main_grad", False) + + +def view_main_grad_as_grouped_buffer( + main_grad: torch.Tensor, + num_groups: int, + weight_shape: tuple[int, ...], + *, + label: str = "", +) -> torch.Tensor: + """Return ``main_grad`` viewed as ``(num_groups, *weight_shape)`` without copy. + Raises if the numel doesn't match or if the existing stride pattern does + not allow a zero-copy view to the grouped layout. + """ + grouped_shape = (num_groups, *weight_shape) + if tuple(main_grad.shape) == grouped_shape: + return main_grad + prefix = f"{label} " if label else "Grouped weight " + if main_grad.numel() != math.prod(grouped_shape): + raise RuntimeError( + f"{prefix}main_grad expected shape {grouped_shape} or matching numel, " + f"but got shape {tuple(main_grad.shape)}" + ) + try: + return main_grad.view(grouped_shape) + except RuntimeError as e: + raise RuntimeError( + f"{prefix}main_grad must be viewable as {grouped_shape} without copy, " + f"but got shape {tuple(main_grad.shape)} and stride " + f"{tuple(main_grad.stride())}" + ) from e + + +def get_dummy_wgrads_for_params( + weight_params: list[torch.nn.Parameter], +) -> list[Optional[torch.Tensor]]: + """Build dummy ``.grad`` placeholders for Megatron-LM wgrad-fusion params. + + For each parameter that exposes ``grad_added_to_main_grad``, set the flag + to ``True`` and return a dummy wgrad tensor (zeroed if + ``zero_out_wgrad`` is also set on the parameter). For parameters without + the flag, the corresponding entry is ``None``. + + The returned list has the same length and order as ``weight_params``. + """ + from ..module.base import get_dummy_wgrad # pylint: disable=import-outside-toplevel + + out: list[Optional[torch.Tensor]] = [] + for wp in weight_params: + if hasattr(wp, "grad_added_to_main_grad"): + wp.grad_added_to_main_grad = True + out.append( + get_dummy_wgrad( + list(wp.size()), + wp.dtype, + zero=getattr(wp, "zero_out_wgrad", False), + ) + ) + else: + out.append(None) + return out + + def validate_grouped_mlp_dims(fc1, glu_op, fc2) -> None: """Validate FC1 / scaled GLU / FC2 dimensions for fused grouped MLP.""" diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 17594726cc..0b0814b112 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -24,7 +24,6 @@ _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, - get_dummy_wgrad, ) from ...tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer @@ -36,7 +35,13 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_dequantize, is_quantized_tensor +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, +) def _wait_async(handle: Optional[Any]) -> None: @@ -1065,16 +1070,9 @@ def op_backward( grad_weight = None if ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = self.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="BasicLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -1104,14 +1102,6 @@ def op_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = self.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([self.weight])[0] return grad_input, [grad_weight] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b0214bb2b0..6c9b7a5f58 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -21,7 +21,6 @@ _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, - get_dummy_wgrad, ) from ...quantization import FP8GlobalStateManager, Recipe from ...tensor import MXFP8Quantizer, MXFP8Tensor, Quantizer @@ -32,7 +31,14 @@ devices_match, round_up_to_nearest_multiple, ) -from .._common import is_quantized_tensor, maybe_dequantize +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, + view_main_grad_as_grouped_buffer, +) from ..op import BasicOperation, OperationContext from ...tensor import GroupedTensor from ...triton.grouped_dbias_dscales import ( @@ -784,36 +790,15 @@ def _get_discrete_weights_for_gemm( out.append(w) return out - def _dummy_main_grad_wgrads(self) -> list[Optional[torch.Tensor]]: - """Return a NEW list of per-output dummy weight gradients for the - ``accumulate_into_main_grad`` (Megatron-LM wgrad fusion) path. - Length: 1 for ``single_grouped_weight=True``, ``num_groups`` - otherwise. + def _get_weight_tensors(self) -> list[torch.nn.Parameter]: + """Return the weight parameters in registration order. + + Length is 1 when ``single_grouped_weight=True`` (one + ``GroupedTensor`` parameter), otherwise ``num_groups``. """ if self.single_grouped_weight: - weight_param = self.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - return [ - get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - ] - return [None] - - out: list[Optional[torch.Tensor]] = [None] * self.num_groups - for group_idx in range(self.num_groups): - wp = getattr(self, f"weight{group_idx}") - if hasattr(wp, "grad_added_to_main_grad"): - wp.grad_added_to_main_grad = True - out[group_idx] = get_dummy_wgrad( - list(wp.size()), - wp.dtype, - zero=getattr(wp, "zero_out_wgrad", False), - ) - return out + return [self.weight] + return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] def _get_grouped_bias_for_gemm( self, @@ -1247,8 +1232,8 @@ def _fuser_backward_split_quantize( ]: num_groups = self.num_groups has_bias = self.has_bias - weight_param = self.weight if self.single_grouped_weight else self.weight0 - device = weight_param.device + weights = self._get_weight_tensors() + device = weights[0].device # Saved tensors from forward pass saved_tensors = ctx.saved_tensors @@ -1305,29 +1290,15 @@ def _fuser_backward_split_quantize( if accumulate_into_main_grad: # Megatron-LM wgrad fusion: GEMM accumulates into the # parameter's ``main_grad`` directly. - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - main_grad = weight_param.main_grad - if isinstance(main_grad, GroupedTensor): - # Legacy path: no contiguous backing; per-group - # quantized tensors as discrete GEMM outputs. - grad_weights = main_grad.quantized_tensors - if grad_weights is None: - grad_weights = main_grad.split_into_quantized_tensors() - else: - if main_grad.shape != grouped_shape: - if main_grad.numel() != math.prod(grouped_shape): - raise RuntimeError( - "GroupedLinear expected grouped weight main_grad to have " - f"shape {grouped_shape} or matching numel, " - f"but got shape {tuple(main_grad.shape)}" - ) - main_grad = main_grad.reshape(grouped_shape) - final_weight_grads[0] = main_grad - grad_weights = [main_grad[idx] for idx in range(num_groups)] - accumulate_into_main_grad = not getattr( - weight_param, "overwrite_main_grad", False + main_grad = get_main_grad_from_param( + weights[0], op_label="GroupedLinear" + ) + main_grad = view_main_grad_as_grouped_buffer( + main_grad, num_groups, weight_shape, label="GroupedLinear weight" ) + final_weight_grads[0] = main_grad + grad_weights = [main_grad[idx] for idx in range(num_groups)] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: final_weight_grads[0] = torch.empty( grouped_shape, dtype=ctx.dtype, device=device @@ -1335,19 +1306,16 @@ def _fuser_backward_split_quantize( grad_weights = [final_weight_grads[0][idx] for idx in range(num_groups)] else: if accumulate_into_main_grad: - for group_idx in range(num_groups): - weight_param = getattr(self, f"weight{group_idx}") - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - grad_weights[group_idx] = weight_param.main_grad - accumulate_into_main_grad = not getattr( - self.weight0, "overwrite_main_grad", False - ) + grad_weights = [ + get_main_grad_from_param(w, op_label="GroupedLinear") + for w in weights + ] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - for group_idx in range(num_groups): - grad_weights[group_idx] = torch.empty( - weight_shape, dtype=ctx.dtype, device=device - ) + grad_weights = [ + torch.empty(weight_shape, dtype=ctx.dtype, device=device) + for _ in range(num_groups) + ] final_weight_grads = list(grad_weights) # Perform dgrad GEMMs @@ -1411,7 +1379,7 @@ def _fuser_backward_split_quantize( # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. if ctx.weight_requires_grad and self._accumulate_into_main_grad: - final_weight_grads = self._dummy_main_grad_wgrads() + final_weight_grads = get_dummy_wgrads_for_params(weights) if not has_bias: grad_params = list(final_weight_grads) @@ -1442,8 +1410,8 @@ def _fuser_backward_grouped_tensor( ]: num_groups = self.num_groups has_bias = self.has_bias - weight_param = self.weight if self.single_grouped_weight else self.weight0 - device = weight_param.device + weights = self._get_weight_tensors() + device = weights[0].device dtype = ctx.dtype with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) @@ -1593,25 +1561,19 @@ def _fuser_backward_grouped_tensor( # Main-grad fusion: GEMM writes directly into ``main_grad``. # ``overwrite_main_grad`` only flips the GEMM's # ``accumulate`` flag. - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - main_grad = weight_param.main_grad - grouped_shape = (num_groups, *weight_shape) - if main_grad.numel() != math.prod(grouped_shape): - raise RuntimeError( - "GroupedLinear expected grouped weight main_grad to have " - f"shape {grouped_shape} or matching numel, " - f"but got shape {tuple(main_grad.shape)}" - ) + main_grad = get_main_grad_from_param( + weights[0], op_label="GroupedLinear" + ) + main_grad = view_main_grad_as_grouped_buffer( + main_grad, num_groups, weight_shape, label="GroupedLinear weight" + ) grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( num_tensors=num_groups, tensor_shape=weight_shape, rowwise_data=main_grad.view(-1), dtype=main_grad.dtype, ) - accumulate_into_main_grad = not getattr( - weight_param, "overwrite_main_grad", False - ) + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_groups, @@ -1624,19 +1586,16 @@ def _fuser_backward_grouped_tensor( wgrad_output = grouped_wgrad else: if accumulate_into_main_grad: - for idx in range(num_groups): - wp = getattr(self, f"weight{idx}") - if hasattr(wp, "__fsdp_param__"): - wp.main_grad = wp.get_main_grad() - final_weight_grads[idx] = wp.main_grad - accumulate_into_main_grad = not getattr( - self.weight0, "overwrite_main_grad", False - ) + final_weight_grads = [ + get_main_grad_from_param(w, op_label="GroupedLinear") + for w in weights + ] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - for idx in range(num_groups): - final_weight_grads[idx] = torch.empty( - weight_shape, dtype=dtype, device=device - ) + final_weight_grads = [ + torch.empty(weight_shape, dtype=dtype, device=device) + for _ in range(num_groups) + ] wgrad_output = final_weight_grads # wgrad GEMM @@ -1662,7 +1621,7 @@ def _fuser_backward_grouped_tensor( # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into # ``main_grad`` again. if ctx.weight_requires_grad and self._accumulate_into_main_grad: - final_weight_grads = self._dummy_main_grad_wgrads() + final_weight_grads = get_dummy_wgrads_for_params(weights) # Assemble grad params in parameter registration order and return. if not has_bias: diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 4c0dd894da..9aa837cbe3 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -8,14 +8,12 @@ from collections.abc import Callable import functools import inspect -import math import os from typing import Optional import torch import transformer_engine_torch as tex -from ...module.base import get_dummy_wgrad from ...quantization import Recipe from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer @@ -27,7 +25,11 @@ from .._common import ( _nvidia_cudnn_frontend_supports_wgrad, fuse_grouped_mlp_ops, + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, maybe_dequantize, + view_main_grad_as_grouped_buffer, validate_grouped_mlp_dims, ) from ...cpp_extensions import general_grouped_gemm_for_grouped_tensor @@ -147,42 +149,27 @@ def _compute_grad_params( accumulate_into_main_grad = False grouped_wgrad = None wgrad_output = None + op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" + weights = fc_op._get_weight_tensors() if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: - weight_param = fc_op.weight if fc_op._accumulate_into_main_grad: # Main-grad fusion: GEMM writes directly into ``main_grad``. # ``overwrite_main_grad`` only flips the GEMM's ``accumulate`` # flag (overwrite vs. accumulate); it does not change the # output buffer. - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - main_grad = weight_param.main_grad - grouped_shape = (num_groups, *weight_shape) - if main_grad.shape != grouped_shape: - if main_grad.numel() != math.prod(grouped_shape): - raise RuntimeError( - f"Grouped MLP fused backward expected {label} main_grad to have " - f"shape {grouped_shape} or matching numel, " - f"but got shape {tuple(main_grad.shape)}" - ) - try: - main_grad = main_grad.view(grouped_shape) - except RuntimeError as e: - raise RuntimeError( - f"Grouped MLP fused backward requires {label} main_grad to be " - f"viewable as {grouped_shape} without copy, but got shape" - f" {tuple(main_grad.shape)} and stride" - f" {tuple(main_grad.stride())}" - ) from e + main_grad = get_main_grad_from_param(weights[0], op_label=op_label) + main_grad = view_main_grad_as_grouped_buffer( + main_grad, num_groups, weight_shape, label=f"{op_label} weight" + ) grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( num_tensors=num_groups, tensor_shape=weight_shape, rowwise_data=main_grad, dtype=main_grad.dtype, ) - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_groups, @@ -192,19 +179,20 @@ def _compute_grad_params( dtype=dtype, ) wgrad_output = grouped_wgrad + w_list = [grouped_wgrad.rowwise_data.view(num_groups, *weight_shape)] else: w_list = [None] * num_groups if ctx.weight_requires_grad: if fc_op._accumulate_into_main_grad: - for idx in range(num_groups): - wp = getattr(fc_op, f"weight{idx}") - if hasattr(wp, "__fsdp_param__"): - wp.main_grad = wp.get_main_grad() - w_list[idx] = wp.main_grad - accumulate_into_main_grad = not getattr(fc_op.weight0, "overwrite_main_grad", False) + w_list = [ + get_main_grad_from_param(w, op_label=op_label) for w in weights + ] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - for idx in range(num_groups): - w_list[idx] = torch.empty(weight_shape, dtype=dtype, device=device) + w_list = [ + torch.empty(weight_shape, dtype=dtype, device=device) + for _ in range(num_groups) + ] wgrad_output = w_list if ctx.weight_requires_grad: @@ -234,41 +222,9 @@ def _compute_grad_params( else: gemm_fn(grouped_x, grouped_dy, wgrad_output) - # Extract results, mark accumulated if needed. Gate the post-GEMM - # ``grad_added_to_main_grad`` + dummy ``.grad`` on the user-requested - # fusion flag, NOT on the (possibly-downgraded) GEMM accumulate flag. - # This matches ``module/linear.py`` and prevents FSDP's post-backward - # hook from re-touching ``main_grad`` when ``overwrite_main_grad=True`` - # (Megatron-FSDP) or when wgrad is delayed. - if fc_op.single_grouped_weight: - packed_wgrad = None - if not delay_wgrad: - packed_wgrad = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - if ( - ctx.weight_requires_grad - and fc_op._accumulate_into_main_grad - and hasattr(weight_param, "grad_added_to_main_grad") - ): - weight_param.grad_added_to_main_grad = True - packed_wgrad = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - w_list = [packed_wgrad] - else: - if delay_wgrad or (ctx.weight_requires_grad and fc_op._accumulate_into_main_grad): - w_list = [None] * num_groups - if ctx.weight_requires_grad and fc_op._accumulate_into_main_grad: - for idx in range(num_groups): - wp = getattr(fc_op, f"weight{idx}") - if hasattr(wp, "grad_added_to_main_grad"): - wp.grad_added_to_main_grad = True - w_list[idx] = get_dummy_wgrad( - list(wp.size()), - wp.dtype, - zero=getattr(wp, "zero_out_wgrad", False), - ) + # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added + if fc_op._accumulate_into_main_grad: + w_list = get_dummy_wgrads_for_params(weights) # Assemble grad_params in parameter registration order. if not fc_op.has_bias: @@ -380,8 +336,7 @@ def fuser_backward( grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) out_shape = list(grad_output.size()) num_groups = fc1_op.num_groups - fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 - device = fc1_weight_param.device + device = fc1_op._get_weight_tensors()[0].device dtype = fc1_ctx.dtype # Saved tensors from FC1 forward diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index c06e212e87..382fecfd07 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -9,9 +9,13 @@ import torch -from ...module.base import get_dummy_wgrad from ...utils import clear_tensor_data from ..basic import BasicLinear, MakeExtraOutput +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, +) from ..op import FusedOperation, FusibleOperation, OperationContext @@ -57,16 +61,9 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="BasicLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -99,15 +96,7 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = linear_op.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([linear_op.weight])[0] return grad_input, [(), (grad_weight,)], [(), ()] diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py index 709073e6f8..b48c2e6d52 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py @@ -9,9 +9,13 @@ import torch -from ...module.base import get_dummy_wgrad from ...utils import clear_tensor_data from ..basic import BasicLinear, ConstantScale +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, +) from ..op import FusedOperation, FusibleOperation, OperationContext @@ -58,16 +62,9 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="BasicLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -99,15 +96,7 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = linear_op.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([linear_op.weight])[0] return grad_input, [(grad_weight,), ()], [(), ()] diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index fbaf69d75d..5c68628bcf 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -17,14 +17,19 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, fill_userbuffers_buffer_for_all_gather, - get_dummy_wgrad, get_ub, ) from ...quantized_tensor import Quantizer from ...tensor.mxfp8_tensor import MXFP8Quantizer from ...utils import canonicalize_device, canonicalize_dtype, clear_tensor_data from ..basic import BasicLinear, Bias, ReduceScatter -from .._common import maybe_dequantize, is_quantized_tensor +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, +) from ..op import FusedOperation, FusibleOperation, OperationContext @@ -519,16 +524,11 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param( + weight_param, op_label="UserbuffersBackwardLinear" + ) + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -563,15 +563,7 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = linear_op.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([linear_op.weight])[0] # Return gradients grad_params = [() for _ in range(len(self.basic_ops))] From e843f1a52e47b6a0f640fcf910c1769d1b941457 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Apr 2026 18:12:27 +0000 Subject: [PATCH 09/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/ops/basic/grouped_linear.py | 14 ++++---------- .../pytorch/ops/fused/backward_grouped_mlp.py | 7 ++----- .../ops/fused/userbuffers_backward_linear.py | 4 +--- 3 files changed, 7 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 6c9b7a5f58..115c3e7b97 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1290,9 +1290,7 @@ def _fuser_backward_split_quantize( if accumulate_into_main_grad: # Megatron-LM wgrad fusion: GEMM accumulates into the # parameter's ``main_grad`` directly. - main_grad = get_main_grad_from_param( - weights[0], op_label="GroupedLinear" - ) + main_grad = get_main_grad_from_param(weights[0], op_label="GroupedLinear") main_grad = view_main_grad_as_grouped_buffer( main_grad, num_groups, weight_shape, label="GroupedLinear weight" ) @@ -1307,8 +1305,7 @@ def _fuser_backward_split_quantize( else: if accumulate_into_main_grad: grad_weights = [ - get_main_grad_from_param(w, op_label="GroupedLinear") - for w in weights + get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: @@ -1561,9 +1558,7 @@ def _fuser_backward_grouped_tensor( # Main-grad fusion: GEMM writes directly into ``main_grad``. # ``overwrite_main_grad`` only flips the GEMM's # ``accumulate`` flag. - main_grad = get_main_grad_from_param( - weights[0], op_label="GroupedLinear" - ) + main_grad = get_main_grad_from_param(weights[0], op_label="GroupedLinear") main_grad = view_main_grad_as_grouped_buffer( main_grad, num_groups, weight_shape, label="GroupedLinear weight" ) @@ -1587,8 +1582,7 @@ def _fuser_backward_grouped_tensor( else: if accumulate_into_main_grad: final_weight_grads = [ - get_main_grad_from_param(w, op_label="GroupedLinear") - for w in weights + get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 9aa837cbe3..e1d1b47db8 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -184,14 +184,11 @@ def _compute_grad_params( w_list = [None] * num_groups if ctx.weight_requires_grad: if fc_op._accumulate_into_main_grad: - w_list = [ - get_main_grad_from_param(w, op_label=op_label) for w in weights - ] + w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: w_list = [ - torch.empty(weight_shape, dtype=dtype, device=device) - for _ in range(num_groups) + torch.empty(weight_shape, dtype=dtype, device=device) for _ in range(num_groups) ] wgrad_output = w_list diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 5c68628bcf..7d67815f9a 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -524,9 +524,7 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - main_grad = get_main_grad_from_param( - weight_param, op_label="UserbuffersBackwardLinear" - ) + main_grad = get_main_grad_from_param(weight_param, op_label="UserbuffersBackwardLinear") accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) grad_weight = main_grad.detach() else: From e55beec751f17f239eedd2bb7a99a2f3afddcbff Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 28 Apr 2026 23:31:12 +0000 Subject: [PATCH 10/18] clean up a but Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 25 ++++++------------- 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 115c3e7b97..de8b877d09 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -703,7 +703,7 @@ def op_backward(self, *args, **kwargs): ) @staticmethod - def _is_grouped_quantize_supported(input_quantizers: Sequence[Optional[Quantizer]]) -> bool: + def _is_graph_safe_path_supported(input_quantizers: Sequence[Optional[Quantizer]]) -> bool: """Whether all input quantizers support the graph-safe grouped-tensor flow. See ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS`` for the gating rationale. @@ -896,10 +896,10 @@ def fuser_forward( # * We just wrap the existing high-precision data as a ``GroupedTensor`` (no quantize call). # * FP32 is excluded because the cublasLt grouped GEMM doesnt support it. # Quantized compute path: - # * Only when the quantizer supports the graph-safe ``tex.group_quantize`` kernel - # (see ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS``; currently MXFP8). + # * Currently only supported for MXFP8. NVFP4 support will be enabled once + # the NVFP4 GGEMM support is added. use_grouped_tensor_path = ( - with_quantized_compute and self._is_grouped_quantize_supported(input_quantizers) + with_quantized_compute and self._is_graph_safe_path_supported(input_quantizers) ) or (not with_quantized_compute and dtype in (torch.bfloat16, torch.float16)) if use_grouped_tensor_path: return self._fuser_forward_grouped_tensor( @@ -1007,8 +1007,6 @@ def _fuser_forward_split_quantize( ) # Add bias * scales when scale_bias is enabled - # Would be done as part of larger refactor for GroupedLinear + GroupedTensor - # integration. if self._scale_bias and has_bias: scales_splits = torch.split(scales, split_sizes_int) out_splits = torch.split(out, split_sizes_int) @@ -1414,15 +1412,6 @@ def _fuser_backward_grouped_tensor( with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) # Saved tensors from forward pass (see _fuser_forward_grouped_tensor). - # ``base_offsets`` is the cumulative-token offset tensor (computed once - # in forward via ``tex.splits_to_offsets``); we scale it by the - # appropriate last dim to get any GroupedTensor's ``tensor_offsets``. - # ``ws`` is either a single ``GroupedTensor`` (single_grouped_weight) - # or a list of per-group tensors; it is passed straight to - # ``general_grouped_gemm_for_grouped_tensor`` in dgrad. - # ``x_data`` / ``x_scale`` carry the saved input for the wgrad pass: - # * Quantized path: columnwise data + scale. - # * Unquantized path: raw rowwise data + ``None`` scale. saved_tensors = ctx.saved_tensors split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] base_offsets, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1542,7 +1531,7 @@ def _fuser_backward_grouped_tensor( ) # params init for wgrad GEMM - accumulate_into_main_grad = self._accumulate_into_main_grad + accumulate_into_main_grad = False weight_shape = (self.out_features, self.in_features) wgrad_output: Any = None grouped_wgrad: Optional[GroupedTensor] = None @@ -1554,7 +1543,7 @@ def _fuser_backward_grouped_tensor( # Can be a GroupedTensor or list of tensors based on single_grouped_weight. if ctx.weight_requires_grad: if self.single_grouped_weight: - if accumulate_into_main_grad: + if self._accumulate_into_main_grad: # Main-grad fusion: GEMM writes directly into ``main_grad``. # ``overwrite_main_grad`` only flips the GEMM's # ``accumulate`` flag. @@ -1580,7 +1569,7 @@ def _fuser_backward_grouped_tensor( final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) wgrad_output = grouped_wgrad else: - if accumulate_into_main_grad: + if self._accumulate_into_main_grad: final_weight_grads = [ get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights ] From 6886e756dc6d498ec2d7b27cda1cc07e7bce1b70 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 29 Apr 2026 22:08:44 +0000 Subject: [PATCH 11/18] fix on l40/hopper to skip Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 2 + .../pytorch/ops/basic/grouped_linear.py | 50 ++++++++++++------- 2 files changed, 34 insertions(+), 18 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index cdc5001a1d..fe98778b58 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2310,6 +2310,8 @@ def test_grouped_linear_cuda_graph_safe( GPU-resident split offsets and is the only flow safe to capture. """ + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") # Skip invalid configurations if quantization is None and quantized_weight: pytest.skip("quantized_weight requires a quantization recipe") diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index de8b877d09..4ecb10180e 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -29,6 +29,7 @@ canonicalize_dtype, clear_tensor_data, devices_match, + get_device_compute_capability, round_up_to_nearest_multiple, ) from .._common import ( @@ -703,15 +704,28 @@ def op_backward(self, *args, **kwargs): ) @staticmethod - def _is_graph_safe_path_supported(input_quantizers: Sequence[Optional[Quantizer]]) -> bool: - """Whether all input quantizers support the graph-safe grouped-tensor flow. - - See ``_GROUPED_QUANTIZE_SUPPORTED_QUANTIZERS`` for the gating rationale. - Currently this is MXFP8 only; every other quantization recipe (fp8 delayed / - current scaling, fp8 block scaling, NVFP4, ...) falls back to the - legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. + def _is_graph_safe_path_supported( + *, + with_quantized_compute: bool, + input_quantizers: Sequence[Optional[Quantizer]], + dtype: torch.dtype, + ) -> bool: + """Whether the graph-safe grouped-tensor flow can be used. + + * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, + which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common + library. That kernel requires Blackwell (SM100) or newer with cuBLAS 13.3+. + * Quantized compute is currently MXFP8-only; every other quantization + recipe (fp8 delayed / current scaling, fp8 block scaling, NVFP4, ...) + falls back to the legacy flow. + * Unquantized compute supports BF16/FP16 only -- FP32 is excluded + because the cublasLt grouped GEMM doesn't support it. """ - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) + if get_device_compute_capability() < (10, 0): + return False + if with_quantized_compute: + return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) + return dtype in (torch.bfloat16, torch.float16) def _get_grouped_weight_for_gemm( self, @@ -891,16 +905,16 @@ def fuser_forward( if self._scale_bias: scales = basic_op_extra_inputs[0][1] - # Dispatch: graph-safe GroupedTensor flow whenever it can be used -- - # Unquantized (bf16/fp16) compute path: - # * We just wrap the existing high-precision data as a ``GroupedTensor`` (no quantize call). - # * FP32 is excluded because the cublasLt grouped GEMM doesnt support it. - # Quantized compute path: - # * Currently only supported for MXFP8. NVFP4 support will be enabled once - # the NVFP4 GGEMM support is added. - use_grouped_tensor_path = ( - with_quantized_compute and self._is_graph_safe_path_supported(input_quantizers) - ) or (not with_quantized_compute and dtype in (torch.bfloat16, torch.float16)) + # Dispatch: graph-safe GroupedTensor flow whenever it can be used. + # See ``_is_graph_safe_path_supported`` for the gating rationale -- + # in short it requires Blackwell (SM100+) plus a supported dtype / + # quantization recipe. Otherwise we fall back to the legacy + # ``tex.split_quantize`` + ``general_grouped_gemm`` flow. + use_grouped_tensor_path = self._is_graph_safe_path_supported( + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + dtype=dtype, + ) if use_grouped_tensor_path: return self._fuser_forward_grouped_tensor( ctx=ctx, From e22c6d92ee979c58bcb0c7e35846ee05b17f9916 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 5 May 2026 08:46:19 +0000 Subject: [PATCH 12/18] address review comments + save activation in backward + common context savings for fused/unfused paths Signed-off-by: Varun Thumbe --- .../pytorch/ops/basic/grouped_linear.py | 274 ++++++++++-------- .../pytorch/ops/fused/backward_grouped_mlp.py | 98 ++----- .../pytorch/ops/fused/forward_grouped_mlp.py | 81 +++--- .../tensor/storage/grouped_tensor_storage.py | 47 +++ 4 files changed, 274 insertions(+), 226 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 35f64695c7..202441e4f1 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -23,6 +23,7 @@ _2X_ACC_WGRAD, ) from ...quantization import FP8GlobalStateManager, Recipe +from ...quantized_tensor import QuantizedTensorStorage from ...tensor import MXFP8Quantizer, MXFP8Tensor, Quantizer from ...utils import ( canonicalize_device, @@ -249,7 +250,7 @@ def backward_dw(self) -> None: else: # Fused MXFP8 grouped MLP saves `GroupedTensor` activations for wgrad. clear_tensor_data( - activations.data, + activations.rowwise_data, activations.columnwise_data, activations.scale_inv, activations.columnwise_scale_inv, @@ -845,7 +846,6 @@ def _get_weight_tensors(self) -> list[torch.nn.Parameter]: def _get_grouped_bias_for_gemm( self, dtype: torch.dtype, - device: torch.device, ) -> Optional[torch.Tensor]: """Build a uniform GroupedTensor of per-group biases for the cublas grouped GEMM. @@ -905,13 +905,11 @@ def fuser_forward( # Quantizers input_quantizers = [None] * num_groups weight_quantizers = [None] * num_groups - grad_output_quantizers = [None] * num_groups with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() if with_quantized_compute: for group_idx in range(num_groups): input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) - grad_output_quantizers[group_idx] = self.get_quantizer("backward", group_idx) # Get autocast dtype if needed if torch.is_autocast_enabled(): @@ -943,36 +941,100 @@ def fuser_forward( input_quantizers=input_quantizers, dtype=dtype, ) + if use_grouped_tensor_path: - return self._fuser_forward_grouped_tensor( - ctx=ctx, + out, tensors_to_save = self._fuser_forward_grouped_tensor( input_=input_, split_sizes=split_sizes, scales=scales, with_quantized_compute=with_quantized_compute, input_quantizers=input_quantizers, weight_quantizers=weight_quantizers, - grad_output_quantizers=grad_output_quantizers, dtype=dtype, input_requires_grad=input_requires_grad, weight_requires_grad=weight_requires_grad, device=device, ) - return self._fuser_forward_split_quantize( - ctx=ctx, + else: + out, tensors_to_save = self._fuser_forward_split_quantize( + input_=input_, + split_sizes=split_sizes, + scales=scales, + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + ) + + # Save tensors and autograd metadata on the basic-op context. + self.fuser_forward_save_ctx( + basic_op_ctxs=basic_op_ctxs, input_=input_, - split_sizes=split_sizes, - scales=scales, - with_quantized_compute=with_quantized_compute, - input_quantizers=input_quantizers, - weight_quantizers=weight_quantizers, - grad_output_quantizers=grad_output_quantizers, - dtype=dtype, - input_requires_grad=input_requires_grad, - weight_requires_grad=weight_requires_grad, - device=device, + tensors_to_save=[tensors_to_save], + requires_grad=[ctx.requires_grad], + basic_op_extra_inputs=basic_op_extra_inputs, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + basic_op_kwargs=basic_op_kwargs, + use_grouped_tensor_path=use_grouped_tensor_path, ) + return out, [()] + + def fuser_forward_save_ctx( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, # pylint: disable=unused-argument + tensors_to_save: list[tuple[Optional[torch.Tensor | QuantizedTensorStorage], ...]], + *, + requires_grad: list[bool], + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], # pylint: disable=unused-argument + prev_op_grad_output_quantizer: Optional[Quantizer], # pylint: disable=unused-argument + next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument + basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument + use_grouped_tensor_path: bool, + ) -> None: + """ + Save tensors and autograd metadata in context. + """ + if not requires_grad[0]: + return + + ctx = basic_op_ctxs[0] + ctx.save_for_backward(*tensors_to_save[0]) + + num_groups = self.num_groups + weight_param = self.weight if self.single_grouped_weight else self.weight0 + + with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + input_quantizers = [None] * num_groups + weight_quantizers = [None] * num_groups + grad_output_quantizers = [None] * num_groups + if with_quantized_compute: + for group_idx in range(num_groups): + input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) + weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) + grad_output_quantizers[group_idx] = self.get_quantizer("backward", group_idx) + + ctx.use_grouped_tensor_path = use_grouped_tensor_path + ctx.with_quantized_compute = with_quantized_compute + ctx.input_quantizers = input_quantizers + ctx.weight_quantizers = weight_quantizers + ctx.grad_output_quantizers = grad_output_quantizers + ctx.grad_input_quantizers = None + # ``split_sizes`` and ``base_split_offsets`` are routed through + # ``save_for_backward`` (see ``_fuser_forward_split_quantize`` and + # ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). + if torch.is_autocast_enabled(): + ctx.dtype = torch.get_autocast_dtype("cuda") + else: + ctx.dtype = weight_param.dtype + ctx.input_requires_grad = requires_grad[0] + ctx.weight_requires_grad = requires_grad[0] and weight_param.requires_grad + # ================================================================== # Legacy `tex.split_quantize` + `general_grouped_gemm` flow. # ``m_splits`` is needed on CPU here, so this flow is NOT cuda-graphable. @@ -980,19 +1042,19 @@ def fuser_forward( def _fuser_forward_split_quantize( self, *, - ctx: OperationContext, input_: torch.Tensor, split_sizes: torch.Tensor, scales: Optional[torch.Tensor], with_quantized_compute: bool, input_quantizers: list[Optional[Quantizer]], weight_quantizers: list[Optional[Quantizer]], - grad_output_quantizers: list[Optional[Quantizer]], dtype: torch.dtype, input_requires_grad: bool, weight_requires_grad: bool, device: torch.device, - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: + """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. + """ num_groups = self.num_groups has_bias = self.has_bias @@ -1070,47 +1132,44 @@ def _fuser_forward_split_quantize( for x in xs: x.update_usage(rowwise_usage=False, columnwise_usage=True) - # Save state for backward pass - if ctx.requires_grad: - saved = [split_sizes] - if self._scale_bias: - saved.append(scales) - saved.extend(xs) - saved.extend(ws) - ctx.save_for_backward(*saved) - ctx.use_grouped_tensor_path = False - ctx.with_quantized_compute = with_quantized_compute - ctx.input_quantizers = input_quantizers - ctx.weight_quantizers = weight_quantizers - ctx.grad_output_quantizers = grad_output_quantizers - ctx.grad_input_quantizers = None - ctx.dtype = dtype - ctx.input_requires_grad = input_requires_grad - ctx.weight_requires_grad = weight_requires_grad - - return out, [()] + # Build the tuple of tensors to save for backward. Layout: + # [split_sizes, base_split_offsets, split_points, + # (scales if scale_bias), *xs, *ws] + # ``base_split_offsets`` and ``split_points`` are unused on the + # split-quantize backward path but are included as ``None`` so the + # saved-tensor layout matches the graph-safe + # ``_fuser_forward_grouped_tensor`` path (and the fused MLP forward). + saved: list[Optional[torch.Tensor]] = [split_sizes, None, None] + if self._scale_bias: + saved.append(scales) + saved.extend(xs) + saved.extend(ws) + return out, tuple(saved) def _fuser_forward_grouped_tensor( self, *, - ctx: OperationContext, input_: torch.Tensor, split_sizes: torch.Tensor, scales: Optional[torch.Tensor], with_quantized_compute: bool, input_quantizers: list[Optional[Quantizer]], weight_quantizers: list[Optional[Quantizer]], - grad_output_quantizers: list[Optional[Quantizer]], dtype: torch.dtype, input_requires_grad: bool, weight_requires_grad: bool, device: torch.device, - ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: - """Graph-safe GroupedTensor forward path.""" + ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: + """Graph-safe GroupedTensor forward path (pure compute). + Returns ``(output, tensors_to_save)``. ``split_sizes``, + ``base_split_offsets`` and ``split_points`` are returned so that + ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them. + """ num_groups = self.num_groups has_bias = self.has_bias - base_offsets = tex.splits_to_offsets(split_sizes, 1) + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_points = base_split_offsets[1:].to(dtype=torch.int) # Flatten to 2D so the first dim is the total token count. original_shape = list(input_.size()) @@ -1132,7 +1191,7 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=x.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_offsets * self.in_features, + tensor_offsets=base_split_offsets * self.in_features, ) # Build the weight GroupedTensor / list. @@ -1165,7 +1224,7 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=out.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_offsets * self.out_features, + tensor_offsets=base_split_offsets * self.out_features, ) # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias @@ -1174,7 +1233,7 @@ def _fuser_forward_grouped_tensor( bias_scale: Optional[torch.Tensor] = None if has_bias: # Bias always needs to be passed as a GroupedTensor for the grouped GEMM. - grouped_bias = self._get_grouped_bias_for_gemm(dtype, device) + grouped_bias = self._get_grouped_bias_for_gemm(dtype) if self._scale_bias: bias_scale = scales.reshape(-1) if bias_scale.dtype != torch.float32: @@ -1197,45 +1256,23 @@ def _fuser_forward_grouped_tensor( if not weight_requires_grad: grouped_x = None - # Save state for backward pass. Following the fused grouped MLP - # pattern, we save the GroupedTensor's component buffers (rather than - # the wrapper) and rebuild it in backward. ``base_offsets`` is saved - # so backward can reuse it without recomputing via ``splits_to_offsets``. - if ctx.requires_grad: - saved: list[Optional[torch.Tensor]] = [split_sizes, base_offsets] - if self._scale_bias: - saved.append(scales) - # For the wgrad input we save (data, scale_inv). - # * Quantized path saves columnwise data + scale. - # * Unquantized path saves the raw rowwise data and a None scale. - if grouped_x is not None: - if with_quantized_compute: - saved.extend( - [ - grouped_x.columnwise_data, - grouped_x.columnwise_scale_inv, - ] - ) - else: - saved.extend([grouped_x.rowwise_data, None]) - else: - saved.extend([None, None]) - if self.single_grouped_weight: - saved.append(grouped_weights) - else: - saved.extend(grouped_weights) - ctx.save_for_backward(*saved) - ctx.use_grouped_tensor_path = True - ctx.with_quantized_compute = with_quantized_compute - ctx.input_quantizers = input_quantizers - ctx.weight_quantizers = weight_quantizers - ctx.grad_output_quantizers = grad_output_quantizers - ctx.grad_input_quantizers = None - ctx.dtype = dtype - ctx.input_requires_grad = input_requires_grad - ctx.weight_requires_grad = weight_requires_grad - - return out, [()] + # Build the tuple of tensors to save for backward. Layout: + # [split_sizes, base_split_offsets, split_points, + # (scales if _scale_bias), grouped_x, *weights] + if grouped_x is not None: + if with_quantized_compute: + # only columnwise data is needed for wgrad + grouped_x.rowwise_data = None + grouped_x.scale_inv = None + saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points] + if self._scale_bias: + saved.append(scales) + saved.append(grouped_x) + if self.single_grouped_weight: + saved.append(grouped_weights) + else: + saved.extend(grouped_weights) + return out, tuple(saved) def fuser_backward( self, @@ -1275,9 +1312,15 @@ def _fuser_backward_split_quantize( weights = self._get_weight_tensors() device = weights[0].device - # Saved tensors from forward pass + # Saved tensors from forward pass. Layout: + # [split_sizes, base_split_offsets, split_points, + # (scales if _scale_bias), *xs, *ws] + # ``base_split_offsets`` and ``split_points`` are unused on this path + # but are present so the saved-tensor layout matches the graph-safe + # path (and the fused MLP forward). saved_tensors = ctx.saved_tensors - split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] + split_sizes = saved_tensors[0] + saved_tensors = saved_tensors[3:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1453,51 +1496,30 @@ def _fuser_backward_grouped_tensor( with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) - # Saved tensors from forward pass (see _fuser_forward_grouped_tensor). + # Saved tensors from forward pass + # Layout: [split_sizes, base_split_offsets, split_points, + # (scales if _scale_bias), grouped_x, *weights] + # ``split_points`` is unused on this path but is present so the + # saved-tensor layout matches the fused MLP forward (which needs it + # for the cuDNN grouped GEMM kernel). saved_tensors = ctx.saved_tensors - split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] - base_offsets, saved_tensors = saved_tensors[0], saved_tensors[1:] + split_sizes = saved_tensors[0] + base_split_offsets = saved_tensors[1] + saved_tensors = saved_tensors[3:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] - x_data, saved_tensors = saved_tensors[0], saved_tensors[1:] - x_scale, saved_tensors = saved_tensors[0], saved_tensors[1:] + grouped_x, saved_tensors = saved_tensors[0], saved_tensors[1:] if self.single_grouped_weight: ws, saved_tensors = saved_tensors[0], saved_tensors[1:] else: ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] # Flatten grad_output to 2D (total_tokens, out_features) - # to figure out total tokens and use it to build the grouped input tensor. + # to figure out total tokens. dy_2d = grad_output.reshape(-1, self.out_features) total_tokens = dy_2d.size(0) - # Rebuild the grouped input tensor used by wgrad. - grouped_x = None - if ctx.weight_requires_grad and x_data is not None: - if with_quantized_compute: - grouped_x = GroupedTensor( - shape=(total_tokens, self.in_features), - dtype=dtype, - num_tensors=num_groups, - quantizer=ctx.input_quantizers[0], - columnwise_data=x_data, - columnwise_scale_inv=x_scale, - first_dims=split_sizes, - tensor_offsets=base_offsets * self.in_features, - with_gemm_swizzled_scales=True, - ) - else: - grouped_x = GroupedTensor( - shape=(total_tokens, self.in_features), - dtype=dtype, - num_tensors=num_groups, - quantizer=None, - data=x_data, - first_dims=split_sizes, - tensor_offsets=base_offsets * self.in_features, - ) - # Build the grad_output GroupedTensor. # Optionally get dbias is fusion available with bgrad_group_quantize dbias_packed = None @@ -1526,7 +1548,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=dy_2d.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_offsets * self.out_features, + tensor_offsets=base_split_offsets * self.out_features, ) # Bias Grads compute if not already computed in bgrad_group_quantize @@ -1540,11 +1562,11 @@ def _fuser_backward_grouped_tensor( dy_2d, scales_f32, bias_packed, - offsets=base_offsets, + offsets=base_split_offsets, ) elif dbias_packed is None: # BF16/FP16 path - dbias_packed = compute_grouped_dbias(dy_2d, base_offsets, num_groups) + dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups) if self.single_grouped_bias: final_bias_grads = [dbias_packed.to(dtype=dtype)] else: @@ -1562,7 +1584,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=grad_input.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_offsets * self.in_features, + tensor_offsets=base_split_offsets * self.in_features, ) general_grouped_gemm_for_grouped_tensor( ws, diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 8a6a316505..fef3460426 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -7,8 +7,6 @@ from __future__ import annotations from collections.abc import Callable import functools -import inspect -import math import os from typing import Optional @@ -341,14 +339,12 @@ def fuser_backward( device = fc1_op._get_weight_tensors()[0].device dtype = fc1_ctx.dtype - # Saved tensors from FC1 forward + # Saved tensors from FC1 forward. + # Layout: [split_sizes, base_split_offsets, split_points, + # grouped_fc1_x, *fc1_weights] saved_tensors = fc1_ctx.saved_tensors - split_sizes, split_points, saved_tensors = ( - saved_tensors[0], - saved_tensors[1], - saved_tensors[2:], - ) - + split_sizes, base_split_offsets, split_points = saved_tensors[:3] + grouped_fc1_x, saved_tensors = saved_tensors[3], saved_tensors[4:] if fc1_op.single_grouped_weight: grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] else: @@ -357,21 +353,21 @@ def fuser_backward( saved_tensors[num_groups:], ) - ( - fc1_x_col_data, - fc1_x_col_scale, - fc1_x_tensor_offsets, - ), saved_tensors = ( - saved_tensors[:3], - saved_tensors[3:], - ) - # Saved tensors from scaled SwiGLU forward swiglu_in, scales = swiglu_ctx.saved_tensors - # Saved tensors from FC2 forward - saved_tensors = fc2_ctx.saved_tensors - _, saved_tensors = saved_tensors[0], saved_tensors[1:] # Assume same split sizes as FC1 + # Saved tensors from FC2 forward. + # Layout: [split_sizes, base_split_offsets, split_points, + # (fc2_scales if _scale_bias), + # grouped_fc2_x, *fc2_weights] + scale_bias = fc2_op._scale_bias and fc2_op.has_bias + saved_tensors = fc2_ctx.saved_tensors[3:] + if fc2_op._scale_bias: + # Saved for the unfused backward path, which reads its own + # per-op scales here. The fused backward below currently reuses + # the SwiGLU ``scales``. + saved_tensors = saved_tensors[1:] + grouped_fc2_x, saved_tensors = saved_tensors[0], saved_tensors[1:] if fc2_op.single_grouped_weight: grouped_fc2_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] else: @@ -380,53 +376,21 @@ def fuser_backward( saved_tensors[num_groups:], ) - ( - fc2_x_col_data, - fc2_x_col_scale, - fc2_x_tensor_offsets, - ), saved_tensors = ( - saved_tensors[:3], - saved_tensors[3:], - ) - # Group splits if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") - scale_bias = fc2_op._scale_bias and fc2_op.has_bias - grouped_fc1_x = None - if fc1_ctx.weight_requires_grad: - grouped_fc1_x = GroupedTensor( - shape=(out_shape[0], fc1_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc1_ctx.input_quantizer, - columnwise_data=fc1_x_col_data, - columnwise_scale_inv=fc1_x_col_scale, - first_dims=split_sizes, - tensor_offsets=fc1_x_tensor_offsets, - with_gemm_swizzled_scales=True, - ) - - grouped_fc2_x = None - if fc2_ctx.weight_requires_grad: - grouped_fc2_x = GroupedTensor( - shape=(out_shape[0], fc2_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc2_ctx.input_quantizer, - columnwise_data=fc2_x_col_data, - columnwise_scale_inv=fc2_x_col_scale, - first_dims=split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - with_gemm_swizzled_scales=True, - ) + if not fc1_ctx.weight_requires_grad: + grouped_fc1_x = None + if not fc2_ctx.weight_requires_grad: + grouped_fc2_x = None # Split grad output tensor and convert dtypes if needed - fc2_ctx.grad_output_quantizer.set_usage( + fc2_grad_output_quantizer = fc2_ctx.grad_output_quantizers[0] + fc2_grad_output_quantizer.set_usage( rowwise=True, columnwise=fc2_ctx.weight_requires_grad ) - fc2_ctx.grad_output_quantizer.optimize_for_gemm = True + fc2_grad_output_quantizer.optimize_for_gemm = True output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None fc2_dy = None @@ -441,14 +405,14 @@ def fuser_backward( if output_fc2_dbias and not scale_bias: grouped_fc2_dy, fc2_dbias_packed = tex.bgrad_group_quantize( fc2_dy, - fc2_ctx.grad_output_quantizer, + fc2_grad_output_quantizer, num_groups, split_sizes, ) else: grouped_fc2_dy = tex.group_quantize( fc2_dy, - fc2_ctx.grad_output_quantizer, + fc2_grad_output_quantizer, num_groups, split_sizes, ) @@ -565,7 +529,7 @@ def fuser_backward( fc2_dy, scales_f32, bias_packed, - offsets=fc1_ctx.base_split_offsets, + offsets=base_split_offsets, dscales=grad_scales, ) fc2_dbias_packed_result = fc2_dbias_packed_result.to(dtype=dtype) @@ -594,12 +558,12 @@ def fuser_backward( fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)] # FC1 grad output for dgrad and wgrad GEMMs - fc1_dy_tensor_offsets = fc1_ctx.base_split_offsets * fc1_weight_shape[0] + fc1_dy_tensor_offsets = base_split_offsets * fc1_weight_shape[0] grouped_fc1_dy = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[0]), dtype=dtype, num_tensors=num_groups, - quantizer=fc1_ctx.grad_output_quantizer, + quantizer=fc1_ctx.grad_output_quantizers[0], data=fc1_dy_row_data, columnwise_data=fc1_dy_col_data, scale_inv=fc1_dy_row_scale, @@ -633,7 +597,7 @@ def fuser_backward( and fc2_op.wgrad_store.delay_wgrad_compute() ): clear_tensor_data( - grouped_fc2_x.data, + grouped_fc2_x.rowwise_data, grouped_fc2_x.columnwise_data, grouped_fc2_x.scale_inv, grouped_fc2_x.columnwise_scale_inv, @@ -729,7 +693,7 @@ def fuser_backward( and fc1_op.wgrad_store.delay_wgrad_compute() ): clear_tensor_data( - grouped_fc1_x.data, + grouped_fc1_x.rowwise_data, grouped_fc1_x.columnwise_data, grouped_fc1_x.scale_inv, grouped_fc1_x.columnwise_scale_inv, diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 599e5f96ae..7f5a52550e 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -160,10 +160,10 @@ def fuser_forward( if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") split_sizes = split_sizes.to(dtype=torch.int64, device=device) - base_offsets = tex.splits_to_offsets(split_sizes, 1) - split_points = base_offsets[1:].to(dtype=torch.int) - fc1_x_tensor_offsets = base_offsets * fc1_weight_shape[1] - fc2_x_tensor_offsets = base_offsets * fc2_weight_shape[1] + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_points = base_split_offsets[1:].to(dtype=torch.int) + fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] + fc2_x_tensor_offsets = base_split_offsets * fc2_weight_shape[1] # Extract post-scales from extra input scales = basic_op_extra_inputs[1][0] @@ -452,27 +452,35 @@ def fuser_forward( # Save state for backward pass if requires_grad: mark_grouped_tensor(grouped_fc1_x, swiglu_in, scales, grouped_fc2_x) - fc1_input_tensors = ( - grouped_fc1_x.columnwise_data, - grouped_fc1_x.columnwise_scale_inv, - fc1_x_tensor_offsets, - ) - # FC1 + + # Save the input ``GroupedTensor``s themselves for the activations. + for grouped_fc_x in (grouped_fc1_x, grouped_fc2_x): + if grouped_fc_x is not None: + grouped_fc_x.rowwise_data = None + grouped_fc_x.scale_inv = None + + # FC1 saved-tensor layout. + # [split_sizes, base_split_offsets, split_points, + # grouped_fc1_x, *fc1_weight_tensors] fc1_weight_tensors = ( [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight ) fc1_ctx.save_for_backward( - split_sizes, split_points, *fc1_weight_tensors, *fc1_input_tensors + split_sizes, + base_split_offsets, + split_points, + grouped_fc1_x, + *fc1_weight_tensors, ) + fc1_ctx.use_grouped_tensor_path = True fc1_ctx.with_quantized_compute = True - fc1_ctx.input_quantizer = fc1_input_quantizer - fc1_ctx.weight_quantizer = fc1_weight_quantizer - fc1_ctx.grad_output_quantizer = fc1_grad_output_quantizer + fc1_ctx.input_quantizers = [fc1_input_quantizer] + fc1_ctx.weight_quantizers = [fc1_weight_quantizer] + fc1_ctx.grad_output_quantizers = [fc1_grad_output_quantizer] fc1_ctx.grad_input_quantizers = None fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad - fc1_ctx.base_split_offsets = base_offsets # Scaled SwiGLU swiglu_ctx.save_for_backward(swiglu_in, scales) @@ -480,25 +488,31 @@ def fuser_forward( swiglu_ctx.extra_input_requires_grad = True swiglu_ctx.dtype = dtype - # FC2 state - if grouped_fc2_x is not None: - fc2_input_tensors = ( - grouped_fc2_x.columnwise_data, - grouped_fc2_x.columnwise_scale_inv, - fc2_x_tensor_offsets, - ) - else: - fc2_input_tensors = (None, None, None) - - if fc2_op.single_grouped_weight: - fc2_ctx.save_for_backward(split_sizes, grouped_fc2_weight, *fc2_input_tensors) - else: - fc2_ctx.save_for_backward(split_sizes, *grouped_fc2_weight, *fc2_input_tensors) - + # FC2 saved-tensor layout. Matches the unfused + # ``GroupedLinear._fuser_forward_grouped_tensor`` layout so the + # unfused backward (basic/grouped_linear.py) can consume the same + # ctx when the fused backward is unavailable. + # [split_sizes, base_split_offsets, split_points, + # (fc2_scales if _scale_bias), + # grouped_fc2_x, *fc2_weight_tensors] + fc2_weight_tensors = ( + [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight + ) + fc2_saved: list[Optional[torch.Tensor]] = [ + split_sizes, + base_split_offsets, + split_points, + ] + if fc2_op._scale_bias: + fc2_saved.append(fc2_scales) + fc2_saved.append(grouped_fc2_x) + fc2_saved.extend(fc2_weight_tensors) + fc2_ctx.save_for_backward(*fc2_saved) + fc2_ctx.use_grouped_tensor_path = True fc2_ctx.with_quantized_compute = True - fc2_ctx.input_quantizer = fc2_input_quantizer - fc2_ctx.weight_quantizer = fc2_weight_quantizer - fc2_ctx.grad_output_quantizer = fc2_grad_output_quantizer + fc2_ctx.input_quantizers = [fc2_input_quantizer] + fc2_ctx.weight_quantizers = [fc2_weight_quantizer] + fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] fc2_ctx.grad_input_quantizers = None fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad @@ -539,3 +553,4 @@ def fuse_forward_ops( # Register fusion if available if ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): register_forward_fusion(fuse_forward_ops, prepend=True) + diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 5f12c3ed8c..92c5ba64de 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -303,6 +303,53 @@ def get_dtype(self) -> torch.dtype: return self.fake_dtype + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], "GroupedTensorStorage"]: + """Prepare the tensor base for saving for backward. + """ + tensors = [ + self.rowwise_data, + self.columnwise_data, + self.scale_inv, + self.columnwise_scale_inv, + self.amax, + self.columnwise_amax, + self.scale, + self.first_dims, + self.last_dims, + self.tensor_offsets, + ] + self.rowwise_data = None + self.columnwise_data = None + self.scale_inv = None + self.columnwise_scale_inv = None + self.amax = None + self.columnwise_amax = None + self.scale = None + self.first_dims = None + self.last_dims = None + self.tensor_offsets = None + self.quantized_tensors = None + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + """Restore the tensor base data from the saved tensors list. + """ + self.rowwise_data = tensors[0] + self.columnwise_data = tensors[1] + self.scale_inv = tensors[2] + self.columnwise_scale_inv = tensors[3] + self.amax = tensors[4] + self.columnwise_amax = tensors[5] + self.scale = tensors[6] + self.first_dims = tensors[7] + self.last_dims = tensors[8] + self.tensor_offsets = tensors[9] + return tensors[10:] + def clear(self) -> None: """ Reset tensor data and clear all buffers. From 7f23de47efae36da5f5d8a3fa63ad9489b10f856 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 5 May 2026 08:53:02 +0000 Subject: [PATCH 13/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/basic/grouped_linear.py | 3 +-- .../pytorch/ops/fused/backward_grouped_mlp.py | 4 +--- transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py | 3 +-- .../pytorch/tensor/storage/grouped_tensor_storage.py | 6 ++---- 4 files changed, 5 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 202441e4f1..4520dabf16 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1053,8 +1053,7 @@ def _fuser_forward_split_quantize( weight_requires_grad: bool, device: torch.device, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: - """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow. - """ + """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" num_groups = self.num_groups has_bias = self.has_bias diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index fef3460426..555bfad88d 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -387,9 +387,7 @@ def fuser_backward( # Split grad output tensor and convert dtypes if needed fc2_grad_output_quantizer = fc2_ctx.grad_output_quantizers[0] - fc2_grad_output_quantizer.set_usage( - rowwise=True, columnwise=fc2_ctx.weight_requires_grad - ) + fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=fc2_ctx.weight_requires_grad) fc2_grad_output_quantizer.optimize_for_gemm = True output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 7f5a52550e..43b42812c4 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -459,7 +459,7 @@ def fuser_forward( grouped_fc_x.rowwise_data = None grouped_fc_x.scale_inv = None - # FC1 saved-tensor layout. + # FC1 saved-tensor layout. # [split_sizes, base_split_offsets, split_points, # grouped_fc1_x, *fc1_weight_tensors] fc1_weight_tensors = ( @@ -553,4 +553,3 @@ def fuse_forward_ops( # Register fusion if available if ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): register_forward_fusion(fuse_forward_ops, prepend=True) - diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 92c5ba64de..485b32328b 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -306,8 +306,7 @@ def get_dtype(self) -> torch.dtype: def prepare_for_saving( self, ) -> Tuple[list[Optional[torch.Tensor]], "GroupedTensorStorage"]: - """Prepare the tensor base for saving for backward. - """ + """Prepare the tensor base for saving for backward.""" tensors = [ self.rowwise_data, self.columnwise_data, @@ -336,8 +335,7 @@ def prepare_for_saving( def restore_from_saved( self, tensors: list[Optional[torch.Tensor]] ) -> list[Optional[torch.Tensor]]: - """Restore the tensor base data from the saved tensors list. - """ + """Restore the tensor base data from the saved tensors list.""" self.rowwise_data = tensors[0] self.columnwise_data = tensors[1] self.scale_inv = tensors[2] From 09ba167281fc2a58a3e34679ec01bf025ef596ce Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 5 May 2026 09:18:42 +0000 Subject: [PATCH 14/18] cleanup Signed-off-by: Varun Thumbe --- tests/pytorch/test_fusible_ops.py | 79 ++----------------- tests/pytorch/utils.py | 72 +++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 6 +- 3 files changed, 81 insertions(+), 76 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 53dfb1b422..d690327fe2 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -7,6 +7,7 @@ from collections.abc import Iterable, Sequence import functools import io +import os import math import random from typing import Optional @@ -42,7 +43,6 @@ ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch.cpp_extensions.gemm import general_grouped_gemm_for_grouped_tensor -from transformer_engine.pytorch.module.base import get_dummy_wgrad import transformer_engine_torch as tex # Import utility functions @@ -51,6 +51,7 @@ assert_close_grads, dtype_tols, make_recipe, + MegatronTrainingHelper, quantization_tols, reset_rng_states, ) @@ -212,76 +213,6 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return out -class MegatronTrainingHelper: - """Test-side stand-in for the Megatron-Core DDP / MegatronFSDP wrapper. - Megatron's DDP wrapper (and MegatronFSDP) owns the per-parameter - ``main_grad`` buffer and the ``overwrite_main_grad`` / - ``grad_added_to_main_grad`` attributes that coordinate - ``fuse_wgrad_accumulation`` with TE modules. These helpers reproduce the - relevant slice of that protocol so TE tests can exercise the - accumulate-into-``main_grad`` code path without pulling in the full - Megatron-Core dependency. - """ - - @staticmethod - def init_main_grad_buffers( - weight_params: Iterable[torch.nn.Parameter], - *, - fill_value: float, - overwrite_main_grad: bool, - zero_out_wgrad: bool = False, - dtype: torch.dtype = torch.float32, - ) -> None: - """Allocate ``main_grad`` and stamp the wrapper attributes on each - param, mirroring what the Megatron DDP/FSDP wrapper does before - backward.""" - for wp in weight_params: - wp.main_grad = torch.full(wp.size(), fill_value, device=wp.device, dtype=dtype) - wp.overwrite_main_grad = overwrite_main_grad - wp.zero_out_wgrad = zero_out_wgrad - wp.grad_added_to_main_grad = False - - @staticmethod - def verify_main_grad_accumulation( - weight_params: Iterable[torch.nn.Parameter], - *, - expected_main_grads: Iterable[torch.Tensor], - rtol: float = 0.0, - atol: float = 0.0, - ) -> None: - """Check that backward produced what the Megatron wrapper expects: - each ``main_grad`` matches ``expected_main_grads``, - ``grad_added_to_main_grad`` was flipped to ``True`` so the wrapper's - post-backward hooks won't double-accumulate, and ``param.grad`` was - replaced by the cached dummy tensor (so a wrapper hook that did - ``main_grad += grad`` would be a no-op rather than double-counting). - """ - for wp, expected in zip(weight_params, expected_main_grads): - torch.testing.assert_close(wp.main_grad.to(expected), expected, rtol=rtol, atol=atol) - - assert wp.grad_added_to_main_grad is True, ( - "weight.grad_added_to_main_grad was not flipped to True; " - "the Megatron DDP/FSDP wrapper hook will double-accumulate." - ) - - # ``.grad`` should be the cached dummy tensor returned by - # ``get_dummy_wgrad`` -- shared storage, not the real wgrad. - expected_dummy = get_dummy_wgrad(list(wp.size()), wp.dtype) - assert ( - wp.grad is not None - ), "weight.grad is None; the Megatron protocol expects a dummy tensor stand-in here." - assert wp.grad.data_ptr() == expected_dummy.data_ptr(), ( - "weight.grad does not share storage with the cached dummy " - "wgrad; downstream wrapper hooks risk double-accumulating." - ) - if getattr(wp, "zero_out_wgrad", False): - assert torch.all(wp.grad == 0), ( - "weight.zero_out_wgrad=True but the dummy weight.grad " - "was not zeroed; downstream hooks reading .grad would " - "see stale bytes from the previous step." - ) - - class TestSequentialContainer: """Tests for sequential container""" @@ -2119,7 +2050,8 @@ def test_grouped_linear( single_grouped_bias: bool, ) -> None: """Grouped GEMM""" - + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and (single_grouped_weight or single_grouped_bias): + pytest.skip("single_grouped_weight/single_grouped_bias requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") # Split sizes split_sizes = [split_alignment * i for i in range(group_size)] random.shuffle(split_sizes) @@ -2321,7 +2253,8 @@ def test_grouped_linear_cuda_graph_safe( Exercises the grouped-tensor / cublas-grouped-gemm path which uses GPU-resident split offsets and is the only flow safe to capture. """ - + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and (single_grouped_weight or single_grouped_bias): + pytest.skip("single_grouped_weight/single_grouped_bias requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if torch.cuda.get_device_capability() < (10, 0): pytest.skip("Grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") # Skip invalid configurations diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 8f8852edc2..c7cbe78a6d 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -8,6 +8,7 @@ import os import random import subprocess +from collections.abc import Iterable from contextlib import contextmanager from typing import Optional, Sequence, Tuple, Dict, Any, List from packaging.version import Version as PkgVersion @@ -27,6 +28,7 @@ check_set_window_size, ) from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend +from transformer_engine.pytorch.module.base import get_dummy_wgrad def str_to_dtype(dtype: str | torch.dtype) -> torch.dtype: @@ -477,3 +479,73 @@ def run_distributed( msg += f"\n--- stderr ---\n{stderr_tail}" raise AssertionError(msg) return result + + +class MegatronTrainingHelper: + """Test-side stand-in for the Megatron-Core DDP / MegatronFSDP wrapper. + Megatron's DDP wrapper (and MegatronFSDP) owns the per-parameter + ``main_grad`` buffer and the ``overwrite_main_grad`` / + ``grad_added_to_main_grad`` attributes that coordinate + ``fuse_wgrad_accumulation`` with TE modules. These helpers reproduce the + relevant slice of that protocol so TE tests can exercise the + accumulate-into-``main_grad`` code path without pulling in the full + Megatron-Core dependency. + """ + + @staticmethod + def init_main_grad_buffers( + weight_params: Iterable[torch.nn.Parameter], + *, + fill_value: float, + overwrite_main_grad: bool, + zero_out_wgrad: bool = False, + dtype: torch.dtype = torch.float32, + ) -> None: + """Allocate ``main_grad`` and stamp the wrapper attributes on each + param, mirroring what the Megatron DDP/FSDP wrapper does before + backward.""" + for wp in weight_params: + wp.main_grad = torch.full(wp.size(), fill_value, device=wp.device, dtype=dtype) + wp.overwrite_main_grad = overwrite_main_grad + wp.zero_out_wgrad = zero_out_wgrad + wp.grad_added_to_main_grad = False + + @staticmethod + def verify_main_grad_accumulation( + weight_params: Iterable[torch.nn.Parameter], + *, + expected_main_grads: Iterable[torch.Tensor], + rtol: float = 0.0, + atol: float = 0.0, + ) -> None: + """Check that backward produced what the Megatron wrapper expects: + each ``main_grad`` matches ``expected_main_grads``, + ``grad_added_to_main_grad`` was flipped to ``True`` so the wrapper's + post-backward hooks won't double-accumulate, and ``param.grad`` was + replaced by the cached dummy tensor (so a wrapper hook that did + ``main_grad += grad`` would be a no-op rather than double-counting). + """ + for wp, expected in zip(weight_params, expected_main_grads): + torch.testing.assert_close(wp.main_grad.to(expected), expected, rtol=rtol, atol=atol) + + assert wp.grad_added_to_main_grad is True, ( + "weight.grad_added_to_main_grad was not flipped to True; " + "the Megatron DDP/FSDP wrapper hook will double-accumulate." + ) + + # ``.grad`` should be the cached dummy tensor returned by + # ``get_dummy_wgrad`` -- shared storage, not the real wgrad. + expected_dummy = get_dummy_wgrad(list(wp.size()), wp.dtype) + assert ( + wp.grad is not None + ), "weight.grad is None; the Megatron protocol expects a dummy tensor stand-in here." + assert wp.grad.data_ptr() == expected_dummy.data_ptr(), ( + "weight.grad does not share storage with the cached dummy " + "wgrad; downstream wrapper hooks risk double-accumulating." + ) + if getattr(wp, "zero_out_wgrad", False): + assert torch.all(wp.grad == 0), ( + "weight.zero_out_wgrad=True but the dummy weight.grad " + "was not zeroed; downstream hooks reading .grad would " + "see stale bytes from the previous step." + ) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 4520dabf16..b150438def 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -43,7 +43,7 @@ view_main_grad_as_grouped_buffer, ) from ..op import BasicOperation, OperationContext -from ...tensor import GroupedTensor +from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, compute_grouped_dbias_dscales, @@ -900,7 +900,7 @@ def fuser_forward( # Check which grads are required ctx = basic_op_ctxs[0] input_requires_grad = ctx.requires_grad - weight_requires_grad = weight_param.requires_grad + weight_requires_grad = ctx.requires_grad and weight_param.requires_grad # Quantizers input_quantizers = [None] * num_groups @@ -988,7 +988,7 @@ def fuser_forward_save_ctx( self, basic_op_ctxs: list[OperationContext], input_: torch.Tensor, # pylint: disable=unused-argument - tensors_to_save: list[tuple[Optional[torch.Tensor | QuantizedTensorStorage], ...]], + tensors_to_save: list[tuple[Optional[torch.Tensor | QuantizedTensorStorage | GroupedTensorStorage], ...]], *, requires_grad: list[bool], basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], # pylint: disable=unused-argument From 56874da2f4ced0f5e8e54bd1b9203fdb510a092d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 5 May 2026 09:19:38 +0000 Subject: [PATCH 15/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 18 ++++++++++++++---- .../pytorch/ops/basic/grouped_linear.py | 4 +++- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index d690327fe2..7691582f97 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2050,8 +2050,13 @@ def test_grouped_linear( single_grouped_bias: bool, ) -> None: """Grouped GEMM""" - if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and (single_grouped_weight or single_grouped_bias): - pytest.skip("single_grouped_weight/single_grouped_bias requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) # Split sizes split_sizes = [split_alignment * i for i in range(group_size)] random.shuffle(split_sizes) @@ -2253,8 +2258,13 @@ def test_grouped_linear_cuda_graph_safe( Exercises the grouped-tensor / cublas-grouped-gemm path which uses GPU-resident split offsets and is the only flow safe to capture. """ - if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and (single_grouped_weight or single_grouped_bias): - pytest.skip("single_grouped_weight/single_grouped_bias requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) if torch.cuda.get_device_capability() < (10, 0): pytest.skip("Grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") # Skip invalid configurations diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b150438def..88a9f7a36f 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -988,7 +988,9 @@ def fuser_forward_save_ctx( self, basic_op_ctxs: list[OperationContext], input_: torch.Tensor, # pylint: disable=unused-argument - tensors_to_save: list[tuple[Optional[torch.Tensor | QuantizedTensorStorage | GroupedTensorStorage], ...]], + tensors_to_save: list[ + tuple[Optional[torch.Tensor | QuantizedTensorStorage | GroupedTensorStorage], ...] + ], *, requires_grad: list[bool], basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], # pylint: disable=unused-argument From 8cad423703321474d6267206be8837e5a4e467b7 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 5 May 2026 09:45:50 +0000 Subject: [PATCH 16/18] address review comments Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/basic/grouped_linear.py | 8 ++++++++ .../pytorch/ops/fused/backward_grouped_mlp.py | 2 ++ 2 files changed, 10 insertions(+) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b150438def..9028d2fa67 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1459,6 +1459,10 @@ def _fuser_backward_split_quantize( # ``main_grad`` again. if ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = get_dummy_wgrads_for_params(weights) + elif ctx.weight_requires_grad and delay_wgrad: + final_weight_grads = ( + [None] if self.single_grouped_weight else [None] * num_groups + ) if not has_bias: grad_params = list(final_weight_grads) @@ -1668,6 +1672,10 @@ def _fuser_backward_grouped_tensor( # ``main_grad`` again. if ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = get_dummy_wgrads_for_params(weights) + elif ctx.weight_requires_grad and delay_wgrad: + final_weight_grads = ( + [None] if self.single_grouped_weight else [None] * num_groups + ) # Assemble grad params in parameter registration order and return. if not has_bias: diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 555bfad88d..b32aa2ec41 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -229,6 +229,8 @@ def _compute_grad_params( # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added if fc_op._accumulate_into_main_grad: w_list = get_dummy_wgrads_for_params(weights) + elif delay_wgrad: + w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups # Assemble grad_params in parameter registration order. if not fc_op.has_bias: From 396070e1a23d187d5d5e4075b9f1351ef0add81a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 5 May 2026 09:47:33 +0000 Subject: [PATCH 17/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/basic/grouped_linear.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index ce18db19a9..e698c2697f 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1462,9 +1462,7 @@ def _fuser_backward_split_quantize( if ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: - final_weight_grads = ( - [None] if self.single_grouped_weight else [None] * num_groups - ) + final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups if not has_bias: grad_params = list(final_weight_grads) @@ -1675,9 +1673,7 @@ def _fuser_backward_grouped_tensor( if ctx.weight_requires_grad and self._accumulate_into_main_grad: final_weight_grads = get_dummy_wgrads_for_params(weights) elif ctx.weight_requires_grad and delay_wgrad: - final_weight_grads = ( - [None] if self.single_grouped_weight else [None] * num_groups - ) + final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups # Assemble grad params in parameter registration order and return. if not has_bias: From 8456ec6dd7c58bb577a08edc8910394cbee8d462 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 5 May 2026 19:01:52 +0000 Subject: [PATCH 18/18] fix linting errors Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py | 2 +- transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index b32aa2ec41..320c7c39e5 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -32,7 +32,7 @@ validate_grouped_mlp_dims, ) from ...cpp_extensions import general_grouped_gemm_for_grouped_tensor -from ...module.base import _2X_ACC_WGRAD, get_dummy_wgrad +from ...module.base import _2X_ACC_WGRAD from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 43b42812c4..91db2ff9b7 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -162,7 +162,6 @@ def fuser_forward( split_sizes = split_sizes.to(dtype=torch.int64, device=device) base_split_offsets = tex.splits_to_offsets(split_sizes, 1) split_points = base_split_offsets[1:].to(dtype=torch.int) - fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] fc2_x_tensor_offsets = base_split_offsets * fc2_weight_shape[1] # Extract post-scales from extra input