diff --git a/tests/pytorch/test_custom_recipe_grouped.py b/tests/pytorch/test_custom_recipe_grouped.py new file mode 100644 index 00000000000..aa14459985a --- /dev/null +++ b/tests/pytorch/test_custom_recipe_grouped.py @@ -0,0 +1,293 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped GEMM with custom (reference) quantizers. + +These drive the two custom-recipe grouped entry points with the reference NVFP4/MXFP4 quantizer, +which is the only in-tree way to reach the Python `qgemm` implementations: the module and ops +grouped paths only ever build custom operands for the discrete (list) form, and the reference +quantizer is what makes a `qgemm` callable without a compiled custom kernel. + +Covered here, because nothing else in the suite reaches it: + * `_custom_grouped_gemm` (list form) — operand roles, m_splits slicing, per-expert bias, and the + `qresult_x`/`qresult_w` hand-off that non-pow-2 NVFP4 needs. + * `_custom_grouped_tensor_gemm` (GroupedTensorStorage form) — per-group offsets and block scales. + * The refusals: dgrad/wgrad, non-TN layouts, malformed bias lists, and the non-pow-2 recipe that + this layout cannot serve. +""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch.cpp_extensions import ( + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) +from transformer_engine.pytorch.custom_recipes import reference_utils +from transformer_engine.pytorch.custom_recipes.gemm import GEMMType, MMParams +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import GroupedTensorStorage + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +SPLITS = [16, 32, 8, 24] +K, N = 256, 128 +TOTAL_M = sum(SPLITS) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="grouped GEMM requires a CUDA device" +) + + +def _quantizer(pow_2_scales: bool) -> NVFP4QuantizerRef: + # MXFP4 (E8M0 scales) is defined on 1x32 tiles, NVFP4 on 1x16. + return NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, + pow_2_scales=pow_2_scales, + quant_tile_shape=(1, 32) if pow_2_scales else (1, 16), + ) + + +def _operands(): + torch.manual_seed(0) + weights = [torch.randn(N, K, device=DEVICE, dtype=DTYPE) for _ in SPLITS] + inputs = [torch.randn(m, K, device=DEVICE, dtype=DTYPE) for m in SPLITS] + biases = [torch.randn(N, device=DEVICE, dtype=DTYPE) for _ in SPLITS] + return weights, inputs, biases + + +def _per_group_reference(quantizer, qweights, qinputs, biases=None) -> torch.Tensor: + """One qgemm per expert, placed at that expert's row offset: what the dispatch must produce.""" + expected = torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE) + offset = 0 + for i, m in enumerate(SPLITS): + expected[offset : offset + m] = quantizer.qgemm( + qinputs[i].data[:m], + qweights[i].data, + MMParams(out_dtype=DTYPE, use_split_accumulator=False), + DTYPE, + qinputs[i].scale[:m], + qweights[i].scale, + biases[i] if biases is not None else None, + gemm_type=GEMMType.FPROP, + qresult_x=qinputs[i], + qresult_w=qweights[i], + ) + offset += m + return expected + + +def _assert_close_to_unquantized(actual: torch.Tensor, weights, inputs, biases=None) -> None: + """4-bit operands are coarse, so only gross (operand-role) mistakes are ruled out here.""" + offset = 0 + for i, m in enumerate(SPLITS): + expected = torch.nn.functional.linear( + inputs[i], weights[i], biases[i] if biases is not None else None + ) + block = actual[offset : offset + m] + deviation = (block - expected).norm() / expected.norm().clamp_min(1e-6) + assert deviation < 0.25, f"expert {i}: quantized result deviates {deviation:.1%}" + offset += m + + +@pytest.mark.parametrize("use_bias", [False, True], ids=["no-bias", "bias"]) +@pytest.mark.parametrize("pow_2_scales", [True, False], ids=["mxfp4", "nvfp4"]) +def test_custom_grouped_gemm_list_form(use_bias: bool, pow_2_scales: bool) -> None: + """`general_grouped_gemm` dispatches per-expert custom operands to the quantizer's qgemm.""" + quantizer = _quantizer(pow_2_scales) + weights, inputs, biases = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + + out = torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE) + general_grouped_gemm( + qweights, + qinputs, + [out], + [None] * len(SPLITS), + DTYPE, + single_output=True, + m_splits=SPLITS, + bias=biases if use_bias else None, + use_bias=use_bias, + ) + + torch.testing.assert_close( + out, _per_group_reference(quantizer, qweights, qinputs, biases if use_bias else None) + ) + _assert_close_to_unquantized(out, weights, inputs, biases if use_bias else None) + + +def test_custom_grouped_tensor_gemm() -> None: + """`general_grouped_gemm_for_grouped_tensor` slices groups out of GroupedTensorStorage.""" + quantizer = _quantizer(pow_2_scales=True) + weights, inputs, _ = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + + row_offsets = [0] + for m in SPLITS: + row_offsets.append(row_offsets[-1] + m) + weight_offsets = [i * N for i in range(len(SPLITS) + 1)] + + grouped_weights = GroupedTensorStorage( + shape=(len(SPLITS) * N, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qweights], dim=0), + scale_inv=torch.cat([t.scale for t in qweights], dim=0), + offsets=weight_offsets, + scale_inv_offsets=weight_offsets, + ) + grouped_inputs = GroupedTensorStorage( + shape=(TOTAL_M, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qinputs], dim=0), + scale_inv=torch.cat([t.scale for t in qinputs], dim=0), + offsets=row_offsets, + scale_inv_offsets=row_offsets, + ) + grouped_out = GroupedTensorStorage( + shape=(TOTAL_M, N), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=None, + data=torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE), + offsets=row_offsets, + ) + + general_grouped_gemm_for_grouped_tensor( + grouped_weights, grouped_inputs, grouped_out, layout="TN" + ) + + torch.testing.assert_close( + grouped_out.rowwise_data, + _per_group_reference(quantizer, qweights, qinputs), + ) + _assert_close_to_unquantized(grouped_out.rowwise_data, weights, inputs) + + +def test_custom_grouped_tensor_gemm_refuses_nvfp4() -> None: + """This layout carries no per-group quantized tensors, so a non-pow-2 recipe must refuse.""" + quantizer = _quantizer(pow_2_scales=False) + weights, inputs, _ = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + row_offsets = [0] + for m in SPLITS: + row_offsets.append(row_offsets[-1] + m) + weight_offsets = [i * N for i in range(len(SPLITS) + 1)] + + grouped_weights = GroupedTensorStorage( + shape=(len(SPLITS) * N, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qweights], dim=0), + scale_inv=torch.cat([t.scale for t in qweights], dim=0), + offsets=weight_offsets, + scale_inv_offsets=weight_offsets, + ) + grouped_inputs = GroupedTensorStorage( + shape=(TOTAL_M, K), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=quantizer, + data=torch.cat([t.data for t in qinputs], dim=0), + scale_inv=torch.cat([t.scale for t in qinputs], dim=0), + offsets=row_offsets, + scale_inv_offsets=row_offsets, + ) + grouped_out = GroupedTensorStorage( + shape=(TOTAL_M, N), + dtype=DTYPE, + num_tensors=len(SPLITS), + quantizer=None, + data=torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE), + offsets=row_offsets, + ) + with pytest.raises(ValueError, match="qresult"): + general_grouped_gemm_for_grouped_tensor( + grouped_weights, grouped_inputs, grouped_out, layout="TN" + ) + + +def test_custom_grouped_gemm_refuses_unsupported_contracts() -> None: + """dgrad/wgrad, non-TN layouts and malformed bias lists are refused before any work.""" + quantizer = _quantizer(pow_2_scales=True) + weights, inputs, biases = _operands() + qweights = [quantizer(w) for w in weights] + qinputs = [quantizer(x) for x in inputs] + out = torch.empty(TOTAL_M, N, device=DEVICE, dtype=DTYPE) + + def call(**kwargs): + general_grouped_gemm( + qweights, + qinputs, + [out], + [None] * len(SPLITS), + DTYPE, + single_output=True, + m_splits=SPLITS, + **kwargs, + ) + + with pytest.raises(NotImplementedError, match="fprop only"): + call(grad=True) + with pytest.raises(NotImplementedError, match="TN layout"): + call(layout="NN") + with pytest.raises(ValueError, match="per-group bias"): + call(use_bias=True) + with pytest.raises(ValueError, match="expected 4 biases"): + call(bias=biases[:2], use_bias=True) + + +def test_reference_quantizer_in_place_update_contract() -> None: + """`quantize_` refreshes every field a later GEMM reads, honours noop_flag, returns self.""" + quantizer = _quantizer(pow_2_scales=False) + tensor = torch.randn(32, K, device=DEVICE, dtype=DTYPE) + qresult = quantizer(tensor) + before_amax = qresult.global_amax_row.clone() + + assert qresult.quantize_(tensor * 4.0) is qresult + assert not torch.equal(before_amax, qresult.global_amax_row), "stale global amax after refresh" + assert tuple(qresult.original_shape) == (32, K) + + snapshot = qresult.data.clone() + qresult.quantize_(tensor * 0.25, noop_flag=torch.tensor([1.0], device=DEVICE)) + assert torch.equal(snapshot, qresult.data), "noop_flag was ignored" + + +def test_reference_qgemm_applies_bias_in_fp32_epilogue() -> None: + """The bias is an epilogue term, so qgemm(bias=x) == qgemm() + x rather than an error.""" + quantizer = _quantizer(pow_2_scales=True) + activations = quantizer(torch.randn(16, K, device=DEVICE, dtype=DTYPE)) + weights = quantizer(torch.randn(N, K, device=DEVICE, dtype=DTYPE)) + bias = torch.randn(N, device=DEVICE, dtype=DTYPE) + params = MMParams(out_dtype=DTYPE, use_split_accumulator=False) + + plain = quantizer.qgemm( + activations.data, + weights.data, + params, + DTYPE, + activations.scale, + weights.scale, + gemm_type=GEMMType.FPROP, + ) + with_bias = quantizer.qgemm( + activations.data, + weights.data, + params, + DTYPE, + activations.scale, + weights.scale, + bias, + gemm_type=GEMMType.FPROP, + ) + torch.testing.assert_close(with_bias, plain + bias, rtol=2e-2, atol=2e-2) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 19b27798c5d..c0f5c80ee48 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -3827,3 +3827,27 @@ def test_swizzle_scales_and_pack_ptrs_for_discrete_weights( swizzled_scales_buffer, expected_swizzled_scales_buffer, ) + + +def test_grouped_tensor_custom_dispatch_classifies_plain_operands(): + """A plain GroupedTensorStorage must not be classified as a custom operand. + + `is_custom` raises on any type that is neither a Quantizer/QuantizedTensorStorage nor a + Tensor, so dispatching on the container itself made every grouped-tensor GEMM raise + before it reached the native implementation. + """ + from transformer_engine.pytorch.cpp_extensions.gemm import _grouped_operand_is_custom + from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import ( + GroupedTensorStorage, + ) + + grouped = GroupedTensorStorage( + shape=(32, 32), + dtype=torch.bfloat16, + num_tensors=2, + quantizer=None, + data=torch.zeros(32, 32, dtype=torch.bfloat16, device="cuda"), + ) + assert not _grouped_operand_is_custom(grouped) + assert not _grouped_operand_is_custom([grouped]) + assert not _grouped_operand_is_custom(None) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 404a6e65613..469f3ea771d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -34,7 +34,7 @@ from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom -from ..custom_recipes.gemm import custom_gemm +from ..custom_recipes.gemm import custom_gemm, GEMMType, MMParams from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer __all__ = [ @@ -1040,6 +1040,162 @@ def strided_batched_gemm( ) +# Custom-recipe dispatch for the LIST-based grouped GEMM. The operand roles are taken from +# GroupedLinear's fprop call site: +# +# out = _validate_or_alloc_output(out, sum(m_splits), weights_fp8[0].size(0), ...) +# general_grouped_gemm(weights_fp8, inputmats, [out], ..., single_output=True, m_splits=m_splits) +# +# A = list of weights (one per expert) +# B = list of inputs (per-expert activation buffers, padded) +# out = [single tensor] of shape [sum(m_splits), A[0].size(0)] +# m_splits[i] = VALID activation rows for expert i +# +# So for each expert: activation rows are B[i][0:m_splits[i]], the weight is A[i], and the result +# is written to out[0][offset : offset+m_splits[i]]. + + +def _custom_grouped_gemm( + A, + B, + out, + quantization_params, + out_dtype, + layout="TN", + m_splits=None, + gelu=False, + grad=False, + accumulate=False, + bias=None, + use_bias=False, + use_split_accumulator=False, + D_dtype=None, + single_output=False, +): + """Grouped GEMM over per-expert weight/activation lists using a custom quantizer's qgemm.""" + if gelu: + raise NotImplementedError("custom grouped GEMM does not define GELU outputs") + # dgrad/wgrad use different operand roles and a transposed weight; only TN fprop is + # implemented here, so they are refused instead of computing the wrong product. + if grad: + raise NotImplementedError("custom grouped GEMM implements fprop only") + if layout != "TN": + raise NotImplementedError( + f"custom grouped GEMM supports the TN layout only, got {layout!r}" + ) + if D_dtype is not None: + raise NotImplementedError("custom grouped GEMM does not define D_dtype") + if quantization_params is not None and any(q is not None for q in quantization_params): + raise NotImplementedError("custom grouped GEMM does not define output quantization") + if not A or not B: + raise ValueError("A and B must both be non-empty") + if not all(is_custom(x) for x in (*A, *B)): + raise TypeError("custom grouped GEMM cannot mix custom and ordinary tensors") + if m_splits is None: + raise ValueError("m_splits is required for the per-expert layout") + # Two output layouts occur: a single concatenated buffer (`single_output=True`) or one tensor + # per group. Both are supported; the earlier version refused the list form, which is the one + # TEGroupedLinear actually uses. + if single_output and len(out) != 1: + raise ValueError("single_output=True requires exactly one output tensor") + if not single_output and len(out) != len(m_splits): + raise ValueError(f"expected {len(m_splits)} outputs, got {len(out)}") + + n = len(m_splits) + if len(A) != n or len(B) != n: + raise ValueError(f"expected {n} weights and {n} inputs, got {len(A)} and {len(B)}") + total_rows = sum(m_splits) + + quantizer = None + for candidate in (*A, *B): + q = getattr(candidate, "_quantizer", None) + if q is not None and callable(getattr(q, "qgemm", None)): + quantizer = q + break + if quantizer is None: + raise TypeError("no quantizer with a callable qgemm on the grouped operands") + + m_params = MMParams(out_dtype=out_dtype, use_split_accumulator=use_split_accumulator) + + dst_full = out[0] if single_output else None + if single_output: + if dst_full.shape[0] < total_rows: + raise ValueError(f"output has {dst_full.shape[0]} rows but sum(m_splits)={total_rows}") + width = dst_full.shape[-1] + else: + width = out[0].shape[-1] + for i, o in enumerate(out): + if o.shape[-1] != width: + raise ValueError(f"out[{i}] width {o.shape[-1]} != {width}") + + if use_bias and bias is None: + raise ValueError("use_bias=True requires a per-group bias list") + if use_bias and len(bias) != n: + raise ValueError(f"expected {n} biases, got {len(bias)}") + + # NOTE (known cost, not a correctness gap): this issues one single-group GEMM per expert group, so a + # step pays ~40x more kernel launches than it has GEMMs of work (~40,000 launches for ~1,000 GEMMs, + # measured on a real training step 2026-11). Per-launch cost is small (~0.05 ms of device time), which + # is why it never shows up in a per-phase average, but the sum is now the largest identified item in + # the custom-recipe path. Replacing this loop with a grouped GEMM kernel (one launch for all groups) + # is the intended fix and is what the surviving difference to the native FP4 path is made of. + offset = 0 + for i in range(n): + rows = m_splits[i] + if rows == 0: + continue + w, act = A[i], B[i] + if w.data is None or act.data is None: + raise ValueError(f"group {i}: missing quantized data (weight or activation)") + if w.scale is None or act.scale is None: + raise ValueError(f"group {i}: missing scale_inv (weight or activation)") + + qw = w.data + sw = w.scale + qx = act.data[:rows] + sx = act.scale + # Activation scales are per-row, so slice the same prefix as the data. + if hasattr(sx, "shape") and sx.ndim >= 1 and sx.shape[0] >= rows: + sx = sx[:rows] + + # The quantized tensors themselves carry the global scales that a non-pow-2 (NVFP4) + # recipe turns into the GEMM alpha, so they go along with the packed data. + result = quantizer.qgemm( + qx, + qw, + m_params, + out_dtype, + sx, + sw, + bias[i] if use_bias else None, + gemm_type=GEMMType.FPROP, + qresult_x=act, + qresult_w=w, + ) + if tuple(result.shape) != (rows, width): + raise ValueError( + f"group {i}: qgemm returned {tuple(result.shape)}, expected {(rows, width)} " + f"(qx={tuple(qx.shape)} qw={tuple(qw.shape)} sx={tuple(sx.shape)} " + f"sw={tuple(sw.shape)} widths={width})" + ) + dst = dst_full[offset : offset + rows] if single_output else out[i] + if tuple(dst.shape) != (rows, width): + raise ValueError( + f"group {i}: output segment is {tuple(dst.shape)}, expected {(rows, width)}" + ) + with torch.no_grad(): + if accumulate: + dst.add_(result) + else: + dst.copy_(result) + offset += rows + + ref = dst_full if single_output else out[0] + bias_output = [torch.empty(0, dtype=ref.dtype, device=ref.device)] + gelu_input = [torch.empty(0, dtype=ref.dtype, device=ref.device)] + return out, bias_output, gelu_input + + def general_grouped_gemm( A: List[torch.Tensor], B: List[torch.Tensor], @@ -1060,6 +1216,25 @@ def general_grouped_gemm( """ TN layout Grouped GEMM with fp8 inputs. """ + if any(is_custom(a) for a in A) or any(is_custom(b) for b in B): + return _custom_grouped_gemm( + A, + B, + out, + quantization_params, + out_dtype, + layout, + m_splits, + gelu, + grad, + accumulate, + bias, + use_bias, + use_split_accumulator, + D_dtype, + single_output, + ) + num_gemms = len(A) transa = layout[0] == "T" @@ -1205,6 +1380,179 @@ def _get_grouped_cublas_workspace(device: int, layout: str) -> torch.Tensor: return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) +# Custom-recipe dispatch for the *grouped-tensor* GEMM, the entry point GroupedLinear's fprop uses. +# A carries the per-expert weights, B the per-expert activations and `out` the grouped result; +# `offsets` / `scale_inv_offsets` slice each group's rows out of `rowwise_data` / +# `rowwise_scale_inv`, and `num_tensors` is the group count. Only TN fprop is implemented. +# +# Unlike the list form, this layout has no per-group QuantizedTensorStorage to hand a quantizer as +# `qresult_x` / `qresult_w`, so a recipe that derives its global scale from them (non-pow-2 NVFP4) +# is refused inside `qgemm` rather than served wrong. + + +def _custom_grouped_tensor_gemm( + A, + B, + out, + *, + layout="TN", + accumulate=False, + use_split_accumulator=False, + bias=None, + bias_scale=None, + grad=False, + alpha=None, + beta=None, +): + """Grouped GEMM over GroupedTensorStorage operands using a custom quantizer's qgemm.""" + import math + + if grad: + raise NotImplementedError("custom grouped-tensor GEMM does not support grad") + if bias is not None or bias_scale is not None: + raise NotImplementedError("custom grouped-tensor GEMM does not support bias") + if alpha is not None or beta is not None: + raise NotImplementedError("custom grouped-tensor GEMM does not support alpha/beta") + if layout != "TN": + raise NotImplementedError( + f"custom grouped-tensor GEMM supports layout TN only, got {layout!r}" + ) + + def expect_grouped(x, name): + if not hasattr(x, "rowwise_data"): + raise TypeError(f"{name} must be a GroupedTensorStorage, got {type(x).__name__}") + if getattr(x, "rowwise_data", None) is None: + raise ValueError( + f"{name}.rowwise_data is None (columnwise-only storage is unsupported)" + ) + return x + + A = expect_grouped(A, "A") + B = expect_grouped(B, "B") + out = expect_grouped(out, "out") + + num = getattr(A, "num_tensors", None) + if not num: + raise ValueError("A.num_tensors is unset; cannot iterate groups") + + def offsets_of(x, name): + offsets = getattr(x, "offsets", None) + if offsets is None: + raise ValueError(f"{name}.offsets is None; per-group slicing is impossible") + offsets = list(offsets) + if len(offsets) != num + 1: + raise ValueError(f"{name}.offsets has {len(offsets)} entries, expected {num + 1}") + return offsets + + a_off = offsets_of(A, "A") + b_off = offsets_of(B, "B") + o_off = offsets_of(out, "out") + + def scale_offsets_of(x, name): + so = getattr(x, "scale_inv_offsets", None) + if so is None: + return None + so = list(so) + if len(so) != num + 1: + raise ValueError(f"{name}.scale_inv_offsets has {len(so)} entries, expected {num + 1}") + return so + + a_s_off = scale_offsets_of(A, "A") + b_s_off = scale_offsets_of(B, "B") + + quantizer = getattr(A, "quantizer", None) + if quantizer is None: + quantizer = getattr(B, "quantizer", None) + if quantizer is None or not callable(getattr(quantizer, "qgemm", None)): + raise TypeError("no quantizer with a callable qgemm on the grouped operands") + + m_params = MMParams(out_dtype=None, use_split_accumulator=use_split_accumulator) + out_dtype = out.rowwise_data.dtype + + # Every caller of this function that reaches the custom path passes `layout="TN"` with + # A = per-expert weight and B = per-expert activation (`module/grouped_linear.py`, + # `ops/basic/grouped_linear.py`), so the operand that feeds qgemm first is B and the row + # count that has to match the output segment is B's, not A's. + for i in range(num): + a0, a1 = a_off[i], a_off[i + 1] + b0, b1 = b_off[i], b_off[i + 1] + rows = b1 - b0 + if rows == 0: + continue + + o0, o1 = o_off[i], o_off[i + 1] + if o1 - o0 != rows: + raise ValueError( + f"group {i}: out segment has {o1 - o0} rows but the activation slice has {rows} " + f"(b_off={b_off[:3]}..., o_off={o_off[:3]}...)" + ) + + qx = B.rowwise_data[b0:b1] + qw = A.rowwise_data[a0:a1] + + # GroupedTensorStorage names its block scales `scale_inv`; `rowwise_scale_inv` is the name + # the per-tensor storages use, so accept either rather than only the one this entry point + # never actually receives. + def block_scales(x): + for field in ("scale_inv", "rowwise_scale_inv"): + buffer = getattr(x, field, None) + if buffer is not None: + return buffer + return None + + sx = None + b_scales = block_scales(B) + if b_scales is not None and b_s_off is not None: + sx = b_scales[b_s_off[i] : b_s_off[i + 1]] + sw = None + a_scales = block_scales(A) + if a_scales is not None and a_s_off is not None: + sw = a_scales[a_s_off[i] : a_s_off[i + 1]] + if sx is None or sw is None: + raise ValueError(f"group {i}: missing rowwise scale_inv for a custom quantizer") + + result = quantizer.qgemm( + qx, + qw, + m_params, + out_dtype, + sx, + sw, + None, + gemm_type=GEMMType.FPROP, + ) + target = out.rowwise_data[o0:o1] + if tuple(result.shape) != tuple(target.shape): + raise ValueError( + f"group {i}: qgemm returned {tuple(result.shape)} but the out segment is " + f"{tuple(target.shape)} (qx={tuple(qx.shape)} qw={tuple(qw.shape)} " + f"sx={tuple(sx.shape)} sw={tuple(sw.shape)})" + ) + with torch.no_grad(): + if accumulate: + target.add_(result) + else: + target.copy_(result) + + return out + + +def _grouped_operand_is_custom(x) -> bool: + """Whether a grouped-GEMM operand is routed to a custom quantizer. + + Operands of `general_grouped_gemm_for_grouped_tensor` are `GroupedTensorStorage` + containers or discrete lists of tensors, and `is_custom` raises on any type that is + neither a `Quantizer`/`QuantizedTensorStorage` nor a `torch.Tensor`. Handing it a + container therefore turned every grouped-tensor GEMM, native ones included, into an + `AssertionError`; the quantizer (or the list members) have to be inspected instead. + """ + if isinstance(x, (list, tuple)): + return any(_grouped_operand_is_custom(item) for item in x) + if isinstance(x, (Quantizer, QuantizedTensorStorage)): + return is_custom(x) + return is_custom(getattr(x, "quantizer", None)) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -1227,6 +1575,21 @@ def general_grouped_gemm_for_grouped_tensor( The caller must ensure that GroupedTensor metadata is already compatible with the underlying GEMM implementation (e.g., aligned offsets and output metadata layout). """ + if any(_grouped_operand_is_custom(x) for x in (A, B, out)): + return _custom_grouped_tensor_gemm( + A, + B, + out, + layout=layout, + accumulate=accumulate, + use_split_accumulator=use_split_accumulator, + bias=bias, + bias_scale=bias_scale, + grad=grad, + alpha=alpha, + beta=beta, + ) + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." if grad: raise NotImplementedError("grad is not supported for grouped_tensor GEMM yet.") diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index 29e525da45d..953994e0547 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -392,6 +392,30 @@ def __repr__(self): ")" ) + def quantize_(self, tensor, *, noop_flag=None): + """Quantise into this tensor in place. + + Same contract as `QuantizedTensorStorage.quantize_`: the packed data, the block scales + and every other field a later GEMM reads are refreshed — including + `global_amax_row`/`global_amax_col`, which a non-pow-2 (NVFP4) recipe turns into the + global scale, so refreshing only the data and block scales would pair fresh values with + a stale scale. A non-zero `noop_flag` leaves the tensor untouched, which is what lets a + replayed graph skip re-quantisation on later microbatches. Returns `self`. + """ + if tensor.dtype not in reference_utils.HIGH_PRECISION_FLOAT_DTYPES: + raise TypeError( + f"Unsupported input dtype {tensor.dtype}, expected one of" + f" {reference_utils.HIGH_PRECISION_FLOAT_DTYPES}" + ) + original_shape = tuple(tensor.shape) + if tensor.ndim < 2: + # `update_quantized` reshapes ndim > 2 only, and `_quantize` below needs a matrix; + # the old `quantize` call folded every non-2-D input into a single row. + tensor = tensor.reshape(1, -1) + self._get_quantizer().update_quantized(tensor, self, noop_flag=noop_flag) + self.original_shape = original_shape + return self + def update_usage( self, rowwise_usage: Optional[bool] = None, @@ -408,6 +432,11 @@ def update_usage( if columnwise_usage is not None: needs_data_transpose = columnwise_usage + # A placeholder (neither direction populated) has nothing to generate from; + # refusing here turns workspace allocation into a hard error. + if not has_data and not has_data_transpose: + return self + # Generate data that is required if needs_data and not has_data: raise RuntimeError("Cannot generate FP4 data, even from FP4 data transpose") @@ -429,16 +458,20 @@ def _create_transpose(self): self.data_t = self.data.t().contiguous() self.scale_t = self.scale - def size(self, *args, **kwargs): # pylint: disable=unused-argument - """Return the original tensor shape, not the internal packed data shape. + def size(self, dim=None): + """Return the logical original shape, honouring `dim`. - FP4 quantization packs two 4-bit values into each 8-bit value, which reduces - the second dimension by half. This method returns the logical shape that - users expect, not the internal packed storage shape. + FP4 packs two 4-bit values into each 8-bit value, halving the second dimension, so + the logical shape is not the storage shape. The no-argument form returns the full + logical shape. Passing `dim` must return that dimension's size: callers such as + `GroupedLinear._validate_or_alloc_output` compute `rows = x.size(0)` and hand it to + `torch.empty`, which raises `type must be tuple of ints, but got torch.Size` when + the whole shape comes back instead. """ if self.original_shape is None: raise RuntimeError("NVFP4TensorRef.size() called but original_shape has not been set") - return torch.Size(self.original_shape) + shape = torch.Size(self.original_shape) + return shape if dim is None else shape[dim] def get_wgrad_sign_vector() -> torch.Tensor: @@ -975,6 +1008,21 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ ) if self.row_scaled_nvfp4: raise ValueError("Row-scaled NVFP4 is only supported for NVFP4 (non-pow2) mode.") + # MXFP4 path: the branch above validated the configuration but never + # prepared the operands, so the common tail raised UnboundLocalError at + # `x_input = row_input`. Prepare them here, and guard the empty splits + # that _split_quantization can produce (torch.max() with no dim raises on + # a 0-element tensor). + row_input = tensor + col_input = tensor.t().contiguous() + if row_input.numel(): + global_amax_row = torch.max(torch.abs(row_input)).to(torch.float32).view(1) + else: + global_amax_row = torch.zeros(1, device=row_input.device, dtype=torch.float32) + if col_input.numel(): + global_amax_col = torch.max(torch.abs(col_input)).to(torch.float32).view(1) + else: + global_amax_col = torch.zeros(1, device=col_input.device, dtype=torch.float32) # TODO(etsykunov): Fix bug where global_amax_row and # global_amax_col are not defined # global_amax = torch.empty(0, device=tensor.device, dtype=torch.float32) @@ -1212,9 +1260,6 @@ def qgemm( qresult_w: QuantizedTensorStorage | None = None, ) -> torch.Tensor: """Python implementation of microblock FP4 GEMM.""" - if bias is not None: - raise ValueError("Bias is not supported in NVFP4QuantizerRef.qgemm") - high_precision_x = cast_from_fp4x2(qx, out_dtype) high_precision_w = cast_from_fp4x2(qw, out_dtype) @@ -1369,5 +1414,10 @@ def qgemm( if out is not None: raise ValueError("Output tensor should be None when accumulate is False.") + # Bias is an epilogue term, as in the native path: add it while the accumulator is still + # float32, so the result is rounded once when it is cast to `out_dtype`. + if bias is not None: + y = y + bias.to(torch.float32) + y = y.to(out_dtype) return y diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index bf4f9bba1d5..0cfa45c2af1 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -963,6 +963,15 @@ def quantize_weight( return tensor, None # Validate workspace + # A custom (reference) quantizer cannot be cached safely. The cached entry is a + # TensorSpec whose storage is driven by `update_usage`, and a gradient-checkpoint + # replay can reach the GEMM after that storage has been released, which surfaces as + # "FPROP GEMM: quantized weight data (B.data) is None" in a model that trains fine + # for a single un-replayed forward. Re-quantising per call is cheap next to a wrong + # result, so the workspace path is skipped for custom quantizers. + if quantizer is not None and getattr(quantizer, "custom", False): + workspace = None + if workspace is not None and quantizer is not None: if not _is_weight_workspace_valid(workspace, quantizer): workspace = None