diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index b561506b31..b85dad1258 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -12,6 +12,7 @@ from torch.nn import Parameter import transformer_engine.pytorch as te +import transformer_engine.pytorch.functional as te_functional from transformer_engine.common import recipe from transformer_engine.pytorch import ( Float8Quantizer, @@ -139,6 +140,50 @@ def dtype_tols(dtype: torch.dtype) -> Dict[str, float]: raise ValueError(f"Unsupported dtype ({dtype})") +def test_functional_grouped_linear_matches_torch_split_path(): + """Functional grouped linear matches a loop of torch linear calls.""" + + dtype = torch.float16 + device = "cuda" + split_sizes = torch.tensor([3, 5, 2], dtype=torch.int64, device=device) + in_features = 16 + out_features = 12 + total_tokens = int(split_sizes.sum().item()) + + torch.manual_seed(seed) + x = torch.randn(total_tokens, in_features, dtype=dtype, device=device) + weights = [ + torch.randn(out_features, in_features, dtype=dtype, device=device) + for _ in range(split_sizes.numel()) + ] + biases = [ + torch.randn(out_features, dtype=dtype, device=device) for _ in range(split_sizes.numel()) + ] + + out, cache = te_functional.grouped_linear( + x, + weights, + split_sizes, + bias=biases, + dtype=dtype, + use_grouped_tensor_path=False, + return_cache=True, + ) + + expected = torch.cat( + [ + torch.nn.functional.linear(x_i, weight, bias) + for x_i, weight, bias in zip(torch.split(x, split_sizes.tolist()), weights, biases) + ], + dim=0, + ) + torch.testing.assert_close(out.float(), expected.float(), **dtype_tols(dtype)) + assert isinstance(cache, dict) + assert cache["path"] == "split" + assert cache["num_groups"] == split_sizes.numel() + assert "saved_tensors" in cache + + param_types = [torch.float32, torch.float16] if is_bf16_available(): param_types.append(torch.bfloat16) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 06db28ee27..ec83de993f 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -66,6 +66,7 @@ mark_not_offload, ManualOffloadSynchronizer, ) +from transformer_engine.pytorch import functional from transformer_engine.pytorch import ops from transformer_engine.pytorch import optimizers from transformer_engine.pytorch.export import onnx_export diff --git a/transformer_engine/pytorch/_functional/__init__.py b/transformer_engine/pytorch/_functional/__init__.py new file mode 100644 index 0000000000..64b3c0764e --- /dev/null +++ b/transformer_engine/pytorch/_functional/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Internal functional implementations shared by PyTorch modules and ops.""" diff --git a/transformer_engine/pytorch/_functional/grouped_linear.py b/transformer_engine/pytorch/_functional/grouped_linear.py new file mode 100644 index 0000000000..83232dedd2 --- /dev/null +++ b/transformer_engine/pytorch/_functional/grouped_linear.py @@ -0,0 +1,717 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Functional grouped-linear compute helpers. + +The functions in this module intentionally avoid depending on either +``TransformerEngineBaseModule`` autograd contexts or operation-fuser contexts. +Callers are responsible for saving/restoring the returned cache in the format +that their autograd surface expects. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import os +from typing import Any, Optional, Sequence + +import torch + +import transformer_engine_torch as tex + +from ..cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor +from ..module.base import _2X_ACC_DGRAD, _2X_ACC_FPROP, _2X_ACC_WGRAD, quantize_weight +from ..ops._common import is_quantized_tensor, maybe_dequantize +from ..quantized_tensor import QuantizedTensorStorage +from ..tensor import Float8CurrentScalingQuantizer, GroupedTensor, GroupedTensorStorage +from ..tensor import MXFP8Quantizer, NVFP4Quantizer, Quantizer +from ..triton.grouped_dbias_dscales import compute_grouped_dbias, compute_grouped_dbias_dscales +from ..utils import get_device_compute_capability + + +@dataclass(slots=True) +class GroupedLinearForwardResult: + """Grouped-linear forward output and backward cache.""" + + output: torch.Tensor + cache: dict[str, Any] + new_weight_workspaces: list[Optional[QuantizedTensorStorage]] + + +def maybe_dequantize_to_dtype( + tensor: torch.Tensor | QuantizedTensorStorage, + dtype: torch.dtype, +) -> torch.Tensor: + """Dequantize quantized tensors or cast regular tensors to ``dtype``.""" + + return maybe_dequantize(tensor, dtype) + + +def canonicalize_split_sizes( + split_sizes: torch.Tensor | Sequence[int], + *, + num_groups: int, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Return split sizes as an int64 tensor. + + ``device`` is optional because legacy split-list paths still need CPU-visible + values, while graph-safe grouped-tensor paths need GPU-resident metadata. + """ + + if not isinstance(split_sizes, torch.Tensor): + split_sizes = torch.tensor(split_sizes, dtype=torch.int64, device=device or "cpu") + elif split_sizes.dtype != torch.int64: + split_sizes = split_sizes.to(dtype=torch.int64) + if device is not None and split_sizes.device != device: + split_sizes = split_sizes.to(device=device) + if split_sizes.size() != (num_groups,): + raise ValueError( + f"Shape of splits tensor ({tuple(split_sizes.size())}) " + f"does not match number of GEMMs ({num_groups})." + ) + return split_sizes + + +def grouped_tensor_path_is_supported( + *, + with_quantized_compute: bool, + input_quantizers: Sequence[Optional[Quantizer]], + dtype: torch.dtype, + output_quantizers: Optional[Sequence[Optional[Quantizer]]] = None, + single_grouped_weight: bool = False, + require_env: bool = False, + fp8_calibration: bool = False, + debug: bool = False, + cpu_offloading: bool = False, + backward_override: Optional[str] = None, + save_original_input: bool = False, +) -> bool: + """Whether the graph-safe GroupedTensor/cuBLASLt path can be used.""" + + if require_env and not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): + return False + if ( + debug + or cpu_offloading + or fp8_calibration + or backward_override is not None + or save_original_input + ): + return False + if output_quantizers is not None and any(q is not None for q in output_quantizers): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if with_quantized_compute: + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + if device_capability < (10, 0) and cublaslt_version < 130500: + return False + return True + if not (10, 0) <= device_capability <= (11, 0): + return False + if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): + return True + if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): + return not single_grouped_weight + return False + + return dtype in (torch.bfloat16, torch.float16) + + +def make_grouped_tensor_from_2d_buffer( + data: torch.Tensor, + *, + num_groups: int, + split_sizes: torch.Tensor, + base_split_offsets: torch.Tensor, + last_dim: int, + dtype: torch.dtype, +) -> GroupedTensorStorage: + """Wrap a packed 2D buffer as a varying-first-dimension GroupedTensor.""" + + return GroupedTensorStorage( + shape=(data.size(0), last_dim), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=data.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * last_dim, + ) + + +def make_grouped_bias( + biases: Sequence[torch.Tensor | QuantizedTensorStorage], + *, + num_groups: int, + out_features: int, + dtype: torch.dtype, +) -> GroupedTensorStorage: + """Pack per-group biases into the GroupedTensor GEMM bias format.""" + + bias_data = torch.stack([maybe_dequantize_to_dtype(bias, dtype) for bias in biases], dim=0) + bias_data = bias_data.contiguous() + return GroupedTensorStorage( + shape=(num_groups, out_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(1, out_features)] * num_groups, + quantizer=None, + data=bias_data.reshape(-1), + ) + + +def get_grouped_tensor_members(grouped: GroupedTensor) -> list[torch.Tensor]: + """Return per-group tensor views from a ``GroupedTensor`` parameter.""" + + members = grouped.quantized_tensors + if members is None: + members = grouped.split_into_quantized_tensors() + return list(members) + + +def prepare_discrete_weights_for_grouped_gemm( + weights: Sequence[torch.Tensor | QuantizedTensorStorage], + weight_quantizers: Sequence[Optional[Quantizer]], + *, + with_quantized_compute: bool, + columnwise_usage: bool, + dtype: torch.dtype, + is_first_microbatch: Optional[bool] = None, + weight_workspaces: Optional[Sequence[Optional[QuantizedTensorStorage]]] = None, + cache_weight: bool = False, + skip_fp8_weight_update: Optional[torch.Tensor] = None, + use_quantize_weight: bool = True, + optimize_for_gemm: bool = True, +) -> tuple[list[torch.Tensor | QuantizedTensorStorage], list[Optional[QuantizedTensorStorage]]]: + """Prepare a list of per-group weights for grouped GEMM.""" + + new_workspaces: list[Optional[QuantizedTensorStorage]] = [None] * len(weights) + if not with_quantized_compute: + return [maybe_dequantize_to_dtype(weight, dtype) for weight in weights], new_workspaces + + out: list[torch.Tensor | QuantizedTensorStorage] = [] + update_ws = is_first_microbatch is None or is_first_microbatch + for idx, weight in enumerate(weights): + if is_quantized_tensor(weight) and not use_quantize_weight: + out.append(weight) + continue + quantizer = weight_quantizers[idx] + if quantizer is None: + raise ValueError("Missing quantizer for grouped-linear weight tensor") + quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + if optimize_for_gemm: + quantizer.optimize_for_gemm = True + if not use_quantize_weight: + out.append(quantizer(weight)) + continue + weight_fp8, new_workspaces[idx] = quantize_weight( + tensor=weight, + quantizer=quantizer, + workspace=weight_workspaces[idx] if weight_workspaces else None, + update_workspace=update_ws, + skip_update_flag=skip_fp8_weight_update, + workspace_dtype=dtype, + cache=cache_weight, + ) + out.append(weight_fp8) + return out, new_workspaces + + +def prepare_grouped_weight_for_grouped_gemm( + weight: GroupedTensor, + weight_quantizers: Sequence[Optional[Quantizer]], + *, + with_quantized_compute: bool, + columnwise_usage: bool, + dtype: torch.dtype, +) -> GroupedTensorStorage: + """Prepare one grouped weight parameter for GroupedTensor GEMM.""" + + is_weight_quantized = weight.quantizer is not None + if is_weight_quantized and with_quantized_compute: + return weight + if is_weight_quantized and not with_quantized_compute: + weight_parts = get_grouped_tensor_members(weight) + weight_data = torch.stack([maybe_dequantize_to_dtype(w, dtype) for w in weight_parts]) + weight_data = weight_data.contiguous() + return GroupedTensorStorage( + shape=(weight.num_tensors * weight.tensor_shapes[0][0], weight.tensor_shapes[0][1]), + dtype=dtype, + num_tensors=weight.num_tensors, + shapes=list(weight.tensor_shapes), + quantizer=None, + data=weight_data.reshape(-1), + ) + if not with_quantized_compute: + if weight.rowwise_data.dtype == dtype: + return weight + weight_data = weight.rowwise_data.to(dtype=dtype) + return GroupedTensorStorage( + shape=weight.logical_shape, + dtype=dtype, + num_tensors=weight.num_tensors, + shapes=list(weight.tensor_shapes), + quantizer=None, + data=weight_data.reshape(-1), + ) + + quantizer = weight_quantizers[0] + if quantizer is None: + raise ValueError("Missing quantizer for grouped-linear grouped weight") + quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + return tex.group_quantize( + weight.rowwise_data.view(weight.logical_shape), + quantizer, + weight.num_tensors, + None, + ) + + +def grouped_linear_forward_grouped_tensor( + input: torch.Tensor, # pylint: disable=redefined-builtin + weights: Sequence[torch.Tensor | QuantizedTensorStorage] | GroupedTensor, + split_sizes: torch.Tensor, + *, + biases: Optional[Sequence[torch.Tensor | QuantizedTensorStorage]] = None, + scales: Optional[torch.Tensor] = None, + input_quantizers: Sequence[Optional[Quantizer]], + weight_quantizers: Sequence[Optional[Quantizer]], + output_quantizers: Optional[Sequence[Optional[Quantizer]]] = None, # pylint: disable=unused-argument + with_quantized_compute: bool, + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + is_first_microbatch: Optional[bool] = None, + weight_workspaces: Optional[Sequence[Optional[QuantizedTensorStorage]]] = None, + cache_weight: bool = False, + skip_fp8_weight_update: Optional[torch.Tensor] = None, + use_quantize_weight: bool = True, + optimize_weight_for_gemm: bool = True, + fprop_use_split_accumulator: bool = _2X_ACC_FPROP, +) -> GroupedLinearForwardResult: + """Graph-safe grouped-linear forward using GroupedTensor metadata.""" + + if isinstance(weights, GroupedTensor): + num_groups = weights.num_tensors + out_features, in_features = weights.tensor_shapes[0] + device = weights.device + else: + num_groups = len(weights) + out_features, in_features = weights[0].size() + device = weights[0].device + + split_sizes = canonicalize_split_sizes(split_sizes, num_groups=num_groups, device=device) + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_points = base_split_offsets[1:].to(dtype=torch.int) + + original_shape = input.shape + x = maybe_dequantize_to_dtype(input, dtype).reshape(-1, in_features) + total_tokens = x.size(0) + + if with_quantized_compute: + input_quantizer = input_quantizers[0] + if input_quantizer is None: + raise ValueError("Missing quantizer for grouped-linear input tensor") + 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: + grouped_x = make_grouped_tensor_from_2d_buffer( + x, + num_groups=num_groups, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=in_features, + dtype=dtype, + ) + + if isinstance(weights, GroupedTensor): + grouped_weights: GroupedTensorStorage | list[torch.Tensor | QuantizedTensorStorage] + grouped_weights = prepare_grouped_weight_for_grouped_gemm( + weights, + weight_quantizers, + with_quantized_compute=with_quantized_compute, + columnwise_usage=input_requires_grad, + dtype=dtype, + ) + new_workspaces = [None] + else: + grouped_weights, new_workspaces = prepare_discrete_weights_for_grouped_gemm( + weights, + weight_quantizers, + with_quantized_compute=with_quantized_compute, + columnwise_usage=input_requires_grad, + dtype=dtype, + is_first_microbatch=is_first_microbatch, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + skip_fp8_weight_update=skip_fp8_weight_update, + use_quantize_weight=use_quantize_weight, + optimize_for_gemm=optimize_weight_for_gemm, + ) + + out = torch.empty((*original_shape[:-1], out_features), dtype=dtype, device=device) + grouped_out = make_grouped_tensor_from_2d_buffer( + out.reshape(-1, out_features), + num_groups=num_groups, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=out_features, + dtype=dtype, + ) + + grouped_bias = None + bias_scale = None + if biases is not None: + grouped_bias = make_grouped_bias( + biases, + num_groups=num_groups, + out_features=out_features, + dtype=dtype, + ) + if scales is not None: + bias_scale = scales.reshape(-1) + if bias_scale.dtype != torch.float32: + bias_scale = bias_scale.to(dtype=torch.float32) + + general_grouped_gemm_for_grouped_tensor( + grouped_weights, + grouped_x, + grouped_out, + layout="TN", + bias=grouped_bias, + bias_scale=bias_scale, + use_split_accumulator=fprop_use_split_accumulator, + ) + + if not input_requires_grad: + grouped_weights = None if isinstance(weights, GroupedTensor) else [None] * num_groups + if not weight_requires_grad: + grouped_x = None + if grouped_x is not None and with_quantized_compute and grouped_x.columnwise_data is not None: + grouped_x.rowwise_data = None + grouped_x.scale_inv = None + + saved: list[Any] = [split_sizes, base_split_offsets, split_points] + if scales is not None: + saved.append(scales) + saved.append(grouped_x) + if isinstance(weights, GroupedTensor): + saved.append(grouped_weights) + else: + saved.extend(grouped_weights) + + cache = { + "path": "grouped_tensor", + "saved_tensors": tuple(saved), + "has_scales": scales is not None, + "num_groups": num_groups, + "in_features": in_features, + "out_features": out_features, + "input_shape": original_shape, + "dtype": dtype, + "device": device, + "single_grouped_weight": isinstance(weights, GroupedTensor), + "has_bias": biases is not None, + } + return GroupedLinearForwardResult(out, cache, new_workspaces) + + +def grouped_linear_forward_split( + input: torch.Tensor, # pylint: disable=redefined-builtin + weights: Sequence[torch.Tensor | QuantizedTensorStorage], + split_sizes: torch.Tensor | Sequence[int], + *, + biases: Optional[Sequence[torch.Tensor | QuantizedTensorStorage]] = None, + scales: Optional[torch.Tensor] = None, + input_quantizers: Sequence[Optional[Quantizer]], + weight_quantizers: Sequence[Optional[Quantizer]], + output_quantizers: Optional[Sequence[Optional[Quantizer]]] = None, + with_quantized_compute: bool, + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + is_first_microbatch: Optional[bool] = None, + weight_workspaces: Optional[Sequence[Optional[QuantizedTensorStorage]]] = None, + cache_weight: bool = False, + skip_fp8_weight_update: Optional[torch.Tensor] = None, + debug_quantize_fn: Optional[Any] = None, + cpu_offloading: bool = False, + use_quantize_weight: bool = True, + optimize_weight_for_gemm: bool = True, + fprop_use_split_accumulator: bool = _2X_ACC_FPROP, +) -> GroupedLinearForwardResult: + """Legacy grouped-linear forward using split tensors and grouped GEMM.""" + + num_groups = len(weights) + out_features, in_features = weights[0].size() + device = weights[0].device + split_sizes = canonicalize_split_sizes(split_sizes, num_groups=num_groups) + split_sizes_int = split_sizes.tolist() + + x = maybe_dequantize_to_dtype(input, dtype).reshape(-1, in_features) + if with_quantized_compute: + if debug_quantize_fn is not None: + xs = debug_quantize_fn(x, input_quantizers, split_sizes_int, dtype) + else: + xs = tex.split_quantize(x, split_sizes_int, input_quantizers) + else: + xs = torch.split(x, split_sizes_int) + + if cpu_offloading: + from ..cpu_offload import start_offload # pylint: disable=import-outside-toplevel + + start_offload(*xs) + + ws, new_workspaces = prepare_discrete_weights_for_grouped_gemm( + weights, + weight_quantizers, + with_quantized_compute=with_quantized_compute, + columnwise_usage=input_requires_grad, + dtype=dtype, + is_first_microbatch=is_first_microbatch, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + skip_fp8_weight_update=skip_fp8_weight_update, + use_quantize_weight=use_quantize_weight, + optimize_for_gemm=optimize_weight_for_gemm, + ) + + bs = None + if biases is not None: + bs = [maybe_dequantize_to_dtype(bias, dtype) for bias in biases] + + out = torch.empty((*input.shape[:-1], out_features), dtype=dtype, device=device) + use_gemm_bias = bs is not None and scales is None + general_grouped_gemm( + ws, + xs, + [out], + output_quantizers or [None] * num_groups, + dtype, + single_output=True, + m_splits=split_sizes_int, + bias=bs if use_gemm_bias else None, + use_bias=use_gemm_bias, + use_split_accumulator=fprop_use_split_accumulator, + ) + + if scales is not None and bs is not None: + scale_splits = torch.split(scales, split_sizes_int) + out_splits = torch.split(out.reshape(-1, out_features), split_sizes_int) + for idx in range(num_groups): + out_splits[idx].add_(bs[idx].unsqueeze(0) * scale_splits[idx].unsqueeze(-1)) + + if not input_requires_grad: + ws = [None] * num_groups + elif with_quantized_compute: + for w, weight in zip(ws, weights): + if w is not weight and is_quantized_tensor(w): + w.update_usage(rowwise_usage=False, columnwise_usage=True) + + if not weight_requires_grad: + xs = [None] * num_groups + elif with_quantized_compute: + for x_i in xs: + if is_quantized_tensor(x_i): + x_i.update_usage(rowwise_usage=False, columnwise_usage=True) + + saved: list[Any] = [split_sizes, None, None] + if scales is not None: + saved.append(scales) + saved.extend(xs) + saved.extend(ws) + + cache = { + "path": "split", + "saved_tensors": tuple(saved), + "has_scales": scales is not None, + "num_groups": num_groups, + "in_features": in_features, + "out_features": out_features, + "input_shape": input.shape, + "dtype": dtype, + "device": device, + "single_grouped_weight": False, + "has_bias": biases is not None, + "split_sizes_int": split_sizes_int, + } + return GroupedLinearForwardResult(out, cache, new_workspaces) + + +def compute_grouped_linear_dbias( + grad_output_2d: torch.Tensor, + split_offsets: torch.Tensor, + *, + num_groups: int, + dtype: torch.dtype, + scales: Optional[torch.Tensor] = None, + biases: Optional[Sequence[torch.Tensor | QuantizedTensorStorage]] = None, +) -> tuple[list[torch.Tensor], Optional[torch.Tensor], torch.Tensor]: + """Compute grouped bias gradients, and optional scale gradients.""" + + grad_scales = None + if scales is not None: + if biases is None: + raise ValueError("Grouped bias tensors are required to compute scale gradients") + bias_packed = torch.stack([maybe_dequantize_to_dtype(bias, dtype) for bias in biases]) + dbias_packed, grad_scales = compute_grouped_dbias_dscales( + grad_output_2d, + scales.to(dtype=torch.float32), + bias_packed, + offsets=split_offsets, + ) + else: + dbias_packed = compute_grouped_dbias(grad_output_2d, split_offsets, num_groups) + return [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)], grad_scales, dbias_packed + + +def grouped_linear_backward_grouped_tensor( + grad_output: torch.Tensor, + grouped_x: Optional[GroupedTensorStorage], + weights: Sequence[torch.Tensor | QuantizedTensorStorage] | GroupedTensorStorage, + split_sizes: torch.Tensor, + base_split_offsets: torch.Tensor, + *, + num_groups: int, + in_features: int, + out_features: int, + dtype: torch.dtype, + device: torch.device, + with_quantized_compute: bool, + grad_output_quantizers: Sequence[Optional[Quantizer]], + input_requires_grad: bool, + weight_requires_grad: bool, + has_bias: bool, + biases: Optional[Sequence[torch.Tensor | QuantizedTensorStorage]] = None, + scales: Optional[torch.Tensor] = None, + wgrad_output: Optional[Sequence[torch.Tensor] | GroupedTensorStorage] = None, + wgrad_store: Optional[Any] = None, + wgrad_store_return: Optional[Any] = None, + accumulate_wgrad: bool = False, + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD, + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD, + cast_grad_output_before_quantize: bool = False, +) -> dict[str, Any]: + """GroupedTensor backward compute for grouped linear. + + Caller owns parameter-order concerns, main-grad/dummy-grad policy, and + whether ``wgrad_output`` points at ordinary grad buffers or parameter + ``main_grad`` buffers. + """ + + dy_2d = grad_output.contiguous().view(-1, out_features) + if cast_grad_output_before_quantize or not with_quantized_compute: + dy_2d = maybe_dequantize_to_dtype(dy_2d, dtype) + + dbias_packed = None + if with_quantized_compute: + grad_output_quantizer = grad_output_quantizers[0] + if grad_output_quantizer is None: + raise ValueError("Missing quantizer for grouped-linear grad output") + grad_output_quantizer.set_usage(rowwise=input_requires_grad, columnwise=weight_requires_grad) + grad_output_quantizer.optimize_for_gemm = True + if has_bias and scales is None and isinstance(grad_output_quantizer, MXFP8Quantizer): + 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 = make_grouped_tensor_from_2d_buffer( + dy_2d, + num_groups=num_groups, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=out_features, + dtype=dtype, + ) + + grad_biases = [None] * num_groups + grad_scales = None + if has_bias: + if dbias_packed is None: + grad_biases, grad_scales, dbias_packed = compute_grouped_linear_dbias( + dy_2d, + base_split_offsets, + num_groups=num_groups, + dtype=dtype, + scales=scales, + biases=biases, + ) + else: + grad_biases = [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)] + + grad_input = None + if input_requires_grad: + weight_iter = weights if isinstance(weights, (list, tuple)) else [weights] + for weight in weight_iter: + if isinstance(weight, QuantizedTensorStorage): + weight.update_usage(columnwise_usage=True) + grad_input = torch.empty( + (dy_2d.size(0), in_features), + dtype=dtype, + device=device, + ) + grouped_grad_input = make_grouped_tensor_from_2d_buffer( + grad_input, + num_groups=num_groups, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=in_features, + dtype=dtype, + ) + general_grouped_gemm_for_grouped_tensor( + weights, + grouped_dy, + grouped_grad_input, + layout="NN", + use_split_accumulator=dgrad_use_split_accumulator, + ) + + if weight_requires_grad: + if grouped_x is None or wgrad_output is None: + raise ValueError("Grouped-linear wgrad requires saved input and output buffers") + def wgrad_gemm(inputmats, grad_output_mats, grad_weights): + general_grouped_gemm_for_grouped_tensor( + inputmats, + grad_output_mats, + grad_weights, + layout="NT", + accumulate=accumulate_wgrad, + use_split_accumulator=wgrad_use_split_accumulator, + ) + return wgrad_store_return + + if wgrad_store is not None and wgrad_store.delay_wgrad_compute(): + wgrad_store.put([grouped_x, grouped_dy, wgrad_output], wgrad_gemm) + else: + wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + + return { + "grad_input": grad_input, + "grad_biases": grad_biases, + "grad_scales": grad_scales, + "dbias_packed": dbias_packed, + "grouped_dy": grouped_dy, + } diff --git a/transformer_engine/pytorch/functional.py b/transformer_engine/pytorch/functional.py new file mode 100644 index 0000000000..89e11eeb8b --- /dev/null +++ b/transformer_engine/pytorch/functional.py @@ -0,0 +1,108 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Functional Transformer Engine PyTorch APIs.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +from ._functional.grouped_linear import ( + GroupedLinearForwardResult, + grouped_linear_forward_grouped_tensor, + grouped_linear_forward_split, + grouped_tensor_path_is_supported, +) +from .quantized_tensor import QuantizedTensorStorage +from .tensor import GroupedTensor, Quantizer + +__all__ = ["grouped_linear"] + + +def grouped_linear( + input: torch.Tensor, # pylint: disable=redefined-builtin + weights: Sequence[torch.Tensor | QuantizedTensorStorage] | GroupedTensor, + split_sizes: torch.Tensor | Sequence[int], + bias: Optional[Sequence[torch.Tensor | QuantizedTensorStorage]] = None, + *, + scales: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + input_quantizers: Optional[Sequence[Optional[Quantizer]]] = None, + weight_quantizers: Optional[Sequence[Optional[Quantizer]]] = None, + output_quantizers: Optional[Sequence[Optional[Quantizer]]] = None, + with_quantized_compute: bool = False, + use_grouped_tensor_path: Optional[bool] = None, + return_cache: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, dict]: + """Apply grouped linear transformations. + + This is equivalent to splitting ``input`` along its first dimension, + applying ``torch.nn.functional.linear`` with one weight and optional bias + per split, and concatenating the results. When ``return_cache=True``, a + dictionary of forward state is returned for advanced callers that implement + a paired custom backward. + """ + + if isinstance(weights, GroupedTensor): + num_groups = weights.num_tensors + default_dtype = weights.dtype + else: + num_groups = len(weights) + default_dtype = weights[0].dtype + dtype = default_dtype if dtype is None else dtype + input_quantizers = input_quantizers or [None] * num_groups + weight_quantizers = weight_quantizers or [None] * num_groups + output_quantizers = output_quantizers or [None] * num_groups + + if use_grouped_tensor_path is None: + use_grouped_tensor_path = grouped_tensor_path_is_supported( + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + output_quantizers=output_quantizers, + dtype=dtype, + single_grouped_weight=isinstance(weights, GroupedTensor), + ) + + if use_grouped_tensor_path: + result = grouped_linear_forward_grouped_tensor( + input, + weights, + split_sizes, + biases=bias, + scales=scales, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + output_quantizers=output_quantizers, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + input_requires_grad=input.requires_grad, + weight_requires_grad=( + weights.requires_grad + if isinstance(weights, GroupedTensor) + else any(weight.requires_grad for weight in weights) + ), + ) + else: + if isinstance(weights, GroupedTensor): + weights = list(weights.quantized_tensors or weights.split_into_quantized_tensors()) + result = grouped_linear_forward_split( + input, + weights, + split_sizes, + biases=bias, + scales=scales, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + output_quantizers=output_quantizers, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + input_requires_grad=input.requires_grad, + weight_requires_grad=any(weight.requires_grad for weight in weights), + ) + + if return_cache: + return result.output, result.cache + return result.output diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 82ee7953c1..e4110c3cec 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -6,7 +6,6 @@ from typing import Union, Optional, Callable, Tuple, List from itertools import chain -import os import warnings import weakref @@ -34,7 +33,6 @@ divide, cast_if_needed, clear_tensor_data, - get_device_compute_capability, init_method_constant, requires_grad, resolve_grouped_linear_single_param_flags, @@ -48,14 +46,21 @@ ) from ..cpp_extensions import ( general_grouped_gemm, - general_grouped_gemm_for_grouped_tensor, +) +from .._functional.grouped_linear import ( + grouped_linear_backward_grouped_tensor, + grouped_linear_forward_grouped_tensor, + grouped_tensor_path_is_supported, + make_grouped_bias, + make_grouped_tensor_from_2d_buffer, + maybe_dequantize_to_dtype, + prepare_discrete_weights_for_grouped_gemm, ) from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload -from ..triton.grouped_dbias_dscales import compute_grouped_dbias -from ..tensor import Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer +from ..tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -79,9 +84,7 @@ def _maybe_dequantize( dtype: torch.dtype, ) -> torch.Tensor: """Dequantize quantized tensors or cast regular tensors to ``dtype``.""" - if isinstance(tensor, QuantizedTensorStorage): - return tensor.dequantize(dtype=dtype) - return cast_if_needed(tensor, dtype) + return maybe_dequantize_to_dtype(tensor, dtype) @staticmethod def _is_grouped_tensor_path_supported( @@ -116,44 +119,18 @@ def _is_grouped_tensor_path_supported( Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - # 1. Filter by environment variable - if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): - return False - # 2. Filter out advanced features - if ( - debug - or cpu_offloading - or fp8_calibration - or backward_override is not None - or save_original_input - ): - return False - # 3. Filter by compute capability and cuBLAS version - device_capability = get_device_compute_capability() - if not (9, 0) <= device_capability <= (11, 0): - return False - cublaslt_version = tex.get_cublasLt_version() - if cublaslt_version < 130300: - return False - if device_capability < (10, 0) and cublaslt_version < 130400: - return False - # 4. Output quantization is not supported. - if any(q is not None for q in output_quantizers): - return False - # 5. Filter by quantization recipes. - if fp8: - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+. - if device_capability < (10, 0) and cublaslt_version < 130500: - return False - return True - # MXFP8 and NVFP4 require Blackwell+. - if not (10, 0) <= device_capability <= (11, 0): - return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) - return activation_dtype in (torch.bfloat16, torch.float16) + return grouped_tensor_path_is_supported( + with_quantized_compute=fp8, + input_quantizers=input_quantizers, + output_quantizers=output_quantizers, + dtype=activation_dtype, + require_env=True, + fp8_calibration=fp8_calibration, + debug=debug, + cpu_offloading=cpu_offloading, + backward_override=backward_override, + save_original_input=save_original_input, + ) @staticmethod def _make_grouped_tensor( @@ -166,14 +143,13 @@ def _make_grouped_tensor( dtype: torch.dtype, ) -> GroupedTensorStorage: """Wrap a packed 2D buffer as a varying-first-dimension GroupedTensorStorage.""" - return GroupedTensorStorage( - shape=(data.size(0), last_dim), + return make_grouped_tensor_from_2d_buffer( + data, + num_groups=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=last_dim, dtype=dtype, - num_tensors=num_gemms, - quantizer=None, - data=data.reshape(-1), - first_dims=split_sizes, - tensor_offsets=base_split_offsets * last_dim, ) @staticmethod @@ -185,17 +161,11 @@ def _make_grouped_bias( dtype: torch.dtype, ) -> GroupedTensorStorage: """Pack per-GEMM biases into the grouped GEMM bias format.""" - bias_data = torch.stack( - [_GroupedLinear._maybe_dequantize(bias, dtype) for bias in biases], - dim=0, - ).contiguous() - return GroupedTensorStorage( - shape=(num_gemms, out_features), + return make_grouped_bias( + biases, + num_groups=num_gemms, + out_features=out_features, dtype=dtype, - num_tensors=num_gemms, - shapes=[(1, out_features)] * num_gemms, - quantizer=None, - data=bias_data.reshape(-1), ) @staticmethod @@ -212,29 +182,17 @@ def _prepare_weights_for_grouped_tensor_gemm( cache_weight: bool, ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]: """Prepare discrete weight tensors for GroupedTensor GEMM.""" - weights_for_gemm: List[torch.Tensor] = [] - new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights) - if not with_quantized_compute: - return ( - [_GroupedLinear._maybe_dequantize(weight, activation_dtype) for weight in weights], - new_workspaces, - ) - - update_ws = is_first_microbatch is None or is_first_microbatch - for idx, weight in enumerate(weights): - weight_quantizer = weight_quantizers[idx] - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - weight_fp8, new_workspaces[idx] = quantize_weight( - tensor=weight, - quantizer=weight_quantizer, - workspace=weight_workspaces[idx] if weight_workspaces else None, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - weights_for_gemm.append(weight_fp8) - return weights_for_gemm, new_workspaces + return prepare_discrete_weights_for_grouped_gemm( + weights, + weight_quantizers, + with_quantized_compute=with_quantized_compute, + columnwise_usage=columnwise_usage, + dtype=activation_dtype, + is_first_microbatch=is_first_microbatch, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + skip_fp8_weight_update=skip_fp8_weight_update, + ) @staticmethod def _forward_grouped_tensor( @@ -267,94 +225,41 @@ def _forward_grouped_tensor( out_features = weights[0].size(0) weight_requires_grad = weights[0].requires_grad - split_sizes = m_splits.to(device=device) - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) - - inp_view = inp.reshape(-1, in_features) - x = cast_if_needed(inp_view, activation_dtype) - if fp8: - input_quantizer = input_quantizers[0] - input_quantizer.set_usage( - rowwise=True, - columnwise=is_grad_enabled and weight_requires_grad, - ) - input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes) - else: - grouped_x = _GroupedLinear._make_grouped_tensor( - x, - num_gemms=num_gemms, - split_sizes=split_sizes, - base_split_offsets=base_split_offsets, - last_dim=in_features, - dtype=activation_dtype, - ) - - columnwise_usage = is_grad_enabled and inp.requires_grad - weights_for_gemm, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( - weights, - weight_quantizers, - weight_workspaces, - with_quantized_compute=fp8, - columnwise_usage=columnwise_usage, - activation_dtype=activation_dtype, - is_first_microbatch=is_first_microbatch, - skip_fp8_weight_update=skip_fp8_weight_update, - cache_weight=cache_weight, - ) - - out = torch.empty( - [x.size(0), out_features], - dtype=activation_dtype, - device=device, - ) - grouped_out = _GroupedLinear._make_grouped_tensor( - out, - num_gemms=num_gemms, - split_sizes=split_sizes, - base_split_offsets=base_split_offsets, - last_dim=out_features, - dtype=activation_dtype, - ) - - grouped_bias = None - if use_bias: - grouped_bias = _GroupedLinear._make_grouped_bias( - biases, - num_gemms=num_gemms, - out_features=out_features, - dtype=activation_dtype, - ) - use_split_accumulator = _2X_ACC_FPROP if fp8: recipe = FP8GlobalStateManager.get_fp8_recipe() if hasattr(recipe, "fp8_gemm_fprop"): use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator - general_grouped_gemm_for_grouped_tensor( - weights_for_gemm, - grouped_x, - grouped_out, - layout="TN", - bias=grouped_bias, - use_split_accumulator=use_split_accumulator, + result = grouped_linear_forward_grouped_tensor( + inp, + weights, + m_splits, + biases=biases if use_bias else None, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + with_quantized_compute=fp8, + dtype=activation_dtype, + input_requires_grad=is_grad_enabled and inp.requires_grad, + weight_requires_grad=is_grad_enabled and weight_requires_grad, + is_first_microbatch=is_first_microbatch, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + skip_fp8_weight_update=skip_fp8_weight_update, + fprop_use_split_accumulator=use_split_accumulator, ) + out = result.output + new_workspaces = result.new_weight_workspaces + saved_tensors = result.cache["saved_tensors"] + split_sizes = saved_tensors[0] + base_split_offsets = saved_tensors[1] + grouped_x = saved_tensors[3] + weights_for_gemm = saved_tensors[4:] if is_grad_enabled: - if weight_requires_grad: - # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data - # in backward pass) - if fp8 and grouped_x.columnwise_data is not None: - grouped_x.rowwise_data = None - grouped_x.scale_inv = None - else: - grouped_x = None - - weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms tensors_to_save, tensor_objects = prepare_for_saving( grouped_x, - *weights_to_save, + *weights_for_gemm, split_sizes, base_split_offsets, ) @@ -763,77 +668,6 @@ def _backward_grouped_tensor( if main_grad is not None: origin_weight.main_grad = main_grad - grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) - dy_2d = cast_if_needed(grad_output_view, ctx.activation_dtype) - dbias_packed = None - if ctx.fp8: - grad_output_quantizer = ctx.grad_output_quantizers[0] - grad_output_quantizer.set_usage( - rowwise=ctx.requires_dgrad, - columnwise=ctx.weights_requires_grad, - ) - grad_output_quantizer.optimize_for_gemm = True - if ctx.use_bias and isinstance(grad_output_quantizer, MXFP8Quantizer): - grouped_dy, dbias_packed = tex.bgrad_group_quantize( - dy_2d, - grad_output_quantizer, - N, - split_sizes, - ) - else: - grouped_dy = tex.group_quantize( - dy_2d, - grad_output_quantizer, - N, - split_sizes, - ) - else: - grouped_dy = _GroupedLinear._make_grouped_tensor( - dy_2d, - num_gemms=N, - split_sizes=split_sizes, - base_split_offsets=base_split_offsets, - last_dim=ctx.weights_shape_0, - dtype=ctx.activation_dtype, - ) - - grad_biases = [None] * N - if ctx.use_bias: - if dbias_packed is None: - dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N) - grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] - - dgrad = None - if ctx.requires_dgrad: - dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - for weight in weights: - if isinstance(weight, QuantizedTensorStorage): - weight.update_usage(columnwise_usage=True) - dgrad = torch.empty( - (dy_2d.size(0), ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, - ) - grouped_dgrad = _GroupedLinear._make_grouped_tensor( - dgrad, - num_gemms=N, - split_sizes=split_sizes, - base_split_offsets=base_split_offsets, - last_dim=ctx.weights_shape_1, - dtype=ctx.activation_dtype, - ) - general_grouped_gemm_for_grouped_tensor( - weights, - grouped_dy, - grouped_dgrad, - layout="NN", - use_split_accumulator=dgrad_gemm_use_split_accumulator, - ) - if ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch @@ -841,12 +675,20 @@ def _backward_grouped_tensor( else: accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD + if ctx.requires_dgrad and ctx.fp8: + recipe = ctx.fp8_recipe + if hasattr(recipe, "fp8_gemm_dgrad"): + dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator + + wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD + if ctx.weights_requires_grad and ctx.fp8: + recipe = ctx.fp8_recipe + if hasattr(recipe, "fp8_gemm_wgrad"): + wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator + + wgrad_list = [None] * N if ctx.weights_requires_grad: - wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator if ctx.fuse_wgrad_accumulation: wgrad_list = main_grads else: @@ -865,22 +707,6 @@ def _backward_grouped_tensor( else False ) - def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights): - general_grouped_gemm_for_grouped_tensor( - inputmats, - grad_output_mats, - grad_weights, - layout="NT", - use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=accumulate, - ) - return None, [None] * N, None - - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad) - else: - grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list) - def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: if ctx.fuse_wgrad_accumulation and hasattr(weight, "grad_added_to_main_grad"): @@ -902,12 +728,41 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): wgrad = None return wgrad + else: + accumulate = False + handle_custom_ddp_from_mcore = None + + backward_result = grouped_linear_backward_grouped_tensor( + grad_output, + grouped_x, + weights, + split_sizes, + base_split_offsets, + num_groups=N, + in_features=ctx.weights_shape_1, + out_features=ctx.weights_shape_0, + dtype=ctx.activation_dtype, + device=ctx.device, + with_quantized_compute=ctx.fp8, + grad_output_quantizers=ctx.grad_output_quantizers, + input_requires_grad=ctx.requires_dgrad, + weight_requires_grad=ctx.weights_requires_grad, + has_bias=ctx.use_bias, + wgrad_output=wgrad_list if ctx.weights_requires_grad else None, + wgrad_store=ctx.wgrad_store, + wgrad_store_return=(None, [None] * N, None), + accumulate_wgrad=accumulate, + dgrad_use_split_accumulator=dgrad_gemm_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_gemm_use_split_accumulator, + cast_grad_output_before_quantize=True, + ) + dgrad = backward_result["grad_input"] + grad_biases = backward_result["grad_biases"] + if ctx.weights_requires_grad: wgrad_list = [ handle_custom_ddp_from_mcore(weight, main_grad, wgrad) for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) ] - else: - wgrad_list = [None] * N if not ctx.use_bias: grad_biases = [None] * N diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 5c96d4658e..3c037ceef5 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -15,7 +15,15 @@ import transformer_engine_torch as tex from ...constants import DType -from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor +from ...cpp_extensions import general_grouped_gemm +from ..._functional.grouped_linear import ( + compute_grouped_linear_dbias, + grouped_linear_backward_grouped_tensor, + grouped_linear_forward_grouped_tensor, + grouped_linear_forward_split, + grouped_tensor_path_is_supported, + prepare_grouped_weight_for_grouped_gemm, +) from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore from ...module.base import ( @@ -27,10 +35,8 @@ from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...quantized_tensor import QuantizedTensorStorage from ...tensor import ( - Float8CurrentScalingQuantizer, MXFP8Quantizer, MXFP8Tensor, - NVFP4Quantizer, Quantizer, ) from ...utils import ( @@ -38,7 +44,6 @@ canonicalize_dtype, clear_tensor_data, devices_match, - get_device_compute_capability, resolve_grouped_linear_single_param_flags, round_up_to_nearest_multiple, ) @@ -52,10 +57,6 @@ ) from ..op import BasicOperation, OperationContext from ...tensor import GroupedTensor, GroupedTensorStorage -from ...triton.grouped_dbias_dscales import ( - compute_grouped_dbias, - compute_grouped_dbias_dscales, -) class GroupedLinear(BasicOperation): @@ -786,30 +787,12 @@ def _is_graph_safe_path_supported( * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - if not (9, 0) <= get_device_compute_capability() <= (11, 0): - return False - if with_quantized_compute: - # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM - # path; the compute-capability range was already checked above. On Hopper it - # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS. - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - if ( - get_device_compute_capability() < (10, 0) - and tex.get_cublasLt_version() < 130500 - ): - return False - return True - # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. - if not (10, 0) <= get_device_compute_capability() <= (11, 0): - return False - if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): - return True - # NVFP4 graph-safe grouped quantization requires RHT and only supports - # discrete weights; otherwise fall back to the split-quantize flow. - if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): - return not single_grouped_weight - return False - return dtype in (torch.bfloat16, torch.float16) + return grouped_tensor_path_is_supported( + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + ) def _get_grouped_weight_for_gemm( self, @@ -822,49 +805,12 @@ def _get_grouped_weight_for_gemm( """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 GroupedTensorStorage( - 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 GroupedTensorStorage( - 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, + return prepare_grouped_weight_for_grouped_gemm( + weight_param, + weight_quantizers, + columnwise_usage=columnwise_usage, + with_quantized_compute=with_quantized_compute, + dtype=dtype, ) def _get_discrete_weights_for_gemm( @@ -1132,10 +1078,6 @@ def _fuser_forward_split_quantize( ) -> 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 - - # 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: @@ -1144,87 +1086,24 @@ def _fuser_forward_split_quantize( weights = self.weight.split_into_quantized_tensors() else: weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] - bs = None - if has_bias: - bs = self._get_bias_tensors(dtype) - - ws = self._get_discrete_weights_for_gemm( + result = grouped_linear_forward_split( + input_, weights, - weight_quantizers, - columnwise_usage=input_requires_grad, + split_sizes, + biases=self._get_bias_tensors(dtype) if self.has_bias else None, + scales=scales if self._scale_bias else None, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + output_quantizers=[None] * num_groups, with_quantized_compute=with_quantized_compute, dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + cpu_offloading=is_cpu_offload_enabled(), + use_quantize_weight=False, + optimize_weight_for_gemm=False, ) - - # Split input tensor and convert dtypes if needed - x = maybe_dequantize(input_, dtype) - xs = None - if with_quantized_compute: - for quantizer in input_quantizers: - quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) - xs = tex.split_quantize(x, split_sizes_int, input_quantizers) - else: - xs = torch.split(x, split_sizes_int) - if is_cpu_offload_enabled(): - live_xs = [t for t in xs if t is not None] - if live_xs: - start_offload(*live_xs) - - # Allocate output tensor - in_shape = list(input_.size()) - out_shape = in_shape[:-1] + [self.out_features] - out = torch.empty(out_shape, dtype=dtype, device=device) - - # Perform GEMMs - use_gemm_bias = has_bias and not self._scale_bias - general_grouped_gemm( - ws, - xs, - [out], - [None] * num_groups, # quantization_params - dtype, - m_splits=split_sizes_int, - bias=bs if use_gemm_bias else None, - use_bias=use_gemm_bias, - use_split_accumulator=_2X_ACC_FPROP, - single_output=True, - ) - - # Add bias * scales when scale_bias is enabled - if self._scale_bias and has_bias: - scales_splits = torch.split(scales, split_sizes_int) - out_splits = torch.split(out, split_sizes_int) - for i in range(num_groups): - out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) - - # Prepare weight tensors for backward pass - if not input_requires_grad: - ws = [None] * num_groups - elif with_quantized_compute: - for w, weight_param in zip(ws, weights): - if w is not weight_param: - w.update_usage(rowwise_usage=False, columnwise_usage=True) - - # Prepare input tensor for backward pass - if not weight_requires_grad: - xs = [None] * num_groups - elif with_quantized_compute: - for x in xs: - x.update_usage(rowwise_usage=False, columnwise_usage=True) - - # 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) + return result.output, result.cache["saved_tensors"] def _fuser_forward_grouped_tensor( self, @@ -1246,117 +1125,34 @@ def _fuser_forward_grouped_tensor( ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them. """ num_groups = self.num_groups - has_bias = self.has_bias - - 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()) - x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) - total_tokens = x.size(0) - - # Build the input GroupedTensor. - if with_quantized_compute: - 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 = GroupedTensorStorage( - 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_split_offsets * self.in_features, - ) - - if is_cpu_offload_enabled() and grouped_x is not None: - start_offload(grouped_x) - - # Build the weight GroupedTensor / list. + biases = self._get_bias_tensors(dtype) if self.has_bias else None if self.single_grouped_weight: - # GroupedTensor - 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, - ) + weights = self.weight else: - # Discrete weights - 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, - ) + weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] - # 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 = GroupedTensorStorage( - shape=(total_tokens, self.out_features), + result = grouped_linear_forward_grouped_tensor( + input_, + weights, + split_sizes, + biases=biases, + scales=scales if self._scale_bias else None, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + with_quantized_compute=with_quantized_compute, dtype=dtype, - num_tensors=num_groups, - quantizer=None, - data=out.reshape(-1), - first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + use_quantize_weight=False, + optimize_weight_for_gemm=False, ) - # 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: - # Bias always needs to be passed as a GroupedTensor for the grouped GEMM. - grouped_bias = self._get_grouped_bias_for_gemm(dtype) - 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. - 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, - ) - - if not input_requires_grad: - grouped_weights = None if self.single_grouped_weight else [None] * num_groups - - if not weight_requires_grad: - grouped_x = None - - # 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: - # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data - # in backward pass) - if with_quantized_compute and grouped_x.columnwise_data is not None: - 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) + saved = result.cache["saved_tensors"] + offset = 4 if self._scale_bias else 3 + grouped_x = saved[offset] + if is_cpu_offload_enabled() and grouped_x is not None: + start_offload(grouped_x) + return result.output, saved def fuser_backward( self, @@ -1431,18 +1227,14 @@ def _fuser_backward_split_quantize( dy_2d = dy.reshape(-1, dy.size(-1)) offsets = torch.zeros(num_groups + 1, dtype=torch.int64, device=device) offsets[1:] = split_sizes.cumsum(0) - if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype)) - scales_f32 = scales.to(dtype=torch.float32) - dbias_packed, grad_scales = compute_grouped_dbias_dscales( - dy_2d, - scales_f32, - bias_packed, - offsets=offsets, - ) - else: - 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)] + grad_biases, grad_scales, _ = compute_grouped_linear_dbias( + dy_2d, + offsets, + num_groups=num_groups, + dtype=ctx.dtype, + scales=scales if self._scale_bias else None, + biases=self._get_bias_tensors(ctx.dtype) if self._scale_bias else None, + ) # Initialize grad weight buffers. accumulate_into_main_grad = self._accumulate_into_main_grad @@ -1603,89 +1395,6 @@ def _fuser_backward_grouped_tensor( 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. - dy_2d = grad_output.reshape(-1, self.out_features) - total_tokens = dy_2d.size(0) - - # Build the grad_output GroupedTensor. - # 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( - rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad - ) - grad_output_quantizer.optimize_for_gemm = True - - if ( - has_bias - and not self._scale_bias - and isinstance(grad_output_quantizer, MXFP8Quantizer) - ): - 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: - dy_2d = maybe_dequantize(dy_2d, dtype) - # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm - grouped_dy = GroupedTensorStorage( - 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_split_offsets * self.out_features, - ) - - # 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)) - scales_f32 = scales.to(dtype=torch.float32) - dbias_packed, grad_scales = compute_grouped_dbias_dscales( - dy_2d, - scales_f32, - bias_packed, - offsets=base_split_offsets, - ) - elif dbias_packed is None: - # BF16/FP16 path - 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: - final_bias_grads = [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)] - - # ---- dgrad GEMM ---------------------------------------------------- - 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 = GroupedTensorStorage( - 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_split_offsets * self.in_features, - ) - general_grouped_gemm_for_grouped_tensor( - ws, - grouped_dy, - grouped_grad_input, - layout="NN", - use_split_accumulator=_2X_ACC_DGRAD, - ) - # params init for wgrad GEMM accumulate_into_main_grad = False weight_shape = (self.out_features, self.in_features) @@ -1737,23 +1446,45 @@ def _fuser_backward_grouped_tensor( ] wgrad_output = final_weight_grads - # wgrad GEMM 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) + + backward_result = grouped_linear_backward_grouped_tensor( + grad_output, + grouped_x, + ws, + split_sizes, + base_split_offsets, + num_groups=num_groups, + in_features=self.in_features, + out_features=self.out_features, + dtype=dtype, + device=device, + with_quantized_compute=with_quantized_compute, + grad_output_quantizers=ctx.grad_output_quantizers, + input_requires_grad=ctx.input_requires_grad, + weight_requires_grad=ctx.weight_requires_grad, + has_bias=has_bias, + biases=self._get_bias_tensors(dtype) if self._scale_bias else None, + scales=scales if self._scale_bias else None, + wgrad_output=wgrad_output, + wgrad_store=self.wgrad_store, + accumulate_wgrad=accumulate_into_main_grad, + ) + grad_input = backward_result["grad_input"] + if grad_input is not None: + grad_input = grad_input.view(*grad_output.size()[:-1], self.in_features) + grad_scales = backward_result["grad_scales"] + dbias_packed = backward_result["dbias_packed"] + final_bias_grads: Optional[list[torch.Tensor]] = None + if has_bias: + if self.single_grouped_bias: + final_bias_grads = [dbias_packed.to(dtype=dtype)] else: - wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + final_bias_grads = backward_result["grad_biases"] # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, # signal that ``main_grad`` already carries the wgrad and replace