diff --git a/docs/envvars.rst b/docs/envvars.rst index 5ff430eabea..56263aaec38 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -367,7 +367,13 @@ Kernel Configuration :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable row-scaled NVFP4 tensors for forward activation quantizers in the ``NVFP4BlockScaling`` recipe. When set to ``1`` (or when ``NVFP4BlockScaling(row_scaled_activation=True)`` is used), rowwise ``amax`` metadata is stored as one FP32 value per tensor row instead of a single scalar. + :Description: Enable row-scaled NVFP4 tensors for the activation and gradient-output quantizers in the ``NVFP4BlockScaling`` recipe; weights are controlled by :envvar:`NVTE_NVFP4_ROW_SCALED_WEIGHT`. When set to ``1`` (or when ``NVFP4BlockScaling(row_scaled_activation=True)`` is used), rowwise ``amax`` metadata is stored as one FP32 value per tensor row and columnwise ``amax`` metadata as one FP32 value per tensor column instead of a single scalar. Random Hadamard transforms and stochastic rounding follow the recipe settings as for per-tensor NVFP4. Sequence parallelism, Userbuffers communication overlap, and ``transformer_engine.pytorch.ops.GroupedLinear`` are not supported. + +.. envvar:: NVTE_NVFP4_ROW_SCALED_WEIGHT + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable row-scaled NVFP4 tensors for the weight quantizers in the ``NVFP4BlockScaling`` recipe. When set to ``1`` (or when ``NVFP4BlockScaling(row_scaled_weight=True)`` is used), weights store one FP32 global ``amax`` per row and per column with 1x16 (1D) block scales, and weight 2D quantization is disabled. When unset, weights use a single per-tensor global ``amax`` with 16x16 (2D) block scales, or 1x16 block scales with ``NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1``. Row-scaled weights require BF16 weight tensors with both dimensions divisible by 32. The restrictions of :envvar:`NVTE_NVFP4_ROW_SCALED_ACTIVATION` apply, and 4over6 on weights and quantized primary weights are not supported. .. envvar:: NVTE_NVFP4_4OVER6 diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 5062cf6b42f..7e31b944b62 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -8,9 +8,14 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm +from transformer_engine.pytorch.cpp_extensions import ( + general_gemm, + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import reference_utils +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -301,6 +306,8 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( single_output: bool, use_4over6: bool = False, nvfp4_4over6_err_mode: str = "MAE", + row_scaled_x: bool = True, + row_scaled_w: bool = False, ): te_dtype = te.DType.kFloat4E2M1 device = "cuda" @@ -317,7 +324,7 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, - row_scaled_nvfp4=True, + row_scaled_nvfp4=row_scaled_x, nvfp4_use_4over6=use_4over6, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) @@ -329,6 +336,7 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + row_scaled_nvfp4=row_scaled_w, nvfp4_use_4over6=use_4over6, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) @@ -622,6 +630,11 @@ def test_nvfp4_gemm_versus_reference( @pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"]) @pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) +@pytest.mark.parametrize( + "row_scaled_x, row_scaled_w", + [(True, False), (False, True), (True, True)], + ids=["row_scaled_x", "row_scaled_w", "row_scaled_xw"], +) def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( m_splits: list[int], k: int, @@ -633,7 +646,17 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( single_output: bool, use_4over6: bool, nvfp4_4over6_err_mode: str, + row_scaled_x: bool, + row_scaled_w: bool, ): + if row_scaled_w: + if use_4over6: + pytest.skip("Row-scaled NVFP4 transpose quantization does not support 4over6.") + if w_dtype != torch.bfloat16 or k % 32 != 0 or n % 32 != 0: + pytest.skip( + "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions " + "divisible by 32." + ) check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( x_dtype=x_dtype, w_dtype=w_dtype, @@ -645,6 +668,8 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( single_output=single_output, use_4over6=use_4over6, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, + row_scaled_x=row_scaled_x, + row_scaled_w=row_scaled_w, ) @@ -838,3 +863,471 @@ def test_nvfp4_ue5m3_gemm_leaves_operands_unswizzled(M: int, K: int, N: int): ref = x_q.dequantize(dtype=torch.float32) @ w_q.dequantize(dtype=torch.float32).t() rel_err = (y1.float() - ref).norm() / ref.norm() assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large" + + +def _row_scaled_grouped_tensor_gemm_supported() -> bool: + if not torch.cuda.is_available(): + return False + if torch.cuda.get_device_capability() < (10, 0): + return False + return tex.get_cublasLt_version() >= 130400 + + +_reason_for_no_grouped_tensor_gemm = ( + "Row-scaled grouped-tensor NVFP4 GEMM requires Blackwell (SM100) and cuBLAS 13.4+." +) + + +def _make_row_scaled_grouped_activation( + m_splits: list[int], + k: int, + *, + columnwise: bool, + tokens: torch.Tensor | None = None, +) -> GroupedTensor: + """Cast tokens to a row-scaled NVFP4 GroupedTensor; splits and k must be multiples of 128.""" + device = "cuda" + if tokens is None: + tokens = torch.randn(sum(m_splits), k, dtype=torch.bfloat16, device=device) + tokens_per_expert = torch.tensor(m_splits, dtype=torch.int64, device=device) + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=columnwise, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + quantizer.optimize_for_gemm = True + return tex.group_quantize(tokens, quantizer, len(m_splits), tokens_per_expert) + + +def _make_output_grouped_tensor( + m_splits: list[int], last_dim: int, device: torch.device, dtype: torch.dtype +) -> GroupedTensor: + return GroupedTensor.make_grouped_tensor( + num_tensors=len(m_splits), + first_dims=torch.tensor(m_splits, device=device, dtype=torch.int64), + last_dims=None, + logical_first_dim=sum(m_splits), + logical_last_dim=last_dim, + quantizer=None, + device=device, + dtype=dtype, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "m_splits, k, n", + [ + ([128, 128], 128, 128), + ([128, 256, 128], 256, 256), + ([256, 128, 384], 512, 256), + ([128, 256, 384, 512], 256, 512), + ], +) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +@pytest.mark.parametrize("layout", ["TN", "NN"], ids=["fprop", "dgrad"]) +@pytest.mark.parametrize("row_scaled_w", [False, True], ids=["per_tensor_w", "row_scaled_w"]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_fprop( + m_splits: list[int], + k: int, + n: int, + out_dtype: torch.dtype, + layout: str, + row_scaled_w: bool, +) -> None: + """Row-scaled grouped-tensor GEMM with discrete weights matches a per-expert dense GEMM.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(7) + device = torch.device("cuda") + num_gemms = len(m_splits) + in_dim, out_dim = (k, n) if layout == "TN" else (n, k) + + activation = _make_row_scaled_grouped_activation(m_splits, in_dim, columnwise=False) + activation_splits = activation.split_into_quantized_tensors() + + w_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=row_scaled_w, + ) + w_quantizer.optimize_for_gemm = True + weights = [] + expected = [] + for i in range(num_gemms): + w = torch.randn(n, k, dtype=torch.bfloat16, device=device) + w_nvfp4 = w_quantizer(w) + weights.append(w_nvfp4) + expected.append( + general_gemm(w_nvfp4, activation_splits[i], out_dtype=out_dtype, layout=layout)[0] + ) + + grouped_out = _make_output_grouped_tensor(m_splits, out_dim, device, out_dtype) + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout=layout) + actual = grouped_out.split_into_quantized_tensors() + + if out_dtype == torch.bfloat16: + tols = dict(atol=0.0, rtol=7.8e-3) + else: + tols = dict(atol=1e-4, rtol=1e-4) + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, **tols) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_cudagraph( + m_splits: list[int], k: int, n: int +) -> None: + """A captured row-scaled grouped-tensor GEMM replays correctly.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(11) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + static_tokens = torch.randn(sum(m_splits), k, dtype=torch.bfloat16, device=device) + activation = _make_row_scaled_grouped_activation( + m_splits, k, columnwise=False, tokens=static_tokens + ) + + w_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_quantizer.optimize_for_gemm = True + weights = [ + w_quantizer(torch.randn(n, k, dtype=torch.bfloat16, device=device)) + for _ in range(num_gemms) + ] + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + + def _run(): + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + + torch.cuda.synchronize() + _run() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run() + + fresh = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + activation.rowwise_data.copy_(fresh.rowwise_data) + activation.scale_inv.copy_(fresh.scale_inv) + activation.amax.copy_(fresh.amax) + graph.replay() + torch.cuda.synchronize() + replayed = [t.clone() for t in grouped_out.split_into_quantized_tensors()] + + _run() + torch.cuda.synchronize() + expected = grouped_out.split_into_quantized_tensors() + for got, ref in zip(replayed, expected): + torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "m_splits, k, n", + [ + ([128, 128], 128, 128), + ([128, 256, 128], 256, 256), + ([256, 128, 384], 512, 256), + ([128, 256, 384, 512], 256, 512), + ], +) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad( + m_splits: list[int], k: int, n: int, out_dtype: torch.dtype +) -> None: + """Row-scaled grouped-tensor NT GEMM matches a per-expert dense GEMM.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(13) + device = torch.device("cuda") + num_gemms = len(m_splits) + + x = _make_row_scaled_grouped_activation(m_splits, k, columnwise=True) + dy = _make_row_scaled_grouped_activation(m_splits, n, columnwise=True) + x_splits = x.split_into_quantized_tensors() + dy_splits = dy.split_into_quantized_tensors() + + expected = [ + general_gemm(x_splits[i], dy_splits[i], out_dtype=out_dtype, layout="NT")[0] + for i in range(num_gemms) + ] + + grouped_out = _make_output_grouped_tensor([n] * num_gemms, k, device, out_dtype) + general_grouped_gemm_for_grouped_tensor(x, dy, grouped_out, layout="NT") + actual = grouped_out.split_into_quantized_tensors() + + if out_dtype == torch.bfloat16: + tols = dict(atol=0.0, rtol=7.8e-3) + else: + tols = dict(atol=1e-4, rtol=1e-4) + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, **tols) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([256, 128, 384], 512, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_wgrad_discrete_out( + m_splits: list[int], k: int, n: int +) -> None: + """Row-scaled NT GEMM into a list output matches a per-expert dense GEMM.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(17) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + x = _make_row_scaled_grouped_activation(m_splits, k, columnwise=True) + dy = _make_row_scaled_grouped_activation(m_splits, n, columnwise=True) + x_splits = x.split_into_quantized_tensors() + dy_splits = dy.split_into_quantized_tensors() + + expected = [ + general_gemm(x_splits[i], dy_splits[i], out_dtype=out_dtype, layout="NT")[0] + for i in range(num_gemms) + ] + + out_list = [torch.empty(n, k, dtype=out_dtype, device=device) for _ in range(num_gemms)] + general_grouped_gemm_for_grouped_tensor(x, dy, out_list, layout="NT") + + for got, ref in zip(out_list, expected): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_uneven_colscale( + m_splits: list[int], k: int, n: int +) -> None: + """Column post-scale with uneven groups matches a per-expert dense GEMM.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(23) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + weights = _make_row_scaled_grouped_activation([n] * num_gemms, k, columnwise=False) + activation = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + w_splits = weights.split_into_quantized_tensors() + act_splits = activation.split_into_quantized_tensors() + + expected = [ + general_gemm(w_splits[i], act_splits[i], out_dtype=out_dtype, layout="TN")[0] + for i in range(num_gemms) + ] + + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + actual = grouped_out.split_into_quantized_tensors() + + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_colscale_cudagraph( + m_splits: list[int], k: int, n: int +) -> None: + """A captured column post-scale with uneven groups replays correctly.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(29) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + static_w = torch.randn(n * num_gemms, k, dtype=torch.bfloat16, device=device) + static_act = torch.randn(sum(m_splits), k, dtype=torch.bfloat16, device=device) + weights = _make_row_scaled_grouped_activation( + [n] * num_gemms, k, columnwise=False, tokens=static_w + ) + activation = _make_row_scaled_grouped_activation( + m_splits, k, columnwise=False, tokens=static_act + ) + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + + def _run(): + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + + torch.cuda.synchronize() + _run() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run() + + fresh_w = _make_row_scaled_grouped_activation([n] * num_gemms, k, columnwise=False) + fresh_act = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + for dst, src in ((weights, fresh_w), (activation, fresh_act)): + dst.rowwise_data.copy_(src.rowwise_data) + dst.scale_inv.copy_(src.scale_inv) + dst.amax.copy_(src.amax) + graph.replay() + torch.cuda.synchronize() + replayed = [t.clone() for t in grouped_out.split_into_quantized_tensors()] + + _run() + torch.cuda.synchronize() + expected = grouped_out.split_into_quantized_tensors() + for got, ref in zip(replayed, expected): + torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("m_splits, k, n", [([128, 256, 128], 256, 256)]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_output_capacity( + m_splits: list[int], k: int, n: int +) -> None: + """Output capacity rows beyond the live region do not affect the result.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(31) + device = torch.device("cuda") + num_gemms = len(m_splits) + out_dtype = torch.float32 + + activation = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False) + activation_splits = activation.split_into_quantized_tensors() + + w_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_quantizer.optimize_for_gemm = True + weights = [] + expected = [] + for i in range(num_gemms): + w_nvfp4 = w_quantizer(torch.randn(n, k, dtype=torch.bfloat16, device=device)) + weights.append(w_nvfp4) + expected.append( + general_gemm(w_nvfp4, activation_splits[i], out_dtype=out_dtype, layout="TN")[0] + ) + + grouped_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + live_numel = grouped_out.rowwise_data.numel() + padded = torch.full( + (live_numel + 128 * n,), + float("nan"), + dtype=grouped_out.rowwise_data.dtype, + device=device, + ) + padded[:live_numel].copy_(grouped_out.rowwise_data) + grouped_out.rowwise_data = padded + + general_grouped_gemm_for_grouped_tensor(weights, activation, grouped_out, layout="TN") + actual = grouped_out.split_into_quantized_tensors() + + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_nvfp4_row_scaled_grouped_tensor_gemm_spare_capacity(accumulate: bool, use_bias: bool): + """Rows past sum(first_dims) stay unchanged; live rows match a full-size output.""" + if not _row_scaled_grouped_tensor_gemm_supported(): + pytest.skip(_reason_for_no_grouped_tensor_gemm) + + torch.manual_seed(37) + device = torch.device("cuda") + m_splits = [128, 0, 256] + num_gemms = len(m_splits) + k, n = 256, 256 + live_rows = sum(m_splits) + capacity = live_rows + 128 + out_dtype = torch.bfloat16 + + weights = _make_row_scaled_grouped_activation([n] * num_gemms, k, columnwise=False) + tokens = torch.randn(capacity, k, dtype=torch.bfloat16, device=device) + activation = _make_row_scaled_grouped_activation(m_splits, k, columnwise=False, tokens=tokens) + reference_activation = _make_row_scaled_grouped_activation( + m_splits, k, columnwise=False, tokens=tokens[:live_rows].clone() + ) + + bias = None + if use_bias: + bias = GroupedTensor.make_grouped_tensor( + num_tensors=num_gemms, + first_dims=None, + last_dims=None, + logical_first_dim=num_gemms, + logical_last_dim=n, + quantizer=None, + device=device, + dtype=torch.bfloat16, + ) + bias.rowwise_data.copy_(torch.randn(num_gemms * n, dtype=torch.bfloat16, device=device)) + + initial = torch.randn(capacity * n, dtype=out_dtype, device=device) + grouped_out = GroupedTensor.make_grouped_tensor( + num_tensors=num_gemms, + first_dims=torch.tensor(m_splits, device=device, dtype=torch.int64), + last_dims=None, + logical_first_dim=capacity, + logical_last_dim=n, + quantizer=None, + device=device, + dtype=out_dtype, + ) + grouped_out.rowwise_data.copy_(initial) + reference_out = _make_output_grouped_tensor(m_splits, n, device, out_dtype) + reference_out.rowwise_data.copy_(initial[: live_rows * n]) + + general_grouped_gemm_for_grouped_tensor( + weights, activation, grouped_out, layout="TN", accumulate=accumulate, bias=bias + ) + general_grouped_gemm_for_grouped_tensor( + weights, + reference_activation, + reference_out, + layout="TN", + accumulate=accumulate, + bias=bias, + ) + + result = grouped_out.rowwise_data.view(capacity, n) + torch.testing.assert_close( + result[:live_rows], reference_out.rowwise_data.view(live_rows, n), atol=0.0, rtol=0.0 + ) + torch.testing.assert_close( + result[live_rows:], initial.view(capacity, n)[live_rows:], atol=0.0, rtol=0.0 + ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py index a740e0ad2cc..96d3b232638 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_row_scaled_quantize.py @@ -74,6 +74,76 @@ def _assert_matches_reference(grouped, x_valid: torch.Tensor, split_sections: li torch.testing.assert_close(g_scale[:, : N // 16], ref_scale, atol=0, rtol=0) +def _row_scaled_rht_quantizer(with_random_sign_mask: bool) -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + row_scaled_nvfp4=True, + ) + + +def _assert_rht_matches_per_expert( + grouped, x_valid: torch.Tensor, split_sections: list[int], with_random_sign_mask: bool +) -> None: + N = x_valid.shape[1] + capacity = int(grouped.logical_shape[0]) + q_row = grouped.rowwise_data.view(torch.uint8).reshape(-1) + s_row = grouped.scale_inv.view(torch.uint8).reshape(capacity, -1) + q_col = grouped.columnwise_data.view(torch.uint8).reshape(-1) + s_col = grouped.columnwise_scale_inv.view(torch.uint8).reshape(-1) + amax_row = grouped.amax.reshape(-1) + amax_col = grouped.columnwise_amax.reshape(-1) + + def assert_equal(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + row_start = 0 + for i, m in enumerate(split_sections): + if m == 0: + continue + ref = _row_scaled_rht_quantizer(with_random_sign_mask)(x_valid[row_start : row_start + m]) + assert_equal(amax_row[row_start : row_start + m], ref._amax_rowwise.reshape(-1)) + assert_equal(amax_col[i * N : (i + 1) * N], ref._amax_columnwise.reshape(-1)) + data_range = slice(row_start * N // 2, (row_start + m) * N // 2) + assert_equal( + q_row[data_range].reshape(m, N // 2), + ref._rowwise_data.view(torch.uint8).reshape(m, N // 2), + ) + assert_equal( + q_col[data_range].reshape(N, m // 2), + ref._columnwise_data.view(torch.uint8).reshape(N, m // 2), + ) + assert_equal( + s_row[row_start : row_start + m, : N // 16], + ref._rowwise_scale_inv.view(torch.uint8)[:m, : N // 16], + ) + assert_equal( + s_col[row_start * N // 16 : (row_start + m) * N // 16].reshape(N, m // 16), + ref._columnwise_scale_inv.view(torch.uint8)[:N, : m // 16], + ) + row_start += m + + +_QUANTIZE_MODES = ("rowwise", "rht", "rht_no_sign_mask") + + +def _mode_quantizer(mode: str) -> NVFP4Quantizer: + if mode == "rowwise": + return _row_scaled_quantizer() + return _row_scaled_rht_quantizer(with_random_sign_mask=mode == "rht") + + +def _assert_mode_matches(grouped, x_valid: torch.Tensor, split_sections: list[int], mode: str): + if mode == "rowwise": + _assert_matches_reference(grouped, x_valid, split_sections) + else: + _assert_rht_matches_per_expert(grouped, x_valid, split_sections, mode == "rht") + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "split_sections, N", @@ -82,17 +152,20 @@ def _assert_matches_reference(grouped, x_valid: torch.Tensor, split_sections: li ([128, 256, 128], 256), ([384, 128, 256], 128), ([256, 128], 512), + ([256, 0, 384], 512), + ([128] * 4, 2048), ], ) -def test_matches_per_expert_quantize(split_sections: list[int], N: int) -> None: +@pytest.mark.parametrize("mode", _QUANTIZE_MODES) +def test_matches_per_expert_quantize(split_sections: list[int], N: int, mode: str) -> None: torch.manual_seed(0) x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") - grouped = _group_quantize(x, len(split_sections), first_dims) + grouped = tex.group_quantize(x, _mode_quantizer(mode), len(split_sections), first_dims) assert grouped.row_scaled_nvfp4 - _assert_matches_reference(grouped, x, split_sections) + _assert_mode_matches(grouped, x, split_sections, mode) @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @@ -127,12 +200,14 @@ def test_paged_stashing_capacity_noop(split_sections: list[int], N: int, capacit @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -def test_capture_once_replay_changed_routing() -> None: +@pytest.mark.parametrize("mode", ("rowwise", "rht")) +def test_capture_once_replay_changed_routing(mode: str) -> None: N = 256 num_experts = 3 capacity = 768 splits_a = [256, 128, 128] - splits_b = [128, 256, 128] + splits_b = [128, 384, 0] + quantizer = _mode_quantizer(mode) first_dims = torch.empty(num_experts, dtype=torch.int64, device="cuda") x = torch.empty((capacity, N), dtype=torch.bfloat16, device="cuda") @@ -151,22 +226,18 @@ def program_routing(splits: list[int], x_src: torch.Tensor) -> None: warmup_stream = torch.cuda.Stream() warmup_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(warmup_stream): - _ = _group_quantize(x, num_experts, first_dims) + _ = tex.group_quantize(x, quantizer, num_experts, first_dims) torch.cuda.current_stream().wait_stream(warmup_stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - grouped = _group_quantize(x, num_experts, first_dims) + grouped = tex.group_quantize(x, quantizer, num_experts, first_dims) - program_routing(splits_a, x_a) - graph.replay() - torch.cuda.synchronize() - _assert_matches_reference(grouped, x_a, splits_a) - - program_routing(splits_b, x_b) - graph.replay() - torch.cuda.synchronize() - _assert_matches_reference(grouped, x_b, splits_b) + for splits, x_src in ((splits_a, x_a), (splits_b, x_b)): + program_routing(splits, x_src) + graph.replay() + torch.cuda.synchronize() + _assert_mode_matches(grouped, x_src, splits, mode) @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index 9fdabf68dd0..9a596bd7635 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -53,23 +53,31 @@ def nvfp4_rht_and_2d_quantization(): return nvfp4_recipe @staticmethod - def nvfp4_row_scaled(): + def nvfp4_row_scaled( + with_rht: bool = False, + with_2d_quantization: bool = False, + row_scaled_weight: bool = False, + ): nvfp4_recipe = recipe.NVFP4BlockScaling() - nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() - nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() - nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams(random_hadamard_transform=with_rht) + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=with_2d_quantization) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams(random_hadamard_transform=with_rht) # Emit row-scaled (per-token) NVFP4 for the forward activation. In a # training Linear this also drives the row-scaled columnwise (transpose) # of the activation that the wgrad GEMM consumes in the backward pass. nvfp4_recipe.row_scaled_activation = True + nvfp4_recipe.row_scaled_weight = row_scaled_weight return nvfp4_recipe @staticmethod def nvfp4_recipe_to_test( - with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False + with_rht: bool = False, + with_2d_quantization: bool = False, + row_scaled: bool = False, + row_scaled_weight: bool = False, ): if row_scaled: - return GetRecipes.nvfp4_row_scaled() + return GetRecipes.nvfp4_row_scaled(with_rht, with_2d_quantization, row_scaled_weight) if with_rht and with_2d_quantization: return GetRecipes.nvfp4_rht_and_2d_quantization() elif with_rht: @@ -81,7 +89,10 @@ def nvfp4_recipe_to_test( def get_nvfp4_quantizer_factory( - with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False + with_rht: bool = False, + with_2d_quantization: bool = False, + row_scaled: bool = False, + row_scaled_weight: bool = False, ): """ Create a quantizer factory for NVFP4 reference implementation. @@ -114,8 +125,6 @@ def factory(role): if role is None: return _default_quantizer() if role.tensor_type == "input": - # Only the forward activation is row-scaled, mirroring the production - # wiring in quantization.py (mode == "forward" and tensor_type != "weight"). return reference_nvfp4.NVFP4QuantizerRef( dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), @@ -129,9 +138,16 @@ def factory(role): quant_tile_shape=(16, 16) if with_2d_quantization else (1, 16), pow_2_scales=False, with_rht=False, + row_scaled_nvfp4=row_scaled_weight, ) if role.tensor_type == "grad_output": - return _default_quantizer() + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, + quant_tile_shape=(1, 16), + pow_2_scales=False, + with_rht=with_rht, + row_scaled_nvfp4=row_scaled, + ) return _default_quantizer() return factory @@ -170,6 +186,7 @@ def check_nvfp4_module_versus_reference( with_rht: bool = False, with_2d_quantization: bool = False, row_scaled: bool = False, + row_scaled_weight: bool = False, rel_tol: float = None, ): """ @@ -184,6 +201,7 @@ def check_nvfp4_module_versus_reference( num_steps: Number of forward/backward steps to test row_scaled: Enable row-scaled (per-token) NVFP4 activation quantization, exercising the row-scaled columnwise (transpose) path consumed by wgrad. + row_scaled_weight: Enable row-scaled NVFP4 weight quantization (requires row_scaled). rel_tol: If set, compare tensors by relative Frobenius-norm error instead of the default bitwise-tight ``assert_close`` (used for the row-scaled path). """ @@ -252,8 +270,12 @@ def check_nvfp4_module_versus_reference( ref_module.layer_norm_bias.copy_(native_module.layer_norm_bias) # Create recipes for native and reference implementations - nvfp4_recipe = GetRecipes.nvfp4_recipe_to_test(with_rht, with_2d_quantization, row_scaled) - nvfp4_ref_factory = get_nvfp4_quantizer_factory(with_rht, with_2d_quantization, row_scaled) + nvfp4_recipe = GetRecipes.nvfp4_recipe_to_test( + with_rht, with_2d_quantization, row_scaled, row_scaled_weight + ) + nvfp4_ref_factory = get_nvfp4_quantizer_factory( + with_rht, with_2d_quantization, row_scaled, row_scaled_weight + ) nvfp4_ref_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_factory) # Training loop comparison @@ -443,10 +465,20 @@ def test_nvfp4_linear_versus_reference( ], ) @pytest.mark.parametrize("num_steps", [1, 3], ids=["single_step", "multi_step"]) +@pytest.mark.parametrize("with_rht", [False, True], ids=["no_rht", "with_rht"]) +@pytest.mark.parametrize( + "weight_scaling", ["per_tensor_1d", "per_tensor_2d", "row_scaled_1d"], ids=lambda s: f"w_{s}" +) +@pytest.mark.parametrize( + "module_class", [te.Linear, te.LayerNormLinear], ids=["Linear", "LayerNormLinear"] +) def test_nvfp4_linear_row_scaled_versus_reference( in_features: int, out_features: int, num_steps: int, + with_rht: bool, + weight_scaling: str, + module_class, ): """End-to-end row-scaled (per-token) NVFP4 Linear forward + backward. @@ -457,19 +489,155 @@ def test_nvfp4_linear_row_scaled_versus_reference( reference quantizer by relative norm (both share the same 4-bit error). """ check_nvfp4_module_versus_reference( - module_class=te.Linear, + module_class=module_class, in_features=in_features, out_features=out_features, bias=False, x_dtype=torch.bfloat16, num_steps=num_steps, - with_rht=False, - with_2d_quantization=False, + with_rht=with_rht, + with_2d_quantization=weight_scaling == "per_tensor_2d", row_scaled=True, + row_scaled_weight=weight_scaling == "row_scaled_1d", rel_tol=2e-2, ) +def _make_default_recipe_module(name: str): + hidden_size, ffn_hidden_size = 512, 1024 + kwargs = dict(params_dtype=torch.bfloat16, device="cuda") + if name == "Linear": + return te.Linear(hidden_size, ffn_hidden_size, **kwargs) + if name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, ffn_hidden_size, **kwargs) + if name == "LayerNormMLP": + return te.LayerNormMLP(hidden_size, ffn_hidden_size, **kwargs) + return te.GroupedLinear( + 4, + hidden_size, + ffn_hidden_size, + use_grouped_tensor=name == "GroupedLinear_grouped_tensor", + **kwargs, + ) + + +def _run_default_recipe_module(name, state, nvfp4_recipe, x, dy): + module = _make_default_recipe_module(name) + module.load_state_dict(state) + x = x.detach().clone().requires_grad_(True) + args = () + if name.startswith("GroupedLinear"): + m_splits = [512, 256, 0, 1280] + if name == "GroupedLinear_grouped_tensor": + m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cuda") + args = (m_splits,) + with te.autocast(enabled=nvfp4_recipe is not None, recipe=nvfp4_recipe): + y = module(x, *args) + y.backward(dy[..., : y.shape[-1]]) + results = {"output": y.detach().float(), "input_grad": x.grad.detach().float()} + for param_name, param in module.named_parameters(): + results[param_name] = param.grad.detach().float() + return results + + +def _relative_error(native: torch.Tensor, ref: torch.Tensor) -> float: + diff = (native - ref).norm().item() + denom = ref.norm().item() + return diff / denom if denom > 0 else diff + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "name", + [ + "Linear", + "LayerNormLinear", + "LayerNormMLP", + "GroupedLinear_split_quantize", + "GroupedLinear_grouped_tensor", + ], +) +@pytest.mark.parametrize( + "row_scaled_activation, row_scaled_weight", + [(True, False), (False, True), (True, True)], + ids=["row_scaled_act", "row_scaled_w", "row_scaled_act_w"], +) +def test_nvfp4_row_scaled_default_recipe_error_versus_per_tensor( + name: str, row_scaled_activation: bool, row_scaled_weight: bool +): + """Row-scaled NVFP4 with default recipe settings is no less accurate than per-tensor.""" + reset_rng_states() + state = _make_default_recipe_module(name).state_dict() + tokens = 2048 + row_magnitude = torch.exp(torch.randn(tokens, 1, device="cuda")) + x = (torch.randn(tokens, 512, device="cuda") * row_magnitude).to(torch.bfloat16) + dy = torch.randn(tokens, 1024, device="cuda", dtype=torch.bfloat16) + + reference = _run_default_recipe_module(name, state, None, x, dy) + per_tensor = _run_default_recipe_module(name, state, recipe.NVFP4BlockScaling(), x, dy) + row_scaled = _run_default_recipe_module( + name, + state, + recipe.NVFP4BlockScaling( + row_scaled_activation=row_scaled_activation, row_scaled_weight=row_scaled_weight + ), + x, + dy, + ) + + for key, ref in reference.items(): + assert torch.isfinite(row_scaled[key]).all(), f"{key} has non-finite values" + err_row_scaled = _relative_error(row_scaled[key], ref) + err_per_tensor = _relative_error(per_tensor[key], ref) + print(f"{key}: row-scaled {err_row_scaled:.4f}, per-tensor {err_per_tensor:.4f}") + assert ( + err_row_scaled <= 1.25 * err_per_tensor + 1e-6 + ), f"{key} relative error {err_row_scaled:.4f} exceeds per-tensor {err_per_tensor:.4f}" + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_nvfp4_row_scaled_weight_rejects_quantized_primary_weights(): + nvfp4_recipe = recipe.NVFP4BlockScaling(row_scaled_weight=True) + with pytest.raises(NotImplementedError, match="row_scaled_weight"): + with te.quantized_model_init(enabled=True, recipe=nvfp4_recipe): + te.Linear(128, 128, params_dtype=torch.bfloat16, device="cuda") + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("name", ["Linear", "GroupedLinear"]) +def test_nvfp4_row_scaled_save_original_input(name: str): + """Row-scaled NVFP4 falls back to saving the quantized input for wgrad.""" + reset_rng_states() + kwargs = dict(params_dtype=torch.bfloat16, device="cuda") + + def make(save_original_input): + if name == "Linear": + return te.Linear(512, 1024, save_original_input=save_original_input, **kwargs) + return te.GroupedLinear(4, 512, 1024, save_original_input=save_original_input, **kwargs) + + state = make(False).state_dict() + x = torch.randn(2048, 512, device="cuda", dtype=torch.bfloat16) + dy = torch.randn(2048, 1024, device="cuda", dtype=torch.bfloat16) + args = ([512, 256, 0, 1280],) if name == "GroupedLinear" else () + nvfp4_recipe = recipe.NVFP4BlockScaling( + row_scaled_activation=True, disable_stochastic_rounding=True + ) + + results = [] + for save_original_input in (False, True): + module = make(save_original_input) + module.load_state_dict(state) + inp = x.clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=nvfp4_recipe): + y = module(inp, *args) + y.backward(dy) + results.append([y, inp.grad] + [p.grad for p in module.parameters()]) + + for native, ref in zip(*results): + assert torch.isfinite(native).all() + torch.testing.assert_close(native, ref, atol=0, rtol=0) + + def check_nvfp4_layernorm_linear_versus_reference( in_features: int, out_features: int, diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index 185ceb2880f..55b26224f63 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -44,6 +44,7 @@ def check_quantization_nvfp4_versus_reference( with_rht: bool = True, with_post_rht_amax: bool = True, with_random_sign_mask: bool = True, + row_scaled: bool = False, ) -> None: assert with_rht and with_post_rht_amax, "RHT and post-RHT amax reduction must be enabled." @@ -70,6 +71,7 @@ def check_quantization_nvfp4_versus_reference( with_rht=with_rht, with_post_rht_amax=with_post_rht_amax, with_random_sign_mask=with_random_sign_mask, + row_scaled_nvfp4=row_scaled, ) if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x) @@ -108,6 +110,7 @@ def check_quantization_nvfp4_versus_reference( quant_tile_shape=(1, 16), with_rht=with_rht, with_random_sign_mask=with_random_sign_mask, + row_scaled_nvfp4=row_scaled, ) x_nvfp4_ref = ref_quantizer.quantize(x) # Extract data from RefNVFP4Tensor @@ -129,12 +132,14 @@ def check_quantization_nvfp4_versus_reference( ref_quantizer._apply_rht(x.t().contiguous()) if with_rht else x.t().contiguous() ) ref_amax_colwise_t = torch.max(torch.abs(x_t_for_amax)).to(torch.float32).view(1) + if row_scaled: + ref_amax_colwise_t = x_nvfp4_ref.global_amax_col # SM120/121 uses TE's native single-K=16 MMA Hadamard arithmetic. cuBLAS may # choose a different reduction order (and applies the random-sign matrix in # the ATen reference orientation), so use the unfused TE kernel as the exact # reference for the fused TE kernel on these architectures. - if torch.cuda.get_device_capability() in ((12, 0), (12, 1)): + if not row_scaled and torch.cuda.get_device_capability() in ((12, 0), (12, 1)): env_name = "NVTE_NVFP4_DISABLE_RHT_CAST_FUSION" old_value = os.environ.get(env_name) os.environ[env_name] = "1" @@ -163,7 +168,9 @@ def check_quantization_nvfp4_versus_reference( ref_amax_colwise_t = None if return_rowwise: - torch.testing.assert_close(amax_rowwise, ref_amax_rowwise, atol=0.0, rtol=0.0) + torch.testing.assert_close( + amax_rowwise.reshape(-1), ref_amax_rowwise.reshape(-1), atol=0.0, rtol=0.0 + ) torch.testing.assert_close(qx, qx_ref, atol=0.0, rtol=0.0) # Compare only the valid portion of scale tensors (reference may not have padding) @@ -172,7 +179,9 @@ def check_quantization_nvfp4_versus_reference( torch.testing.assert_close(sx_valid, sx_ref, atol=0.0, rtol=0.0) if return_transpose: - torch.testing.assert_close(amax_colwise, ref_amax_colwise_t, atol=0.0, rtol=0.0) + torch.testing.assert_close( + amax_colwise.reshape(-1), ref_amax_colwise_t.reshape(-1), atol=0.0, rtol=0.0 + ) torch.testing.assert_close(qx_t, qx_t_ref, atol=0.0, rtol=0.0) @@ -224,6 +233,7 @@ def check_quantization_nvfp4_versus_reference( @pytest.mark.parametrize( "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] ) +@pytest.mark.parametrize("row_scaled", [False, True], ids=["tensor_scaled", "row_scaled"]) def test_rht_with_quantization_block_tiling_versus_reference( x_dtype: torch.dtype, M: int, @@ -231,6 +241,7 @@ def test_rht_with_quantization_block_tiling_versus_reference( quantize_mode: str, use_cpp_allocator: bool, with_random_sign_mask: bool, + row_scaled: bool, ) -> None: if quantize_mode == "rowwise_only": @@ -240,6 +251,8 @@ def test_rht_with_quantization_block_tiling_versus_reference( return_rowwise = True return_transpose = True elif quantize_mode == "columnwise_only": + if row_scaled: + pytest.skip("Row-scaled NVFP4 quantization requires rowwise usage") return_rowwise = False return_transpose = True else: @@ -254,6 +267,7 @@ def test_rht_with_quantization_block_tiling_versus_reference( return_transpose=return_transpose, use_cpp_allocator=use_cpp_allocator, with_random_sign_mask=with_random_sign_mask, + row_scaled=row_scaled, ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py index 2ed496c7304..49ffe75cb2e 100755 --- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -from typing import List, Tuple +from typing import List, Optional, Tuple import pytest import torch @@ -128,9 +128,19 @@ def _build_hadamard_matrix( return out.view(original_shape) +def row_amaxes(x_nvfp4) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + if not x_nvfp4._row_scaled_nvfp4: + return None, None + return x_nvfp4._amax_rowwise.reshape(-1, 1), x_nvfp4._amax_columnwise.reshape(-1, 1) + + def quantize_fp4( - x: torch.Tensor, use_stochastic_rounding: bool, use_2D: bool, use_RHT: bool -) -> torch.Tensor: + x: torch.Tensor, + use_stochastic_rounding: bool, + use_2D: bool, + use_RHT: bool, + row_scaled: bool = False, +) -> Tuple[torch.Tensor, ...]: nvfp4_quantizer = NVFP4Quantizer( rowwise=True, columnwise=True, @@ -140,6 +150,7 @@ def quantize_fp4( with_post_rht_amax=True, stochastic_rounding=use_stochastic_rounding, with_2d_quantization=use_2D, + row_scaled_nvfp4=row_scaled, ) x_nvfp4_sut = nvfp4_quantizer(x) @@ -153,7 +164,7 @@ def quantize_fp4( assert x_nvfp4_sut._columnwise_scale_inv is not None sx_t: torch.Tensor = x_nvfp4_sut._columnwise_scale_inv - return qx, sx, qx_t, sx_t + return qx, sx, qx_t, sx_t, *row_amaxes(x_nvfp4_sut) def group_quantize_fp4( @@ -163,7 +174,8 @@ def group_quantize_fp4( use_RHT: bool, split_sections: list[int], use_tex_split_quantize: bool = True, -) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]: + row_scaled: bool = False, +) -> Tuple[List, ...]: """ Group quantize function with toggle between tex.split_quantize and manual split/call methods. @@ -173,7 +185,9 @@ def group_quantize_fp4( use_2D (bool): Use 2D quantization. use_RHT (bool): Use RHT. split_sections (list[int]): Split sizes for inputs. - use_tex_split_quantize (bool): Toggle method. If True, use tex.split_quantize, else use manual split and per-quantizer invocation. + use_tex_split_quantize (bool): Toggle method. If True, use tex.split_quantize (or + tex.group_quantize for row-scaled), else use manual split and per-quantizer invocation. + row_scaled (bool): Use row-scaled NVFP4. Returns: tuple: Lists of quantized tensors and scale tensors for all sections. @@ -189,12 +203,19 @@ def group_quantize_fp4( with_post_rht_amax=True, stochastic_rounding=use_stochastic_rounding, with_2d_quantization=use_2D, + row_scaled_nvfp4=row_scaled, ) for _ in range(num_tensors) ] if use_tex_split_quantize: - outputs = tex.split_quantize(x, split_sections, nvfp4_quantizers) + if row_scaled: + first_dims = torch.tensor(split_sections, dtype=torch.int64, device=x.device) + grouped = tex.group_quantize(x, nvfp4_quantizers[0], num_tensors, first_dims) + outputs = grouped.split_into_quantized_tensors() + else: + outputs = tex.split_quantize(x, split_sections, nvfp4_quantizers) + amax_list = [row_amaxes(output) for output in outputs] qx_list = [output._rowwise_data.view(dtype=torch.uint8) for output in outputs] sx_list = [output._rowwise_scale_inv for output in outputs] qx_t_list = [output._columnwise_data.view(dtype=torch.uint8) for output in outputs] @@ -205,6 +226,7 @@ def group_quantize_fp4( sx_list = [] qx_t_list = [] sx_t_list = [] + amax_list = [] for i in range(num_tensors): x_chunk = x_chunks[i] x_nvfp4_sut = nvfp4_quantizers[i](x_chunk) @@ -220,12 +242,13 @@ def group_quantize_fp4( sx_list.append(sx) qx_t_list.append(qx_t) sx_t_list.append(sx_t) + amax_list.append(row_amaxes(x_nvfp4_sut)) - return qx_list, sx_list, qx_t_list, sx_t_list + return qx_list, sx_list, qx_t_list, sx_t_list, amax_list def check_quantization_nvfp4_versus_reference( - x_dtype: torch.dtype, M: int, N: int, use_2D: bool, use_RHT: bool + x_dtype: torch.dtype, M: int, N: int, use_2D: bool, use_RHT: bool, row_scaled: bool = False ) -> None: device = "cuda" torch.manual_seed(seed) @@ -236,11 +259,14 @@ def check_quantization_nvfp4_versus_reference( if use_RHT: y = RHT(y) amax = torch.max(torch.abs(x)).float() - q_rn, s_rn, q_t_rn, s_t_rn = quantize_fp4( - x, use_stochastic_rounding=False, use_2D=use_2D, use_RHT=use_RHT + q_rn, s_rn, q_t_rn, s_t_rn, amax_rn, amax_t_rn = quantize_fp4( + x, use_stochastic_rounding=False, use_2D=use_2D, use_RHT=use_RHT, row_scaled=row_scaled ) + amax_t = amax + if row_scaled: + amax, amax_t = amax_rn, amax_t_rn dq_rn = dequantize_fp4(q_rn, s_rn, amax) - dq_t_rn = dequantize_fp4(q_t_rn, s_t_rn, amax) + dq_t_rn = dequantize_fp4(q_t_rn, s_t_rn, amax_t) error_rn = (dq_rn - x).float() me_rn = torch.sqrt((error_rn * error_rn).mean()) error_t_rn = (dq_t_rn - y).float() @@ -248,12 +274,12 @@ def check_quantization_nvfp4_versus_reference( sr_result = torch.zeros_like(x).float() sr_t_result = torch.zeros_like(x).float().t().contiguous() for i in range(n_iters): - q_sr, s_sr, q_t_sr, s_t_sr = quantize_fp4( - x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=use_RHT + q_sr, s_sr, q_t_sr, s_t_sr, _, _ = quantize_fp4( + x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=use_RHT, row_scaled=row_scaled ) dq_sr = dequantize_fp4(q_sr, s_sr, amax) - dq_t_sr = dequantize_fp4(q_t_sr, s_t_sr, amax) + dq_t_sr = dequantize_fp4(q_t_sr, s_t_sr, amax_t) sr_result += dq_sr.float() sr_t_result += dq_t_sr.float() @@ -290,6 +316,7 @@ def check_group_quantization_nvfp4_versus_reference( use_RHT: bool, num_splits: int, use_tex_split_quantize: bool = True, + row_scaled: bool = False, ) -> None: device = "cuda" torch.manual_seed(seed) @@ -299,23 +326,25 @@ def check_group_quantization_nvfp4_versus_reference( x_total = torch.randn((M, N), dtype=x_dtype, device=device) * 2 - 1 x_splits = torch.split(x_total, split_sections) - q_rn_list, s_rn_list, q_t_rn_list, s_t_rn_list = group_quantize_fp4( + q_rn_list, s_rn_list, q_t_rn_list, s_t_rn_list, amax_rn_list = group_quantize_fp4( x_total, use_stochastic_rounding=False, use_2D=use_2D, use_RHT=use_RHT, split_sections=split_sections, use_tex_split_quantize=use_tex_split_quantize, + row_scaled=row_scaled, ) sr_n_iter_results = [] for i in range(n_iters): - q_sr_list, s_sr_list, q_t_sr_list, s_t_sr_list = group_quantize_fp4( + q_sr_list, s_sr_list, q_t_sr_list, s_t_sr_list, _ = group_quantize_fp4( x_total, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=use_RHT, split_sections=split_sections, use_tex_split_quantize=use_tex_split_quantize, + row_scaled=row_scaled, ) sr_n_iter_results.append((q_sr_list, s_sr_list, q_t_sr_list, s_t_sr_list)) @@ -323,7 +352,9 @@ def check_group_quantization_nvfp4_versus_reference( y = x.t().contiguous() if use_RHT: y = RHT(y) - amax = torch.max(torch.abs(x)).float() + amax = amax_t = torch.max(torch.abs(x)).float() + if row_scaled: + amax, amax_t = amax_rn_list[i] # fetch q_rn, s_rn, q_t_rn, s_t_rn q_rn = q_rn_list[i] @@ -332,7 +363,7 @@ def check_group_quantization_nvfp4_versus_reference( s_t_rn = s_t_rn_list[i] dq_rn = dequantize_fp4(q_rn, s_rn, amax) - dq_t_rn = dequantize_fp4(q_t_rn, s_t_rn, amax) + dq_t_rn = dequantize_fp4(q_t_rn, s_t_rn, amax_t) error_rn = (dq_rn - x).float() me_rn = torch.sqrt((error_rn * error_rn).mean()) error_t_rn = (dq_t_rn - y).float() @@ -347,7 +378,7 @@ def check_group_quantization_nvfp4_versus_reference( s_t_sr = result_sr[3][i] dq_sr = dequantize_fp4(q_sr, s_sr, amax) - dq_t_sr = dequantize_fp4(q_t_sr, s_t_sr, amax) + dq_t_sr = dequantize_fp4(q_t_sr, s_t_sr, amax_t) sr_result += dq_sr.float() sr_t_result += dq_t_sr.float() @@ -379,19 +410,24 @@ def check_group_quantization_nvfp4_versus_reference( @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("use_2D", [False, True], ids=str) @pytest.mark.parametrize("use_RHT", [False, True], ids=str) +@pytest.mark.parametrize("row_scaled", [False, True], ids=["tensor_scaled", "row_scaled"]) def test_quantization_block_tiling_versus_reference( x_dtype: torch.dtype, use_2D: bool, use_RHT: bool, + row_scaled: bool, M: int, N: int, ) -> None: if x_dtype == torch.float32 and use_RHT: pytest.skip("RHT is only supported with bfloat16") + if row_scaled and (use_2D or x_dtype != torch.bfloat16): + pytest.skip("Row-scaled NVFP4 requires 1D quantization and bfloat16 input") check_quantization_nvfp4_versus_reference( x_dtype=x_dtype, use_2D=use_2D, use_RHT=use_RHT, + row_scaled=row_scaled, M=M, N=N, ) @@ -411,12 +447,14 @@ def test_quantization_block_tiling_versus_reference( @pytest.mark.parametrize("use_RHT", [True], ids=str) @pytest.mark.parametrize("num_splits", [4, 8], ids=str) @pytest.mark.parametrize("use_tex_split_quantize", [True, False], ids=str) +@pytest.mark.parametrize("row_scaled", [False, True], ids=["tensor_scaled", "row_scaled"]) def test_group_stochastic_rounding_quantization_versus_reference( x_dtype: torch.dtype, use_2D: bool, use_RHT: bool, num_splits: int, use_tex_split_quantize: bool, + row_scaled: bool, M: int, N: int, ) -> None: @@ -430,29 +468,39 @@ def test_group_stochastic_rounding_quantization_versus_reference( N=N, num_splits=num_splits, use_tex_split_quantize=use_tex_split_quantize, + row_scaled=row_scaled, ) @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("use_2D", [False, True], ids=str) +@pytest.mark.parametrize("row_scaled", [False, True], ids=["tensor_scaled", "row_scaled"]) def test_stochastic_rounding_picks_an_adjacent_fp4_value( - x_dtype: torch.dtype, use_2D: bool + x_dtype: torch.dtype, use_2D: bool, row_scaled: bool ) -> None: + if row_scaled and (use_2D or x_dtype != torch.bfloat16): + pytest.skip("Row-scaled NVFP4 requires 1D quantization and bfloat16 input") device = "cuda" torch.manual_seed(seed) M, N = 512, 1024 x = torch.randn((M, N), dtype=x_dtype, device=device) - amax = torch.max(torch.abs(x)).float() + amax = amax_t = torch.max(torch.abs(x)).float() - q, s, q_t, s_t = quantize_fp4(x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False) - q_redraw, _, q_t_redraw, _ = quantize_fp4( - x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False + q, s, q_t, s_t, amax_row, amax_col = quantize_fp4( + x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False, row_scaled=row_scaled + ) + if row_scaled: + amax, amax_t = amax_row, amax_col + q_redraw, _, q_t_redraw, _, _, _ = quantize_fp4( + x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False, row_scaled=row_scaled ) assert not torch.equal(q, q_redraw), "two stochastic rounding draws are identical" assert not torch.equal(q_t, q_t_redraw), "two stochastic rounding draws are identical" - def check_adjacent(qx: torch.Tensor, sx: torch.Tensor, reference: torch.Tensor) -> None: + def check_adjacent( + qx: torch.Tensor, sx: torch.Tensor, amax: torch.Tensor, reference: torch.Tensor + ) -> None: # The kernel rounds reference / block_scale onto the E2M1 grid, so every # output has to be one of the two grid values that bracket it, i.e. within # one grid step. An all-zero output fails this everywhere the input is not @@ -465,5 +513,30 @@ def check_adjacent(qx: torch.Tensor, sx: torch.Tensor, reference: torch.Tensor) gap = (fp4_to_fp32(unpack_fp4(qx)) - exact).abs() assert torch.all(gap <= step * 1.0001), f"largest gap to the exact value is {gap.max()}" - check_adjacent(q, s, x) - check_adjacent(q_t, s_t, x.t().contiguous()) + check_adjacent(q, s, amax, x) + check_adjacent(q_t, s_t, amax_t, x.t().contiguous()) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_grouped_row_scaled_stochastic_rounding_capture_redraws() -> None: + split_sections = [256, 128] + N = 512 + x = torch.randn((sum(split_sections), N), dtype=torch.bfloat16, device="cuda") + first_dims = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + quantizer = NVFP4Quantizer(stochastic_rounding=True, row_scaled_nvfp4=True) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + tex.group_quantize(x, quantizer, len(split_sections), first_dims) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + grouped = tex.group_quantize(x, quantizer, len(split_sections), first_dims) + + graph.replay() + first = grouped.rowwise_data.clone() + graph.replay() + torch.cuda.synchronize() + assert not torch.equal(first, grouped.rowwise_data), "graph replays reuse the same SR draw" diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 0a0c42b03eb..887d777b38d 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -33,6 +33,7 @@ general_grouped_gemm_for_grouped_tensor, ) from transformer_engine.pytorch.constants import TE_DType +import transformer_engine.pytorch.module.grouped_linear as grouped_linear_module from transformer_engine.pytorch.module.grouped_linear import ( _GroupedLinear, is_module_grouped_tensor_path_supported, @@ -104,20 +105,20 @@ def nvfp4_row_scaled(): return nvfp4_recipe -def nvfp4_row_scaled_quantized_backward(): +def nvfp4_row_scaled_quantized_backward(with_rht: bool = False): # Same row-scaled activation recipe as nvfp4_row_scaled(), but with # backward_override=None so the backward runs in NVFP4 instead of falling back # to high precision. nvfp4_recipe = recipe.NVFP4BlockScaling( - disable_rht=True, + disable_rht=not with_rht, disable_stochastic_rounding=True, disable_2d_quantization=True, row_scaled_activation=True, backward_override=None, ) - nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams(random_hadamard_transform=with_rht) nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() - nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams(random_hadamard_transform=with_rht) return nvfp4_recipe @@ -406,16 +407,12 @@ def test_grouped_linear_accuracy( @pytest.mark.parametrize("num_gemms", [1, 3]) @pytest.mark.parametrize("bs", [2]) @pytest.mark.parametrize("bias", all_boolean) -def test_grouped_linear_row_scaled_quantized_backward(dtype, num_gemms, bs, bias, model="126m"): - """Row-scaled NVFP4 GroupedLinear with quantized (non-fallback) NVFP4 backward. - - With ``backward_override=None`` the wgrad is computed in NVFP4: the row-scaled - activation becomes operand A of the ``NT`` grouped GEMM, which this PR routes - through the per-expert dense ``general_gemm`` loop. GroupedLinear must then - match a stack of independent dense ``Linear`` layers bit-for-bit, since both - execute the exact same per-expert quantize + GEMM kernels. - """ - recipe_row_scaled = nvfp4_row_scaled_quantized_backward() +@pytest.mark.parametrize("with_rht", all_boolean) +def test_grouped_linear_row_scaled_quantized_backward( + dtype, num_gemms, bs, bias, with_rht, model="126m" +): + """Row-scaled NVFP4 GroupedLinear with quantized backward matches per-expert Linear.""" + recipe_row_scaled = nvfp4_row_scaled_quantized_backward(with_rht) config = model_configs[model] if dtype not in get_nvfp4_inp_supported_dtypes(recipe_row_scaled, dtype): pytest.skip(f"Input dtype {dtype} not supported for row-scaled NVFP4.") @@ -428,6 +425,7 @@ def test_grouped_linear_row_scaled_quantized_backward(dtype, num_gemms, bs, bias bias=bias, params_dtype=dtype, device="cuda", + use_grouped_tensor=False, ) .cuda() .eval() @@ -473,8 +471,6 @@ def test_grouped_linear_row_scaled_quantized_backward(dtype, num_gemms, bs, bias fuse_wgrad_accumulation=False, ) - # GroupedLinear is a per-expert loop over the same dense kernels, so the - # forward output, dgrad, and (row-scaled) wgrad must match bit-for-bit. for o, o_ref in zip(outputs, outputs_ref): torch.testing.assert_close(o, o_ref, rtol=0, atol=0) @@ -2063,6 +2059,17 @@ def test_grouped_parameter_layout_selected_by_constructor( assert isinstance(parameters["bias"], GroupedTensor) +@pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) +def test_op_grouped_linear_rejects_row_scaled_nvfp4(): + op = te.ops.GroupedLinear(2, 128, 128, device="cuda", dtype=torch.bfloat16) + x = torch.randn(256, 128, device="cuda", dtype=torch.bfloat16) + split_sizes = torch.tensor([128, 128], device="cuda", dtype=torch.int64) + fp8_recipe = recipe.NVFP4BlockScaling(row_scaled_activation=True) + with autocast(enabled=True, recipe=fp8_recipe): + with pytest.raises(NotImplementedError, match="row_scaled_activation"): + op(x, split_sizes) + + @pytest.mark.parametrize( "m_splits,exception", [([256, 256], ValueError), (torch.tensor([256, 256]), ValueError)], @@ -2760,6 +2767,10 @@ def _run_grouped_linear_path( recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), ), + pytest.param( + recipe.NVFP4BlockScaling(disable_stochastic_rounding=True, row_scaled_activation=True), + marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), + ), pytest.param( recipe.Float8BlockScaling(), marks=pytest.mark.skipif( @@ -2773,6 +2784,7 @@ def _run_grouped_linear_path( "mxfp8", "mxfp8_2d", "nvfp4", + "nvfp4_row_scaled", "fp8_block_scaling", ], ) @@ -2848,6 +2860,219 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( torch.testing.assert_close(grouped_tensor_out.float(), legacy_out.float(), **tols) +@pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) +@pytest.mark.parametrize("bias", _ALL_BOOLEAN) +def test_grouped_linear_row_scaled_nvfp4_selects_grouped_tensor_path(bias, monkeypatch): + fp8_recipe = recipe.NVFP4BlockScaling( + disable_stochastic_rounding=True, row_scaled_activation=True + ) + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported(fp8_recipe, dtype): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + + m_splits = [128, 0, 384] + num_gemms = len(m_splits) + in_features = 128 + out_features = 256 + torch.manual_seed(1234) + x_base = (0.1 * torch.randn(sum(m_splits), in_features, device="cuda")).to(dtype) + dy = (0.1 * torch.randn(sum(m_splits), out_features, device="cuda")).to(dtype) + weights = [ + (0.1 * torch.randn(out_features, in_features, device="cuda")).to(dtype) + for _ in range(num_gemms) + ] + biases = [(0.1 * torch.randn(out_features, device="cuda")).to(dtype) for _ in range(num_gemms)] + + expected = _run_grouped_linear_path( + enable_grouped_tensor_path=True, + fp8_recipe=fp8_recipe, + bias=bias, + fp8_model_params=False, + delay_wgrad_compute=False, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + monkeypatch=monkeypatch, + ) + + FP8GlobalStateManager.reset() + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) + grouped_gemm_calls = [] + grouped_gemm = grouped_linear_module.general_grouped_gemm_for_grouped_tensor + + def counting_grouped_gemm(*args, **kwargs): + grouped_gemm_calls.append(1) + return grouped_gemm(*args, **kwargs) + + monkeypatch.setattr( + grouped_linear_module, "general_grouped_gemm_for_grouped_tensor", counting_grouped_gemm + ) + module = GroupedLinear( + num_gemms, in_features, out_features, bias=bias, params_dtype=dtype, device="cuda" + ) + with torch.no_grad(): + for i in range(num_gemms): + getattr(module, f"weight{i}").copy_(weights[i]) + if bias: + getattr(module, f"bias{i}").copy_(biases[i]) + x = x_base.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=fp8_recipe): + y = module(x, m_splits) + y.backward(dy) + + assert grouped_gemm_calls + outputs = [y, x.grad] + for i in range(num_gemms): + outputs.append(getattr(module, f"weight{i}").grad) + if bias: + outputs.append(getattr(module, f"bias{i}").grad) + for actual, ref in zip(outputs, expected): + torch.testing.assert_close(actual, ref, rtol=0, atol=0) + + +@pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) +def test_grouped_linear_row_scaled_nvfp4_4over6_keeps_split_quantize_path(monkeypatch): + fp8_recipe = recipe.NVFP4BlockScaling(row_scaled_activation=True, nvfp4_4over6="weights") + dtype = torch.bfloat16 + assert not is_module_grouped_tensor_path_supported(fp8_recipe, dtype) + + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) + grouped_gemm_calls = [] + grouped_gemm = grouped_linear_module.general_grouped_gemm_for_grouped_tensor + + def counting_grouped_gemm(*args, **kwargs): + grouped_gemm_calls.append(1) + return grouped_gemm(*args, **kwargs) + + monkeypatch.setattr( + grouped_linear_module, "general_grouped_gemm_for_grouped_tensor", counting_grouped_gemm + ) + m_splits = [128, 0, 384] + module = GroupedLinear(len(m_splits), 128, 256, params_dtype=dtype, device="cuda") + x = torch.randn(sum(m_splits), 128, dtype=dtype, device="cuda", requires_grad=True) + with autocast(enabled=True, recipe=fp8_recipe): + y = module(x, m_splits) + y.backward(torch.randn_like(y)) + + assert not grouped_gemm_calls + assert torch.isfinite(y).all() and torch.isfinite(x.grad).all() + + +@pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) +@pytest.mark.parametrize( + "m_splits, in_features, out_features", + [ + ([96, 160], 128, 128), + ([64, 64], 128, 128), + ([32, 32], 128, 128), + ([128, 128], 128, 192), + ([128, 128], 192, 128), + ], +) +def test_grouped_linear_row_scaled_nvfp4_unaligned_shapes_keep_split_quantize_path( + m_splits, in_features, out_features, monkeypatch +): + fp8_recipe = recipe.NVFP4BlockScaling( + disable_stochastic_rounding=True, row_scaled_activation=True + ) + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported(fp8_recipe, dtype): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + + num_gemms = len(m_splits) + torch.manual_seed(1234) + x_base = (0.1 * torch.randn(sum(m_splits), in_features, device="cuda")).to(dtype) + dy = (0.1 * torch.randn(sum(m_splits), out_features, device="cuda")).to(dtype) + weights = [ + (0.1 * torch.randn(out_features, in_features, device="cuda")).to(dtype) + for _ in range(num_gemms) + ] + + expected = _run_grouped_linear_path( + enable_grouped_tensor_path=False, + fp8_recipe=fp8_recipe, + bias=False, + fp8_model_params=False, + delay_wgrad_compute=False, + x_base=x_base, + dy=dy, + weights=weights, + biases=None, + m_splits=m_splits, + monkeypatch=monkeypatch, + ) + + FP8GlobalStateManager.reset() + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) + grouped_gemm_calls = [] + grouped_gemm = grouped_linear_module.general_grouped_gemm_for_grouped_tensor + + def counting_grouped_gemm(*args, **kwargs): + grouped_gemm_calls.append(1) + return grouped_gemm(*args, **kwargs) + + monkeypatch.setattr( + grouped_linear_module, "general_grouped_gemm_for_grouped_tensor", counting_grouped_gemm + ) + module = GroupedLinear( + num_gemms, in_features, out_features, bias=False, params_dtype=dtype, device="cuda" + ) + with torch.no_grad(): + for i in range(num_gemms): + getattr(module, f"weight{i}").copy_(weights[i]) + x = x_base.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=fp8_recipe): + y = module(x, m_splits) + y.backward(dy) + + assert not grouped_gemm_calls + outputs = [y, x.grad] + [getattr(module, f"weight{i}").grad for i in range(num_gemms)] + for actual, ref in zip(outputs, expected): + torch.testing.assert_close(actual, ref, rtol=0, atol=0) + + +@pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) +@pytest.mark.parametrize("bias", _ALL_BOOLEAN) +def test_grouped_linear_row_scaled_nvfp4_weight_only_backward(bias): + fp8_recipe = recipe.NVFP4BlockScaling( + disable_stochastic_rounding=True, row_scaled_activation=True + ) + dtype = torch.bfloat16 + if not is_module_grouped_tensor_path_supported(fp8_recipe, dtype): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + + m_splits = torch.tensor([128, 0, 384], dtype=torch.int64, device="cuda") + num_gemms = m_splits.numel() + torch.manual_seed(1234) + x_base = (0.1 * torch.randn(512, 128, device="cuda")).to(dtype) + dy = (0.1 * torch.randn(512, 256, device="cuda")).to(dtype) + module = GroupedLinear( + num_gemms, + 128, + 256, + bias=bias, + params_dtype=dtype, + device="cuda", + use_grouped_tensor=True, + ) + + def run(input_requires_grad): + FP8GlobalStateManager.reset() + module.zero_grad(set_to_none=True) + x = x_base.detach().clone().requires_grad_(input_requires_grad) + with autocast(enabled=True, recipe=fp8_recipe): + y = module(x, m_splits) + y.backward(dy) + return [p.grad.clone() for p in module.parameters()] + + expected = run(True) + actual = run(False) + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monkeypatch): if not is_module_grouped_tensor_path_supported( None, diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 3b95bd9bb9f..6d6225bccf8 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -624,7 +624,7 @@ def expected_e4m3_max(tensor_type): QuantizerRole(module_type="linear", tensor_type="grad_input"), ], ).make_quantizers() - assert [q.row_scaled_nvfp4 for q in backward_quantizers] == [False, False] + assert [q.row_scaled_nvfp4 for q in backward_quantizers] == [True, False] assert [q.nvfp4_use_4over6 for q in backward_quantizers] == [False, False] assert [q.nvfp4_e4m3_max for q in backward_quantizers] == [0, 0] assert [q.nvfp4_4over6_err_mode for q in backward_quantizers] == [nvfp4_4over6_err_mode] * 2 @@ -658,6 +658,34 @@ def expected_e4m3_max(tensor_type): ] +@pytest.mark.skipif(not fp4_available, reason=reason_for_no_fp4) +@pytest.mark.parametrize("row_scaled_activation", [False, True], ids=["act", "row_scaled_act"]) +@pytest.mark.parametrize("row_scaled_weight", [False, True], ids=["w", "row_scaled_w"]) +@pytest.mark.parametrize("disable_2d_quantization", [False, True], ids=["2d", "1d"]) +def test_nvfp4_row_scaled_weight_quantizer_roles( + row_scaled_activation, row_scaled_weight, disable_2d_quantization +): + recipe = NVFP4BlockScaling( + row_scaled_activation=row_scaled_activation, + row_scaled_weight=row_scaled_weight, + disable_2d_quantization=disable_2d_quantization, + ) + weight_quantizer, input_quantizer = NVFP4BlockScalingRecipeState( + recipe, + mode="forward", + num_quantizers=2, + roles=[ + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="input"), + ], + ).make_quantizers() + assert weight_quantizer.row_scaled_nvfp4 == row_scaled_weight + assert weight_quantizer.with_2d_quantization == ( + not row_scaled_weight and not disable_2d_quantization + ) + assert input_quantizer.row_scaled_nvfp4 == row_scaled_activation + + @pytest.mark.skipif(not fp4_available, reason=reason_for_no_fp4) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 13c3694d040..2d7603f463f 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -564,8 +564,6 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for grouped NVTE_NVFP4_1D_SCALING."); NVTE_CHECK(output_tensor->row_scaled_nvfp4, "Grouped NVFP4 quantization is only implemented for row-scaled NVFP4."); - NVTE_CHECK(!quant_config_cpp.stochastic_rounding, - "Grouped row-scaled NVFP4 quantization does not support stochastic rounding."); NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); NVTE_CHECK(quant_config_cpp.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index cfc5dd42c81..2942978c7c1 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -26,6 +26,7 @@ #include "../core/grouped_layout.cuh" #include "core_nvfp4.cuh" #include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" +#include "swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -102,18 +103,31 @@ static __global__ void group_fused_amax_zero_kernel(GroupedFusedAmaxArgs args, c } // Shared amax body; host-split and graph-safe kernels differ only in the prologue. -template +template __device__ __forceinline__ void group_compute_fused_amax_body(const CUtensorMap &tensor_map_input, float *const row_amax_out, float *const col_amax_out, - const int split_start) { + const int split_start, + const uint16_t rht_sign_mask_t = 0) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using IType = __nv_bfloat16; using IType2 = ptx::FPx2; using IType3D = IType[FA_BUFFS_NUM][FA_TILE_DIM_Y][FA_TILE_DIM_X]; + constexpr int RHT_DIM = quantize_transpose_tuned_kernel::SCALE_DIM; + static_assert(FA_TILE_DIM_Y % RHT_DIM == 0, "Amax tile must hold whole RHT blocks."); const bool leading_thread = (threadIdx.x == 0); const int tid = threadIdx.x; + + __shared__ float sRHT[APPLY_COL_RHT ? RHT_DIM * RHT_DIM : 1]; + if constexpr (APPLY_COL_RHT) { + for (int i = tid; i < RHT_DIM * RHT_DIM; i += FA_THREADS_NUM) { + const int r = i / RHT_DIM; + const int c = i % RHT_DIM; + const bool negative = ((__popc(r & c) ^ (rht_sign_mask_t >> r)) & 1) != 0; + sRHT[i] = negative ? -0.25f : 0.25f; + } + } const int block_offset_Y = blockIdx.y * FA_CHUNK_DIM_Y; const int block_offset_X = blockIdx.x * FA_CHUNK_DIM_X; const int local_row = block_offset_Y - split_start + tid; @@ -217,10 +231,26 @@ __device__ __forceinline__ void group_compute_fused_amax_body(const CUtensorMap if constexpr (DO_COL) { if (stage_X == my_col_stage_X) { float local_max = col_partial; + if constexpr (APPLY_COL_RHT) { +#pragma unroll + for (int blk = 0; blk < FA_TILE_DIM_Y / RHT_DIM; ++blk) { + IType v[RHT_DIM]; +#pragma unroll + for (int e = 0; e < RHT_DIM; ++e) { + v[e] = sIn[buff_in][blk * RHT_DIM + e][my_col_in_subtile]; + } + quantize_transpose_tuned_kernel::apply_rht_16(v, sRHT); +#pragma unroll + for (int e = 0; e < RHT_DIM; ++e) { + local_max = fmaxf(local_max, fabsf(static_cast(v[e]))); + } + } + } else { #pragma unroll - for (int e = 0; e < FA_TILE_DIM_Y; ++e) { - const IType v = sIn[buff_in][e][my_col_in_subtile]; - local_max = fmaxf(local_max, fabsf(static_cast(v))); + for (int e = 0; e < FA_TILE_DIM_Y; ++e) { + const IType v = sIn[buff_in][e][my_col_in_subtile]; + local_max = fmaxf(local_max, fabsf(static_cast(v))); + } } col_partial = local_max; } @@ -299,13 +329,13 @@ static __global__ void group_fused_amax_zero_graph_safe_kernel(float *row_amax_a // Graph-safe: per-expert outputs and row start derived on device from offsets; // CTAs past the live row sum noop. -template +template __global__ void __launch_bounds__(FA_THREADS_NUM) group_compute_fused_amax_graph_safe_kernel(const __grid_constant__ CUtensorMap tensor_map_input, float *const row_amax_all, float *const col_amax_all, const int64_t *const __restrict__ offsets_ptr, const size_t num_tensors, const float *noop, - const size_t cols) { + const size_t cols, const uint16_t rht_sign_mask_t) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -323,8 +353,8 @@ __global__ void __launch_bounds__(FA_THREADS_NUM) float *const row_amax_out = DO_ROW ? row_amax_all + split_start : nullptr; float *const col_amax_out = DO_COL ? col_amax_all + tensor_id * cols : nullptr; - group_compute_fused_amax_body(tensor_map_input, row_amax_out, col_amax_out, - split_start); + group_compute_fused_amax_body( + tensor_map_input, row_amax_out, col_amax_out, split_start, rht_sign_mask_t); #else NVTE_DEVICE_ERROR("Grouped fused amax kernel requires SM 10.0+ (Blackwell)."); #endif // __CUDA_ARCH__ >= 1000 @@ -420,7 +450,9 @@ inline void group_row_amax(const Tensor &input, const Tensor *noop, // output->tensor_offsets, amax goes to the contiguous rowwise [sum_M] / columnwise // [num_tensors * K] buffers. BF16-only, VARYING_FIRST_DIM. inline void group_row_amax_graph_safe(const GroupedTensor *input, const Tensor *noop, - GroupedTensor *output, cudaStream_t stream) { + GroupedTensor *output, cudaStream_t stream, + const bool apply_col_rht = false, + const uint16_t rht_sign_mask_t = 0) { #if FP4_TYPE_SUPPORTED using namespace group_row_scaled_amax_kernel; @@ -479,12 +511,16 @@ inline void group_row_amax_graph_safe(const GroupedTensor *input, const Tensor * const dim3 block(FA_THREADS_NUM, 1, 1); TRANSFORMER_ENGINE_SWITCH_CONDITION( - do_row, DO_ROW, TRANSFORMER_ENGINE_SWITCH_CONDITION(do_col, DO_COL, { - auto kernel = group_compute_fused_amax_graph_safe_kernel; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - kernel<<>>(tensor_map_input, row_amax_all, col_amax_all, - offsets_ptr, num_tensors, noop_ptr, cols); - })); + do_row, DO_ROW, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + do_col, DO_COL, TRANSFORMER_ENGINE_SWITCH_CONDITION(apply_col_rht, APPLY_COL_RHT, { + auto kernel = group_compute_fused_amax_graph_safe_kernel < DO_ROW, DO_COL, + APPLY_COL_RHT && DO_COL > ; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); + kernel<<>>(tensor_map_input, row_amax_all, + col_amax_all, offsets_ptr, num_tensors, + noop_ptr, cols, rht_sign_mask_t); + }))); NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); @@ -1390,22 +1426,31 @@ using namespace ptx; #if FP4_TYPE_SUPPORTED -template +template __device__ __forceinline__ void group_row_scaled_cast_body( const CUtensorMap &tensor_map_input, uint8_t *const q_row_base, ScaleType *const s_dec_row_base, const float *const row_amax_base, uint8_t *const q_col_base, ScaleType *const s_dec_col_base, const float *const col_amax_base, const size_t scale_stride_col, const int M_t, const int local_row_base, const int ctaid_X, const int ctaid_Y, const size_t cols, - const size_t scale_stride_row) { + const size_t scale_stride_row, const size_t *rng_state, const IType *rht_matrix) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - constexpr bool USE_STOCHASTIC_ROUNDING = false; // row-scaled rejects SR constexpr bool ROW_SCALED_NVFP4 = true; - // SR off: the RNG stays inert (random_uint4 is never consumed); it exists only - // to satisfy the shared encode-helper signatures. + __shared__ float sRHT[APPLY_COLWISE_RHT ? SCALE_DIM * SCALE_DIM : 1]; + if constexpr (APPLY_COLWISE_RHT) { + for (int i = threadIdx.x; i < SCALE_DIM * SCALE_DIM; i += THREADS_NUM) { + sRHT[i] = static_cast(rht_matrix[i]); + } + } + + const size_t rng_sequence = + threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; RNG_t rng; - rng.init(0, 0, 0); - uint4 random_uint4 = uint4{0, 0, 0, 0}; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; int rnd_idx = 0; const bool leading_thread = (threadIdx.x == 0); @@ -1516,10 +1561,11 @@ __device__ __forceinline__ void group_row_scaled_cast_body( random_uint4, rnd_idx); if constexpr (RETURN_TRANSPOSE) { - colwise_scaling( - sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, /*S_enc_colwise=*/1.0f, stage_Y, stage_X, buff_in, - buff_out_tr, col_amax_base, static_cast(block_offset_X), cols, rng, random_uint4, - rnd_idx); + colwise_scaling( + sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, + /*S_enc_colwise=*/1.0f, stage_Y, stage_X, buff_in, buff_out_tr, col_amax_base, + static_cast(block_offset_X), cols, rng, random_uint4, rnd_idx, sRHT); } // Make the encode-helper SMEM writes visible before the scatter. @@ -1570,8 +1616,18 @@ __device__ __forceinline__ void group_row_scaled_cast_body( row += THREADS_NUM) { ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); const size_t local_row = static_cast(local_row_base) + row; - const size_t scale_idx = local_row * scale_stride_row + scales_block_offset_X_rowwise; - scales_vec.store_to_elts(&s_dec_row_base[scale_idx], 0, SCALES_PER_CHUNK_X); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + const ScaleType *scales = reinterpret_cast(&scales_vec); +#pragma unroll + for (int k = 0; k < SCALES_PER_CHUNK_X; ++k) { + const size_t col = scales_block_offset_X_rowwise + k; + s_dec_row_base[swizzle::gemm_swizzled_scale_idx(local_row, col, scale_stride_row)] = + scales[k]; + } + } else { + const size_t scale_idx = local_row * scale_stride_row + scales_block_offset_X_rowwise; + scales_vec.store_to_elts(&s_dec_row_base[scale_idx], 0, SCALES_PER_CHUNK_X); + } } } if constexpr (RETURN_TRANSPOSE) { @@ -1582,8 +1638,17 @@ __device__ __forceinline__ void group_row_scaled_cast_body( row_tr += THREADS_NUM) { ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); const size_t global_col = static_cast(block_offset_X) + row_tr; - const size_t scale_idx = global_col * scale_stride_col + local_block_offset_M; - scales_vec.store_to_elts(&s_dec_col_base[scale_idx], 0, SCALES_PER_CHUNK_Y); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + const ScaleType *scales = reinterpret_cast(&scales_vec); +#pragma unroll + for (int k = 0; k < SCALES_PER_CHUNK_Y; ++k) { + s_dec_col_base[swizzle::gemm_swizzled_scale_idx(global_col, local_block_offset_M + k, + scale_stride_col)] = scales[k]; + } + } else { + const size_t scale_idx = global_col * scale_stride_col + local_block_offset_M; + scales_vec.store_to_elts(&s_dec_col_base[scale_idx], 0, SCALES_PER_CHUNK_Y); + } } } @@ -1599,13 +1664,15 @@ __device__ __forceinline__ void group_row_scaled_cast_body( } // Grid covers the input capacity; CTAs past the live row sum noop. -template +template __global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_graph_safe_kernel( const __grid_constant__ CUtensorMap tensor_map_input, uint8_t *const q_row_all, ScaleType *const s_dec_row_all, const float *const row_amax_all, uint8_t *const q_col_all, ScaleType *const s_dec_col_all, const float *const col_amax_all, const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, - const size_t num_tensors, const float *noop, const size_t cols, const size_t scale_stride_row) { + const size_t num_tensors, const float *noop, const size_t cols, const size_t scale_stride_row, + const size_t *rng_state, const IType *rht_matrix) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -1639,10 +1706,11 @@ __global__ void __launch_bounds__(THREADS_NUM) group_row_scaled_cast_nvfp4_graph RETURN_TRANSPOSE ? s_dec_col_all + (cols / SCALE_DIM) * row_start : nullptr; const float *const col_amax_base = RETURN_TRANSPOSE ? col_amax_all + tensor_id * cols : nullptr; - group_row_scaled_cast_body( + group_row_scaled_cast_body( tensor_map_input, q_row_base, s_dec_row_base, row_amax_base, q_col_base, s_dec_col_base, col_amax_base, scale_stride_col, M_t, local_row_base, ctaid_X, ctaid_Y, cols, - scale_stride_row); + scale_stride_row, rng_state, rht_matrix); #else NVTE_DEVICE_ERROR("Grouped row-scaled NVFP4 cast kernel requires SM 10.0+ (Blackwell)."); #endif // __CUDA_ARCH__ >= 1000 @@ -1656,22 +1724,46 @@ template inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const Tensor *noop, GroupedTensor *output, const QuantizationConfig *quant_config, - cudaStream_t stream) { + cudaStream_t stream, + const Tensor *rht_matrix = nullptr) { #if FP4_TYPE_SUPPORTED using namespace group_row_scaled_cast_kernel; using namespace quantize_transpose_tuned_kernel; NVTE_CHECK(input->dtype() == DType::kBFloat16, "Grouped row-scaled NVFP4 cast requires BF16 input."); + const bool apply_colwise_rht = rht_matrix != nullptr; + if (apply_colwise_rht) { + NVTE_CHECK(output->has_columnwise_data(), "Columnwise RHT requires columnwise output."); + NVTE_CHECK(rht_matrix->dtype() == DType::kBFloat16, "Columnwise RHT matrix must be BF16."); + NVTE_CHECK( + rht_matrix->dim() == 2 && rht_matrix->shape()[0] == 16 && rht_matrix->shape()[1] == 16, + "Columnwise RHT matrix must have shape [16, 16]."); + } + const IType *rht_matrix_ptr = + apply_colwise_rht ? reinterpret_cast(rht_matrix->data.dptr) : nullptr; NVTE_CHECK(output->all_same_last_dim() || output->all_same_shape(), "Graph-safe grouped row-scaled NVFP4 cast only supports a common last dimension " "(VARYING_FIRST_DIM)."); const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; + const bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; checkCuDriverContext(stream); CheckNoopTensor(*noop, "cast_noop"); + const size_t *rng_state = nullptr; + if (use_stochastic_rounding) { + NVTE_CHECK(quant_config->rng_state != nullptr, + "Stochastic rounding requires an RNG state in the quantization config."); + const Tensor &rng_state_tensor = *convertNVTETensor(quant_config->rng_state); + NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64, + "RNG state should contain 2 64-bit values."); + NVTE_CHECK(rng_state_tensor.data.shape == Shape{2}, + "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); + rng_state = reinterpret_cast(rng_state_tensor.data.dptr); + } + const size_t num_tensors = input->num_tensors; const size_t capacity_rows = input->logical_shape.data[0]; const size_t cols = input->logical_shape.data[1]; @@ -1684,6 +1776,7 @@ inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const T TunableConfig::CHUNK_DIM_Y, " (got ", capacity_rows, ")."); const bool return_transpose = output->has_columnwise_data(); + const bool with_gemm_swizzled = output->with_gemm_swizzled_scales; NVTE_CHECK(output->scale_inv.dptr != nullptr, "Graph-safe cast requires rowwise scales."); NVTE_CHECK(output->amax.dptr != nullptr, "Graph-safe cast requires rowwise amax."); // Per-row scale stride, from a 2D [rows, stride] scale or a flat 1D buffer. @@ -1743,22 +1836,30 @@ inline void group_row_scaled_cast_graph_safe(const GroupedTensor *input, const T const dim3 block(THREADS_NUM, 1, 1); TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_fast_math, USE_FAST_MATH, - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - const int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; - const int out_scales_transpose_mem = RETURN_TRANSPOSE ? buff_size_scales_transpose : 0; - const int dshmem_size = buff_size_aligned_in + buff_size_aligned_out + - out_mem_colwise_data + buff_size_scales + out_scales_transpose_mem + - TMA_SHMEM_ALIGNMENT; - auto kernel = group_row_scaled_cast_nvfp4_graph_safe_kernel; - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - kernel<<>>(tensor_map_input, q_row_all, s_dec_row_all, - row_amax_all, q_col_all, s_dec_col_all, - col_amax_all, offsets_ptr, first_dims_ptr, - num_tensors, noop_ptr, cols, scale_stride_row); - })); + use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_fast_math, USE_FAST_MATH, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + return_transpose, RETURN_TRANSPOSE, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled, WITH_GEMM_SWIZZLED_SCALES, + TRANSFORMER_ENGINE_SWITCH_CONDITION(apply_colwise_rht, APPLY_COLWISE_RHT, { + const int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + const int out_scales_transpose_mem = + RETURN_TRANSPOSE ? buff_size_scales_transpose : 0; + const int dshmem_size = buff_size_aligned_in + buff_size_aligned_out + + out_mem_colwise_data + buff_size_scales + + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + auto kernel = group_row_scaled_cast_nvfp4_graph_safe_kernel < ScaleType, + USE_STOCHASTIC_ROUNDING, USE_FAST_MATH, RETURN_TRANSPOSE, + WITH_GEMM_SWIZZLED_SCALES, APPLY_COLWISE_RHT && RETURN_TRANSPOSE > ; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>( + tensor_map_input, q_row_all, s_dec_row_all, row_amax_all, q_col_all, + s_dec_col_all, col_amax_all, offsets_ptr, first_dims_ptr, num_tensors, + noop_ptr, cols, scale_stride_row, rng_state, rht_matrix_ptr); + }))))); NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 3c24f9a4111..709ead2db82 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -184,13 +184,33 @@ compute_nvfp4_scaling_coefficient(const ScaleType decode_scale, const float glob } } +// Rounds to BF16 to match the unfused RHT path. +__device__ __forceinline__ void apply_rht_16(IType (&v)[SCALE_DIM], + const float *__restrict__ rht_matrix) { + float x[SCALE_DIM]; +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + x[i] = static_cast(v[i]); + } +#pragma unroll + for (int c = 0; c < SCALE_DIM; ++c) { + float acc = 0.0f; +#pragma unroll + for (int r = 0; r < SCALE_DIM; ++r) { + acc = fmaf(x[r], rht_matrix[r * SCALE_DIM + c], acc); + } + v[c] = static_cast(acc); + } +} + template + bool ROW_SCALED_NVFP4, bool APPLY_COLWISE_RHT = false> __device__ __forceinline__ void colwise_scaling( const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_tr_ptr, ScaleType *__restrict__ sSFcolwise_ptr, const float S_enc_colwise, const int stage_Y, const int stage_X, const int buff_in, const int buff_out_tr, const float *amax_colwise_ptr, - const size_t col_offset, const size_t cols, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { + const size_t col_offset, const size_t cols, RNG_t &rng, uint4 &random_uint4, int &rnd_idx, + const float *rht_matrix = nullptr) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn2x = *reinterpret_cast(sIn_ptr); @@ -224,7 +244,18 @@ __device__ __forceinline__ void colwise_scaling( ptx::ld_shared_b32(&sIn2x[buff_in][in_thread_offset_Y + i][in_thread_offset_X]); rIn[0][i] = elt_pair.x; rIn[1][i] = elt_pair.y; - ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, elt_pair); + if constexpr (!APPLY_COLWISE_RHT) { + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, elt_pair); + } + } + if constexpr (APPLY_COLWISE_RHT) { + apply_rht_16(rIn[0], rht_matrix); + apply_rht_16(rIn[1], rht_matrix); +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + const IType2 elt_pair = {rIn[0][i], rIn[1][i]}; + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, elt_pair); + } } const float block_amax[2] = {static_cast(__habs(thread_amax_2x.x)), static_cast(__habs(thread_amax_2x.y))}; diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 0528ac0927a..cd8a1b670eb 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -543,6 +543,28 @@ void group_hadamard_transform_amax_graph_safe(const GroupedTensor* input, Groupe #endif // CUDA_VERSION >= 12080 } +void group_row_scaled_hadamard_transform_cast_graph_safe(const GroupedTensor* input, + GroupedTensor* output, + const Tensor& hadamard_matrix, + const QuantizationConfig& quant_config, + cudaStream_t stream) { + NVTE_CHECK(!quant_config.nvfp4_2d_quantization, + "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(quant_config.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Grouped row-scaled NVFP4 quantization does not support 4over6."); + Tensor dummy_tensor; + const Tensor* noop_tensor = &dummy_tensor; + if (quant_config.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config.noop_tensor); + } + const DType scale_dtype = + output->scale_inv.has_data() ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + dispatch::nvfp4::group_row_scaled_cast_graph_safe( + input, noop_tensor, output, &quant_config, stream, &hadamard_matrix);); +} + } // namespace transformer_engine void nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input, @@ -558,6 +580,15 @@ void nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input return; } + if (output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING && output_tensor->row_scaled_nvfp4) { + NVTE_CHECK(random_sign_mask == 0, + "Row-scaled NVFP4 does not support rowwise RHT; random_sign_mask must be 0."); + dispatch::nvfp4::group_row_amax_graph_safe(input_tensor, /*noop=*/nullptr, output_tensor, + stream, /*apply_col_rht=*/true, + static_cast(random_sign_mask_t)); + return; + } + // Call the group tensor Hadamard transform amax implementation. group_hadamard_transform_amax_graph_safe( input_tensor, output_tensor, static_cast(random_sign_mask), diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu index df19e907fd0..a1f6da4acc8 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -1313,6 +1313,12 @@ void group_row_col_rht_gemm_ntt_w_sfc_graph_safe( } // namespace } // namespace detail +void group_row_scaled_hadamard_transform_cast_graph_safe(const GroupedTensor *input, + GroupedTensor *output, + const Tensor &hadamard_matrix, + const QuantizationConfig &quant_config, + cudaStream_t stream); + void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, GroupedTensor *output, const Tensor &hadamard_matrix_, @@ -1477,8 +1483,6 @@ void nvte_group_hadamard_transform_cast_fusion_graph_safe( GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); - Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace); - QuantizationConfig quant_config_cpp; if (quant_config != nullptr) { quant_config_cpp = *reinterpret_cast(quant_config); @@ -1488,6 +1492,15 @@ void nvte_group_hadamard_transform_cast_fusion_graph_safe( return; } + if (output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING && output_tensor->row_scaled_nvfp4) { + group_row_scaled_hadamard_transform_cast_graph_safe(input_tensor, output_tensor, + *convertNVTETensorCheck(hadamard_matrix), + quant_config_cpp, stream); + return; + } + + Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace); + // Call the multi-tensor Hadamard transform amax implementation. group_hadamard_transform_cast_fusion_graph_safe( input_tensor, output_tensor, *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index 8f1a213cec1..b1d2ad89a34 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -109,6 +109,9 @@ void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outp * therefore, this is a graph safe API and the grouped-tensor argument is passed as a single device structure. * * \param[in] input NVTEGroupedTensor representing grouped input tensors. + * For row-scaled NVFP4 output, the row-wise amax is per row of the input and the column-wise + * amax is per column of the transposed RHT output; random_sign_mask must be 0. + * * \param[in,out] output NVTEGroupedTensor for output amax (row/col). Only the row-wise and * column-wise amaxes are updated. * \param[in] random_sign_mask 16-bit sign mask for RHT. @@ -166,6 +169,9 @@ void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETenso * \param[in,out] output NVTEGroupedTensor for output (row/column-wise quantized results). * \param[in] hadamard_matrix Hadamard matrix to use for transformation. * \param[in] quant_config Quantization configuration. + * For row-scaled NVFP4 output, the row-wise data is quantized without RHT using per-row amax, + * the column-wise data uses RHT with per-column amax, and quant_workspace may be null. + * * \param[in] quant_workspace Workspace buffer. Must be at least 4 bytes. * \param[in] stream CUDA stream used for the operation. */ diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 258504f9443..4cdafbbc1bb 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -576,9 +576,26 @@ class NVFP4BlockScaling(Recipe): disable_2d_quantization : bool, default = False If set to `True`, 1D block scaling with block size 16 is used for all tensors. row_scaled_activation : bool, default = False - If set to `True`, forward activation quantizers emit row-scaled - NVFP4 tensors. In this mode, rowwise ``amax`` metadata is stored - as a vector with one FP32 value per tensor row. + If set to `True`, activation and gradient-output quantizers emit + row-scaled NVFP4 tensors. Weights follow ``row_scaled_weight``. + In this mode, rowwise ``amax`` metadata is stored as a vector with + one FP32 value per tensor row and columnwise ``amax`` metadata as a + vector with one FP32 value per tensor column. Random Hadamard + transforms and stochastic rounding follow the same settings as + per-tensor NVFP4. Sequence parallelism, Userbuffers + communication overlap, and ``te.ops.GroupedLinear`` are not + supported. + row_scaled_weight : bool, default = False + If set to `True`, weight quantizers emit row-scaled NVFP4 tensors + with 1D (1x16) block scales, i.e. one global scale per output + channel for the forward GEMM and per input channel for the dgrad + GEMM. Weight 2D quantization is disabled in this mode. If `False`, + weights use a single per-tensor global scale with 16x16 (2D) + block scales, or 1x16 (1D) block scales when + ``disable_2d_quantization=True``. Row-scaled weights require BF16 + weight tensors with both dimensions divisible by 32. The same + restrictions as ``row_scaled_activation`` apply, and 4over6 on + weights and quantized primary weights are not supported. nvfp4_4over6 : {'none', 'weights', 'activations', 'all'}, default = 'none' Enable 4over6 adaptive NVFP4 block scaling for selected tensor scopes. For each selected FP4 block, quantization compares @@ -606,6 +623,7 @@ class NVFP4BlockScaling(Recipe): ) disable_2d_quantization: bool = os.getenv("NVTE_NVFP4_DISABLE_2D_QUANTIZATION", "0") == "1" row_scaled_activation: bool = os.getenv("NVTE_NVFP4_ROW_SCALED_ACTIVATION", "0") == "1" + row_scaled_weight: bool = os.getenv("NVTE_NVFP4_ROW_SCALED_WEIGHT", "0") == "1" nvfp4_4over6: str = os.getenv("NVTE_NVFP4_4OVER6", "none") nvfp4_4over6_e4m3_use_256: str = os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") nvfp4_4over6_err_mode: str = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").upper() @@ -639,6 +657,9 @@ def __post_init__(self) -> None: assert ( self.nvfp4_4over6_err_mode in _NVFP4_4OVER6_ERR_MODES ), "NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'." + assert not ( + self.row_scaled_weight and self.nvfp4_4over6 in ("weights", "all") + ), "Row-scaled NVFP4 weights (row_scaled_weight=True) do not support 4over6." # Quantization params # Note: RHT is currently only applied to column-wise usage so that @@ -651,7 +672,7 @@ def __post_init__(self) -> None: self.fp4_quant_fwd_weight = QParams( random_hadamard_transform=False, stochastic_rounding=False, - fp4_2d_quantization=not self.disable_2d_quantization, + fp4_2d_quantization=not (self.disable_2d_quantization or self.row_scaled_weight), ) self.fp4_quant_bwd_grad = QParams( random_hadamard_transform=not self.disable_rht, @@ -669,6 +690,7 @@ def _make_repr(self) -> str: f"backward_override={self.backward_override}, " f"disable_rht={self.disable_rht}, " f"row_scaled_activation={self.row_scaled_activation}, " + f"row_scaled_weight={self.row_scaled_weight}, " f"nvfp4_4over6={self.nvfp4_4over6}, " f"nvfp4_4over6_e4m3_use_256={self.nvfp4_4over6_e4m3_use_256}, " f"nvfp4_4over6_err_mode={self.nvfp4_4over6_err_mode}, " diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 8acc547c821..4fb76d6a2f1 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -85,6 +85,18 @@ def _is_nvfp4_row_scaled_tensor(tensor: torch.Tensor) -> bool: return isinstance(tensor, NVFP4TensorStorage) and tensor._row_scaled_nvfp4 +def _nvfp4_unit_amax_alias( + tensor: NVFP4TensorStorage, amax_key: str +) -> Tuple[NVFP4TensorStorage, torch.Tensor]: + """Return a per-tensor alias with unit ``amax_key`` and the replaced FP32 amax.""" + metadata = tensor.get_metadata() + amax = metadata[amax_key] + assert amax is not None and amax.dtype == torch.float32 + metadata[amax_key] = amax.new_ones(1) + metadata["row_scaled_nvfp4"] = False + return NVFP4TensorStorage(**metadata), amax + + def _nvfp4_row_scaled_gemm_inputs( A: NVFP4TensorStorage, B: NVFP4TensorStorage, @@ -93,27 +105,171 @@ def _nvfp4_row_scaled_gemm_inputs( transb: bool, ) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor, torch.Tensor]: """Return per-tensor GEMM aliases and row/column FP32 output scales.""" - A_metadata = A.get_metadata() - a_amax_key = "amax_rowwise" if transa else "amax_columnwise" - output_col_scales = A_metadata[a_amax_key] - assert output_col_scales is not None - A_metadata[a_amax_key] = output_col_scales.new_ones(1) - A_metadata["row_scaled_nvfp4"] = False - - B_metadata = B.get_metadata() - b_amax_key = "amax_columnwise" if transb else "amax_rowwise" - output_row_scales = B_metadata[b_amax_key] - assert output_row_scales is not None - B_metadata[b_amax_key] = output_row_scales.new_ones(1) - B_metadata["row_scaled_nvfp4"] = False - - assert output_row_scales.dtype == torch.float32 and output_col_scales.dtype == torch.float32 - return ( - NVFP4TensorStorage(**A_metadata), - NVFP4TensorStorage(**B_metadata), - output_row_scales.view(-1, 1), - output_col_scales.view(1, -1), + gemm_A, output_col_scales = _nvfp4_unit_amax_alias( + A, "amax_rowwise" if transa else "amax_columnwise" + ) + gemm_B, output_row_scales = _nvfp4_unit_amax_alias( + B, "amax_columnwise" if transb else "amax_rowwise" ) + return gemm_A, gemm_B, output_row_scales.view(-1, 1), output_col_scales.view(1, -1) + + +def _nvfp4_row_scaled_grouped_gemm( + A: List[NVFP4TensorStorage], + B: List[NVFP4TensorStorage], + out: List[torch.Tensor], + *, + transa: bool, + transb: bool, + m_splits: Optional[List[int]], + bias: Optional[List[torch.Tensor]], + grad: bool, + accumulate: bool, + use_split_accumulator: bool, + single_output: bool, + workspaces: List[torch.Tensor], + math_sm_count: int, +) -> bool: + """Run row-scaled NVFP4 grouped GEMM as one per-tensor grouped GEMM plus output scaling. + + Returns ``False`` without side effects when the configuration is not supported. + """ + num_gemms = len(A) + if torch.cuda.is_current_stream_capturing(): + return False + if not all(isinstance(t, NVFP4TensorStorage) for t in itertools.chain(A, B)): + return False + if any(t._scale_dtype == DType.kFloat8UE5M3 for t in itertools.chain(A, B)): + return False + a_amax_attr = "_amax_rowwise" if transa else "_amax_columnwise" + b_amax_attr = "_amax_columnwise" if transb else "_amax_rowwise" + col_scales = [getattr(a, a_amax_attr) for a in A] + row_scales = [getattr(b, b_amax_attr) for b in B] + if any(s is None or s.dtype != torch.float32 for s in itertools.chain(col_scales, row_scales)): + return False + row_scales = [s.view(-1) for s in row_scales] + col_scales = [s.view(-1) for s in col_scales] + + out_dtype = out[0].dtype + if any(o.dtype != out_dtype for o in out): + return False + if single_output: + if m_splits is None or len(m_splits) != num_gemms: + return False + col_numel = col_scales[0].numel() + if col_numel not in (1, out[0].shape[-1]): + return False + if any(s.numel() != col_numel for s in col_scales): + return False + per_row_scales = all(s.numel() == m for s, m in zip(row_scales, m_splits)) + if not per_row_scales and (col_numel == 1 or any(s.numel() != 1 for s in row_scales)): + return False + fp32_out = torch.empty_like(out[0], dtype=torch.float32) + gemm_out = [fp32_out] + else: + shape = out[0].shape + if len(shape) != 2 or any(o.shape != shape for o in out): + return False + row_numel = row_scales[0].numel() + col_numel = col_scales[0].numel() + if row_numel not in (1, shape[0]) or col_numel not in (1, shape[1]): + return False + if any(s.numel() != row_numel for s in row_scales): + return False + if any(s.numel() != col_numel for s in col_scales): + return False + fp32_out = torch.empty((num_gemms, *shape), dtype=torch.float32, device=out[0].device) + gemm_out = list(fp32_out.unbind(0)) + + ones = get_cached_ones_tensor(1, torch.float32, out[0].device) + saved = [(t, t._amax_rowwise, t._amax_columnwise, t._row_scaled_nvfp4) for t in A + B] + try: + for a in A: + setattr(a, a_amax_attr, ones) + a._row_scaled_nvfp4 = False + for b in B: + setattr(b, b_amax_attr, ones) + b._row_scaled_nvfp4 = False + empty_tensors = [_empty_tensor()] * num_gemms + tex.te_general_grouped_gemm( + A, + transa, + B, + transb, + gemm_out, + TE_DType[torch.float32], + m_splits, + empty_tensors, + TE_DType[torch.bfloat16], + single_output, + empty_tensors, + grad, + workspaces, + workspaces[0].shape[0], + False, + use_split_accumulator, + math_sm_count, + ) + finally: + for t, amax_rowwise, amax_columnwise, row_scaled in saved: + t._amax_rowwise = amax_rowwise + t._amax_columnwise = amax_columnwise + t._row_scaled_nvfp4 = row_scaled + + if single_output: + fp32_2d = fp32_out.view(-1, fp32_out.shape[-1]) + total_rows = fp32_2d.shape[0] + repeats = torch.tensor(m_splits, dtype=torch.int64, pin_memory=True).to( + fp32_out.device, non_blocking=True + ) + if col_numel == 1: + col_factor = torch.cat(col_scales).repeat_interleave(repeats, output_size=total_rows) + fp32_2d.mul_((torch.cat(row_scales) * col_factor).view(-1, 1)) + elif per_row_scales: + col_factor = torch.stack(col_scales).repeat_interleave( + repeats, dim=0, output_size=total_rows + ) + fp32_2d.mul_(torch.cat(row_scales).view(-1, 1)) + fp32_2d.mul_(col_factor) + else: + factor = torch.stack(col_scales) * torch.cat(row_scales).view(-1, 1) + fp32_2d.mul_(factor.repeat_interleave(repeats, dim=0, output_size=total_rows)) + if bias is not None: + fp32_2d.add_( + torch.stack(bias).repeat_interleave(repeats, dim=0, output_size=total_rows) + ) + if accumulate: + out[0].add_(fp32_out.to(dtype=out_dtype)) + else: + out[0].copy_(fp32_out) + return True + + row_factor = torch.stack(row_scales).view(num_gemms, -1, 1) + col_factor = torch.stack(col_scales).view(num_gemms, 1, -1) + + def _is_empty(t: NVFP4TensorStorage) -> bool: + data = t._rowwise_data if t._rowwise_data is not None else t._columnwise_data + return data is None or data.numel() == 0 + + empty_groups = [i for i in range(num_gemms) if _is_empty(A[i]) or _is_empty(B[i])] + for i in empty_groups: + row_factor[i].zero_() + col_factor[i].zero_() + if col_factor.shape[-1] == 1: + fp32_out.mul_(row_factor * col_factor) + elif row_factor.shape[1] == 1: + fp32_out.mul_(col_factor * row_factor) + else: + fp32_out.mul_(row_factor) + fp32_out.mul_(col_factor) + if bias is not None: + fp32_out.add_(torch.stack(bias).unsqueeze(1)) + converted = fp32_out if out_dtype == torch.float32 else fp32_out.to(dtype=out_dtype) + if accumulate: + torch._foreach_add_(list(out), list(converted.unbind(0))) + else: + torch._foreach_copy_(list(out), list(converted.unbind(0))) + return True _NATIVE_GEMM_INPUT_STORAGES = ( @@ -1000,6 +1156,27 @@ def general_grouped_gemm( if isinstance(quantization_params[0], DebugQuantizer): use_general_gemm_impl = True elif any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in itertools.chain(A, B)): + if ( + not gelu + and not (grad and use_bias) + and all(q is None for q in quantization_params) + and _nvfp4_row_scaled_grouped_gemm( + A, + B, + out, + transa=transa, + transb=transb, + m_splits=m_splits, + bias=bias if use_bias else None, + grad=grad, + accumulate=accumulate, + use_split_accumulator=use_split_accumulator, + single_output=single_output, + workspaces=workspaces, + math_sm_count=sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))), + ) + ): + return (out[0] if single_output else out), grad_bias if grad else bias, gelu_input use_general_gemm_impl = True elif any( isinstance(t, NVFP4TensorStorage) and t._scale_dtype == DType.kFloat8UE5M3 @@ -1141,12 +1318,55 @@ def general_grouped_gemm_for_grouped_tensor( if is_discrete_in and is_discrete_out: raise ValueError("Both A and out are discrete. This is not supported yet.") - if isinstance(A, GroupedTensorStorage) and A.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") - if isinstance(B, GroupedTensorStorage) and B.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") - if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: - raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + def _is_row_scaled(operand) -> bool: + return isinstance(operand, GroupedTensorStorage) and operand.row_scaled_nvfp4 + + def _list_has_row_scaled(operand) -> bool: + return isinstance(operand, (list, tuple)) and any( + getattr(t, "_row_scaled_nvfp4", False) for t in operand + ) + + if _list_has_row_scaled(B): + raise NotImplementedError( + "Row-scaled NVFP4 B operands passed as a list are not supported; pass a GroupedTensor." + ) + row_scaled_A_list = _list_has_row_scaled(A) + if row_scaled_A_list and not all(_is_nvfp4_row_scaled_tensor(a) for a in A): + raise NotImplementedError("Row-scaled NVFP4 A operands must all be row-scaled.") + + row_scaled = row_scaled_A_list or _is_row_scaled(A) or _is_row_scaled(B) + row_scaled_accumulate = accumulate + post_scale_rows = None # per-row outer scale -> output rows (from B) + post_scale_cols = None # per-group per-col outer scale -> output cols (from A) + if row_scaled: + if beta is not None: + raise NotImplementedError( + "Row-scaled NVFP4 grouped GEMM does not support explicit beta; use accumulate." + ) + if row_scaled_A_list: + amax_key = "amax_rowwise" if transa else "amax_columnwise" + A, col_scales = zip(*(_nvfp4_unit_amax_alias(a, amax_key) for a in A)) + A = list(A) + post_scale_cols = torch.cat([s.view(-1) for s in col_scales]) + elif _is_row_scaled(A): + A = A.copy() + if transa: + post_scale_cols = A.amax + A.amax = post_scale_cols.new_ones(A.num_tensors) + else: + post_scale_cols = A.columnwise_amax + A.columnwise_amax = post_scale_cols.new_ones(A.num_tensors) + A.row_scaled_nvfp4 = False + if _is_row_scaled(B): + B = B.copy() + if transb: + post_scale_rows = B.columnwise_amax + B.columnwise_amax = post_scale_rows.new_ones(B.num_tensors) + else: + post_scale_rows = B.amax + B.amax = post_scale_rows.new_ones(B.num_tensors) + B.row_scaled_nvfp4 = False + accumulate = False def _is_fp8_blockwise(operand) -> bool: if isinstance(operand, (list, tuple)): @@ -1213,14 +1433,38 @@ def _is_fp8_blockwise(operand) -> bool: sm_count = get_sm_count() sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) - return grouped_gemm_impl( + if not row_scaled: + return grouped_gemm_impl( + A, + transa, + B, + transb, + out, + bias, + bias_scale, + alpha, + beta, + workspace_setup, + workspace_cublas, + use_split_accumulator, + sm_count, + ) + + # Groups with an empty contraction dimension must read back as zeros. + if is_discrete_out: + scratch = [torch.zeros_like(o, dtype=torch.float32) for o in out] + else: + scratch = out.copy() + scratch.rowwise_data = torch.zeros_like(out.rowwise_data, dtype=torch.float32) + scratch.fake_dtype = torch.float32 + grouped_gemm_impl( A, transa, B, transb, - out, - bias, - bias_scale, + scratch, + None, + None, alpha, beta, workspace_setup, @@ -1228,3 +1472,65 @@ def _is_fp8_blockwise(operand) -> bool: use_split_accumulator, sm_count, ) + + def _store(dst: torch.Tensor, src: torch.Tensor) -> None: + if row_scaled_accumulate: + dst.add_(src.to(dst.dtype)) + else: + dst.copy_(src.to(dst.dtype)) + + if is_discrete_out: + row_offset = 0 + col_offset = 0 + for i in range(num_tensors): + scaled = scratch[i] + rows = scaled.shape[-2] + cols = scaled.shape[-1] + if post_scale_rows is not None: + row_scale_i = post_scale_rows[row_offset : row_offset + rows] + scaled = scaled.mul_(row_scale_i.reshape((1,) * (scaled.dim() - 2) + (rows, 1))) + if post_scale_cols is not None: + col_scale_i = post_scale_cols[col_offset : col_offset + cols] + scaled = scaled.mul_(col_scale_i.reshape((1,) * (scaled.dim() - 1) + (cols,))) + _store(out[i], scaled) + row_offset += rows + col_offset += cols + else: + if out.last_dims is not None: + raise NotImplementedError( + "Row-scaled grouped GEMM does not support a varying per-expert output width." + ) + last_dim = out.logical_shape[-1] + live_rows = out.logical_shape[0] + flat = scratch.rowwise_data.view(-1, last_dim)[:live_rows] + dst = out.rowwise_data.view(-1, last_dim)[:live_rows] + live_mask = None + if out.first_dims is None: + rows = live_rows // num_tensors + repeats = torch.full((num_tensors,), rows, dtype=torch.int64, device=flat.device) + else: + # Spare capacity rows past sum(first_dims) form a trailing group so that + # repeat_interleave can use output_size without a host sync. + used_rows = out.first_dims.sum() + repeats = torch.cat((out.first_dims, (live_rows - used_rows).view(1))) + live_mask = (torch.arange(live_rows, device=flat.device) < used_rows).view(-1, 1) + + def _per_row(per_group: torch.Tensor) -> torch.Tensor: + if live_mask is not None: + per_group = torch.cat((per_group, per_group.new_zeros(1, last_dim))) + return per_group.repeat_interleave(repeats, dim=0, output_size=live_rows) + + if post_scale_rows is not None: + flat.mul_(post_scale_rows.view(-1, 1)) + if post_scale_cols is not None: + flat.mul_(_per_row(post_scale_cols.view(num_tensors, last_dim))) + if bias is not None: + bias_rows = _per_row(bias.rowwise_data.view(num_tensors, last_dim)) + if bias_scale is not None: + bias_rows = bias_rows * bias_scale.view(-1, 1)[:live_rows] + flat.add_(bias_rows) + if live_mask is not None: + fill = 0.0 if row_scaled_accumulate else dst.to(flat.dtype) + flat = torch.where(live_mask, flat, fill) + _store(dst, flat) + return out diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 3d61dc525f8..a441a94b562 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -216,12 +216,10 @@ void group_quantize_nvfp4_row_scaled_impl(const GroupedTensorWrapper &grouped_in const NVFP4Quantizer *nvfp4_quantizer_cpp, const std::optional &noop_flag, cudaStream_t stream) { - NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht, - "Grouped row-scaled NVFP4 quantization does not support RHT."); + NVTE_CHECK(!nvfp4_quantizer_cpp->with_rht || nvfp4_quantizer_cpp->with_post_rht_amax, + "Grouped row-scaled NVFP4 RHT quantization expects post-RHT amax buffers."); NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization, "Grouped row-scaled NVFP4 quantization does not support 2D quantization."); - NVTE_CHECK(!nvfp4_quantizer_cpp->stochastic_rounding, - "Grouped row-scaled NVFP4 quantization does not support stochastic rounding."); NVTE_CHECK(!nvfp4_quantizer_cpp->with_amax_reduction, "Grouped row-scaled NVFP4 quantization does not support amax reduction."); NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, @@ -231,9 +229,38 @@ void group_quantize_nvfp4_row_scaled_impl(const GroupedTensorWrapper &grouped_in if (noop_flag.has_value()) { quant_config_cpp.set_noop_tensor(noop_flag->data()); } + + at::Tensor rng_state; + TensorWrapper te_rng_state; + if (nvfp4_quantizer_cpp->stochastic_rounding) { + constexpr size_t rng_elts_per_thread = 1024; + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); + rng_state = torch::empty({2}, at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA)); + philox_unpack(philox_args, static_cast(rng_state.data_ptr())); + te_rng_state = makeTransformerEngineTensor(rng_state); + quant_config_cpp.set_rng_state(te_rng_state.data()); + quant_config_cpp.set_stochastic_rounding(true); + } if (transformer_engine::getenv("NVTE_USE_FAST_MATH")) { quant_config_cpp.set_use_fast_math(true); } + if (nvfp4_quantizer_cpp->with_rht && + grouped_output_tensor.get_columnwise_data().data_ptr != nullptr) { + NVTE_CHECK(!noop_flag.has_value(), + "Grouped row-scaled NVFP4 RHT quantization does not support a noop flag."); + auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_amax_graph_safe( + grouped_input_tensor.data(), grouped_output_tensor.data(), 0, + nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream); + nvte_group_hadamard_transform_cast_fusion_graph_safe( + grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(), + quant_config_cpp, nullptr, stream); + }); + return; + } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor.data(), quant_config_cpp, stream); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 960e0fe46f4..60ec0319a7b 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1967,7 +1967,7 @@ bool NVFP4Quantizer::meets_alignment(const std::vector& shape, NVFP4Quantizer::ShapeAlignment NVFP4Quantizer::gemm_swizzled_scales_alignment( bool optimize_for_gemm, bool with_rht, bool with_2d_quantization, bool row_scaled_nvfp4, bool nvfp4_use_4over6, int sm_arch, std::optional disable_rht_cast_fusion) { - if (!optimize_for_gemm) { + if (!optimize_for_gemm || row_scaled_nvfp4) { return std::nullopt; } if (with_rht) { @@ -2063,7 +2063,11 @@ std::pair NVFP4Quantizer::create_tensor( // nvte_compute_amax_with_config will zero out the pointer if needed const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; if (!disable_second_level_scale) { - amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); + // No kernel writes the per-column amaxes of an empty row-scaled tensor, but the GEMM + // output scaling still reads them. + amax_columnwise = (row_scaled_nvfp4 && flat_first_dim == 0) + ? at::zeros({amax_cols}, bit32_tensor_opts) + : at::empty({amax_cols}, bit32_tensor_opts); } } @@ -2205,8 +2209,6 @@ std::pair NVFP4Quantizer::create_grouped_tenso const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 grouped quantization requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 grouped quantization does not support columnwise usage."); } const int64_t total_data_elements = total_elements / 2; @@ -2229,7 +2231,11 @@ std::pair NVFP4Quantizer::create_grouped_tenso const int64_t total_scale_elements = static_cast(product(scale_shape)); columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); if (!disable_second_level_scale) { - columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + const int64_t col_amax_elements = + row_scaled_nvfp4 + ? static_cast(num_tensors) * static_cast(logical_last_dim) + : static_cast(num_tensors); + columnwise_amax = at::empty({col_amax_elements}, float_opts); } } @@ -2529,6 +2535,7 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( out_identity_scale_inv.shape); out_identity.set_amax(out_identity_amax.data_ptr, static_cast(out_identity_amax.dtype), out_identity_amax.shape); + out_identity.set_row_scaled_nvfp4(out.get_row_scaled_nvfp4()); NVTE_SCOPED_GIL_RELEASE( { nvte_quantize_v2(input.data(), out_identity.data(), quant_config, stream); }); @@ -2560,6 +2567,7 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( out_transpose.set_amax(out_columnwise_amax.data_ptr, static_cast(out_columnwise_amax.dtype), out_columnwise_amax.shape); + out_transpose.set_row_scaled_nvfp4(out.get_row_scaled_nvfp4()); // Invoking fallback RHT kernel unfused. NVTE_SCOPED_GIL_RELEASE({ @@ -2642,19 +2650,17 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou const bool row_scaled_nvfp4 = out.get_row_scaled_nvfp4(); if (row_scaled_nvfp4) { - NVTE_CHECK(!this->with_rht, "Row-scaled NVFP4 quantization does not support RHT."); NVTE_CHECK(!this->with_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); - NVTE_CHECK(!this->stochastic_rounding, - "Row-scaled NVFP4 quantization does not support stochastic rounding."); NVTE_CHECK(!this->with_amax_reduction, - "Row-scaled NVFP4 quantization does not support amax reduction."); + "Row-scaled NVFP4 quantization does not support amax reduction " + "(e.g. sequence parallelism)."); NVTE_CHECK(cols % 16 == 0, "Row-scaled NVFP4 quantization requires last dim divisible by 16."); } // Restriction for the RHT cast fusion kernel because we are using MMA hardware for computing RHT const bool eligible_for_rht_cast_fusion = - input.dtype() == DType::kBFloat16 && + !row_scaled_nvfp4 && input.dtype() == DType::kBFloat16 && NVFP4Quantizer::is_eligible_for_rht_cast_fusion(convertShape(input.shape())); // Stochastic rounding @@ -2711,7 +2717,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // We need: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) - if (compute_amax && !this->disable_second_level_scale) { + if (compute_amax && !row_scaled_nvfp4 && !this->disable_second_level_scale) { NVTE_SCOPED_GIL_RELEASE({ nvte_hadamard_transform_amax(input.data(), out.data(), 0, this->rht_matrix_random_sign_mask_t, stream); diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 3f074df177e..0c6b0b9d7f8 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -263,7 +263,7 @@ def nvfp4_row_scaled_fwd_mxfp8_bwd_factory( * ``grad_output`` -> MXFP8 * everything else -> MXFP8 - Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid RHT, + NVFP4 is used only for fprop here, so the forward NVFP4 quantizers avoid RHT, stochastic rounding, and 2D scaling. The backward input and weight representations are quantized to MXFP8 from the dequantized NVFP4 forward representations rather than directly from the original high-precision @@ -315,7 +315,7 @@ def nvfp4_row_scaled_fwd_high_precision_bwd_factory( * ``grad_output`` -> ``IdentityQuantizer`` * everything else -> plain NVFP4 - Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, + NVFP4 is used only for fprop here, so the forward quantizers avoid RHT, stochastic rounding, and 2D scaling. The backward input and weight representations are high-precision values diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index 29e525da45d..e2b7bc55740 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -1292,23 +1292,20 @@ def qgemm( fp8_max_w = _nvfp4_effective_scale_max(qresult_w_scale_dtype, qresult_w_e4m3_max) factor = NVFP4_FP4_MAX * NVFP4_FP4_MAX * fp8_max_x * fp8_max_w - if gemm_type == gemm.GEMMType.WGRAD: - partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col - # A row-scaled operand contributes a per-output-column (N) vector - # here, so broadcast along the last axis. Selecting the axis from - # gemm_type (rather than matching numel against M) avoids the - # square-matrix ambiguity where M == N. - if partial_alpha.numel() > 1: - partial_alpha = partial_alpha.reshape(1, -1) - else: - partial_alpha = partial_alpha.squeeze(-1) + amax_x = ( + qresult_x.global_amax_col + if gemm_type == gemm.GEMMType.WGRAD + else qresult_x.global_amax_row + ) + amax_w = ( + qresult_w.global_amax_row + if gemm_type == gemm.GEMMType.FPROP + else qresult_w.global_amax_col + ) + if amax_x.numel() > 1 or amax_w.numel() > 1: + partial_alpha = amax_x.reshape(-1, 1) * amax_w.reshape(1, -1) else: - partial_alpha = qresult_x.global_amax_row * qresult_w.global_amax_row - # A row-scaled operand contributes a per-output-row (M) vector. - if partial_alpha.numel() > 1: - partial_alpha = partial_alpha.reshape(-1, 1) - else: - partial_alpha = partial_alpha.squeeze(-1) + partial_alpha = (amax_x * amax_w).squeeze(-1) alpha = torch.div(partial_alpha, factor) M, K = high_precision_x.shape diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 3d04882698a..d3df7d23003 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -69,6 +69,9 @@ def can_reconstruct_wgrad_input_from_original(quantizer) -> bool: target = getattr(quantizer, "parent_quantizer", quantizer) if target is None: target = quantizer + if getattr(target, "row_scaled_nvfp4", False): + # Row-scaled NVFP4 quantization cannot produce a columnwise-only tensor. + return False if isinstance(target, HybridQuantizer): if target.columnwise_source == "original": return True diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 9e90bcd12c9..85e5b9fed67 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1963,6 +1963,11 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: quantizer = self.quantizers["scaling_fwd"][fp8_meta_index] if quantizer is None: raise RuntimeError("Weight quantizer has not been initialized") + if getattr(quantizer, "row_scaled_nvfp4", False): + raise NotImplementedError( + "Row-scaled NVFP4 weights (row_scaled_weight=True) do not support " + "quantized primary weights." + ) quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) quantizer.internal = False # HybridQuantizer is included so its current-scaling / NVFP4 diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index a2e1c4210b7..d2f0d4231e5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -76,6 +76,7 @@ HybridQuantizer, IdentityQuantizer, MXFP8Quantizer, + NVFP4Quantizer, ) from ..quantized_tensor import ( QuantizedTensorStorage, @@ -115,9 +116,9 @@ def is_module_grouped_tensor_path_supported( 13.5+ required for FP8 per-tensor current scaling on Hopper, and 13.6+ required for FP8 block scaling on Hopper. * FP32 is unsupported by the cuBLASLt grouped GEMM. - * With MXFP8 or FP8 block scaling, every split size must be a multiple of - 128. This is checked only on the device, so other split sizes can produce - incorrect results. + * With MXFP8, FP8 block scaling, or NVFP4, every split size must be a + multiple of 128. This is checked only on the device, so other split sizes + can produce incorrect results. Runtime-only restrictions such as debug mode, CPU offloading, calibration, output quantization, and backend selection are checked separately by @@ -135,7 +136,9 @@ def is_module_grouped_tensor_path_supported( getattr(recipe, "backward_override", None), getattr(recipe, "disable_rht", None), getattr(recipe, "row_scaled_activation", None), + getattr(recipe, "row_scaled_weight", None), getattr(recipe, "fp8_format", None), + getattr(recipe, "nvfp4_4over6", None), dtype, device_index, ) @@ -147,7 +150,9 @@ def _is_module_grouped_tensor_path_supported( backward_override: Optional[str], disable_rht: Optional[bool], row_scaled_activation: Optional[bool], + row_scaled_weight: Optional[bool], fp8_format: Optional[RecipeFormat], + nvfp4_4over6: Optional[str], dtype: torch.dtype, device_index: int, ) -> bool: @@ -183,11 +188,18 @@ def _is_module_grouped_tensor_path_supported( if recipe_type.mxfp8(): return device_capability >= (10, 0) if recipe_type.nvfp4(): + if row_scaled_activation: + return ( + device_capability >= (10, 0) + and cublaslt_version >= 130400 + and fp8_format == RecipeFormat.E4M3 + and nvfp4_4over6 == "none" + ) return ( device_capability >= (10, 0) and not disable_rht - and not row_scaled_activation and fp8_format == RecipeFormat.E4M3 + and (not row_scaled_weight or cublaslt_version >= 130400) ) return False @@ -1113,7 +1125,9 @@ def _grouped_linear_graph_safe_backward( if args.fp8: grad_output_quantizer = args.grad_output_quantizers[0] grad_output_quantizer.set_usage( - rowwise=args.requires_dgrad, + rowwise=( + args.requires_dgrad or getattr(grad_output_quantizer, "row_scaled_nvfp4", False) + ), columnwise=args.weights_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True @@ -1416,6 +1430,11 @@ def _prepare_weights_for_grouped_tensor_gemm( weight_quantizer = weight_quantizers[0] if weight_quantizer is None: raise RuntimeError("Quantized grouped compute requires a weight quantizer.") + if isinstance(weight_quantizer, NVFP4Quantizer) and not weight_quantizer.with_rht: + raise NotImplementedError( + "single_grouped_weight is not supported for NVFP4 without RHT; " + "use per-expert weights." + ) weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) # forward() already applied _enable_weight_preswizzle(); preserve that decision # because not every quantizer and weight shape supports fused quantize-swizzle. @@ -1713,7 +1732,11 @@ class GroupedLinear(TransformerEngineBaseModule): The native path requires CUDA ``m_splits``. ``None`` preserves the deprecated ``NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM`` environment-variable selection for compatibility. New callers should pass a boolean explicitly. - With MXFP8 or FP8 block scaling, the native path requires every + When neither is set, row-scaled NVFP4 recipes select the native path when + it is supported and the feature dimensions and token count are multiples + of 128. Host ``m_splits`` must also be multiples of 128 and are copied to + the device outside CUDA Graph capture. + With MXFP8, FP8 block scaling, or NVFP4, the native path requires every ``m_splits`` entry to be a multiple of 128. This is checked only on the device, so other split sizes can produce incorrect results. @@ -1765,6 +1788,10 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + self._use_grouped_tensor_unset = ( + use_grouped_tensor is None + and os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM") is None + ) if use_grouped_tensor is None: use_grouped_tensor_env = os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM") if use_grouped_tensor_env is not None: @@ -2426,8 +2453,27 @@ def forward( save_original_input = False cpu_offloading = is_cpu_offload_enabled() + use_grouped_tensor = self.use_grouped_tensor + auto_grouped_tensor = ( + not use_grouped_tensor + and self._use_grouped_tensor_unset + and _recipe is not None + and _recipe.nvfp4() + and _recipe.row_scaled_activation + and self.in_features % 128 == 0 + and self.out_features % 128 == 0 + and (inp.numel() // self.in_features) % 128 == 0 + and ( + m_splits.device.type == te_device_type() + or ( + not torch.cuda.is_current_stream_capturing() + and bool((m_splits % 128 == 0).all()) + ) + ) + ) + use_grouped_tensor = use_grouped_tensor or auto_grouped_tensor use_grouped_tensor_path = False - if self.use_grouped_tensor and not ( + if use_grouped_tensor and not ( self.fp8_calibration or debug or cpu_offloading @@ -2439,7 +2485,7 @@ def forward( device=inp.device, ) if ( - self.use_grouped_tensor + use_grouped_tensor and not use_grouped_tensor_path and (self.single_grouped_weight or use_grouped_bias) ): @@ -2450,6 +2496,8 @@ def forward( " single_grouped_bias to allow the split-quantize fallback." ) if use_grouped_tensor_path: + if auto_grouped_tensor and m_splits.device.type != te_device_type(): + m_splits = m_splits.to(device=inp.device) if m_splits.device.type != te_device_type(): raise ValueError( "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 86fa740c5b6..1ccf096a393 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1111,6 +1111,16 @@ def fuser_forward( # quantization recipe. Otherwise we fall back to the legacy # ``tex.split_quantize`` + ``general_grouped_gemm`` flow. recipe = FP8GlobalStateManager.get_fp8_recipe() if with_quantized_compute else None + if ( + recipe is not None + and recipe.nvfp4() + and (recipe.row_scaled_activation or recipe.row_scaled_weight) + ): + raise NotImplementedError( + "Row-scaled NVFP4 (row_scaled_activation=True or row_scaled_weight=True) is not " + "supported by te.ops.GroupedLinear. Use transformer_engine.pytorch.GroupedLinear " + "instead." + ) use_grouped_tensor_path = is_op_fuser_grouped_tensor_path_supported( recipe, dtype, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index bba320aad17..13d586bd544 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1103,7 +1103,11 @@ def fuse_grouped_mlp_ops( if recipe.disable_rht: # Graph-safe grouped quantize is only supported with RHT return ops - if recipe.row_scaled_activation or recipe.nvfp4_4over6 != "none": + if ( + recipe.row_scaled_activation + or recipe.row_scaled_weight + or recipe.nvfp4_4over6 != "none" + ): # 4over6 doesn't used fused kernels return ops if recipe.fp8_format == RecipeFormat.UE5M3: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 355e13f0989..9dd6c9a64e2 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1807,9 +1807,12 @@ def _make(tensor_type: str) -> NVFP4Quantizer: with_2d_quantization=qparams.fp4_2d_quantization, stochastic_rounding=qparams.stochastic_rounding, row_scaled_nvfp4=( - self.mode == "forward" - and tensor_type != "weight" - and self.recipe.row_scaled_activation + self.recipe.row_scaled_weight + if tensor_type == "weight" + else ( + self.recipe.row_scaled_activation + and tensor_type in ("input", "output", "grad_output") + ) ), nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 6a48f85c6c7..334f158cd8b 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -860,13 +860,12 @@ def make_grouped_tensor( raise ValueError( "Row-scaled NVFP4 grouped quantization requires rowwise usage." ) - if columnwise_usage: - raise ValueError( - "Row-scaled NVFP4 grouped quantization does not support columnwise usage." - ) total_amax_elements = ( sum(math.prod(s[:-1]) for s in shape) if row_scaled_nvfp4 else num_tensors ) + total_columnwise_amax_elements = ( + sum(s[-1] for s in shape) if row_scaled_nvfp4 else num_tensors + ) if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8, but FP4 packs 2 values per byte) @@ -899,7 +898,9 @@ def make_grouped_tensor( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) if not disable_second_level_scale: - columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + columnwise_amax = torch.empty( + total_columnwise_amax_elements, dtype=torch.float32, device=device + ) elif compatible_recipe.float8_block_scaling(): scale_inv_dtype = DType.kFloat32 @@ -1116,6 +1117,7 @@ def split_into_quantized_tensors( columnwise_scale_inv_offsets.append(cum) self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets nvfp4_rowwise_amax_offsets = None + nvfp4_columnwise_amax_offsets = None row_scaled_nvfp4 = self.row_scaled_nvfp4 nvfp4_use_4over6 = self.nvfp4_use_4over6 nvfp4_e4m3_max = self.nvfp4_e4m3_max @@ -1123,9 +1125,13 @@ def split_into_quantized_tensors( if recipe.nvfp4() and row_scaled_nvfp4: cum = 0 nvfp4_rowwise_amax_offsets = [0] + ccum = 0 + nvfp4_columnwise_amax_offsets = [0] for i in range(self.num_tensors): cum += math.prod(self.tensor_shapes[i][:-1]) nvfp4_rowwise_amax_offsets.append(cum) + ccum += self.tensor_shapes[i][-1] + nvfp4_columnwise_amax_offsets.append(ccum) for i in range(self.num_tensors): quantizer = self.quantizer @@ -1329,7 +1335,12 @@ def split_into_quantized_tensors( amax_rowwise = self.amax[i : i + 1] if self.columnwise_amax is not None: - amax_columnwise = self.columnwise_amax[i : i + 1] + if nvfp4_columnwise_amax_offsets is not None: + camax_start = nvfp4_columnwise_amax_offsets[i] + camax_end = nvfp4_columnwise_amax_offsets[i + 1] + amax_columnwise = self.columnwise_amax[camax_start:camax_end] + else: + amax_columnwise = self.columnwise_amax[i : i + 1] if quantizer.internal: nvfp4_tensor_class = NVFP4TensorStorage