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 3e3164e4fa..09babf7767 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 @@ -239,11 +239,28 @@ def test_row_scaled_matches_per_row_quantization(): @_skip_no_cuda -def test_row_scaled_rejects_16x16(): +def test_per_row_amax_16x16_matches_per_tile_slices(): + """A tile-uniform per-row amax vector with 16x16 blocks == quantizing + each 16-row slice with its scalar amax (the stacked-expert-weight + contract).""" + torch.manual_seed(0) 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") + slice_amax = x.float().abs().view(4, 16 * 256).amax(dim=1) + expanded = slice_amax.repeat_interleave(16) + codes, scales = four_over_six_quantize(x, expanded, block="16x16") + for s in range(4): + ref_codes, ref_scales = four_over_six_quantize( + x[s * 16 : (s + 1) * 16].contiguous(), slice_amax[s], block="16x16" + ) + torch.testing.assert_close( + codes[s * 16 : (s + 1) * 16], ref_codes, atol=0, rtol=0 + ) + torch.testing.assert_close( + scales[s * 16 : (s + 1) * 16].view(torch.uint8), + ref_scales.view(torch.uint8), + atol=0, + rtol=0, + ) @_skip_no_cuda @@ -262,11 +279,16 @@ def test_dequant_sqnr(block): @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) + if row_scaled and block == "16x16": + # Per-row amaxes with 16x16 tiles follow the stacked-expert-weight + # contract: constant within every 16-row tile. + amax = x.float().abs().view(8, 16 * 512).amax(dim=1).repeat_interleave(16) + elif row_scaled: + amax = x.abs().amax(dim=1).to(torch.float32) + else: + amax = 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 @@ -560,8 +582,6 @@ def test_cutedsl_bitwise_matches_reference( cases: rounding-boundary straddles, dtype-max saturation, bf16 subnormals, and negative zeros. """ - if row_scaled and block == "16x16": - pytest.skip("row-scaled is 1x16 only") shapes = [(128, 256), (64, 1024), (384, 256)] if block == "1x16": shapes.append((100, 320)) diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py new file mode 100644 index 0000000000..56aeef646b --- /dev/null +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six_grouped.py @@ -0,0 +1,625 @@ +# 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.utils import is_sm_at_least_100, torch_version_at_least + +if not torch_version_at_least("2.10.0"): + pytest.skip( + "four_over_six_grouped reads FP4 scaled_grouped_mm scale/swizzle " + "types at import time (torch 2.10+)", + allow_module_level=True, + ) + +import torchao.prototype.moe_training.nvfp4_training.four_over_six as four_over_six_module +from torchao.float8.float8_utils import compute_error +from torchao.prototype.moe_training.config import NVFP4FourOverSixTrainingOpConfig +from torchao.prototype.moe_training.nvfp4_training import four_over_six_grouped +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + four_over_six_linear, + four_over_six_quantize, + nvfp4_dequantize, +) +from torchao.prototype.moe_training.nvfp4_training.four_over_six_grouped import ( + four_over_six_grouped_mm, +) +from torchao.prototype.moe_training.utils import _quantize_then_scaled_grouped_mm + +_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_grouped_mm)", +) + + +def _make_grouped_inputs(group_sizes, K, N, seed=0, device="cuda"): + """Packed activations, stacked expert weights, and end offsets.""" + torch.manual_seed(seed) + M = sum(group_sizes) + E = len(group_sizes) + A = torch.randn(M, K, dtype=torch.bfloat16, device=device) + B = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.1 + offs = torch.tensor(group_sizes, dtype=torch.int32, device=device).cumsum( + 0, dtype=torch.int32 + ) + return A, B, offs + + +@_skip_no_sm100 +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +def test_group_expanded_amax_matches_per_split_quantize(err_mode, e4m3_scale_bound): + """One quantize call with group-expanded amaxes == a per-split loop.""" + group_sizes = [128, 384, 256] + A, _, offs = _make_grouped_inputs(group_sizes, K=256, N=128) + group_amax = torch.stack( + [ + A[start:end].abs().amax().to(torch.float32) + for start, end in zip([0, *offs.tolist()[:-1]], offs.tolist()) + ] + ) + expanded = group_amax.repeat_interleave(torch.tensor(group_sizes, device=A.device)) + codes, scales = four_over_six_quantize( + A, expanded, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + start = 0 + for g, end in enumerate(offs.tolist()): + split_codes, split_scales = four_over_six_quantize( + A[start:end].contiguous(), + group_amax[g], + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + torch.testing.assert_close(codes[start:end], split_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales[start:end].view(torch.uint8), + split_scales.view(torch.uint8), + atol=0, + rtol=0, + ) + start = end + + +@_skip_no_sm100 +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_per_tensor_grouped_forward_matches_dense_loop(weight_block): + """Grouped forward vs dense four_over_six GEMMs per 128-aligned group. + + The quantized operands are bitwise-identical by construction (pinned by + the amax-expansion test above); the GEMM outputs are compared bitwise + and fall back to an SQNR bound if the grouped and dense kernels reduce + in different orders. + """ + group_sizes = [128, 256, 128] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N) + y = four_over_six_grouped_mm(A, B, offs, weight_block=weight_block) + + start = 0 + refs = [] + for e, end in enumerate(offs.tolist()): + refs.append( + four_over_six_linear( + A[start:end].contiguous(), + B[e], + None, + "mae", + 256, + False, + "high_precision", + weight_block, + ) + ) + start = end + y_ref = torch.cat(refs) + if not torch.equal(y, y_ref): + sqnr = compute_error(y_ref.float(), y.float()) + assert sqnr > 85.0, f"grouped vs dense-loop forward SQNR {sqnr:.1f} dB" + print(f"\ngrouped GEMM reduction differs from dense: SQNR {sqnr:.1f} dB") + + +@_skip_no_sm100 +@pytest.mark.parametrize( + "group_sizes, pad", + [ + pytest.param([128, 256, 128], False, id="uniform-128-aligned"), + pytest.param([128] * 64, False, id="uniform-128-rows-per-expert"), + pytest.param( + [128 if e % 8 == 0 else 0 for e in range(64)], + False, + id="decode-like-8-active-56-empty", + ), + pytest.param([1, 220, 77], True, id="ragged-padded"), + ], +) +def test_row_scaled_grouped_forward_matches_loop_oracle(group_sizes, pad, monkeypatch): + """Row-scaled fused single-GEMM forward vs per-group dense-GEMM oracles. + + torch's grouped GEMM only emits bf16 high-precision output — one + rounding ahead of the fp32 row scale that a loop of dense FP32-output + GEMMs does not have. The exactness + check compares against the dense loop with that rounding emulated on + its GEMM outputs (bitwise, with the reduction-order SQNR fallback); + the raw-loop comparison bounds the rounding cost. The quantized + operands are identical by construction. + """ + A, B, offs = _make_grouped_inputs(group_sizes, K=2048, N=768, seed=5) + N = B.shape[1] + y_fused = four_over_six_grouped_mm( + A, B, offs, row_scaled_activation=True, pad_token_groups_for_grouped_mm=pad + ) + + def _dense_loop(): + refs = [] + start = 0 + for e, end in enumerate(offs.tolist()): + rows = A[start:end] + start = end + if rows.shape[0] == 0: + refs.append(A.new_zeros(0, N)) + continue + padded_rows = 128 * ((rows.shape[0] + 127) // 128) + padded = torch.zeros( + padded_rows, A.shape[1], dtype=A.dtype, device=A.device + ) + padded[: rows.shape[0]] = rows + ref = four_over_six_linear(padded, B[e], None, "mae", 256, True) + refs.append(ref[: rows.shape[0]]) + return torch.cat(refs) + + y_loop = _dense_loop() + dense_gemm = four_over_six_module._scaled_mm_nvfp4 + + def _bf16_rounded_gemm(*args): + return dense_gemm(*args).to(torch.bfloat16).to(torch.float32) + + monkeypatch.setattr(four_over_six_module, "_scaled_mm_nvfp4", _bf16_rounded_gemm) + y_emul = _dense_loop() + + assert y_fused.shape == y_loop.shape + if not torch.equal(y_fused, y_emul): + sqnr = compute_error(y_emul.float(), y_fused.float()) + assert sqnr > 85.0, f"fused vs rounding-emulated oracle SQNR {sqnr:.1f} dB" + print(f"\nfused grouped GEMM reduction differs from dense: SQNR {sqnr:.1f} dB") + rounding_sqnr = compute_error(y_loop.float(), y_fused.float()) + assert rounding_sqnr > 45.0, f"fused vs loop-oracle SQNR {rounding_sqnr:.1f} dB" + print(f"\nbf16 GEMM-output rounding cost vs loop: SQNR {rounding_sqnr:.1f} dB") + + +@_skip_no_sm100 +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +@pytest.mark.parametrize("weight_block", ["1x16", "16x16"]) +def test_batched_weight_quantize_matches_per_expert_loop( + err_mode, e4m3_scale_bound, weight_block +): + """The one-call flattened weight quantize == the per-expert loop.""" + torch.manual_seed(2) + E, N, K = 5, 128, 256 + B = torch.randn(E, N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + weight_amax = B.abs().amax(dim=(1, 2)).to(torch.float32) + codes, scales = four_over_six_grouped._quantize_expert_weights( + B, weight_amax, weight_block, err_mode, e4m3_scale_bound + ) + for e in range(E): + ref_codes, ref_scales = four_over_six_quantize( + B[e], + weight_amax[e], + block=weight_block, + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + torch.testing.assert_close(codes[e], ref_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales[e].view(torch.uint8), + ref_scales.view(torch.uint8), + atol=0, + rtol=0, + ) + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_grouped_backward_high_precision(row_scaled_activation): + """dx/dw are bf16 grouped GEMMs on the original operands.""" + group_sizes = [128, 256, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, B, offs, row_scaled_activation=row_scaled_activation + ) + dy = torch.randn_like(y) + y.backward(dy) + + dx_ref = torch._grouped_mm(dy, B.detach(), offs=offs, out_dtype=torch.bfloat16) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), A.detach(), offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_grouped_backward_dequantized(row_scaled_activation, weight_block): + """dx/dw are bf16 grouped GEMMs on dequantized fprop operands.""" + group_sizes = [128, 256, 128] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + err_mode="mse", + row_scaled_activation=row_scaled_activation, + weight_block=weight_block, + backward_override="dequantized", + ) + dy = torch.randn_like(y) + y.backward(dy) + + A_hp, B_hp = A.detach(), B.detach() + if row_scaled_activation: + x_amax = A_hp.abs().amax(dim=1).to(torch.float32) + else: + group_amax = [] + start = 0 + for end in offs.tolist(): + group_amax.append(A_hp[start:end].abs().amax().to(torch.float32)) + start = end + x_amax = torch.stack(group_amax).repeat_interleave( + torch.tensor(group_sizes, device=A.device) + ) + x_codes, x_scales = four_over_six_quantize(A_hp, x_amax, err_mode="mse") + x_dq = nvfp4_dequantize(x_codes, x_scales, x_amax) + w_dq = [] + for e in range(B.shape[0]): + w_amax = B_hp[e].abs().amax().to(torch.float32) + w_codes, w_scales = four_over_six_quantize( + B_hp[e], w_amax, block=weight_block, err_mode="mse" + ) + w_dq.append(nvfp4_dequantize(w_codes, w_scales, w_amax)) + w_dq = torch.stack(w_dq) + + dx_ref = torch._grouped_mm(dy, w_dq, offs=offs, out_dtype=torch.bfloat16) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), x_dq, offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("backward_override", ["high_precision", "dequantized"]) +def test_grouped_backward_empty_groups(backward_override): + """Decode-like offsets with zero-size groups, forward and backward. + + dx/dw are bitwise vs the same grouped GEMMs on the reference operands; + experts that received no tokens get all-zero weight gradients. + """ + group_sizes = [128, 0, 256, 0, 0, 128] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N, seed=9) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + row_scaled_activation=True, + backward_override=backward_override, + ) + dy = torch.randn_like(y) + y.backward(dy) + + A_hp, B_hp = A.detach(), B.detach() + if backward_override == "high_precision": + x_ref, w_ref = A_hp, B_hp + else: + x_amax = A_hp.abs().amax(dim=1).to(torch.float32) + x_codes, x_scales = four_over_six_quantize(A_hp, x_amax) + x_ref = nvfp4_dequantize(x_codes, x_scales, x_amax) + w_dq = [] + for e in range(B.shape[0]): + w_amax = B_hp[e].abs().amax().to(torch.float32) + w_codes, w_scales = four_over_six_quantize(B_hp[e], w_amax, block="16x16") + w_dq.append(nvfp4_dequantize(w_codes, w_scales, w_amax)) + w_ref = torch.stack(w_dq) + + dx_ref = torch._grouped_mm(dy, w_ref, offs=offs, out_dtype=torch.bfloat16) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), x_ref, offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + empty = [e for e, size in enumerate(group_sizes) if size == 0] + assert (B.grad[empty] == 0).all() + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_grouped_padding_matches_aligned(row_scaled_activation, monkeypatch): + """Unaligned groups with padding == an aligned construction, per group. + + The row-scaled dense references emulate the fused grouped GEMM's bf16 + output rounding (see the loop-oracle test) so the comparison stays + bitwise; padding semantics are identical on every path. + """ + if row_scaled_activation: + dense_gemm = four_over_six_module._scaled_mm_nvfp4 + monkeypatch.setattr( + four_over_six_module, + "_scaled_mm_nvfp4", + lambda *args: dense_gemm(*args).to(torch.bfloat16).to(torch.float32), + ) + K, N = 256, 384 + aligned_sizes = [128, 256, 128] + ragged_sizes = [100, 220, 77] + A_al, B, offs_al = _make_grouped_inputs(aligned_sizes, K=K, N=N, seed=3) + # Ragged view: the first rows of each aligned group, so every ragged + # group's rows (and hence its amax and quantization) exist verbatim in + # the aligned run. + ragged_rows = [] + start = 0 + for size, ragged in zip(aligned_sizes, ragged_sizes): + ragged_rows.append(A_al[start : start + ragged]) + start += size + A_rg = torch.cat(ragged_rows).contiguous() + offs_rg = torch.tensor(ragged_sizes, dtype=torch.int32, device=A_al.device).cumsum( + 0, dtype=torch.int32 + ) + + y_rg = four_over_six_grouped_mm( + A_rg, + B, + offs_rg, + row_scaled_activation=row_scaled_activation, + pad_token_groups_for_grouped_mm=True, + ) + assert y_rg.shape == (sum(ragged_sizes), N) + + # Reference: dense per-group forward on the ragged rows padded to 128. + start = 0 + for e, ragged in enumerate(ragged_sizes): + rows = A_rg[start : start + ragged] + padded = torch.zeros( + 128 * ((ragged + 127) // 128), K, dtype=rows.dtype, device=rows.device + ) + padded[:ragged] = rows + if row_scaled_activation: + ref = four_over_six_linear(padded, B[e], None, "mae", 256, True) + else: + # Per-tensor group scale comes from the real rows' amax; the + # zero padding rows cannot change it. + ref = four_over_six_linear( + padded, B[e], None, "mae", 256, False, "high_precision" + ) + torch.testing.assert_close( + y_rg[start : start + ragged], ref[:ragged], atol=0, rtol=0 + ) + start += ragged + + +@_skip_no_sm100 +def test_grouped_validation(): + group_sizes = [128, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=128) + with pytest.raises(ValueError, match="no quantized backward"): + four_over_six_grouped_mm(A, B, offs, backward_override="quantized") + with pytest.raises(ValueError, match="1D int32"): + four_over_six_grouped_mm(A, B, offs.to(torch.int64)) + with pytest.raises(ValueError, match="one group-end offset per expert"): + four_over_six_grouped_mm(A, B, offs[:1]) + with pytest.raises(ValueError, match="divisible by 128"): + four_over_six_grouped_mm(A[:, :144], B[:, :, :144].contiguous(), offs) + + +@_skip_no_sm100 +def test_grouped_rl_rollout_recipe_point(): + """The RL rollout recipe point: row-scaled + MSE + bound 256 + 1x16 + weights + dequantized backward, on ragged token groups.""" + group_sizes = [100, 220, 77] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384, seed=7) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + err_mode="mse", + e4m3_scale_bound=256, + row_scaled_activation=True, + weight_block="1x16", + backward_override="dequantized", + pad_token_groups_for_grouped_mm=True, + ) + assert y.shape == (sum(group_sizes), 384) + y.backward(torch.randn_like(y)) + assert A.grad is not None and A.grad.shape == A.shape + assert B.grad is not None and B.grad.shape == B.shape + sqnr = compute_error( + torch._grouped_mm(A.detach(), B.detach().transpose(-2, -1), offs=offs).float(), + y.float(), + ) + assert sqnr > 14.0, f"quantization noise floor too high: {sqnr:.1f} dB" + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +@pytest.mark.parametrize("weight_block", ["16x16", "1x16"]) +def test_grouped_backward_dequantized_ragged(row_scaled_activation, weight_block): + """Ragged groups + padding + dequantized backward, value-checked. + + This is the composition the torchtitan grouped-experts hook ships; + the padded rows quantize to zeros and are unpadded away before the + backward GEMMs, so the reference can quantize the ragged rows directly. + """ + group_sizes = [100, 220, 77] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N, seed=7) + A.requires_grad_(True) + B.requires_grad_(True) + y = four_over_six_grouped_mm( + A, + B, + offs, + err_mode="mse", + row_scaled_activation=row_scaled_activation, + weight_block=weight_block, + backward_override="dequantized", + pad_token_groups_for_grouped_mm=True, + ) + dy = torch.randn_like(y) + y.backward(dy) + + A_hp, B_hp = A.detach(), B.detach() + if row_scaled_activation: + x_amax = A_hp.abs().amax(dim=1).to(torch.float32) + else: + # Group amaxes come from the real rows; zero padding cannot raise them. + group_amax = [] + start = 0 + for end in offs.tolist(): + group_amax.append(A_hp[start:end].abs().amax().to(torch.float32)) + start = end + x_amax = torch.stack(group_amax).repeat_interleave( + torch.tensor(group_sizes, device=A.device) + ) + x_codes, x_scales = four_over_six_quantize(A_hp, x_amax, err_mode="mse") + x_dq = nvfp4_dequantize(x_codes, x_scales, x_amax) + w_dq = [] + for e in range(B.shape[0]): + w_amax = B_hp[e].abs().amax().to(torch.float32) + w_codes, w_scales = four_over_six_quantize( + B_hp[e], w_amax, block=weight_block, err_mode="mse" + ) + w_dq.append(nvfp4_dequantize(w_codes, w_scales, w_amax)) + w_dq = torch.stack(w_dq) + + dx_ref = torch._grouped_mm(dy, w_dq, offs=offs, out_dtype=torch.bfloat16) + dw_ref = torch._grouped_mm( + dy.transpose(-2, -1), x_dq, offs=offs, out_dtype=torch.bfloat16 + ) + torch.testing.assert_close(A.grad, dx_ref, atol=0, rtol=0) + torch.testing.assert_close(B.grad, dw_ref, atol=0, rtol=0) + + +@_skip_no_sm100 +@pytest.mark.parametrize("tail_rows", [0, 128]) +def test_dispatcher_grouped_mm_four_over_six(tail_rows): + """NVFP4FourOverSixTrainingOpConfig drives this op through the grouped + GEMM dispatcher, bitwise vs a direct call. + + The dispatcher hands weights over as B_t = (E, K, N). Activation + buffers over-allocated past offs[-1] (padded token dispatchers + over-allocate to worst-case capacity) come back zero-extended with + zero tail gradients, and the logical rows match the exact-shape + reference — proof the garbage tail cannot feed the per-group amaxes. + """ + group_sizes = [128, 256, 128] + K, N = 256, 384 + A, B, offs = _make_grouped_inputs(group_sizes, K=K, N=N, seed=17) + M_logical = A.shape[0] + if tail_rows: + # Garbage tail: any leak into the last group's amax would flip its + # scale chain and break the bitwise comparison below. + tail = torch.full((tail_rows, K), 123.0, dtype=A.dtype, device=A.device) + A = torch.cat([A, tail]) + kwargs = dict( + err_mode="mse", + e4m3_scale_bound=256, + row_scaled_activation=False, + weight_block="1x16", + backward_override="dequantized", + pad_token_groups_for_grouped_mm=False, + ) + config = NVFP4FourOverSixTrainingOpConfig(**kwargs) + + A_d = A.clone().requires_grad_(True) + B_d = B.clone().requires_grad_(True) + y_d = _quantize_then_scaled_grouped_mm( + A_d, B_d.transpose(-2, -1), config=config, offs=offs + ) + assert y_d.shape == (A.shape[0], N) + dy = torch.randn_like(y_d) + y_d.backward(dy) + + A_r = A[:M_logical].clone().requires_grad_(True) + B_r = B.clone().requires_grad_(True) + y_r = four_over_six_grouped_mm(A_r, B_r, offs, **kwargs) + y_r.backward(dy[:M_logical]) + + torch.testing.assert_close(y_d[:M_logical], y_r, atol=0, rtol=0) + torch.testing.assert_close(A_d.grad[:M_logical], A_r.grad, atol=0, rtol=0) + torch.testing.assert_close(B_d.grad, B_r.grad, atol=0, rtol=0) + if tail_rows: + assert (y_d[M_logical:] == 0).all() + assert (A_d.grad[M_logical:] == 0).all() + + +@_skip_no_sm100 +@pytest.mark.parametrize("backward_override", ["high_precision", "dequantized"]) +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +def test_grouped_compile(backward_override, row_scaled_activation): + """fullgraph compile of the grouped op, forward and backward. + + The op is nonstrict-traced under compile, so eager numerics carry over + bitwise. Both scale granularities compile: the row-scaled forward is a + single fused grouped GEMM with no host reads. + """ + group_sizes = [128, 256, 128] + A, B, offs = _make_grouped_inputs(group_sizes, K=256, N=384, seed=11) + + # Compile the decorated op directly (the mxfp8 grouped test's pattern); + # calling a nonstrict-traced function from a compiled frame is rejected. + A_e = A.clone().requires_grad_(True) + B_e = B.clone().requires_grad_(True) + y_eager = four_over_six_grouped_mm( + A_e, + B_e, + offs, + err_mode="mse", + row_scaled_activation=row_scaled_activation, + backward_override=backward_override, + ) + dy = torch.randn_like(y_eager) + y_eager.backward(dy) + + A_c = A.clone().requires_grad_(True) + B_c = B.clone().requires_grad_(True) + try: + y_compiled = torch.compile(four_over_six_grouped_mm, fullgraph=True)( + A_c, + B_c, + offs, + err_mode="mse", + row_scaled_activation=row_scaled_activation, + backward_override=backward_override, + ) + except torch._dynamo.exc.Unsupported as e: + if "nonstrict_trace" in str(e): + pytest.skip( + "this torch build rejects autograd.Function outputs from " + "nonstrict_trace-ed functions (the mxfp8 grouped compile " + "test's pattern); coverage resumes on builds that accept it" + ) + raise + y_compiled.backward(dy) + + torch.testing.assert_close(y_compiled, y_eager, atol=0, rtol=0) + torch.testing.assert_close(A_c.grad, A_e.grad, atol=0, rtol=0) + torch.testing.assert_close(B_c.grad, B_e.grad, atol=0, rtol=0) diff --git a/torchao/prototype/moe_training/README.md b/torchao/prototype/moe_training/README.md index f0858fb946..b6dce584a0 100644 --- a/torchao/prototype/moe_training/README.md +++ b/torchao/prototype/moe_training/README.md @@ -13,6 +13,7 @@ - [End-to-end training benchmark with TorchTitan: Llama4 Scout vs bfloat16 baseline](#end-to-end-training-benchmark-with-torchtitan-llama4-scout-vs-bfloat16-baseline) - [Implementation details for developers](#implementation-details-for-developers) - [Limitations](#limitations) +- [NVFP4 four-over-six MoE training](#nvfp4-four-over-six-moe-training) ## Overview This prototype provides: @@ -317,3 +318,29 @@ For all other ops, these training tensor subclasses behave like regular torch.Te ## Limitations - The new CUDA kernel for MXFP8 quantization of the non-transposed expert weights in the backwards pass does not support TP yet. + +## NVFP4 four-over-six MoE training +NVFP4 four-over-six is a prototype adaptive block-scaling recipe: per +16-value block, the standard map-to-6 encoding competes with a 1.5x-scale +map-to-4 candidate and the lower-error one is kept. Forward grouped +GEMMs run in NVFP4 (each token group quantized as its own tensor); backwards +are high-precision or dequantized grouped GEMMs only (no quantized backward). Requires SM100+, PyTorch 2.10+, the CuTe DSL runtime packages +(`nvidia-cutlass-dsl`, `cuda-python`, `apache-tvm-ffi`) for the quantize +fast path, and K/N % 128 == 0 with 128-row-aligned token groups (or +`pad_token_groups_for_grouped_mm=True`). +The grouped op is traceable under `torch.compile` (nonstrict trace) in +both scale granularities; the grouped GEMM dispatcher branch is eager-only +(it reads `offs[-1]` on the host for the over-allocation tail slice). +```python +import torch +from torchao.prototype.moe_training.nvfp4_training.four_over_six_grouped import ( + four_over_six_grouped_mm, +) + +# A: (M, K) packed token groups, B: (E, N, K) expert weights, +# offs: cumulative int32 group-end offsets. +out = four_over_six_grouped_mm(A, B, offs, row_scaled_activation=True) +``` +Framework integrations drive the same op through the grouped GEMM +dispatcher with `NVFP4FourOverSixTrainingOpConfig` (see the torchtitan +converters); `quantize_` model conversion for this config is future work. diff --git a/torchao/prototype/moe_training/config.py b/torchao/prototype/moe_training/config.py index e60c7f7682..f72b1bedbd 100644 --- a/torchao/prototype/moe_training/config.py +++ b/torchao/prototype/moe_training/config.py @@ -227,6 +227,67 @@ def __hash__(self): ) +# register as pytree constant so we can use dynamo nonstrict trace in torchao.prototype.moe_training.ep +@register_as_pytree_constant +@dataclass +class NVFP4FourOverSixTrainingOpConfig(TrainingOpBaseConfig): + """ + The NVFP4FourOverSixTrainingOpConfig defines the NVFP4 four-over-six + training config for grouped GEMM ops. + + Four-over-six scores two candidate encodings per 16-value block and keeps + the lower-error one; the knobs mirror + ``torchao.prototype.moe_training.nvfp4_training.four_over_six_grouped.four_over_six_grouped_mm``, + which the grouped GEMM dispatcher drives from this config. + """ + + # Candidate-selection error metric for four-over-six, "mae" or "mse". + err_mode: str = "mae" + + # Global E4M3 scale bound; 256 leaves map-to-4 headroom, 448 uses the full range. + e4m3_scale_bound: int = 256 + + # Whether to derive one FP32 global scale per activation row instead of per token group. + row_scaled_activation: bool = False + + # Weight quantization block shape, "16x16" or "1x16". + weight_block: str = "16x16" + + # Backward computation override for the grouped GEMM. None or "high_precision" + # computes both gradients with plain grouped GEMMs on the saved high-precision + # operands. "dequantized" computes them from the dequantized forward operands. + # Grouped four-over-six has no quantized backward. + backward_override: Optional[str] = None + + # Whether to pad the token group sizes to multiples of 128 (the grouped GEMM alignment). + pad_token_groups_for_grouped_mm: bool = False + + def __eq__(self, other): + if isinstance(other, NVFP4FourOverSixTrainingOpConfig): + return ( + self.err_mode == other.err_mode + and self.e4m3_scale_bound == other.e4m3_scale_bound + and self.row_scaled_activation == other.row_scaled_activation + and self.weight_block == other.weight_block + and self.backward_override == other.backward_override + and self.pad_token_groups_for_grouped_mm + == other.pad_token_groups_for_grouped_mm + ) + return NotImplemented + + def __hash__(self): + return hash( + ( + self.err_mode, + self.e4m3_scale_bound, + self.row_scaled_activation, + self.weight_block, + self.backward_override, + self.pad_token_groups_for_grouped_mm, + ) + ) + + @register_quantize_module_handler(Float8TrainingOpConfig) @register_quantize_module_handler(MXFP8TrainingOpConfig) def _moe_training_transform( 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 bf563c241f..0b82298890 100644 --- a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -188,8 +188,13 @@ def four_over_six_quantize( 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). + global_amax: scalar FP32 amax, or a (R,) per-row amax vector — the + row-scaled activation variant for 1x16 blocks, or per-expert + amaxes expanded over expert rows for stacked expert weights. + With 16x16 blocks the vector must be constant within every + 16-row tile (each row of a tile derives the tile's scale chain + from its own amax entry; the grouped path's expert boundaries + keep tiles amax-uniform). 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. @@ -212,8 +217,6 @@ def four_over_six_quantize( 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, " diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py index 81ecc3258a..58da39a34a 100644 --- a/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py @@ -606,7 +606,8 @@ def four_over_six_quantize_cutedsl( Args: x: (R, C) bfloat16 or float32, contiguous, C % 64 == 0. - global_amax: scalar FP32 amax, or (R,) per-row amax (1x16 only). + global_amax: scalar FP32 amax, or (R,) per-row amax (with 16x16 + blocks, constant within every 16-row tile; see the quantizer). block: "1x16" or "16x16". err_mode: "mae" or "mse". e4m3_scale_bound: 256 or 448. diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py new file mode 100644 index 0000000000..3f6d494ea4 --- /dev/null +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six_grouped.py @@ -0,0 +1,468 @@ +# 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. + +"""Differentiable NVFP4 four-over-six grouped GEMM for MoE training. + +Grouped counterpart of ``four_over_six_mm`` for routed-expert layers: +``A`` holds token groups packed along dim 0, ``B`` holds one weight matrix +per expert, and ``offs`` marks each group's end row. Every token group is +quantized as its own tensor: + +* per-tensor activations (the default): each token group gets its own + global scale from that group's amax. The group amaxes are expanded to a + per-row amax vector so the whole packed tensor quantizes in one + ``four_over_six_quantize`` call — bitwise identical to quantizing each + group separately, because the quantizer derives every row's scale chain + from that row's amax entry. The forward GEMM is one + ``F.scaled_grouped_mm`` with per-group second-level scales. +* row-scaled activations: one global scale per token row. The forward is a + single ``F.scaled_grouped_mm`` carrying the constant per-tensor factor in + every group's slot, its bf16 output upcast and scaled by the raw per-row + amaxes (torch's grouped GEMM only emits bf16 high-precision output). That + is one GEMM-output rounding away from a per-group loop of dense + four-over-six GEMMs, and the tests pin the fused output against a + rounding-emulated loop oracle. + +Weights always quantize per expert with per-tensor scales +(``weight_block`` selects 16x16 tiles or 1x16 blocks, as in the dense op). + +Gradients never quantize with four-over-six, so the grouped backward +supports only the high-precision and dequantized overrides of +``four_over_six_mm``: + +* ``"high_precision"`` (the default): bf16 grouped GEMMs on the saved + original operands; +* ``"dequantized"``: bf16 grouped GEMMs on dequantizations of the rowwise + operands the forward consumed — the RL train/inference-consistency mode. + +``"quantized"`` raises. Requires K % 128 == 0 and N % 128 == 0; token +groups must be 128-row aligned unless ``pad_token_groups_for_grouped_mm`` +is set, which zero-pads each group to the next 128 multiple before +quantization (zero rows quantize to zero codes and are sliced away from the +output). +""" + +from typing import Optional + +import torch +import torch.nn.functional as F + +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + FP4_E2M1_MAX, + _global_decode_scale, + four_over_six_quantize, + nvfp4_dequantize, +) +from torchao.prototype.moe_training.nvfp4_training.group_hadamard_utils import ( + _DEVICE_ASSERTS, +) +from torchao.prototype.moe_training.utils import ( + conditional_nostrict_trace, + pad_token_groups, + unpad_token_groups, +) +from torchao.prototype.mx_formats.utils import to_blocked +from torchao.quantization.quantize_.common import KernelPreference +from torchao.utils import is_sm_at_least_100 + +_ALIGNMENT = 128 +_SCALE_RECIPE = [F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise] +_SWIZZLE = [F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE] + +__all__ = ["four_over_six_grouped_mm"] + + +@conditional_nostrict_trace +def four_over_six_grouped_mm( + A: torch.Tensor, + B: torch.Tensor, + offs: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + err_mode: str = "mae", + e4m3_scale_bound: int = 256, + row_scaled_activation: bool = False, + weight_block: str = "16x16", + backward_override: Optional[str] = None, + pad_token_groups_for_grouped_mm: bool = False, +) -> torch.Tensor: + """Quantize grouped activations and expert weights with four-over-six. + + ``A`` has shape ``(M, K)``, ``B`` has shape ``(E, N, K)``, and ``offs`` + contains the cumulative row-end offset for each expert. Knobs match + ``four_over_six_mm``; see the module docstring for the grouped-specific + backward and alignment semantics. + """ + output = _FourOverSixGroupedMM.apply( + A, + B, + offs, + err_mode, + e4m3_scale_bound, + row_scaled_activation, + weight_block, + backward_override, + pad_token_groups_for_grouped_mm, + ) + if bias is not None: + output = output + bias.to(output.dtype) + return output + + +def _expand_group_amax( + row_amax: torch.Tensor, group_end_offsets: torch.Tensor, num_experts: int +) -> torch.Tensor: + """Per-row amax vector holding each row's group amax. + + Rows past the final offset (the pad-helper's over-allocated tail) take + the last group's amax; they are all-zero and never enter the GEMM. + """ + group_idx = torch.searchsorted( + group_end_offsets, + torch.arange(row_amax.shape[0], device=row_amax.device, dtype=torch.int32), + right=True, + ).clamp_(max=num_experts - 1) + group_amax = torch.zeros( + num_experts, dtype=torch.float32, device=row_amax.device + ).scatter_reduce_( + 0, group_idx, row_amax.to(torch.float32), reduce="amax", include_self=True + ) + return group_amax[group_idx], group_amax + + +def _quantize_expert_weights( + weight: torch.Tensor, + weight_amax: torch.Tensor, + weight_block: str, + err_mode: str, + e4m3_scale_bound: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-expert four-over-six quantization of a stacked (E, N, K) weight. + + Flattening experts along rows and expanding each expert's amax over its + rows quantizes the whole stack in one call — bitwise identical to a + per-expert loop, because the quantizer derives every row's scale chain + from that row's amax entry, 1x16 blocks never cross rows, and 16x16 + tiles never cross experts (N % 128 == 0 keeps expert boundaries 16-row + aligned, and every row of a tile carries the same expert's amax). + """ + num_experts, N, K = weight.shape + flat_codes, flat_scales = four_over_six_quantize( + weight.reshape(num_experts * N, K), + weight_amax.repeat_interleave(N), + block=weight_block, + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + return ( + flat_codes.view(num_experts, N, K // 2), + flat_scales.view(num_experts, N, K // 16), + ) + + +def _dequantize_expert_weights( + codes: torch.Tensor, + scales: torch.Tensor, + weight_amax: torch.Tensor, + e4m3_scale_bound: int, +) -> torch.Tensor: + """Dequantize stacked per-expert codes back to a bf16 (E, N, K) weight. + + Flattening experts along rows and expanding each expert's amax over its + rows reproduces the per-expert scalar dequantization exactly — the + decode chain reads one amax entry per row either way. + """ + num_experts, N = codes.shape[0], codes.shape[1] + row_amax = weight_amax.to(torch.float32).repeat_interleave(N) + flat = nvfp4_dequantize( + codes.reshape(num_experts * N, -1), + scales.reshape(num_experts * N, -1), + row_amax, + e4m3_scale_bound=e4m3_scale_bound, + ) + return flat.view(num_experts, N, -1) + + +def _row_scaled_single_grouped_gemm( + x_codes: torch.Tensor, + x_scales: torch.Tensor, + x_amax: torch.Tensor, + x_global: torch.Tensor, + w_codes: torch.Tensor, + w_scales: torch.Tensor, + w_global: torch.Tensor, + padded_group_end_offsets: torch.Tensor, +) -> torch.Tensor: + """One ``F.scaled_grouped_mm`` covering every group of the row-scaled + forward. + + The GEMM epilogue applies the E4M3 block scales and the per-tensor + factors — the constant 1/(6*bound) in every group's activation slot and + each expert's amax/(6*bound) — and the raw per-row amax multiply plus + the final bf16 cast happen on the upcast output. The GEMM emits bf16 + (torch's grouped GEMM has no FP32 output mode), which costs one + rounding before the row scale relative to a loop of dense FP32-output + GEMMs per group. Rows past the final offset may hold garbage; the pad + helper's unpad drops them. + """ + num_experts = w_codes.shape[0] + output = F.scaled_grouped_mm( + x_codes.view(torch.float4_e2m1fn_x2), + w_codes.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + # scaled_grouped_mm consumes swizzled scale bytes viewed at the + # logical 2D shape, as in the per-tensor forward; the view needs + # the 128-row alignment the forward enforces. One to_blocked over + # the row-flattened expert scales equals the per-expert stack + # bitwise because N % 128 == 0 keeps expert boundaries on swizzle + # row-block boundaries. + scale_a=[ + to_blocked(x_scales).view(x_scales.shape), + x_global.expand(num_experts).contiguous(), + ], + scale_recipe_a=_SCALE_RECIPE, + scale_b=[ + to_blocked(w_scales.reshape(-1, w_scales.shape[-1])).view(num_experts, -1), + w_global, + ], + scale_recipe_b=_SCALE_RECIPE, + swizzle_a=_SWIZZLE, + swizzle_b=_SWIZZLE, + offs=padded_group_end_offsets, + output_dtype=torch.bfloat16, + ) + return (output.to(torch.float32) * x_amax.view(-1, 1)).to(torch.bfloat16) + + +class _FourOverSixGroupedMM(torch.autograd.Function): + """NVFP4 four-over-six grouped forward with override-only backward.""" + + @staticmethod + def forward( + ctx, + input_act: torch.Tensor, + weight: torch.Tensor, + group_end_offsets: torch.Tensor, + err_mode: str, + e4m3_scale_bound: int, + row_scaled_activation: bool, + weight_block: str, + backward_override: Optional[str], + pad_token_groups_for_grouped_mm: bool, + ) -> torch.Tensor: + if group_end_offsets.ndim != 1 or group_end_offsets.dtype != torch.int32: + raise ValueError("offs must be a 1D int32 tensor") + if not group_end_offsets.is_contiguous(): + raise ValueError("offs must be contiguous") + if group_end_offsets.numel() != weight.shape[0]: + raise ValueError("offs must contain one group-end offset per expert") + if not is_sm_at_least_100(): + raise NotImplementedError( + "NVFP4 four-over-six grouped GEMM requires SM100+" + ) + if backward_override is None: + backward_override = "high_precision" + if backward_override not in ("high_precision", "dequantized"): + if backward_override == "quantized": + raise ValueError( + "grouped four-over-six has no quantized backward; use " + "'high_precision' or 'dequantized'" + ) + raise ValueError( + f"backward_override must be 'high_precision' or 'dequantized', " + f"got {backward_override!r}" + ) + + num_tokens, K = input_act.shape + num_experts, N, _ = weight.shape + if K % _ALIGNMENT != 0 or N % _ALIGNMENT != 0: + raise ValueError( + f"K and N must be divisible by {_ALIGNMENT}; got K={K}, N={N}" + ) + if _DEVICE_ASSERTS: + group_sizes = torch.diff( + group_end_offsets, prepend=group_end_offsets.new_zeros(1) + ) + torch.ops.aten._assert_async.msg( + torch.all(group_sizes >= 0), "offs must be non-decreasing" + ) + torch.ops.aten._assert_async.msg( + group_end_offsets[-1] == num_tokens, + "the final group-end offset must equal A.shape[0]", + ) + if not pad_token_groups_for_grouped_mm: + torch.ops.aten._assert_async.msg( + torch.all(group_sizes % _ALIGNMENT == 0), + "every token group must be 128-row aligned when padding is disabled", + ) + + input_act = input_act.to(torch.bfloat16).contiguous() + weight = weight.to(torch.bfloat16).contiguous() + original_input = input_act + + padded_group_start_offsets = None + if pad_token_groups_for_grouped_mm: + # The fused pad/unpad CUDA kernels only accept alignment_size 32 + # and at most 32 groups; this op needs 128-row alignment with any + # expert count, so it pins the pure-torch path. + input_act, padded_group_start_offsets, padded_group_end_offsets = ( + pad_token_groups( + input_act, + group_end_offsets, + alignment_size=_ALIGNMENT, + kernel_preference=KernelPreference.EMULATED, + ) + ) + else: + padded_group_end_offsets = group_end_offsets + + row_amax = input_act.abs().amax(dim=1) + group_amax = None + if row_scaled_activation: + x_amax = row_amax.to(torch.float32) + else: + x_amax, group_amax = _expand_group_amax( + row_amax, padded_group_end_offsets, num_experts + ) + weight_amax = weight.abs().amax(dim=(1, 2)).to(torch.float32) + + x_codes, x_scales = four_over_six_quantize( + input_act, + x_amax, + block="1x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_codes, w_scales = _quantize_expert_weights( + weight, weight_amax, weight_block, err_mode, e4m3_scale_bound + ) + w_global = _global_decode_scale(weight_amax, e4m3_scale_bound) + + if row_scaled_activation: + # The row-scaled forward carries the constant 1/(6*bound) factor + # in every group's per-tensor slot and scales the GEMM output by + # the raw per-row amaxes; one F.scaled_grouped_mm covers all + # groups. + x_global = torch.full( + (), + 1.0 / (FP4_E2M1_MAX * float(e4m3_scale_bound)), + dtype=torch.float32, + device=input_act.device, + ) + output = _row_scaled_single_grouped_gemm( + x_codes, + x_scales, + x_amax, + x_global, + w_codes, + w_scales, + w_global, + padded_group_end_offsets, + ) + else: + output = F.scaled_grouped_mm( + x_codes.view(torch.float4_e2m1fn_x2), + w_codes.view(torch.float4_e2m1fn_x2).transpose(-2, -1), + # scaled_grouped_mm consumes swizzled scale bytes viewed at the + # logical 2D shape (the layout the group quantize kernels + # return); the view needs the 128-row alignment enforced above. + scale_a=[ + to_blocked(x_scales).view(x_scales.shape), + _global_decode_scale(group_amax, e4m3_scale_bound), + ], + scale_recipe_a=_SCALE_RECIPE, + # One flattened to_blocked covers all experts, bitwise equal + # to a per-expert stack (see _row_scaled_single_grouped_gemm). + scale_b=[ + to_blocked(w_scales.reshape(-1, w_scales.shape[-1])).view( + num_experts, -1 + ), + w_global, + ], + scale_recipe_b=_SCALE_RECIPE, + swizzle_a=_SWIZZLE, + swizzle_b=_SWIZZLE, + offs=padded_group_end_offsets, + output_dtype=torch.bfloat16, + ) + + if pad_token_groups_for_grouped_mm: + output = unpad_token_groups( + output, + group_end_offsets, + padded_group_start_offsets, + num_tokens, + alignment_size=_ALIGNMENT, + kernel_preference=KernelPreference.EMULATED, + ) + + if backward_override == "high_precision": + ctx.save_for_backward(original_input, weight, group_end_offsets) + else: + if padded_group_start_offsets is None: + padded_group_start_offsets = group_end_offsets.new_zeros(0) + ctx.save_for_backward( + x_codes, + x_scales, + x_amax, + w_codes, + w_scales, + weight_amax, + group_end_offsets, + padded_group_start_offsets, + ) + ctx.backward_override = backward_override + ctx.e4m3_scale_bound = e4m3_scale_bound + ctx.pad_token_groups_for_grouped_mm = pad_token_groups_for_grouped_mm + ctx.num_tokens = num_tokens + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_output = grad_output.to(torch.bfloat16).contiguous() + + if ctx.backward_override == "high_precision": + input_act, weight, group_end_offsets = ctx.saved_tensors + else: + ( + x_codes, + x_scales, + x_amax, + w_codes, + w_scales, + weight_amax, + group_end_offsets, + padded_group_start_offsets, + ) = ctx.saved_tensors + input_act = nvfp4_dequantize( + x_codes, x_scales, x_amax, e4m3_scale_bound=ctx.e4m3_scale_bound + ) + if ctx.pad_token_groups_for_grouped_mm: + input_act = unpad_token_groups( + input_act, + group_end_offsets, + padded_group_start_offsets, + ctx.num_tokens, + alignment_size=_ALIGNMENT, + kernel_preference=KernelPreference.EMULATED, + ) + weight = _dequantize_expert_weights( + w_codes, w_scales, weight_amax, ctx.e4m3_scale_bound + ) + + grad_input = torch._grouped_mm( + grad_output, + weight, + offs=group_end_offsets, + out_dtype=torch.bfloat16, + ) + grad_weight = torch._grouped_mm( + grad_output.transpose(-2, -1), + input_act, + offs=group_end_offsets, + out_dtype=torch.bfloat16, + ) + return grad_input, grad_weight, None, None, None, None, None, None, None diff --git a/torchao/prototype/moe_training/utils.py b/torchao/prototype/moe_training/utils.py index 9ba6fe78a3..09d8f56f5a 100644 --- a/torchao/prototype/moe_training/utils.py +++ b/torchao/prototype/moe_training/utils.py @@ -10,6 +10,7 @@ from torchao.prototype.moe_training.config import ( Float8TrainingOpConfig, MXFP8TrainingOpConfig, + NVFP4FourOverSixTrainingOpConfig, TrainingOpBaseConfig, ) from torchao.prototype.mx_formats.mx_tensor import to_mx @@ -405,6 +406,36 @@ def _quantize_then_scaled_grouped_mm( offs, **kwargs, ) + elif isinstance(config, NVFP4FourOverSixTrainingOpConfig): + from torchao.prototype.moe_training.nvfp4_training.four_over_six_grouped import ( + four_over_six_grouped_mm, + ) + + # The dispatcher hands expert weights over as B_t with shape (E, K, N); + # the four-over-six op takes them in their stored (E, N, K) layout. + # Callers with padded token dispatchers over-allocate A past offs[-1], + # while the op requires offs[-1] == A.shape[0] and the unwritten tail + # rows must not feed its per-group amaxes: slice to the logical rows + # and zero-extend the output, which also routes zero gradients to the + # tail. The host read of offs[-1] keeps this dispatcher branch + # eager-only (the op itself is nonstrict-traced under compile). + num_tokens = int(offs[-1]) + tail_rows = A.shape[0] - num_tokens + output = four_over_six_grouped_mm( + A[:num_tokens] if tail_rows > 0 else A, + B_t.transpose(-2, -1), + offs, + bias, + err_mode=config.err_mode, + e4m3_scale_bound=config.e4m3_scale_bound, + row_scaled_activation=config.row_scaled_activation, + weight_block=config.weight_block, + backward_override=config.backward_override, + pad_token_groups_for_grouped_mm=config.pad_token_groups_for_grouped_mm, + ) + if tail_rows > 0: + output = torch.nn.functional.pad(output, (0, 0, 0, tail_rows)) + return output else: raise ValueError(f"Unsupported config type: {type(config)}")