diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py index 96f6272f78..1e37a8adb3 100644 --- a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py @@ -12,6 +12,7 @@ from torchao.float8.float8_utils import compute_error from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( NVFP4FourOverSixLinear, + four_over_six_dequantize, four_over_six_global_encode_scale, four_over_six_linear, four_over_six_quantize, @@ -170,6 +171,44 @@ def test_dequant_sqnr(block): assert compute_error(x.float(), dq).item() > 14.0 +@_skip_no_cuda +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_dequantize_roundtrip(block, row_scaled): + """four_over_six_dequantize reconstructs the quantized values.""" + if row_scaled and block == "16x16": + pytest.skip("row-scaled is 1x16 only") + torch.manual_seed(0) + x = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda") + amax = (x.abs().amax(dim=1) if row_scaled else x.abs().amax()).to(torch.float32) + codes, scales = four_over_six_quantize(x, amax, block=block) + dq = four_over_six_dequantize(codes, scales, amax, out_dtype=torch.float32) + assert compute_error(x.float(), dq).item() > 14.0 + # Zero blocks reconstruct exactly: scale byte 0x00 makes the decode + # scale exactly zero regardless of the global amax. + x[:, :16] = 0.0 + codes, scales = four_over_six_quantize(x, amax, block=block) + dq = four_over_six_dequantize(codes, scales, amax, out_dtype=torch.float32) + assert (dq[:, :16] == 0.0).all() + + +@_skip_no_cuda +def test_dequantize_validation(): + codes = torch.zeros(32, 128, dtype=torch.uint8, device="cuda") + scales = torch.zeros(32, 16, dtype=torch.uint8, device="cuda").view( + torch.float8_e4m3fn + ) + amax = torch.ones((), dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="e4m3_scale_bound"): + four_over_six_dequantize(codes, scales, amax, e4m3_scale_bound=128) + with pytest.raises(ValueError, match="scales must have shape"): + four_over_six_dequantize(codes, scales[:, :8], amax) + with pytest.raises(ValueError, match="row vector"): + four_over_six_dequantize( + codes, scales, torch.ones(7, dtype=torch.float32, device="cuda") + ) + + @_skip_no_sm100 @pytest.mark.parametrize("row_scaled_activation", [False, True]) @pytest.mark.parametrize("bias", [False, True]) @@ -222,6 +261,138 @@ def test_linear_rejects_unaligned_dims(): four_over_six_linear(x, w, None, "mae", 256, False) +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_backward_override_high_precision(row_scaled_activation): + """dx/dw are the plain bf16 GEMMs on the original operands.""" + torch.manual_seed(0) + M, K, N = 256, 512, 384 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = (torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1).requires_grad_( + True + ) + y = four_over_six_linear( + x, w, None, "mae", 256, row_scaled_activation, "high_precision" + ) + dy = torch.randn_like(y) + y.backward(dy) + torch.testing.assert_close(x.grad, dy @ w.detach(), atol=0, rtol=0) + torch.testing.assert_close(w.grad, dy.t() @ x.detach(), atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_backward_override_dequantized(row_scaled_activation, weight_block): + """dx/dw are bf16 GEMMs on dequantizations of the rowwise fprop operands.""" + torch.manual_seed(0) + M, K, N = 256, 512, 384 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = (torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1).requires_grad_( + True + ) + y = four_over_six_linear( + x, w, None, "mae", 256, row_scaled_activation, "dequantized", weight_block + ) + dy = torch.randn_like(y) + y.backward(dy) + + x_hp, w_hp = x.detach(), w.detach() + x_amax = ( + x_hp.abs().amax(dim=1) if row_scaled_activation else x_hp.abs().amax() + ).to(torch.float32) + w_amax = w_hp.abs().amax().to(torch.float32) + x_codes, x_scales = four_over_six_quantize(x_hp, x_amax) + w_codes, w_scales = four_over_six_quantize(w_hp, w_amax, block=weight_block) + x_dq = four_over_six_dequantize(x_codes, x_scales, x_amax) + w_dq = four_over_six_dequantize(w_codes, w_scales, w_amax) + torch.testing.assert_close(x.grad, dy @ w_dq, atol=0, rtol=0) + torch.testing.assert_close(w.grad, dy.t() @ x_dq, atol=0, rtol=0) + + +@_skip_no_sm100 +def test_row_scaled_default_backward_is_high_precision(): + """row_scaled + backward_override=None keeps the pre-override behavior.""" + torch.manual_seed(0) + M, K, N = 256, 512, 384 + x_hp = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w_hp = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + dy = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") + + def run(override): + x = x_hp.clone().requires_grad_(True) + w = w_hp.clone().requires_grad_(True) + y = four_over_six_linear(x, w, None, "mae", 256, True, override) + y.backward(dy) + return y.detach(), x.grad, w.grad + + y0, dx0, dw0 = run(None) + y1, dx1, dw1 = run("high_precision") + torch.testing.assert_close(y0, y1, atol=0, rtol=0) + torch.testing.assert_close(dx0, dx1, atol=0, rtol=0) + torch.testing.assert_close(dw0, dw1, atol=0, rtol=0) + + +@_skip_no_sm100 +def test_weight_block_1x16_forward(): + """weight_block='1x16' quantizes the fprop weight with 1x16 blocks.""" + from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + _global_decode_scale, + _scaled_mm_nvfp4, + ) + + torch.manual_seed(0) + M, K, N = 256, 512, 384 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + y = four_over_six_linear(x, w, None, "mae", 256, False, None, "1x16") + + x_amax = x.abs().amax().to(torch.float32) + w_amax = w.abs().amax().to(torch.float32) + x_codes, x_scales = four_over_six_quantize(x, x_amax) + w_codes, w_scales = four_over_six_quantize(w, w_amax, block="1x16") + y_ref = _scaled_mm_nvfp4( + x_codes, + x_scales, + _global_decode_scale(x_amax, 256), + w_codes.t(), + w_scales, + _global_decode_scale(w_amax, 256), + torch.bfloat16, + ) + torch.testing.assert_close(y, y_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +def test_backward_override_validation(): + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="no quantized backward"): + four_over_six_linear(x, w, None, "mae", 256, True, "quantized") + with pytest.raises(ValueError, match="backward_override"): + four_over_six_linear(x, w, None, "mae", 256, False, "bf16") + with pytest.raises(ValueError, match="weight_block"): + four_over_six_linear(x, w, None, "mae", 256, False, None, "8x8") + + +@_skip_no_sm100 +def test_linear_module_backward_override(): + lin = NVFP4FourOverSixLinear( + 512, + 384, + backward_override="dequantized", + weight_block="1x16", + device="cuda", + dtype=torch.bfloat16, + ) + x = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda", requires_grad=True) + y = lin(x) + y.sum().backward() + assert y.shape == (128, 384) + assert lin.weight.grad is not None + assert x.grad is not None + + @_skip_no_cuda @pytest.mark.parametrize("err_mode", ["mae", "mse"]) @pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) @@ -277,6 +448,49 @@ def test_bitwise_parity_with_transformer_engine( ) +@_skip_no_cuda +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +@pytest.mark.parametrize("row_scaled", [False, True]) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32]) +def test_dequantize_bitwise_parity_with_transformer_engine( + e4m3_scale_bound, row_scaled, out_dtype +): + """four_over_six_dequantize matches TE's NVFP4 dequantize kernel bitwise.""" + te = pytest.importorskip("transformer_engine.pytorch") + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + if not te.is_nvfp4_available(): + pytest.skip("NVFP4 not available in this TransformerEngine build") + + torch.manual_seed(0) + M, N = 256, 512 + x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + row_scaled_nvfp4=row_scaled, + nvfp4_use_4over6=True, + nvfp4_e4m3_max=e4m3_scale_bound, + ) + t = quantizer(x) + te_dq = t.dequantize(dtype=out_dtype) + amax = (x.abs().amax(dim=1) if row_scaled else x.abs().amax()).to(torch.float32) + codes, scales = four_over_six_quantize( + x, amax, e4m3_scale_bound=e4m3_scale_bound + ) + dq = four_over_six_dequantize( + codes, + scales, + amax, + e4m3_scale_bound=e4m3_scale_bound, + out_dtype=out_dtype, + ) + torch.testing.assert_close(dq, te_dq, atol=0, rtol=0) + + @_skip_no_cutedsl @pytest.mark.parametrize("err_mode", ["mae", "mse"]) @pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py index 4330f91099..a0ea577766 100644 --- a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -51,6 +51,22 @@ * backward with row-scaled activations: high-precision (bf16) GEMMs. A row-scaled four-over-six tensor has no columnwise form — the per-row scales do not transpose — so the quantized wgrad operand cannot be produced. + +Those backward defaults can be overridden with ``backward_override``, +mirroring TransformerEngine's ``NVTE_BACKWARD_OVERRIDE`` recipe field: + +* ``"quantized"``: the standard-NVFP4-gradient backward above (the + per-tensor default; rejected for row-scaled activations); +* ``"high_precision"``: bf16 GEMMs on the saved original operands (the + row-scaled default); +* ``"dequantized"``: bf16 GEMMs on dequantizations of the rowwise operands + the forward GEMM consumed, so the gradients differentiate the + quantized-forward function itself — the RL train/inference-consistency + mode. Only 4-bit codes and scales are saved for backward, the same + activation-memory win TransformerEngine PR #3141 documents. + +Weights quantize with 16x16 tiles by default; ``weight_block="1x16"`` mirrors +TransformerEngine's ``disable_2d_quantization``. """ from typing import Optional @@ -63,6 +79,7 @@ f4_unpacked_to_f32, f32_to_f4_unpacked, pack_uint4, + unpack_uint4, ) from torchao.prototype.mx_formats.utils import to_blocked @@ -78,6 +95,7 @@ __all__ = [ "four_over_six_global_encode_scale", "four_over_six_quantize", + "four_over_six_dequantize", "four_over_six_mm", "four_over_six_linear", "NVFP4FourOverSixLinear", @@ -255,6 +273,63 @@ def four_over_six_quantize( return pack_uint4(codes.view(rows, cols)), scales +def four_over_six_dequantize( + codes: torch.Tensor, + scales: torch.Tensor, + global_amax: torch.Tensor, + *, + e4m3_scale_bound: int = 256, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize packed FP4 codes and block scales back to high precision. + + Inverse of :func:`four_over_six_quantize`, transcribed operation for + operation from TransformerEngine's ``dequantize_nvfp4.cuh``: the per-block + decode scale is ``(f32(scale) * amax) * factor_inv`` with + ``factor_inv = 1 / (6 * bound)`` a correctly-rounded FP32 reciprocal, and + each element is ``f32(code) * decode_scale`` cast to ``out_dtype``. + Scales from either block granularity dequantize identically (a 16x16 tile + stores its scale byte on every row). + + Args: + codes: (R, C//2) uint8 packed FP4 codes. + scales: (R, C//16) float8_e4m3fn block scales. + global_amax: scalar FP32 amax, or a (R,) per-row amax vector for the + row-scaled variant. + e4m3_scale_bound: the bound the codes were quantized with. + out_dtype: output dtype (the kernel's OType cast). + """ + if e4m3_scale_bound not in (256, 448): + raise ValueError(f"e4m3_scale_bound must be 256 or 448, got {e4m3_scale_bound}") + rows, packed_cols = codes.shape + cols = packed_cols * 2 + if scales.shape != (rows, cols // 16): + raise ValueError( + f"scales must have shape ({rows}, {cols // 16}), " + f"got {tuple(scales.shape)}" + ) + row_scaled = global_amax.dim() == 1 and global_amax.numel() == rows + if not row_scaled and global_amax.numel() != 1: + raise ValueError( + f"global_amax must be a scalar or a ({rows},) row vector, " + f"got shape {tuple(global_amax.shape)}" + ) + values = f4_unpacked_to_f32(unpack_uint4(codes)).view(rows, cols // 16, 16) + amax = global_amax.to(torch.float32) + if row_scaled: + amax = amax.view(rows, 1) + # The reciprocal must come from a true FP32 division (see _candidate_error + # on why a python-scalar denominator double-rounds). + factor_inv = torch.ones((), dtype=torch.float32, device=codes.device) / torch.full( + (), + FP4_E2M1_MAX * float(e4m3_scale_bound), + dtype=torch.float32, + device=codes.device, + ) + decode_scale = (scales.to(torch.float32) * amax) * factor_inv + return (values * decode_scale.unsqueeze(-1)).to(out_dtype).view(rows, cols) + + def _standard_rtne_quantize( x: torch.Tensor, global_amax: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -323,6 +398,10 @@ class four_over_six_mm(torch.autograd.Function): With row-scaled activations the backward runs in bf16 instead (see the module docstring), saving the high-precision operands. + ``backward_override`` selects among the quantized, high-precision, and + dequantized backwards described in the module docstring; ``None`` keeps + the defaults above. ``weight_block`` selects the weight tile granularity. + Requires: M % 128 == 0, K % 128 == 0, N % 128 == 0. """ @@ -335,6 +414,8 @@ def forward( err_mode: str = "mae", e4m3_scale_bound: int = 256, row_scaled_activation: bool = False, + backward_override: Optional[str] = None, + weight_block: str = "16x16", ): M = input_hp.shape[:-1].numel() K = input_hp.shape[-1] @@ -348,6 +429,24 @@ def forward( f"four_over_six_mm requires M, K, N all divisible by 128; " f"got M={M}, K={K}, N={N}" ) + if backward_override is None: + backward_override = ( + "high_precision" if row_scaled_activation else "quantized" + ) + if backward_override not in ("quantized", "high_precision", "dequantized"): + raise ValueError( + f"backward_override must be 'quantized', 'high_precision', or " + f"'dequantized', got {backward_override!r}" + ) + if backward_override == "quantized" and row_scaled_activation: + raise ValueError( + "row-scaled four-over-six has no quantized backward; use " + "'high_precision' or 'dequantized'" + ) + if weight_block not in ("1x16", "16x16"): + raise ValueError( + f"weight_block must be '1x16' or '16x16', got {weight_block!r}" + ) input_2d = input_hp.reshape(-1, K).contiguous() if row_scaled_activation: @@ -366,7 +465,7 @@ def forward( w_codes, w_scales = four_over_six_quantize( weight_hp, w_amax, - block="16x16", + block=weight_block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound, ) @@ -407,8 +506,19 @@ def forward( if bias is not None: output = output + bias - if row_scaled_activation: + if backward_override == "high_precision": ctx.save_for_backward(input_2d, weight_hp) + elif backward_override == "dequantized": + # The rowwise operands the forward GEMM just consumed; backward + # dequantizes them, differentiating the quantized-forward function. + ctx.save_for_backward( + x_codes, + x_scales, + x_amax, + w_codes, + w_scales, + w_amax, + ) else: x_col_codes, x_col_scales = four_over_six_quantize( input_2d.t().contiguous(), @@ -420,7 +530,7 @@ def forward( w_col_codes, w_col_scales = four_over_six_quantize( weight_hp.t().contiguous(), w_amax, - block="16x16", + block=weight_block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound, ) @@ -432,7 +542,7 @@ def forward( w_col_scales, w_amax, ) - ctx.row_scaled_activation = row_scaled_activation + ctx.backward_override = backward_override ctx.e4m3_scale_bound = e4m3_scale_bound ctx.input_orig_shape = input_hp.shape ctx.has_bias = bias is not None @@ -443,10 +553,27 @@ def backward(ctx, grad_output: torch.Tensor): grad_output = grad_output.contiguous() grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]) - if ctx.row_scaled_activation: + if ctx.backward_override == "high_precision": input_2d, weight_hp = ctx.saved_tensors grad_input = (grad_output_2d @ weight_hp).reshape(ctx.input_orig_shape) grad_weight = grad_output_2d.t() @ input_2d + elif ctx.backward_override == "dequantized": + ( + x_codes, + x_scales, + x_amax, + w_codes, + w_scales, + w_amax, + ) = ctx.saved_tensors + weight_dq = four_over_six_dequantize( + w_codes, w_scales, w_amax, e4m3_scale_bound=ctx.e4m3_scale_bound + ) + input_dq = four_over_six_dequantize( + x_codes, x_scales, x_amax, e4m3_scale_bound=ctx.e4m3_scale_bound + ) + grad_input = (grad_output_2d @ weight_dq).reshape(ctx.input_orig_shape) + grad_weight = grad_output_2d.t() @ input_dq else: ( x_col_codes, @@ -488,7 +615,7 @@ def backward(ctx, grad_output: torch.Tensor): if ctx.has_bias else None ) - return grad_input, grad_weight, grad_bias, None, None, None + return grad_input, grad_weight, grad_bias, None, None, None, None, None four_over_six_linear = four_over_six_mm.apply @@ -501,6 +628,8 @@ class NVFP4FourOverSixLinear(nn.Linear): forward GEMM operands use four-over-six NVFP4, gradients use standard NVFP4 (or bf16 when ``row_scaled_activation`` is set — see the module docstring for why row-scaled has no quantized backward). + ``backward_override`` and ``weight_block`` pass through to + :class:`four_over_six_mm`. """ def __init__( @@ -511,6 +640,8 @@ def __init__( err_mode: str = "mae", e4m3_scale_bound: int = 256, row_scaled_activation: bool = False, + backward_override: Optional[str] = None, + weight_block: str = "16x16", device=None, dtype=None, ): @@ -518,6 +649,8 @@ def __init__( self.err_mode = err_mode self.e4m3_scale_bound = e4m3_scale_bound self.row_scaled_activation = row_scaled_activation + self.backward_override = backward_override + self.weight_block = weight_block def forward(self, x: torch.Tensor) -> torch.Tensor: return four_over_six_linear( @@ -527,6 +660,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.err_mode, self.e4m3_scale_bound, self.row_scaled_activation, + self.backward_override, + self.weight_block, ) @classmethod @@ -536,6 +671,8 @@ def from_linear( err_mode: str = "mae", e4m3_scale_bound: int = 256, row_scaled_activation: bool = False, + backward_override: Optional[str] = None, + weight_block: str = "16x16", ) -> "NVFP4FourOverSixLinear": new = cls( mod.in_features, @@ -544,6 +681,8 @@ def from_linear( err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound, row_scaled_activation=row_scaled_activation, + backward_override=backward_override, + weight_block=weight_block, device=mod.weight.device, dtype=mod.weight.dtype, )