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 new file mode 100644 index 0000000000..00fa7e7d3f --- /dev/null +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py @@ -0,0 +1,496 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026, NVIDIA CORPORATION. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + + +import pytest +import torch + +from torchao.float8.float8_utils import compute_error +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + NVFP4FourOverSixLinear, + four_over_six_global_encode_scale, + four_over_six_linear, + four_over_six_quantize, + nvfp4_dequantize, +) +from torchao.prototype.moe_training.nvfp4_training.nvfp4_training import ( + NVFP4Linear, + NVFP4TrainingConfig, +) +from torchao.prototype.mx_formats.kernels import f4_unpacked_to_f32, unpack_uint4 +from torchao.quantization import quantize_ +from torchao.utils import is_sm_at_least_100, torch_version_at_least + +_skip_no_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) +_skip_no_sm100 = pytest.mark.skipif( + not ( + torch.cuda.is_available() + and is_sm_at_least_100() + and torch_version_at_least("2.10.0") + ), + reason="requires SM100+ and PyTorch 2.10+ (FP4 scaled_mm)", +) + + +def _dequantize(codes, scales, global_amax, e4m3_scale_bound): + """Reconstruct FP32 values from packed codes, block scales, and global amax.""" + rows = codes.shape[0] + values = f4_unpacked_to_f32(unpack_uint4(codes)).view(rows, -1, 16) + s_dec = 1.0 / four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) + if s_dec.dim() == 1: + s_dec = s_dec.view(rows, 1, 1) + return (values * scales.to(torch.float32).unsqueeze(-1) * s_dec).view(rows, -1) + + +def _map6_reference(x, global_amax, e4m3_scale_bound): + """Standard (map-to-6 only) encoding with the four-over-six scale chain.""" + from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + _FP32_MAX, + FP4_E2M1_MAX, + FP8_E4M3_MAX, + _fp4_rtne, + ) + + rows, cols = x.shape + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) + fp4_max = torch.full((), FP4_E2M1_MAX, dtype=torch.float32, device=x.device) + base = (xf.abs().amax(dim=-1) / fp4_max) * s_enc + scale6 = base.clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + inv6 = (1.0 / (scale6.to(torch.float32) * (1.0 / s_enc))).clamp(max=_FP32_MAX) + _, values6 = _fp4_rtne(xf * inv6.unsqueeze(-1)) + s_dec = (1.0 / s_enc).view(-1, 1, 1) if s_enc.dim() == 1 else 1.0 / s_enc + dequant6 = values6 * scale6.to(torch.float32).unsqueeze(-1) * s_dec + return dequant6.view(rows, cols) + + +@_skip_no_cuda +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +def test_scales_are_candidate_scales(err_mode, e4m3_scale_bound, block): + """Every stored block scale is one of the two candidate scales.""" + torch.manual_seed(0) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + _, scales = four_over_six_quantize( + x, amax, block=block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + + xf = x.float().view(128, 16, 16) + if block == "16x16": + tiles = x.float().abs().view(8, 16, 16, 16) + block_amax = tiles.amax(dim=(1, 3)).repeat_interleave(16, dim=0) + else: + block_amax = xf.abs().amax(dim=-1) + s_enc = four_over_six_global_encode_scale(amax, e4m3_scale_bound) + fp4_max = torch.full((), 6.0, dtype=torch.float32, device="cuda") + base = (block_amax / fp4_max) * s_enc + scale6 = base.clamp(max=448.0).to(torch.float8_e4m3fn).view(torch.uint8) + scale4 = (base * 1.5).clamp(max=448.0).to(torch.float8_e4m3fn).view(torch.uint8) + got = scales.view(torch.uint8) + assert ((got == scale6) | (got == scale4)).all() + + +@_skip_no_cuda +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +def test_selection_not_worse_than_map6(e4m3_scale_bound): + """Per-block MAE of the stored encoding <= the map-to-6-only encoding.""" + torch.manual_seed(0) + x = torch.randn(256, 512, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize( + x, amax, block="1x16", err_mode="mae", e4m3_scale_bound=e4m3_scale_bound + ) + dq = _dequantize(codes, scales, amax, e4m3_scale_bound) + dq6 = _map6_reference(x, amax, e4m3_scale_bound) + xf = x.float() + err = (dq - xf).abs().view(256, -1, 16).sum(dim=-1).double() + err6 = (dq6 - xf).abs().view(256, -1, 16).sum(dim=-1).double() + # Selection minimizes the FP32 sequential-sum error; allow FP32-vs-FP64 + # summation slack on ties. + assert (err <= err6 + 1e-4).all() + # And the recipe must actually engage: some blocks pick map-to-4. + assert (err < err6 - 1e-4).any() + + +@_skip_no_cuda +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +def test_selection_minimizes_err_mode(err_mode, e4m3_scale_bound): + """The stored encoding's per-block error is the minimum over both + candidates under the configured metric (the map-to-6 comparison above + only bounds the mae side).""" + from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + _FP32_MAX, + FP4_E2M1_MAX, + FP8_E4M3_MAX, + _candidate_error, + _fp4_rtne, + ) + + torch.manual_seed(0) + rows, cols = 256, 512 + x = torch.randn(rows, cols, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + _, scales = four_over_six_quantize( + x, amax, block="1x16", err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + + # Recompute both candidate encodings with the quantizer's own chain. + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(amax, e4m3_scale_bound) + fp4_max = torch.full((), FP4_E2M1_MAX, dtype=torch.float32, device=x.device) + base = (xf.abs().amax(dim=-1) / fp4_max) * s_enc + scale6 = base.clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + scale4 = (base * 1.5).clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + s_dec = 1.0 / s_enc + inv6 = (1.0 / (scale6.to(torch.float32) * s_dec)).clamp(max=_FP32_MAX) + inv4 = (1.0 / (scale4.to(torch.float32) * s_dec)).clamp(max=_FP32_MAX) + _, values6 = _fp4_rtne(xf * inv6.unsqueeze(-1)) + _, values4 = _fp4_rtne(xf * inv4.unsqueeze(-1)) + err6 = _candidate_error( + values6, scale6.unsqueeze(-1), xf, amax, err_mode, e4m3_scale_bound + ) + err4 = _candidate_error( + values4, scale4.unsqueeze(-1), xf, amax, err_mode, e4m3_scale_bound + ) + + # Every stored scale is one of the candidates, and its error is the + # candidate minimum (equal candidate bytes encode identically, so the + # attribution below is unambiguous). + stored = scales.view(torch.uint8) + scale6_u8 = scale6.view(torch.uint8) + scale4_u8 = scale4.view(torch.uint8) + assert ((stored == scale6_u8) | (stored == scale4_u8)).all() + stored_err = torch.where(stored == scale4_u8, err4, err6) + min_err = torch.where(err4 < err6, err4, err6) + torch.testing.assert_close(stored_err, min_err, atol=0, rtol=0) + # Both candidates must win somewhere for the check to bite. + assert (err4 < err6).any() and (err6 < err4).any() + + +@_skip_no_cuda +def test_row_scaled_matches_per_row_quantization(): + """Row-scaled output == each row quantized alone with its own scalar amax.""" + torch.manual_seed(0) + x = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") + row_amax = x.abs().amax(dim=1).to(torch.float32) + codes, scales = four_over_six_quantize(x, row_amax, block="1x16") + for r in range(0, 64, 17): + codes_r, scales_r = four_over_six_quantize(x[r : r + 1], row_amax[r].view(())) + torch.testing.assert_close(codes[r : r + 1], codes_r, atol=0, rtol=0) + torch.testing.assert_close( + scales[r : r + 1].view(torch.uint8), + scales_r.view(torch.uint8), + atol=0, + rtol=0, + ) + + +@_skip_no_cuda +def test_row_scaled_rejects_16x16(): + x = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") + row_amax = x.abs().amax(dim=1).to(torch.float32) + with pytest.raises(ValueError, match="1x16 blocks only"): + four_over_six_quantize(x, row_amax, block="16x16") + + +@_skip_no_cuda +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +def test_dequant_sqnr(block): + torch.manual_seed(0) + x = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize(x, amax, block=block) + dq = _dequantize(codes, scales, amax, 256) + 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): + """nvfp4_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 = nvfp4_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 = nvfp4_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"): + nvfp4_dequantize(codes, scales, amax, e4m3_scale_bound=128) + with pytest.raises(ValueError, match="scales must have shape"): + nvfp4_dequantize(codes, scales[:, :8], amax) + with pytest.raises(ValueError, match="row vector"): + nvfp4_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]) +def test_linear_forward_backward(row_scaled_activation, bias): + 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 + ) + b = ( + torch.randn(N, dtype=torch.bfloat16, device="cuda", requires_grad=True) + if bias + else None + ) + y = four_over_six_linear(x, w, b, "mae", 256, row_scaled_activation) + assert y.shape == (M, N) + dy = torch.randn_like(y) + y.backward(dy) + + y_ref = x.detach().float() @ w.detach().float().t() + if bias: + y_ref = y_ref + b.detach().float() + dx_ref = dy.float() @ w.detach().float() + dw_ref = dy.float().t() @ x.detach().float() + assert compute_error(y_ref, y.float()).item() > 14.0 + assert compute_error(dx_ref, x.grad.float()).item() > 14.0 + assert compute_error(dw_ref, w.grad.float()).item() > 14.0 + if bias: + # grad_bias is reduced in bf16, matching nvfp4_linear. + torch.testing.assert_close(b.grad, dy.sum(dim=0)) + + +@_skip_no_sm100 +def test_linear_module(): + torch.manual_seed(0) + lin = NVFP4FourOverSixLinear(512, 384, 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 + + +@_skip_no_cuda +def test_linear_rejects_unaligned_dims(): + x = torch.randn(100, 512, dtype=torch.bfloat16, device="cuda") + w = torch.randn(384, 512, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="divisible by 128"): + 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 = nvfp4_dequantize(x_codes, x_scales, x_amax) + w_dq = nvfp4_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_cuda +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") + + +@_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 + + +def test_training_config_four_over_six_recipe_swap(): + model = torch.nn.Sequential( + torch.nn.Linear(512, 384, bias=True), torch.nn.Linear(384, 512, bias=False) + ) + weight = model[0].weight + quantize_( + model, + NVFP4TrainingConfig( + recipe="four_over_six", + err_mode="mse", + e4m3_scale_bound=448, + row_scaled_activation=True, + backward_override="dequantized", + weight_block="1x16", + ), + ) + for mod in model: + assert type(mod) is NVFP4FourOverSixLinear + assert mod.err_mode == "mse" + assert mod.e4m3_scale_bound == 448 + assert mod.row_scaled_activation is True + assert mod.backward_override == "dequantized" + assert mod.weight_block == "1x16" + assert model[0].weight is weight + assert model[0].bias is not None + assert model[1].bias is None + # Re-quantizing leaves already-converted modules alone — under the same + # recipe and under the other one (no silent cross-recipe rewrap). + converted = model[0] + quantize_(model, NVFP4TrainingConfig(recipe="four_over_six")) + assert model[0] is converted + quantize_(model, NVFP4TrainingConfig()) + assert model[0] is converted + + +def test_training_config_default_recipe_swap(): + model = torch.nn.Sequential(torch.nn.Linear(512, 384, bias=False)) + quantize_(model, NVFP4TrainingConfig()) + assert type(model[0]) is NVFP4Linear + converted = model[0] + quantize_(model, NVFP4TrainingConfig(recipe="four_over_six")) + assert model[0] is converted + + +def test_training_config_recipe_validation(): + with pytest.raises(ValueError, match="recipe must be"): + NVFP4TrainingConfig(recipe="4over6") + with pytest.raises(ValueError, match="err_mode configures the 'four_over_six'"): + NVFP4TrainingConfig(err_mode="mse") + with pytest.raises(ValueError, match="stay at its default under recipe='default'"): + NVFP4TrainingConfig(backward_override="dequantized") + with pytest.raises(ValueError, match="err_mode must be"): + NVFP4TrainingConfig(recipe="four_over_six", err_mode="rmse") + with pytest.raises(ValueError, match="e4m3_scale_bound must be"): + NVFP4TrainingConfig(recipe="four_over_six", e4m3_scale_bound=384) + with pytest.raises(ValueError, match="backward_override must be"): + NVFP4TrainingConfig(recipe="four_over_six", backward_override="bf16") + with pytest.raises(ValueError, match="weight_block must be"): + NVFP4TrainingConfig(recipe="four_over_six", weight_block="32x32") + with pytest.raises(ValueError, match="world_size configures the 'default'"): + NVFP4TrainingConfig(recipe="four_over_six", world_size=2) diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py new file mode 100644 index 0000000000..72fa422d5e --- /dev/null +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -0,0 +1,713 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2026, NVIDIA CORPORATION. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +"""NVFP4 four-over-six quantization and the linear layer that consumes it. + +Four-over-six is an adaptive NVFP4 block-scaling recipe: every quantization +block is encoded twice and the candidate with the lower dequantization error +is stored. + +* The **map-to-6** candidate is the standard NVFP4 encoding: the E4M3 block + scale maps the block amax to FP4 value 6. +* The **map-to-4** candidate expands the E4M3 block scale by 1.5x, so FP4 + value 4 reaches the range that value 6 reaches in the standard encoding. + The FP4 grid is denser around 4 than around 6, which lowers error for + blocks whose mass sits below the amax. + +Errors are compared per block with a configurable metric (mean-absolute or +mean-squared, computed in the input domain); ties select map-to-6. To leave +E4M3 headroom for the 1.5x scale expansion, the global (per-tensor) scale is +derived from a reduced E4M3 bound of 256 by default instead of 448. + +Two global-scale granularities are supported for activations: + +* per-tensor: one FP32 scale for the whole tensor (the default), and +* row-wise: one FP32 scale per tensor row, derived from that row's amax. + +The reference arithmetic pins every rounding step, so results are +reproducible bit for bit across implementations. Two details are +load-bearing: + +* The block-scale association is ``(block_amax / 6) * S_enc`` — one division + then one multiply. The standard NVFP4 path uses + ``block_amax * (S_enc * (1/6))``, which rounds differently on a fraction of + blocks. +* The per-block error is accumulated sequentially in element order with FP32 + round-to-nearest adds, and 16x16 tiles reduce their 16 row-group errors in + a pairwise halving tree. Both orders affect candidate selection on ties + near the FP32 rounding boundary. + +``four_over_six_linear`` mirrors the recipe's training semantics: + +* forward GEMM: activations quantized 1x16 four-over-six (optionally + row-scaled), weights quantized 16x16 four-over-six; +* backward with per-tensor activations: gradients use standard NVFP4 + round-to-nearest-even (four-over-six never applies to gradients), and the + saved columnwise activation/weight codes are four-over-six; +* 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``: + +* ``"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, which also + cuts activation memory. + +Weights quantize with 16x16 tiles by default; ``weight_block="1x16"`` selects +rowwise blocks instead. +""" + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from torchao.prototype.mx_formats.kernels import ( + f4_unpacked_to_f32, + f32_to_f4_unpacked, + pack_uint4, + unpack_uint4, +) +from torchao.prototype.mx_formats.utils import to_blocked + +FP4_E2M1_MAX = 6.0 +FP8_E4M3_MAX = 448.0 +_FP32_MAX = torch.finfo(torch.float32).max + +__all__ = [ + "four_over_six_global_encode_scale", + "four_over_six_quantize", + "nvfp4_dequantize", + "four_over_six_mm", + "four_over_six_linear", + "NVFP4FourOverSixLinear", +] + + +def four_over_six_global_encode_scale( + global_amax: torch.Tensor, e4m3_scale_bound: int = 256 +) -> torch.Tensor: + """Global encode scale: bound * 6 / amax. + + ``global_amax`` may be a scalar (per-tensor) or a 1-D per-row vector. + ``amax == 0`` gives inf and an enormous amax underflows the scale to + zero; both fall back to the identity scale. + """ + amax = global_amax.to(torch.float32) + candidate = torch.full_like(amax, float(e4m3_scale_bound) * FP4_E2M1_MAX) / amax + candidate = candidate.clamp(max=_FP32_MAX) + return torch.where( + (amax == 0.0) | (candidate == 0.0), torch.ones_like(candidate), candidate + ) + + +def _fp4_rtne(scaled: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Round-to-nearest-even FP4 codes and their exact FP32 values. + + Reproduces ``cvt.rn.satfinite.e2m1x2.f32`` followed by + ``cvt.rn.f16x2.e2m1x2`` (E2M1 values are exact in FP16 and FP32). + """ + clamped = scaled.clamp(-FP4_E2M1_MAX, FP4_E2M1_MAX) + codes = f32_to_f4_unpacked(clamped) + return codes, f4_unpacked_to_f32(codes) + + +def _candidate_error( + values: torch.Tensor, + scale_fp8: torch.Tensor, + xf: torch.Tensor, + global_amax: torch.Tensor, + err_mode: str, + e4m3_scale_bound: int, +) -> torch.Tensor: + """Per-block dequantization error: FP32 adds in element order, per 1x16 group. + + values/xf: (rows, num_groups, 16); scale_fp8: (rows, num_groups, 1); + global_amax broadcastable against (rows, num_groups). + """ + sf = scale_fp8.to(torch.float32)[..., 0] + # The denominator must be a tensor: dividing by a python scalar lowers to a + # multiply by its (inexact) reciprocal, which double-rounds and flips + # candidate picks near error ties. Tensor-tensor division is a true + # correctly-rounded FP32 division. + err_denom = torch.full( + (), + FP4_E2M1_MAX * float(e4m3_scale_bound), + dtype=torch.float32, + device=xf.device, + ) + err = torch.zeros_like(xf[..., 0]) + for idx in range(16): + val = ((values[..., idx] * sf) * global_amax) / err_denom + diff = val - xf[..., idx] + if err_mode == "mse": + err = err + diff * diff + else: + err = err + diff.abs() + return err + + +def _tile_error_tree_sum(err: torch.Tensor) -> torch.Tensor: + """Reduce 16 row-group errors per 16x16 tile in the warp-shuffle tree order. + + err: (rows, num_groups) with rows % 16 == 0 -> (rows // 16, num_groups). + """ + rows = err.view(err.shape[0] // 16, 16, err.shape[1]) + rows = rows[:, 0:8] + rows[:, 8:16] + rows = rows[:, 0:4] + rows[:, 4:8] + rows = rows[:, 0:2] + rows[:, 2:4] + return rows[:, 0] + rows[:, 1] + + +def four_over_six_quantize( + x: torch.Tensor, + global_amax: torch.Tensor, + *, + block: str = "1x16", + err_mode: str = "mae", + e4m3_scale_bound: int = 256, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2-D tensor to NVFP4 with four-over-six block selection. + + Args: + x: (R, C) bfloat16 or float32, C % 16 == 0 (R % 16 == 0 for 16x16). + global_amax: scalar FP32 amax, or a (R,) per-row amax vector for the + row-scaled variant (1x16 blocks only). + block: "1x16" (activations/gradient operands) or "16x16" (weights). + err_mode: "mae" or "mse" candidate-selection error metric. + e4m3_scale_bound: 256 (default, leaves map-to-4 headroom) or 448. + + Returns: + (codes, scales): (R, C//2) uint8 packed FP4 codes (low nibble = even + element) and (R, C//16) float8_e4m3fn block scales. + """ + if block not in ("1x16", "16x16"): + raise ValueError(f"block must be '1x16' or '16x16', got {block!r}") + if err_mode not in ("mae", "mse"): + raise ValueError(f"err_mode must be 'mae' or 'mse', got {err_mode!r}") + if e4m3_scale_bound not in (256, 448): + raise ValueError(f"e4m3_scale_bound must be 256 or 448, got {e4m3_scale_bound}") + if x.dim() != 2: + raise ValueError(f"x must be 2D, got {x.dim()}D") + rows, cols = x.shape + if cols % 16 != 0: + raise ValueError(f"C must be divisible by 16, got C={cols}") + if block == "16x16" and rows % 16 != 0: + raise ValueError(f"16x16 blocks require R divisible by 16, got R={rows}") + row_scaled = global_amax.dim() == 1 and global_amax.numel() == rows + if row_scaled and block != "1x16": + raise ValueError("row-scaled four-over-six supports 1x16 blocks only") + 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)}" + ) + + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) + if row_scaled: + s_enc = s_enc.view(rows, 1) + err_amax = global_amax.to(torch.float32).view(rows, 1) + else: + err_amax = global_amax.to(torch.float32) + + if block == "16x16": + tiles = xf.abs().view(rows // 16, 16, cols // 16, 16) + block_amax = tiles.amax(dim=(1, 3)).repeat_interleave(16, dim=0) + else: + block_amax = xf.abs().amax(dim=-1) + + # Scale-pair construction: base = (block_amax / 6) * S_enc, then the 1.5x + # map-to-4 expansion; both capped at the full E4M3 range. The divisor is a + # tensor for a true correctly-rounded FP32 division (a python-scalar + # divisor lowers to a reciprocal multiply, which double-rounds). + fp4_max = torch.full((), FP4_E2M1_MAX, dtype=torch.float32, device=xf.device) + base = (block_amax / fp4_max) * s_enc + scale6 = base.clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + scale4 = (base * 1.5).clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + s_dec = 1.0 / s_enc + inv6 = (1.0 / (scale6.to(torch.float32) * s_dec)).clamp(max=_FP32_MAX) + inv4 = (1.0 / (scale4.to(torch.float32) * s_dec)).clamp(max=_FP32_MAX) + + codes6, values6 = _fp4_rtne(xf * inv6.unsqueeze(-1)) + codes4, values4 = _fp4_rtne(xf * inv4.unsqueeze(-1)) + err6 = _candidate_error( + values6, scale6.unsqueeze(-1), xf, err_amax, err_mode, e4m3_scale_bound + ) + err4 = _candidate_error( + values4, scale4.unsqueeze(-1), xf, err_amax, err_mode, e4m3_scale_bound + ) + if block == "16x16": + pick4 = ( + _tile_error_tree_sum(err4) < _tile_error_tree_sum(err6) + ).repeat_interleave(16, dim=0) + else: + pick4 = err4 < err6 + + codes = torch.where(pick4.unsqueeze(-1), codes4, codes6) + scales = torch.where(pick4, scale4, scale6) + return pack_uint4(codes.view(rows, cols)), scales + + +def nvfp4_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 NVFP4 codes and block scales back to high precision. + + This is the standard NVFP4 decode — four-over-six changes only the + encode-side scale selection. It is a pure-PyTorch correctness helper, + not an optimized kernel; a fused decode kernel is future work (TODO). + 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 (this + recipe family defaults to 256; standard NVFP4 uses 448). + 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}), 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)}" + ) + return _nvfp4_dequantize_op(codes, scales, global_amax, e4m3_scale_bound, out_dtype) + + +# Registered as a custom op so torch.compile keeps the eager decode: inductor +# codegen of the fused unpack + broadcast-scale graph miscompiles the +# low-nibble lane (torch 2.14 nightly), and the dequantized backward's whole +# contract is bitwise parity with the fprop operands. +# TODO: file the upstream pytorch inductor issue and cite it here as the +# removal trigger for this workaround. +@torch.library.custom_op("torchao::nvfp4_dequantize", mutates_args=()) +def _nvfp4_dequantize_op( + codes: torch.Tensor, + scales: torch.Tensor, + global_amax: torch.Tensor, + e4m3_scale_bound: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + rows, packed_cols = codes.shape + cols = packed_cols * 2 + row_scaled = global_amax.dim() == 1 and global_amax.numel() == rows + 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) + + +@_nvfp4_dequantize_op.register_fake +def _(codes, scales, global_amax, e4m3_scale_bound, out_dtype): + return codes.new_empty((codes.shape[0], codes.shape[1] * 2), dtype=out_dtype) + + +def _standard_rtne_quantize( + x: torch.Tensor, global_amax: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Standard NVFP4 1x16 round-to-nearest-even quantize for gradient operands. + + The gradient scale chain keeps the standard association + ``block_amax * (S_enc * (1/6))`` and the full 448 E4M3 bound; only + non-gradient four-over-six tensors use the ``(block_amax / 6) * S_enc`` + association above. + """ + rows, cols = x.shape + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound=448) + block_amax = xf.abs().amax(dim=-1) + scales = ( + (block_amax * (s_enc * (1.0 / FP4_E2M1_MAX))) + .clamp(max=FP8_E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + enc = (1.0 / (scales.to(torch.float32) * (1.0 / s_enc))).clamp(max=_FP32_MAX) + codes, _ = _fp4_rtne(xf * enc.unsqueeze(-1)) + return pack_uint4(codes.view(rows, cols)), scales + + +def _global_decode_scale(amax: torch.Tensor, e4m3_scale_bound: int) -> torch.Tensor: + """Per-tensor decode scale consumed by the GEMM: amax / (bound * 6). + + The scalar divisor (a reciprocal-multiply lowering) is fine here, unlike + the encode chain's tensor divisors: this factor never enters the encode + arithmetic — consumers reconstruct it for the GEMM's per-tensor scale + slot — so it sits outside the div.rn encode contract, and the + linear-level tests pin the resulting GEMM outputs. + """ + return amax.to(torch.float32) / (float(e4m3_scale_bound) * FP4_E2M1_MAX) + + +def _scaled_mm_nvfp4( + a_codes: torch.Tensor, + a_scales: torch.Tensor, + a_global: torch.Tensor, + b_codes_t: torch.Tensor, + b_scales: torch.Tensor, + b_global: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Block-scaled FP4 GEMM with per-tensor second-level scales. + + a_codes: (M, K//2) uint8; b_codes_t: (K//2, N) transposed uint8 view; + a_scales/b_scales: plain (rows, K//16) float8 block scales (swizzled here). + """ + return F.scaled_mm( + a_codes.view(torch.float4_e2m1fn_x2), + b_codes_t.view(torch.float4_e2m1fn_x2), + scale_a=[to_blocked(a_scales).flatten(), a_global], + scale_recipe_a=[F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise], + scale_b=[to_blocked(b_scales).flatten(), b_global], + scale_recipe_b=[F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise], + swizzle_a=[F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE], + swizzle_b=[F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE], + output_dtype=out_dtype, + ) + + +@torch._dynamo.allow_in_graph +class four_over_six_mm(torch.autograd.Function): + """NVFP4 four-over-six quantized matmul. + + 3 GEMMs: + forward: x_row @ W.T = output (1x16 four-over-six x, 16x16 four-over-six W) + backward: dy_row @ W.T = grad_input (standard-NVFP4 dy; saved columnwise W) + backward: dy_col.T @ x_col = grad_weight (standard-NVFP4 dy; saved columnwise x) + + 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. Non-bf16 inputs are + cast to bf16 and gradients are always bf16, so leaves that require grad + must be bf16 (matching ``nvfp4_mm_triton``). + """ + + @staticmethod + def forward( + ctx, + input_hp: torch.Tensor, + weight_hp: torch.Tensor, + bias: Optional[torch.Tensor], + 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] + N = weight_hp.shape[0] + if input_hp.dtype != torch.bfloat16: + input_hp = input_hp.to(torch.bfloat16) + if weight_hp.dtype != torch.bfloat16: + weight_hp = weight_hp.to(torch.bfloat16) + if M % 128 != 0 or K % 128 != 0 or N % 128 != 0: + raise ValueError( + 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'" + ) + input_2d = input_hp.reshape(-1, K).contiguous() + + if row_scaled_activation: + x_amax = input_2d.abs().amax(dim=1).to(torch.float32) + else: + x_amax = input_2d.abs().amax().to(torch.float32) + w_amax = weight_hp.abs().amax().to(torch.float32) + + x_codes, x_scales = four_over_six_quantize( + input_2d, + x_amax, + block="1x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_codes, w_scales = four_over_six_quantize( + weight_hp, + w_amax, + block=weight_block, + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_global = _global_decode_scale(w_amax, e4m3_scale_bound) + + if row_scaled_activation: + # The GEMM's per-tensor slot cannot hold a per-row scale: run it + # with the constant 1/(6*bound) factor, then apply the raw per-row + # amaxes on the FP32 output before the bf16 cast. + x_global = torch.full( + (), + 1.0 / (FP4_E2M1_MAX * float(e4m3_scale_bound)), + dtype=torch.float32, + device=input_2d.device, + ) + output = _scaled_mm_nvfp4( + x_codes, + x_scales, + x_global, + w_codes.t(), + w_scales, + w_global, + torch.float32, + ) + output = (output * x_amax.view(-1, 1)).to(torch.bfloat16) + else: + x_global = _global_decode_scale(x_amax, e4m3_scale_bound) + output = _scaled_mm_nvfp4( + x_codes, + x_scales, + x_global, + w_codes.t(), + w_scales, + w_global, + torch.bfloat16, + ) + output = output.reshape(*input_hp.shape[:-1], N) + if bias is not None: + output = output + bias.to(output.dtype) + + 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(), + x_amax, + block="1x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_col_codes, w_col_scales = four_over_six_quantize( + weight_hp.t().contiguous(), + w_amax, + block=weight_block, + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + ctx.save_for_backward( + x_col_codes, + x_col_scales, + x_amax, + w_col_codes, + w_col_scales, + w_amax, + ) + 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 + return output + + @staticmethod + 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.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 = nvfp4_dequantize( + w_codes, w_scales, w_amax, e4m3_scale_bound=ctx.e4m3_scale_bound + ) + input_dq = nvfp4_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, + x_col_scales, + x_amax, + w_col_codes, + w_col_scales, + w_amax, + ) = ctx.saved_tensors + dy_amax = grad_output_2d.abs().amax().to(torch.float32) + dy_row_codes, dy_row_scales = _standard_rtne_quantize( + grad_output_2d, dy_amax + ) + dy_col_codes, dy_col_scales = _standard_rtne_quantize( + grad_output_2d.t().contiguous(), dy_amax + ) + dy_global = _global_decode_scale(dy_amax, 448) + grad_input = _scaled_mm_nvfp4( + dy_row_codes, + dy_row_scales, + dy_global, + w_col_codes.t(), + w_col_scales, + _global_decode_scale(w_amax, ctx.e4m3_scale_bound), + torch.bfloat16, + ).reshape(ctx.input_orig_shape) + grad_weight = _scaled_mm_nvfp4( + dy_col_codes, + dy_col_scales, + dy_global, + x_col_codes.t(), + x_col_scales, + _global_decode_scale(x_amax, ctx.e4m3_scale_bound), + torch.bfloat16, + ) + + grad_bias = ( + grad_output.sum(dim=tuple(range(grad_output.dim() - 1))) + if ctx.has_bias + else None + ) + return grad_input, grad_weight, grad_bias, None, None, None, None, None + + +four_over_six_linear = four_over_six_mm.apply + + +class NVFP4FourOverSixLinear(nn.Linear): + """Linear layer with NVFP4 four-over-six quantized GEMMs. + + Drop-in replacement for nn.Linear implementing the four-over-six recipe: + 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`. ``bias`` defaults off, matching the recipe's + usage (unlike nn.Linear). + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + 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, + ): + super().__init__(in_features, out_features, bias, device=device, dtype=dtype) + 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( + x, + self.weight, + self.bias, + self.err_mode, + self.e4m3_scale_bound, + self.row_scaled_activation, + self.backward_override, + self.weight_block, + ) + + @classmethod + def from_linear( + cls, + mod: nn.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, + mod.out_features, + mod.bias is not None, + 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, + ) + if mod.weight.device != torch.device("meta"): + new.weight = mod.weight + if mod.bias is not None: + new.bias = mod.bias + return new diff --git a/torchao/prototype/moe_training/nvfp4_training/nvfp4_training.py b/torchao/prototype/moe_training/nvfp4_training/nvfp4_training.py index ba0a93a29e..bd3c0afdbd 100644 --- a/torchao/prototype/moe_training/nvfp4_training/nvfp4_training.py +++ b/torchao/prototype/moe_training/nvfp4_training/nvfp4_training.py @@ -71,12 +71,12 @@ def _make_rht_sign_vector( class NVFP4TrainingConfig(AOBaseConfig): """Configuration for NVFP4 quantized training. - When passed to quantize_(), replaces nn.Linear modules with - NVFP4Linear, which quantizes all three GEMMs (forward - and backward) to NVFP4. + When passed to quantize_(), replaces nn.Linear modules with the module + implementing the selected recipe, which quantizes all three GEMMs + (forward and backward) to NVFP4. Args: - kernel_preference: Backend for quantization kernels. + kernel_preference: ("default" recipe) Backend for quantization kernels. TRITON: Pure-Triton RHT + stochastic rounding path. CUTEDSL: CuteDSL kernels for the full quantize path (amax, forward RTNE quantize, SR backward quantize, and 2D weight quantize). @@ -84,11 +84,13 @@ class NVFP4TrainingConfig(AOBaseConfig): by 256. Under tensor parallel the same constraints apply to each per-rank shard, and the per-rank M shard must be divisible by 256. Default: TRITON. - process_group: Optional ProcessGroup for tensor-parallel TP. + process_group: ("default" recipe) Optional ProcessGroup for + tensor-parallel TP. When set, forward dispatches to the selected NVFP4 tensor-parallel path (TRITON or CUTEDSL). - world_size: TP world size. Inferred from process_group if None. - rht_sign_vector: Optional {-1, 1} sign vector of length 16 for the + world_size: ("default" recipe) TP world size. Inferred from + process_group if None. + rht_sign_vector: ("default" recipe) Optional {-1, 1} sign vector of length 16 for the randomized Hadamard transform. When None, each NVFP4Linear draws its own random vector. In multi-rank settings (FSDP) replicas will therefore have different bases — harmless for convergence but @@ -96,12 +98,90 @@ class NVFP4TrainingConfig(AOBaseConfig): consistency should broadcast a single vector before calling quantize_() and pass it here. The TP path always enforces consistency via _replicate_rht_sign_vector regardless of this field. + recipe: Which NVFP4 training recipe to install. + "default": NVFP4Linear — randomized Hadamard transform + stochastic + rounding, configured by the fields above. + "four_over_six": NVFP4FourOverSixLinear — adaptive per-block + candidate selection between the standard map-to-6 encoding and + a 1.5x-scale map-to-4 encoding, configured by the fields below. + Stateless (no RHT/SR buffers, no TP support). + Each recipe's fields must stay at their defaults under the other + recipe. + err_mode: ("four_over_six" recipe) Candidate-selection error metric, + "mae" or "mse". + e4m3_scale_bound: ("four_over_six" recipe) Global E4M3 scale bound; + 256 leaves map-to-4 headroom, 448 uses the full range. + row_scaled_activation: ("four_over_six" recipe) Derive one FP32 global + scale per activation row instead of per tensor. + backward_override: ("four_over_six" recipe) Backward mode — None + (recipe default), "quantized", "high_precision", or "dequantized". + weight_block: ("four_over_six" recipe) Weight block granularity, + "16x16" or "1x16". """ kernel_preference: KernelPreference = KernelPreference.TRITON process_group: Optional[object] = field(default=None, compare=False) world_size: Optional[int] = None rht_sign_vector: Optional[object] = field(default=None, compare=False) + recipe: str = "default" + err_mode: str = "mae" + e4m3_scale_bound: int = 256 + row_scaled_activation: bool = False + backward_override: Optional[str] = None + weight_block: str = "16x16" + + def __post_init__(self): + if self.recipe not in ("default", "four_over_six"): + raise ValueError( + f"recipe must be 'default' or 'four_over_six', got {self.recipe!r}" + ) + if self.recipe == "default": + four_over_six_defaults = ( + ("err_mode", self.err_mode, "mae"), + ("e4m3_scale_bound", self.e4m3_scale_bound, 256), + ("row_scaled_activation", self.row_scaled_activation, False), + ("backward_override", self.backward_override, None), + ("weight_block", self.weight_block, "16x16"), + ) + for name, value, default in four_over_six_defaults: + if value != default: + raise ValueError( + f"{name} configures the 'four_over_six' recipe and must " + f"stay at its default under recipe='default', got {value!r}" + ) + return + if self.err_mode not in ("mae", "mse"): + raise ValueError(f"err_mode must be 'mae' or 'mse', got {self.err_mode!r}") + if self.e4m3_scale_bound not in (256, 448): + raise ValueError( + f"e4m3_scale_bound must be 256 or 448, got {self.e4m3_scale_bound}" + ) + if self.backward_override not in ( + None, + "quantized", + "high_precision", + "dequantized", + ): + raise ValueError( + f"backward_override must be None, 'quantized', 'high_precision', " + f"or 'dequantized', got {self.backward_override!r}" + ) + if self.weight_block not in ("1x16", "16x16"): + raise ValueError( + f"weight_block must be '1x16' or '16x16', got {self.weight_block!r}" + ) + default_recipe_defaults = ( + ("kernel_preference", self.kernel_preference, KernelPreference.TRITON), + ("process_group", self.process_group, None), + ("world_size", self.world_size, None), + ("rht_sign_vector", self.rht_sign_vector, None), + ) + for name, value, default in default_recipe_defaults: + if value is not default and value != default: + raise ValueError( + f"{name} configures the 'default' recipe and must stay at " + f"its default under recipe='four_over_six', got {value!r}" + ) class NVFP4Linear(nn.Linear): @@ -242,15 +322,31 @@ def _nvfp4_training_transform( config: NVFP4TrainingConfig, parameter_name: Optional[str] = None, ) -> nn.Module: - """Handler for quantize_(): replaces nn.Linear with NVFP4Linear.""" - if isinstance(module, NVFP4Linear): + """Handler for quantize_(): replaces nn.Linear with config.recipe's module. + + Modules already converted to either recipe are left alone. + """ + from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + NVFP4FourOverSixLinear, + ) + + if isinstance(module, (NVFP4Linear, NVFP4FourOverSixLinear)): + return module + if not isinstance(module, nn.Linear): return module - if isinstance(module, nn.Linear): - return NVFP4Linear.from_linear( + if config.recipe == "four_over_six": + return NVFP4FourOverSixLinear.from_linear( module, - kernel_preference=config.kernel_preference, - process_group=config.process_group, - world_size=config.world_size, - rht_sign_vector=config.rht_sign_vector, + err_mode=config.err_mode, + e4m3_scale_bound=config.e4m3_scale_bound, + row_scaled_activation=config.row_scaled_activation, + backward_override=config.backward_override, + weight_block=config.weight_block, ) - return module + return NVFP4Linear.from_linear( + module, + kernel_preference=config.kernel_preference, + process_group=config.process_group, + world_size=config.world_size, + rht_sign_vector=config.rht_sign_vector, + )