diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 76550db5f8..0d97288e98 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -15,6 +15,7 @@ import torch import transformer_engine.pytorch as te +from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch.ops.fused.grouped_mlp import ( _cudnn_frontend_supports_grouped_gemm_srelu, _cudnn_frontend_version_supported, @@ -229,6 +230,23 @@ def make_reference_and_test_tensors( return ref, test +class _InjectGrad(torch.autograd.Function): + """Replace the gradient flowing into ``x`` with ``grad``. + + Mirrors how an FP8 token dispatch delivers a pre-quantized ``GroupedTensor`` + grad output: the downstream op simply returns one as its grad input. + """ + + @staticmethod + def forward(ctx, x, grad): # pylint: disable=arguments-differ + ctx.injected_grad = grad + return x + + @staticmethod + def backward(ctx, grad_output): # pylint: disable=arguments-differ + return ctx.injected_grad, None + + class TestGroupedLinearOp: """Tests for advanced features with grouped linear basic op""" @@ -436,6 +454,189 @@ def test_grouped_linear( else: assert b_test.grad is None + @staticmethod + def _make_rowwise_mxfp8_wire_input( + x_hp: torch.Tensor, + group_size: int, + split_sizes: torch.Tensor, + ) -> "GroupedTensor": + """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format).""" + wire_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False + ) + x_wire = tex.group_quantize( + x_hp, wire_quantizer, group_size, split_sizes.to(dtype=torch.int64) + ) + # Sanity: the wire tensor is rowwise-only with unswizzled scales. + assert x_wire.columnwise_data is None + assert not x_wire._with_gemm_swizzled_scales + return x_wire + + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear_prequantized_mxfp8_input( + self, + *, + group_size: int = 4, + weight_shape: tuple[int, int] = (256, 256), + split_alignment: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + weight_requires_grad: bool, + ) -> None: + """Rowwise-only MXFP8 GroupedTensor input (FP8 token dispatch wire format). + + The input arrives already rowwise-quantized with compact scales. The op + must feed the rowwise data to the forward GEMM as-is and manufacture the + columnwise copy for the wgrad GEMM. The reference run consumes the + *dequantized* wire tensor (the only data a layer can see after FP8 + dispatch) on the normal quantize-from-BF16 path; because MXFP8 requant + is idempotent along the rowwise axis and both paths derive the + columnwise copy from the same dequantized data, the two runs must match + bit-for-bit. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype) + + # Split sizes (including an empty group) + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + out_features, in_features = weight_shape + total_tokens = int(split_sizes.sum().item()) + in_shape = (total_tokens, in_features) + + # Wire-format input and its exact dequantization (the reference input). + x_hp = torch.rand(in_shape, dtype=dtype, device=device) - 0.5 + x_wire = self._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes) + x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(in_shape) + + dy = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5 + recipe = make_recipe("mxfp8") + op = te.ops.GroupedLinear( + group_size, in_features, out_features, bias=False, device=device, dtype=dtype + ) + with torch.no_grad(): + for param in op.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + + def _run(x): + with te.autocast(enabled=True, recipe=recipe): + y = op(x, split_sizes) + wgrads = [] + if weight_requires_grad: + y.backward(dy) + for group_idx in range(group_size): + weight = getattr(op, f"weight{group_idx}") + wgrads.append(weight.grad.detach().clone()) + weight.grad = None + return y.detach(), wgrads + + y_ref, wgrads_ref = _run(x_ref) + y_test, wgrads_test = _run(x_wire) + + # Bit-exact match expected (identical quantized inputs and kernels). + torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(wgrads_test, wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + + @pytest.mark.parametrize("bias_mode", ("none", "plain", "scaled")) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear_prequantized_mxfp8_grad( + self, + *, + group_size: int = 4, + weight_shape: tuple[int, int] = (256, 256), + split_alignment: int = 128, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + bias_mode: str, + weight_requires_grad: bool, + ) -> None: + """Rowwise-only MXFP8 GroupedTensor grad output (FP8 token dispatch, backward). + + Mirrors ``test_grouped_linear_prequantized_mxfp8_input`` on the backward + side: the rowwise data feeds the dgrad GEMM and the columnwise copy is + manufactured for wgrad. Covers all three bias gradient sources: none, + ``plain`` (fused into the columnwise stage of the quantize kernel, or + reduced from the dequantized grad when frozen weights leave no columnwise + stage), and ``scaled`` (``scale_bias``, whose dbias/dscales need the + dequantized grad because they depend on the routing probabilities). + + With frozen weights TE also requires frozen biases, so bias gradients are + not observable there; ``dscales`` still is, since the probabilities are an + input rather than a parameter. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype) + + has_bias = bias_mode != "none" + use_scale_bias = bias_mode == "scaled" + + # Split sizes (including an empty group) + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + out_features, in_features = weight_shape + total_tokens = int(split_sizes.sum().item()) + + x = torch.rand((total_tokens, in_features), dtype=dtype, device=device) - 0.5 + dy_hp = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + + # Wire-format grad output and its exact dequantization (the reference grad). + dy_wire = self._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes) + dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, out_features + ) + + recipe = make_recipe("mxfp8") + op = te.ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=has_bias, + device=device, + dtype=dtype, + scale_bias=use_scale_bias, + ) + if not weight_requires_grad: + # TE requires bias.requires_grad to match weight.requires_grad. + for param in op.parameters(): + param.requires_grad_(False) + + def _run(grad): + x_in = x.detach().clone().requires_grad_() + probs_in = probs.detach().clone().requires_grad_(use_scale_bias) + extra_inputs = (split_sizes, probs_in) if use_scale_bias else (split_sizes,) + with te.autocast(enabled=True, recipe=recipe): + y = op(x_in, *extra_inputs) + # Deliver ``grad`` as the op's grad output, as FP8 dispatch would. + _InjectGrad.apply(y, grad).backward(torch.ones_like(y)) + grads = [("dx", x_in.grad)] + if use_scale_bias: + grads.append(("dprobs", probs_in.grad)) + if weight_requires_grad: + for group_idx in range(group_size): + weight = getattr(op, f"weight{group_idx}") + grads.append((f"w{group_idx}", weight.grad.detach().clone())) + weight.grad = None + if has_bias: + bias_param = getattr(op, f"bias{group_idx}") + grads.append((f"b{group_idx}", bias_param.grad.detach().clone())) + bias_param.grad = None + return grads + + grads_ref = _run(dy_ref) + grads_test = _run(dy_wire) + + # Bit-exact match expected (identical quantized grads and kernels). + for (_, grad_test), (_, grad_ref) in zip(grads_test, grads_ref): + torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0) + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) @pytest.mark.parametrize( "quantization", @@ -1076,6 +1277,208 @@ def _make_module(): assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols) assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_mlp_prequantized_mxfp8_input( + self, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + weight_requires_grad: bool, + ) -> None: + """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor input. + + Production (FP8 token dispatch) path: FC1 receives an already + rowwise-quantized input with compact scales. The fused op must feed the + rowwise data to the forward GEMM and manufacture FC1's columnwise copy + for wgrad. Compared bit-for-bit against a run on the dequantized wire + input (see ``test_grouped_linear_prequantized_mxfp8_input``). + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + maybe_skip_quantization( + "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype + ) + + # Split sizes (including an empty group); sum is a multiple of 128. + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + glu_interleave_size = 32 + + # Wire-format FC1 input and its exact dequantization (reference input). + x_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + x_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes) + x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, hidden_size + ) + + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype + ) + fc2 = te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + with torch.no_grad(): + for param in module.parameters(): + param.requires_grad_(requires_grad=weight_requires_grad) + + def _run(x): + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + fc1_wgrads, fc2_wgrads = [], [] + if weight_requires_grad: + y.backward(dy) + for group_idx in range(group_size): + fc1_w = getattr(fc1, f"weight{group_idx}") + fc2_w = getattr(fc2, f"weight{group_idx}") + fc1_wgrads.append(fc1_w.grad.detach().clone()) + fc2_wgrads.append(fc2_w.grad.detach().clone()) + fc1_w.grad = None + fc2_w.grad = None + return y.detach(), fc1_wgrads, fc2_wgrads + + y_ref, fc1_wgrads_ref, fc2_wgrads_ref = _run(x_ref) + y_test, fc1_wgrads_test, fc2_wgrads_test = _run(x_wire) + + # Confirm the CuTeDSL fused op was actually formed (not the fallback). + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Bit-exact match expected (identical quantized inputs and kernels). + torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(fc1_wgrads_test, fc1_wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + for wgrad_test, wgrad_ref in zip(fc2_wgrads_test, fc2_wgrads_ref): + torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0) + + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_mlp_prequantized_mxfp8_grad( + self, + *, + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + bias: bool, + weight_requires_grad: bool, + ) -> None: + """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor grad output. + + FC2 receives the pre-quantized grad, as an FP8 token dispatch delivers it + on the backward pass. With ``bias`` FC2 uses ``scale_bias``, whose + dbias/dscales need the dequantized grad rather than the fused dbias. + """ + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system") + if not weight_requires_grad: + # Independent of pre-quantization: the fused forward saves the input + # activations whenever anything requires grad, then asserts they carry + # columnwise data -- which is only built when the weights need grads. + pytest.skip("Fused grouped MLP does not support frozen weights") + maybe_skip_quantization( + "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype + ) + + # Split sizes (including an empty group); sum is a multiple of 128. + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + total_tokens = int(split_sizes.sum().item()) + glu_interleave_size = 32 + + x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + probs = torch.rand((total_tokens,), dtype=dtype, device=device) + dy_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5 + + # Wire-format grad output and its exact dequantization (the reference grad). + dy_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes) + dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view( + total_tokens, hidden_size + ) + + recipe = make_recipe("mxfp8") + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te.ops.GroupedLinear( + group_size, hidden_size, 2 * hidden_size, bias=bias, device=device, dtype=dtype + ) + fc2 = te.ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device=device, + dtype=dtype, + scale_bias=bias, + ) + module = te.ops.Sequential( + fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2 + ) + + # Frozen experts (weights) with a still-training bias/router is the case + # where the dequantized grad is needed but is not a byproduct of wgrad. + if not weight_requires_grad: + for fc in (fc1, fc2): + for group_idx in range(group_size): + getattr(fc, f"weight{group_idx}").requires_grad_(False) + + def _run(grad): + x_in = x.detach().clone().requires_grad_() + fc2_extra = (split_sizes, probs) if bias else (split_sizes,) + with te.autocast(enabled=True, recipe=recipe): + y = module(x_in, split_sizes, probs, *fc2_extra) + _InjectGrad.apply(y, grad).backward(torch.ones_like(y)) + grads = [("dx", x_in.grad)] + for name, fc in (("fc1", fc1), ("fc2", fc2)): + for group_idx in range(group_size): + if weight_requires_grad: + weight = getattr(fc, f"weight{group_idx}") + grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone())) + weight.grad = None + if bias: + bias_param = getattr(fc, f"bias{group_idx}") + grads.append((f"{name}_b{group_idx}", bias_param.grad.detach().clone())) + bias_param.grad = None + return grads + + grads_ref = _run(dy_ref) + grads_test = _run(dy_wire) + + # Confirm the CuTeDSL fused op was actually formed (not the fallback). + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU) + + # Bit-exact match expected (identical quantized grads and kernels), except + # bias gradients: the fused kernels generate them with an accumulation + # that is not reproducible run to run (two runs on identical inputs differ + # by one BF16 ulp). Same tolerances as + # ``test_grouped_mlp_single_weight_numerics``. + bias_tols = {"rtol": 0.05, "atol": 0.015625} + for (name, grad_test), (_, grad_ref) in zip(grads_test, grads_ref): + if "_b" in name: + torch.testing.assert_close(grad_test, grad_ref, **bias_tols) + else: + torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0) + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize( diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 607346ce30..7d2589bee3 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -10,11 +10,17 @@ import torch +import transformer_engine_torch as tex from transformer_engine_torch import FP8TensorMeta +from ..constants import MXFP8_BLOCK_SCALING_SIZE, TE_DType from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..tensor.float8_tensor import Float8Tensor +from ..tensor.grouped_tensor import GroupedTensor +from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..quantized_tensor import QuantizedTensorStorage +from ..triton.grouped_dbias_dscales import compute_grouped_dbias from ..utils import canonicalize_dtype @@ -66,6 +72,199 @@ def maybe_dequantize( return tensor +def grouped_storage_from_grouped_tensor(tensor: GroupedTensor) -> GroupedTensorStorage: + """Repack a ``GroupedTensor`` into a ``GroupedTensorStorage``. + + ``GroupedTensor`` is a ``torch.Tensor`` subclass, so the CPU offload + infrastructure's ``prepare_for_saving`` treats it as a plain tensor and + does not decompose it into its component data tensors. By repacking into + a ``GroupedTensorStorage`` (not a ``torch.Tensor``), the fuser's + ``prepare_for_saving`` call correctly decomposes the activation before + ``save_for_backward``. + """ + return GroupedTensorStorage( + shape=tensor.logical_shape, + dtype=tensor.fake_dtype, + num_tensors=tensor.num_tensors, + shapes=tensor.tensor_shapes, + quantizer=tensor.quantizer, + data=tensor.rowwise_data, + columnwise_data=tensor.columnwise_data, + scale_inv=tensor.scale_inv, + columnwise_scale_inv=tensor.columnwise_scale_inv, + amax=tensor.amax, + columnwise_amax=tensor.columnwise_amax, + scale=tensor.scale, + first_dims=tensor.first_dims, + last_dims=tensor.last_dims, + tensor_offsets=tensor.tensor_offsets, + offsets=tensor.offsets, + scale_inv_offsets=tensor.scale_inv_offsets, + columnwise_scale_inv_offsets=tensor.columnwise_scale_inv_offsets, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + row_scaled_nvfp4=tensor.row_scaled_nvfp4, + nvfp4_use_4over6=tensor.nvfp4_use_4over6, + nvfp4_e4m3_max=tensor.nvfp4_e4m3_max, + ) + + +def prepare_prequantized_mxfp8_input_for_gemm( + grouped_x: GroupedTensorStorage, + quantizer: MXFP8Quantizer, + num_groups: int, + split_sizes: torch.Tensor, + dtype: torch.dtype, + *, + with_columnwise: bool, + with_dbias: bool = False, + with_dequantized: bool = False, + tensor_offsets: Optional[torch.Tensor] = None, +) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Make an already-quantized MXFP8 grouped input GEMM-ready (in place). + + For inputs that arrive rowwise-quantized (e.g. FP8 token dispatch), where the + high-precision tensor no longer exists. The rowwise data feeds the GEMM as-is + and its scales are swizzled; the columnwise copy cannot be derived from it + (the two directions scale along perpendicular axes) so it is manufactured by + dequantize + columnwise-only requantize. + + The input must be rowwise-only with unswizzled scales, and each group's token + count must be a multiple of 128 so per-group scales start on a swizzle-tile + boundary. + + TODO: optimize and fuse the round-trips requant + + Parameters + ---------- + grouped_x : GroupedTensorStorage + Rowwise-quantized input, updated in place. + quantizer : MXFP8Quantizer + The op's input quantizer. Supplies the FP8 dtype for the columnwise copy. + num_groups : int + Number of groups. + split_sizes : torch.Tensor + Per-group row counts, on device. + dtype : torch.dtype + High-precision dtype to dequantize to. + with_columnwise : bool + Whether to build the columnwise copy that the wgrad GEMM consumes. + with_dbias : bool, default = ``False`` + Whether to also produce the per-group bias gradient. When a columnwise + copy is built the quantize kernel accumulates it in the stage it is + already running, so it costs no extra pass; otherwise it is reduced from + the dequantized tensor. + with_dequantized : bool, default = ``False`` + Whether to dequantize even when no columnwise copy is needed. It must be + requested here rather than recovered later, since the rowwise scales are + swizzled before returning and dequantization requires unswizzled ones. + tensor_offsets : torch.Tensor, optional + Per-group element offsets for the columnwise requantize. + + Returns + ------- + torch.Tensor or None + Per-group bias gradient, when ``with_dbias``. + torch.Tensor or None + Dequantized input, when it was materialized (``with_columnwise``, where it + is a byproduct, or ``with_dequantized``). Used by callers that cannot use + the fused ``dbias``, e.g. ``scale_bias``, whose dbias/dscales depend on + the routing probabilities. + """ + if grouped_x.rowwise_data is None: + raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise data.") + if grouped_x.scale_inv is None: + raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise scales.") + if grouped_x._with_gemm_swizzled_scales: + raise NotImplementedError("Pre-quantized MXFP8 grouped input must have unswizzled scales.") + if grouped_x.columnwise_data is not None: + # Columnwise grouped scales have a per-group layout, so the global + # single-tensor swizzle below cannot convert them. + raise NotImplementedError( + "Pre-quantized MXFP8 grouped input with unswizzled scales must be rowwise-only." + ) + if grouped_x.quantizer is not None and grouped_x.quantizer.dtype != quantizer.dtype: + # The forward GEMM consumes the input's rowwise data verbatim while the + # wgrad GEMM consumes the columnwise copy we manufacture with ``quantizer``. + # A dtype mismatch would make the two directions disagree numerically. + raise ValueError( + f"Pre-quantized MXFP8 grouped input has FP8 dtype {grouped_x.quantizer.dtype}, " + f"but the op's input quantizer expects {quantizer.dtype}." + ) + + # Manufacture columnwise data for the wgrad GEMM: dequantize the rowwise + # wire data and requantize columnwise-only. + dbias = None + dequantized = None + if with_columnwise or with_dbias or with_dequantized: + dequantized = tex.group_dequantize(grouped_x, TE_DType[dtype]).rowwise_data.view( + grouped_x.logical_shape + ) + if with_columnwise: + colwise_quantizer = quantizer.copy() + colwise_quantizer.set_usage(rowwise=False, columnwise=True) + colwise_quantizer.optimize_for_gemm = True + colwise_quantizer.internal = True + if with_dbias: + colwise_x, dbias = tex.bgrad_group_quantize( + dequantized, + colwise_quantizer, + num_groups, + split_sizes, + tensor_offsets=tensor_offsets, + ) + else: + colwise_x = tex.group_quantize( + dequantized, + colwise_quantizer, + num_groups, + split_sizes, + tensor_offsets=tensor_offsets, + ) + grouped_x.columnwise_data = colwise_x.columnwise_data + grouped_x.columnwise_scale_inv = colwise_x.columnwise_scale_inv + elif with_dbias: + # No columnwise stage to accumulate into (e.g. frozen weights need no + # wgrad), so reduce the dequantized grad directly. + dbias = compute_grouped_dbias( + dequantized, tex.splits_to_offsets(split_sizes, 1), num_groups + ) + + # Convert rowwise scales to the GEMM-swizzled layout. The grouped GEMM + # reads activation scales as one (total_tokens, cols) matrix, so the + # single-tensor swizzle applies. Swizzling allocates a new scale buffer; + # the original unswizzled scales are left untouched. + # 128-alignment (see docstring) means the scale array needs no padding. + total_tokens, cols = grouped_x.logical_shape + if total_tokens % 128 != 0 or cols % 128 != 0: + raise ValueError( + "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, " + f"but got ({total_tokens}, {cols})." + ) + scale_shape = (total_tokens, cols // MXFP8_BLOCK_SCALING_SIZE) + if grouped_x.scale_inv.numel() != math.prod(scale_shape): + raise ValueError( + f"Pre-quantized MXFP8 grouped input has {grouped_x.scale_inv.numel()} rowwise " + f"scales, but expected {math.prod(scale_shape)} for shape {scale_shape}." + ) + tmp = MXFP8Tensor( + shape=(total_tokens, cols), + dtype=dtype, + fp8_dtype=quantizer.dtype, + rowwise_data=grouped_x.rowwise_data.view(total_tokens, cols), + rowwise_scale_inv=grouped_x.scale_inv.view(scale_shape), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=quantizer, + requires_grad=False, + with_gemm_swizzled_scales=False, + ) + tex.swizzle_scales_for_gemm_(tmp) + grouped_x.scale_inv = tmp._rowwise_scale_inv.view(-1) + grouped_x._with_gemm_swizzled_scales = True + + return dbias, dequantized + + def maybe_autocast_dtype( *, device_type: str = "cuda", diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index d3e2dc0a13..eb113a5fc4 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -47,8 +47,10 @@ get_accumulate_flag_in_param, get_dummy_wgrads_for_params, get_main_grad_from_param, + grouped_storage_from_grouped_tensor, is_quantized_tensor, maybe_dequantize, + prepare_prequantized_mxfp8_input_for_gemm, validate_or_alloc_output, view_main_grad_as_grouped_buffer, ) @@ -1180,6 +1182,11 @@ def _fuser_forward_split_quantize( out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" + if isinstance(input_, GroupedTensor): + raise NotImplementedError( + "Pre-quantized GroupedTensor input is only supported on the " + "graph-safe grouped-tensor path." + ) num_groups = self.num_groups has_bias = self.has_bias @@ -1304,11 +1311,42 @@ def _fuser_forward_grouped_tensor( # 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) + prequantized_mxfp8_input = ( + with_quantized_compute + and isinstance(input_, GroupedTensor) + and isinstance(input_quantizers[0], MXFP8Quantizer) + and isinstance(input_.quantizer, MXFP8Quantizer) + ) + if prequantized_mxfp8_input: + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, in_features) layout; just validate the shape. + if input_.dim() != 2 or input_.size(-1) != self.in_features: + raise ValueError( + "GroupedTensor input must have shape (total_tokens, " + f"{self.in_features}), but got {tuple(input_.size())}." + ) + total_tokens = input_.size(0) + else: + x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) + total_tokens = x.size(0) # Build the input GroupedTensor. - if with_quantized_compute: + if prequantized_mxfp8_input: + # Rowwise-only MXFP8 input (e.g. FP8 token dispatch): feed the + # rowwise data to the forward GEMM as-is, manufacture the + # columnwise copy needed by the wgrad GEMM, and swizzle the + # rowwise scales for the GEMM. + grouped_x = grouped_storage_from_grouped_tensor(input_) + prepare_prequantized_mxfp8_input_for_gemm( + grouped_x, + input_quantizers[0], + num_groups, + split_sizes, + dtype, + with_columnwise=weight_requires_grad, + tensor_offsets=base_split_offsets * self.in_features, + ) + elif with_quantized_compute: input_quantizer = input_quantizers[0] input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) input_quantizer.optimize_for_gemm = True @@ -1661,8 +1699,25 @@ def _fuser_backward_grouped_tensor( # 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) + prequantized_mxfp8_grad = ( + with_quantized_compute + and isinstance(grad_output, GroupedTensor) + and isinstance(ctx.grad_output_quantizers[0], MXFP8Quantizer) + and isinstance(grad_output.quantizer, MXFP8Quantizer) + ) + if prequantized_mxfp8_grad: + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, out_features) layout; just validate the shape. + if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features: + raise ValueError( + "GroupedTensor grad output must have shape (total_tokens, " + f"{self.out_features}), but got {tuple(grad_output.size())}." + ) + dy_2d = None + total_tokens = grad_output.size(0) + else: + 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 @@ -1679,7 +1734,28 @@ def _fuser_backward_grouped_tensor( fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad ) - if has_bias and not self._scale_bias and fuse_bgrad: + if prequantized_mxfp8_grad: + # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse the + # rowwise data for the dgrad GEMM and manufacture the columnwise copy + # for wgrad. ``scale_bias`` needs the high-precision grad below, so it + # takes the dequantized tensor instead of the fused dbias. + grouped_dy = grouped_storage_from_grouped_tensor(grad_output) + dbias_packed, dy_2d = prepare_prequantized_mxfp8_input_for_gemm( + grouped_dy, + grad_output_quantizer, + num_groups, + split_sizes, + dtype, + with_columnwise=ctx.weight_requires_grad, + with_dbias=has_bias and not self._scale_bias, + with_dequantized=has_bias and self._scale_bias, + tensor_offsets=base_split_offsets * self.out_features, + ) + if not (has_bias and self._scale_bias): + # Only scale_bias reads the dequantized grad; drop it so the + # buffer is freed instead of living until backward ends. + dy_2d = None + elif has_bias and not self._scale_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( dy_2d, grad_output_quantizer, num_groups, split_sizes ) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 83954a9b3d..a5645ff488 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -52,8 +52,10 @@ get_accumulate_flag_in_param, get_dummy_wgrads_for_params, get_main_grad_from_param, + grouped_storage_from_grouped_tensor, is_quantized_tensor, maybe_dequantize, + prepare_prequantized_mxfp8_input_for_gemm, validate_or_alloc_output, view_main_grad_as_grouped_buffer, ) @@ -919,7 +921,16 @@ def fuser_forward( # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) - input_ = input_.reshape(-1, fc1_weight_shape[1]) + if isinstance(input_, GroupedTensor): + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, in_features) layout; just validate the shape. + if input_.dim() != 2 or input_.size(-1) != fc1_weight_shape[1]: + raise ValueError( + "GroupedTensor input must have shape (total_tokens, " + f"{fc1_weight_shape[1]}), but got {tuple(input_.size())}." + ) + else: + input_ = input_.reshape(-1, fc1_weight_shape[1]) in_shape = list(input_.size()) if in_shape[0] % 128 != 0: raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).") @@ -1080,36 +1091,23 @@ def fuser_forward( or isinstance(fc1_input_quantizer, NVFP4Quantizer) and isinstance(input_quantizer, NVFP4Quantizer) ): - # GroupedTensor is a torch.Tensor subclass, so the CPU offload - # infrastructure's prepare_for_saving treats it as a plain tensor - # and does not decompose it into its component data tensors. By - # repacking into a GroupedTensorStorage (not a torch.Tensor), we - # ensure the fuser's prepare_for_saving call correctly decomposes - # the activation before save_for_backward. - grouped_fc1_x = GroupedTensorStorage( - shape=input_.logical_shape, - dtype=input_.fake_dtype, - num_tensors=input_.num_tensors, - shapes=input_.tensor_shapes, - quantizer=input_.quantizer, - data=input_.rowwise_data, - columnwise_data=input_.columnwise_data, - scale_inv=input_.scale_inv, - columnwise_scale_inv=input_.columnwise_scale_inv, - amax=input_.amax, - columnwise_amax=input_.columnwise_amax, - scale=input_.scale, - first_dims=input_.first_dims, - last_dims=input_.last_dims, - tensor_offsets=input_.tensor_offsets, - offsets=input_.offsets, - scale_inv_offsets=input_.scale_inv_offsets, - columnwise_scale_inv_offsets=input_.columnwise_scale_inv_offsets, - with_gemm_swizzled_scales=input_._with_gemm_swizzled_scales, - row_scaled_nvfp4=input_.row_scaled_nvfp4, - nvfp4_use_4over6=input_.nvfp4_use_4over6, - nvfp4_e4m3_max=input_.nvfp4_e4m3_max, - ) + grouped_fc1_x = grouped_storage_from_grouped_tensor(input_) + if ( + isinstance(fc1_input_quantizer, MXFP8Quantizer) + and not grouped_fc1_x._with_gemm_swizzled_scales + ): + # Rowwise-only MXFP8 input (e.g. FP8 token dispatch): + # manufacture the columnwise copy needed by the wgrad GEMM + # and swizzle the rowwise scales for the forward GEMM. + prepare_prequantized_mxfp8_input_for_gemm( + grouped_fc1_x, + fc1_input_quantizer, + num_groups, + split_sizes, + dtype, + with_columnwise=weight_requires_grad, + tensor_offsets=fc1_x_tensor_offsets, + ) else: fc1_x = maybe_dequantize(input_, dtype) grouped_fc1_x = _group_quantize_for_grouped_mlp( @@ -1591,7 +1589,16 @@ def fuser_backward( # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) - grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) + if isinstance(grad_output, GroupedTensor): + # GroupedTensor forbids reshape and is already in the canonical + # (total_tokens, out_features) layout; just validate the shape. + if grad_output.dim() != 2 or grad_output.size(-1) != fc2_weight_shape[0]: + raise ValueError( + "GroupedTensor grad output must have shape (total_tokens, " + f"{fc2_weight_shape[0]}), but got {tuple(grad_output.size())}." + ) + else: + 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 @@ -1659,12 +1666,36 @@ def fuser_backward( isinstance(fc2_grad_output_quantizer, NVFP4Quantizer) and isinstance(grad_output_quantizer, NVFP4Quantizer) ) + prequantized_mxfp8_grad = isinstance(fc2_grad_output_quantizer, MXFP8Quantizer) if ( - not output_fc2_dbias + (not output_fc2_dbias or prequantized_mxfp8_grad) and isinstance(grad_output, GroupedTensor) and fc2_grad_output_quantizer_matches ): - grouped_fc2_dy = grad_output + if prequantized_mxfp8_grad: + # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse + # the rowwise data for the dgrad GEMM and manufacture FC2's + # columnwise copy for wgrad. ``scale_bias`` needs the + # high-precision grad below, so it takes the dequantized tensor + # instead of the fused dbias. + grouped_fc2_dy = grouped_storage_from_grouped_tensor(grad_output) + fc2_dbias_packed, fc2_dy = prepare_prequantized_mxfp8_input_for_gemm( + grouped_fc2_dy, + fc2_grad_output_quantizer, + num_groups, + split_sizes, + dtype, + with_columnwise=fc2_ctx.weight_requires_grad, + with_dbias=output_fc2_dbias and not scale_bias, + with_dequantized=scale_bias, + tensor_offsets=base_split_offsets * fc2_weight_shape[0], + ) + if not scale_bias: + # Only scale_bias reads the dequantized grad; drop it so the + # buffer is freed instead of living until backward ends. + fc2_dy = None + else: + grouped_fc2_dy = grad_output else: fc2_dy = maybe_dequantize(grad_output, dtype) if output_fc2_dbias and not scale_bias: