From e36b0952f0e628cecbe57091b5c2d95b36e0d471 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Mon, 20 Jul 2026 15:53:05 -0700 Subject: [PATCH 1/9] improve device grouped linear Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_grouped_linear.py | 421 ++++++++++++++++++ tests/pytorch/test_grouped_tensor.py | 141 ++++++ transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/cast.cpp | 35 +- .../pytorch/csrc/extensions/pybind.cpp | 2 +- .../pytorch/module/grouped_linear.py | 413 +++++++++++++---- 6 files changed, 917 insertions(+), 97 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 10d0dbd227..3ebda27cfc 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1504,6 +1504,20 @@ def test_fp8_grouped_gemm(shape, accumulate): ) +def _require_native_grouped_tensor_gemm(*, mxfp8: bool = False) -> None: + """Skip unless the native GroupedTensor grouped GEMM is available.""" + device_capability = torch.cuda.get_device_capability() + if not (9, 0) <= device_capability <= (11, 0): + pytest.skip("Native GroupedTensor grouped GEMM requires Hopper or Blackwell.") + if mxfp8 and device_capability < (10, 0): + pytest.skip("MXFP8 native GroupedTensor grouped GEMM requires Blackwell.") + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + pytest.skip("Native GroupedTensor grouped GEMM requires cuBLASLt 13.3+.") + if device_capability < (10, 0) and cublaslt_version < 130400: + pytest.skip("Native GroupedTensor grouped GEMM on Hopper requires cuBLASLt 13.4+.") + + @pytest.fixture(autouse=True) def _reset_fp8_state(monkeypatch): monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "0") @@ -1512,10 +1526,417 @@ def _reset_fp8_state(monkeypatch): monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) +@pytest.mark.parametrize( + "m_splits,exception", + [([256, 256], TypeError), (torch.tensor([256, 256]), ValueError)], + ids=["python-list", "cpu-tensor"], +) +def test_single_grouped_weight_rejects_host_m_splits(monkeypatch, m_splits, exception): + """A single parent parameter must never fall back to host-split per-expert GEMMs.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + grouped_linear = GroupedLinear( + 2, + 64, + 64, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + grouped_gemm_backend="grouped_tensor", + ) + x = torch.randn(512, 64, dtype=torch.bfloat16, device="cuda") + with pytest.raises(exception, match="requires.*CUDA"): + grouped_linear(x, m_splits) + + +def test_single_grouped_weight_native_bf16_matches_discrete(monkeypatch): + """Native BF16 grouped storage produces one parent wgrad and no member split calls.""" + _require_native_grouped_tensor_gemm() + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + + num_gemms = 3 + in_features = 128 + out_features = 128 + m_splits = torch.tensor([256, 512, 256], dtype=torch.int64, device="cuda") + total_tokens = int(m_splits.sum()) + weights = torch.randn( + num_gemms, + out_features, + in_features, + dtype=torch.bfloat16, + device="cuda", + ) + + discrete = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + grouped_gemm_backend="grouped_tensor", + ) + single = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + grouped_gemm_backend="grouped_tensor", + ) + with torch.no_grad(): + for idx in range(num_gemms): + getattr(discrete, f"weight{idx}").copy_(weights[idx]) + single.weight.rowwise_data.view_as(weights).copy_(weights) + + x = torch.randn(total_tokens, in_features, dtype=torch.bfloat16, device="cuda") + dy = torch.randn(total_tokens, out_features, dtype=torch.bfloat16, device="cuda") + x_discrete = x.detach().clone().requires_grad_(True) + x_single = x.detach().clone().requires_grad_(True) + + y_discrete = discrete(x_discrete, m_splits) + y_single = single(x_single, m_splits) + y_discrete.backward(dy) + y_single.backward(dy) + + tolerances = dict(rtol=1e-2, atol=5e-3) + torch.testing.assert_close(y_single.float(), y_discrete.float(), **tolerances) + torch.testing.assert_close(x_single.grad.float(), x_discrete.grad.float(), **tolerances) + discrete_wgrad = torch.stack( + [getattr(discrete, f"weight{idx}").grad for idx in range(num_gemms)] + ) + torch.testing.assert_close(single.weight.grad.float(), discrete_wgrad.float(), **tolerances) + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_single_grouped_weight_mxfp8_workspace_cache(monkeypatch): + """BF16 primary weights update one persistent MXFP8 grouped workspace per iteration.""" + _require_native_grouped_tensor_gemm(mxfp8=True) + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + mxfp8_recipe = recipe.MXFP8BlockScaling() + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + grouped_gemm_backend="grouped_tensor", + ) + + x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + + with autocast(enabled=True, recipe=mxfp8_recipe): + grouped_linear(x, m_splits, is_first_microbatch=True) + workspace = grouped_linear._fp8_workspaces["weight"] + pointers = ( + workspace.rowwise_data.data_ptr(), + workspace.columnwise_data.data_ptr(), + workspace.scale_inv.data_ptr(), + workspace.columnwise_scale_inv.data_ptr(), + ) + old_data = workspace.rowwise_data.clone() + + with torch.no_grad(): + grouped_linear.weight.rowwise_data.add_(1) + grouped_linear(x, m_splits, is_first_microbatch=False) + assert torch.equal(workspace.rowwise_data, old_data) + + grouped_linear(x, m_splits, is_first_microbatch=True) + assert pointers == ( + workspace.rowwise_data.data_ptr(), + workspace.columnwise_data.data_ptr(), + workspace.scale_inv.data_ptr(), + workspace.columnwise_scale_inv.data_ptr(), + ) + assert not torch.equal(workspace.rowwise_data, old_data) + + +@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8) +def test_single_grouped_primary_mxfp8_bypasses_weight_workspace(monkeypatch): + """An MXFP8 primary grouped parameter is already GEMM-ready and is not requantized.""" + _require_native_grouped_tensor_gemm(mxfp8=True) + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + mxfp8_recipe = recipe.MXFP8BlockScaling() + with quantized_model_init(enabled=True, recipe=mxfp8_recipe): + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + grouped_gemm_backend="grouped_tensor", + ) + + x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=mxfp8_recipe): + grouped_linear(x, m_splits, is_first_microbatch=True) + assert grouped_linear.weight.quantizer is not None + assert "weight" not in grouped_linear._fp8_workspaces + + def _clone_outputs(outputs): return [None if out is None else out.detach().clone() for out in outputs] +def _grouped_linear_weight_params(module): + if module.single_grouped_weight: + return [module.weight] + return [getattr(module, f"weight{i}") for i in range(module.num_gemms)] + + +def _grouped_linear_bias_params(module): + if not module.use_bias: + return [] + if module.single_grouped_bias: + return [module.bias] + return [getattr(module, f"bias{i}") for i in range(module.num_gemms)] + + +def _run_grouped_parameter_layout( + *, + grouped_gemm_backend, + fp8_recipe, + single_grouped_weight, + single_grouped_bias, + use_bias, + delay_wgrad_compute, + fuse_wgrad_accumulation, + x_base, + dy, + weights, + biases, + m_splits, +): + """Run one layout and return all numerically observable forward/backward results.""" + FP8GlobalStateManager.reset() + num_gemms, out_features, in_features = weights.shape + module = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=torch.bfloat16, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + delay_wgrad_compute=delay_wgrad_compute, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + grouped_gemm_backend=grouped_gemm_backend, + ) + + with torch.no_grad(): + if single_grouped_weight: + module.weight.rowwise_data.view_as(weights).copy_(weights) + else: + for i in range(num_gemms): + getattr(module, f"weight{i}").copy_(weights[i]) + if use_bias: + if single_grouped_bias: + module.bias.rowwise_data.view_as(biases).copy_(biases) + else: + for i in range(num_gemms): + getattr(module, f"bias{i}").copy_(biases[i]) + + flat_main_grad = None + initial_main_grad = None + weight_params = _grouped_linear_weight_params(module) + if fuse_wgrad_accumulation: + # MCore owns one flat FP32 grad buffer. Discrete parameters receive per-expert + # views, while a single grouped parameter receives one view over the full range. + flat_main_grad = torch.full( + (weights.numel(),), + 0.25, + dtype=torch.float32, + device="cuda", + ) + packed_main_grad = flat_main_grad.view_as(weights) + main_grad_views = ( + [packed_main_grad] + if single_grouped_weight + else [packed_main_grad[i] for i in range(num_gemms)] + ) + for param, main_grad in zip(weight_params, main_grad_views): + param.main_grad = main_grad + param.overwrite_main_grad = False + param.zero_out_wgrad = False + param.grad_added_to_main_grad = False + assert ( + param.main_grad.untyped_storage().data_ptr() + == flat_main_grad.untyped_storage().data_ptr() + ) + initial_main_grad = flat_main_grad.clone() + + x = x_base.detach().clone().requires_grad_(True) + m_splits_arg = ( + torch.tensor(m_splits, dtype=torch.int64, device="cuda") + if grouped_gemm_backend == "grouped_tensor" + else m_splits + ) + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + y = module( + x, + m_splits_arg, + is_first_microbatch=False if fuse_wgrad_accumulation else None, + ) + y.backward(dy) + + if fuse_wgrad_accumulation and delay_wgrad_compute: + torch.testing.assert_close(flat_main_grad, initial_main_grad, rtol=0, atol=0) + + # The grouped-tensor path computes dbias during the main backward even when dW is delayed. + if use_bias and grouped_gemm_backend == "grouped_tensor": + assert all(param.grad is not None for param in _grouped_linear_bias_params(module)) + + if delay_wgrad_compute: + module.backward_dw() + + if fuse_wgrad_accumulation: + assert not torch.equal(flat_main_grad, initial_main_grad) + for param in weight_params: + assert param.grad_added_to_main_grad + assert ( + param.main_grad.untyped_storage().data_ptr() + == flat_main_grad.untyped_storage().data_ptr() + ) + packed_wgrad = flat_main_grad.view_as(weights) + elif single_grouped_weight: + packed_wgrad = module.weight.grad.view_as(weights) + else: + packed_wgrad = torch.stack([param.grad for param in weight_params]) + + packed_dbias = None + if use_bias: + bias_params = _grouped_linear_bias_params(module) + if single_grouped_bias: + packed_dbias = bias_params[0].grad.view_as(biases) + else: + packed_dbias = torch.stack([param.grad for param in bias_params]) + + return { + "output": y.detach().clone(), + "dgrad": x.grad.detach().clone(), + "wgrad": packed_wgrad.detach().clone(), + "dbias": None if packed_dbias is None else packed_dbias.detach().clone(), + } + + +_GROUPED_PARAMETER_LAYOUTS = [ + pytest.param(False, False, False, id="no-bias-discrete-weight"), + pytest.param(False, True, False, id="no-bias-single-weight"), + pytest.param(True, False, False, id="bias-discrete-weight-discrete-bias"), + pytest.param(True, True, False, id="bias-single-weight-discrete-bias"), + pytest.param(True, False, True, id="bias-discrete-weight-single-bias"), + pytest.param(True, True, True, id="bias-single-weight-single-bias"), +] + + +@pytest.mark.parametrize( + "use_bias,single_grouped_weight,single_grouped_bias", _GROUPED_PARAMETER_LAYOUTS +) +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + ], +) +@pytest.mark.parametrize("delay_wgrad_compute", _ALL_BOOLEAN) +@pytest.mark.parametrize("fuse_wgrad_accumulation", _ALL_BOOLEAN) +def test_grouped_parameter_layout_matches_cpu_m_splits( + monkeypatch, + use_bias, + single_grouped_weight, + single_grouped_bias, + fp8_recipe, + delay_wgrad_compute, + fuse_wgrad_accumulation, +): + """Match CUDA m_splits and all meaningful parameter layouts against the legacy path.""" + use_mxfp8 = fp8_recipe is not None + _require_native_grouped_tensor_gemm(mxfp8=use_mxfp8) + FP8GlobalStateManager.reset() + + torch.manual_seed(1234) + # MXFP8 grouped quantization requires each expert problem to be 256-aligned. Use the same + # aligned shapes for BF16 so both precision modes exercise an identical problem layout. + num_gemms = 2 + in_features = 256 + out_features = 256 + m_splits = [256, 512] + total_tokens = sum(m_splits) + x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16) + dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16) + weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to( + torch.bfloat16 + ) + biases = None + if use_bias: + biases = (0.1 * torch.randn(num_gemms, out_features, device="cuda")).to(torch.bfloat16) + + # The CPU m_splits baseline is explicitly the legacy, discrete-parameter contract. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") + reference = _run_grouped_parameter_layout( + grouped_gemm_backend="legacy", + fp8_recipe=fp8_recipe, + single_grouped_weight=False, + single_grouped_bias=False, + use_bias=use_bias, + delay_wgrad_compute=delay_wgrad_compute, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + ) + + # Enable single parameters only for the CUDA m_splits target. The explicit layout flags + # below still decide whether this particular case uses discrete or grouped parameters. + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + result = _run_grouped_parameter_layout( + grouped_gemm_backend="grouped_tensor", + fp8_recipe=fp8_recipe, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + use_bias=use_bias, + delay_wgrad_compute=delay_wgrad_compute, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + ) + + tolerances = dict(rtol=1e-5, atol=1e-5) + + for name in ("output", "dgrad", "wgrad", "dbias"): + if reference[name] is None: + assert result[name] is None + else: + torch.testing.assert_close( + result[name].float(), + reference[name].float(), + **tolerances, + msg=f"Mismatch for {name}", + ) + + def _run_grouped_linear_path( *, enable_grouped_tensor_path: bool, diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 4dd52ee2bd..48196c0f39 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -1204,6 +1204,147 @@ def test_group_dequantize_cudagraph_capturable(self) -> None: for exp, got in zip(expected_tensors, static_tensors): assert torch.equal(got, exp) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_quantize_reuses_destination_and_honors_noop(self) -> None: + """Verify the in-place and graph-safe contracts of grouped MXFP8 quantization. + + A captured training graph records the addresses of all MXFP8 weight storage, not just + the Python ``GroupedTensor`` object. Re-quantizing an updated BF16 weight must therefore + refresh the existing FP8 payloads and scales without replacing any allocation. During + graph replay, ``noop_flag`` must additionally allow the captured quantization kernel to + execute as a no-op on microbatches that should reuse the cached weight. + """ + num_tensors = 2 + # Treat two contiguous [256, 256] matrices as one grouped BF16 source. These dimensions + # satisfy the MXFP8 block-alignment requirements in both rowwise and columnwise layouts. + shape = (num_tensors * 256, 256) + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + # Forward GEMM consumes rowwise weight storage, while backward dgrad consumes columnwise + # storage. Both representations and both scale tensors must be updated and kept stable. + quantizer.set_usage(rowwise=True, columnwise=True) + # Store the quantized output in GEMM-ready (swizzled) scale layout. This makes the test + # cover the exact destination format used by cached MXFP8 model weights. + quantizer.optimize_for_gemm = True + + # The first call owns allocation: it creates the persistent destination that subsequent + # optimizer updates and CUDA graph replays must modify in place. + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + output = tex.group_quantize(source, quantizer, num_tensors, None) + + # CUDA graphs retain these raw addresses. Object identity alone is insufficient because + # replacing one internal payload or scale allocation would leave the graph with a stale + # pointer even if ``output`` remained the same Python object. + pointers = ( + output.rowwise_data.data_ptr(), + output.columnwise_data.data_ptr(), + output.scale_inv.data_ptr(), + output.columnwise_scale_inv.data_ptr(), + ) + + # Eager in-place update: quantize new BF16 values into the previously allocated output. + # An independently allocated quantization provides the numerical reference. + updated_source = source + 1 + updated = tex.group_quantize( + updated_source, + quantizer, + num_tensors, + None, + output=output, + ) + reference = tex.group_quantize(updated_source, quantizer, num_tensors, None) + + # The API must return the caller-provided destination and preserve every captured address. + assert updated is output + assert pointers == ( + output.rowwise_data.data_ptr(), + output.columnwise_data.data_ptr(), + output.scale_inv.data_ptr(), + output.columnwise_scale_inv.data_ptr(), + ) + + # In-place quantization must still produce exactly the same FP8 bytes and scales as a + # normal allocating call, for both GEMM orientations. + assert torch.equal(output.rowwise_data, reference.rowwise_data) + assert torch.equal(output.columnwise_data, reference.columnwise_data) + assert torch.equal(output.scale_inv, reference.scale_inv) + assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv) + + # Eager no-op: a nonzero device flag tells the kernel to preserve the cached destination. + # Clone all four buffers so this checks payloads and metadata independently. + old_buffers = ( + output.rowwise_data.clone(), + output.columnwise_data.clone(), + output.scale_inv.clone(), + output.columnwise_scale_inv.clone(), + ) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + + # The source is deliberately different. If the kernel ignores ``noop_flag``, at least + # one payload or scale tensor below will change and expose the failure. + tex.group_quantize( + updated_source * 2, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + assert torch.equal(output.rowwise_data, old_buffers[0]) + assert torch.equal(output.columnwise_data, old_buffers[1]) + assert torch.equal(output.scale_inv, old_buffers[2]) + assert torch.equal(output.columnwise_scale_inv, old_buffers[3]) + + # Capture one fixed launch. ``graph_source``, ``noop``, and ``output`` keep stable device + # addresses; replay changes only their contents. This is how weight caching works when + # different microbatches replay the same CUDA graph. + graph_source = updated_source.clone() + # zero means "perform quantization"; nonzero means "leave destination untouched". + noop.zero_() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + tex.group_quantize( + graph_source, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + + # Replay with new source contents and noop=0. The captured kernel must refresh the same + # destination allocations, and the result must equal a fresh eager quantization. + graph_source.copy_(source - 1) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + reference = tex.group_quantize(graph_source, quantizer, num_tensors, None) + assert torch.equal(output.rowwise_data, reference.rowwise_data) + assert torch.equal(output.columnwise_data, reference.columnwise_data) + assert torch.equal(output.scale_inv, reference.scale_inv) + assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv) + + # Snapshot the successfully refreshed cache before testing the captured no-op branch. + old_buffers = ( + output.rowwise_data.clone(), + output.columnwise_data.clone(), + output.scale_inv.clone(), + output.columnwise_scale_inv.clone(), + ) + + # Replay the exact same captured graph with different source data but noop=1. The launch + # still occurs, which is required for CUDA graph structural consistency, but the kernel + # must not mutate any cached FP8 payload or scale buffer. + graph_source.copy_(source + 2) + noop.fill_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(output.rowwise_data, old_buffers[0]) + assert torch.equal(output.columnwise_data, old_buffers[1]) + assert torch.equal(output.scale_inv, old_buffers[2]) + assert torch.equal(output.columnwise_scale_inv, old_buffers[3]) + @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"]) @pytest.mark.parametrize("direction", ["rowwise", "columnwise"]) @pytest.mark.skipif( diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..9d6cbaa32b 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -350,7 +350,7 @@ py::object dequantize(const py::handle &input, DType otype); py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, - std::optional noop_flag); + std::optional noop_flag, const py::object &output); py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..fd11ef1496 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -276,7 +276,7 @@ void compute_grouped_fp8_current_scaling_amax_and_scale( py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, - std::optional noop_flag) { + std::optional noop_flag, const py::object &output) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -305,11 +305,34 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const GetTransformerEngineDType(tensor.scalar_type()), std::vector{static_cast(tensor.numel())}); - // Create output GroupedTensor. - auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( - num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), - py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, - logical_first_dim, logical_last_dim); + // Create a GroupedTensor or reuse an existing destination. Reusing the destination is + // required for weight caching and CUDA graph replay because captured GEMMs retain the + // original data and scale pointers. + py::object grouped_output_py; + auto grouped_output_tensor_cpp = [&]() -> GroupedTensorWrapper { + if (!output.is_none()) { + NVTE_CHECK(!first_dims.has_value() && !last_dims.has_value() && !tensor_offsets.has_value(), + "group_quantize: output reuse currently requires uniform tensor shapes."); + NVTE_CHECK(output.attr("num_tensors").cast() == num_tensors, + "group_quantize: output has a different number of tensors."); + NVTE_CHECK(output.attr("logical_shape").cast>() == logical_shape, + "group_quantize: output has an incompatible logical shape."); + py::object output_quantizer = output.attr("quantizer"); + NVTE_CHECK(!output_quantizer.is_none(), + "group_quantize: output must have quantized storage."); + NVTE_CHECK(output_quantizer.is(quantizer), + "group_quantize: output must have been created by the same quantizer."); + grouped_output_py = output; + return GroupedTensorFromPyTorchGroupedTensor(output); + } + + auto result = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, last_dims, tensor_offsets, + logical_first_dim, logical_last_dim); + grouped_output_py = std::move(result.second); + return std::move(result.first); + }(); // dispatch to scaling methods enum class GroupedQuantizationMode { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7e9d114be8..2dd1e1b971 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -208,7 +208,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), - py::arg("noop_flag") = py::none()); + py::arg("noop_flag") = py::none(), py::arg("output") = py::none()); transformer_engine::pytorch::bind_quantize_with_amax_extensions(m); m.def("group_dequantize", transformer_engine::pytorch::group_dequantize, "Dequantize group tensor", py::arg("input"), py::arg("otype")); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 32963c07ca..4702ca84f2 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -98,6 +98,7 @@ def _maybe_dequantize( @staticmethod def _is_grouped_tensor_path_supported( *, + grouped_gemm_backend: str, fp8: bool, fp8_calibration: bool, debug: bool, @@ -132,8 +133,8 @@ def _is_grouped_tensor_path_supported( Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - # 1. Filter by environment variable - if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): + # 1. Filter by the explicitly selected backend. + if grouped_gemm_backend != "grouped_tensor": return False # 2. Filter out advanced features if ( @@ -235,14 +236,101 @@ def _prepare_weights_for_grouped_tensor_gemm( weight_quantizers: List[Optional[Quantizer]], weight_workspaces: List[Optional[QuantizedTensorStorage]], *, + num_gemms: int, + single_grouped_weight: bool, with_quantized_compute: bool, columnwise_usage: bool, activation_dtype: torch.dtype, is_first_microbatch: Optional[bool], skip_fp8_weight_update: Optional[torch.Tensor], cache_weight: bool, - ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]: - """Prepare discrete weight tensors for GroupedTensor GEMM.""" + ) -> Tuple[ + Union[GroupedTensorStorage, List[torch.Tensor]], + List[Optional[QuantizedTensorStorage]], + ]: + """Prepare a grouped parameter or discrete weights for GroupedTensor GEMM.""" + if single_grouped_weight: + weight = weights[0] + if not isinstance(weight, GroupedTensorStorage): + raise TypeError( + "single_grouped_weight requires the weight parameter to be a GroupedTensor." + ) + + new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] + if weight.quantizer is not None: + if not with_quantized_compute: + raise RuntimeError( + "Quantized single grouped weights require quantized grouped GEMM compute." + ) + return weight, new_workspaces + + if not with_quantized_compute: + if weight.rowwise_data is None: + raise RuntimeError("Single grouped weight has no rowwise storage.") + if weight.rowwise_data.dtype == activation_dtype: + return weight, new_workspaces + data = weight.rowwise_data.to(dtype=activation_dtype) + return ( + GroupedTensorStorage( + shape=weight.logical_shape, + dtype=activation_dtype, + num_tensors=num_gemms, + shapes=weight.tensor_shapes, + quantizer=None, + data=data, + ), + new_workspaces, + ) + + weight_quantizer = weight_quantizers[0] + if weight_quantizer is None: + raise RuntimeError("MXFP8 grouped compute requires a weight quantizer.") + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + weight_quantizer.optimize_for_gemm = True + + workspace = weight_workspaces[0] if weight_workspaces else None + + if workspace is not None and ( + workspace.quantizer is not weight_quantizer + or (columnwise_usage and workspace.columnwise_data is None) + ): + workspace = None + + if weight.rowwise_data is None: + raise RuntimeError("Single grouped weight has no rowwise storage to quantize.") + source = weight.rowwise_data.view(weight.logical_shape) + update_workspace = is_first_microbatch is None or is_first_microbatch + if workspace is None: + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + ) + elif skip_fp8_weight_update is not None: + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + noop_flag=skip_fp8_weight_update, + output=workspace, + ) + elif update_workspace: + grouped_weight = tex.group_quantize( + source, + weight_quantizer, + num_gemms, + None, + output=workspace, + ) + else: + grouped_weight = workspace + + if cache_weight: + new_workspaces[0] = grouped_weight + return grouped_weight, new_workspaces + weights_for_gemm: List[torch.Tensor] = [] new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights) if not with_quantized_compute: @@ -282,13 +370,10 @@ def _validate_or_alloc_output( """ if buffer is None: return torch.empty((rows, cols), dtype=dtype, device=device) - if buffer.dim() != 2: - raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.") - if buffer.size(0) != rows: - raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.") - if buffer.size(1) != cols: + expected_shape = (rows, cols) + if buffer.shape != expected_shape: raise ValueError( - f"Output buffer last dim {buffer.size(1)} does not match required {cols}." + f"Output buffer shape {tuple(buffer.shape)} must match required {expected_shape}." ) if buffer.dtype != dtype: raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.") @@ -302,6 +387,38 @@ def _validate_or_alloc_output( raise ValueError("Output buffer must not require gradient.") return buffer + @staticmethod + def _prepare_bias_for_grouped_tensor_gemm( + biases: Tuple[torch.Tensor, ...], + *, + single_grouped_bias: bool, + num_gemms: int, + out_features: int, + dtype: torch.dtype, + ) -> GroupedTensorStorage: + """Prepare grouped or discrete bias storage for grouped GEMM.""" + if not single_grouped_bias: + return _GroupedLinear._make_grouped_bias( + biases, + num_gemms=num_gemms, + out_features=out_features, + dtype=dtype, + ) + + bias = biases[0] + if not isinstance(bias, GroupedTensorStorage): + raise TypeError("single_grouped_bias requires a GroupedTensor parameter.") + if bias.rowwise_data.dtype == dtype: + return bias + return GroupedTensorStorage( + shape=bias.logical_shape, + dtype=dtype, + num_tensors=num_gemms, + shapes=bias.tensor_shapes, + quantizer=None, + data=bias.rowwise_data.to(dtype=dtype), + ) + @staticmethod def _forward_grouped_tensor( ctx, @@ -323,6 +440,8 @@ def _forward_grouped_tensor( weight_workspaces: List[Optional[QuantizedTensorStorage]], cache_weight: bool, skip_fp8_weight_update: Optional[torch.Tensor], + single_grouped_weight: bool, + single_grouped_bias: bool, weights: Tuple[torch.Tensor, ...], biases: Tuple[torch.Tensor, ...], out: Optional[torch.Tensor] = None, @@ -332,7 +451,7 @@ def _forward_grouped_tensor( num_gemms = len(m_splits) device = inp.device in_features = weights[0].size(-1) - out_features = weights[0].size(0) + out_features = weights[0].size(-2) weight_requires_grad = weights[0].requires_grad split_sizes = m_splits.to(device=device) @@ -363,6 +482,8 @@ def _forward_grouped_tensor( weights, weight_quantizers, weight_workspaces, + num_gemms=num_gemms, + single_grouped_weight=single_grouped_weight, with_quantized_compute=fp8, columnwise_usage=columnwise_usage, activation_dtype=activation_dtype, @@ -389,8 +510,9 @@ def _forward_grouped_tensor( grouped_bias = None if use_bias: - grouped_bias = _GroupedLinear._make_grouped_bias( + grouped_bias = _GroupedLinear._prepare_bias_for_grouped_tensor_gemm( biases, + single_grouped_bias=single_grouped_bias, num_gemms=num_gemms, out_features=out_features, dtype=activation_dtype, @@ -421,7 +543,9 @@ def _forward_grouped_tensor( else: grouped_x = None - weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms + weights_to_save = [weights_for_gemm] if single_grouped_weight else weights_for_gemm + if not inp.requires_grad: + weights_to_save = [None] * len(weights_to_save) tensors_to_save, tensor_objects = prepare_for_saving( grouped_x, *weights_to_save, @@ -439,16 +563,18 @@ def _forward_grouped_tensor( ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_weight_quantizers = grad_weight_quantizers ctx.weights_requires_grad = weight_requires_grad + ctx.single_grouped_weight = single_grouped_weight + ctx.single_grouped_bias = single_grouped_bias if fuse_wgrad_accumulation and ctx.weights_requires_grad: ctx.origin_weight_refs = [weakref.ref(w) for w in weights] ctx.origin_weights_overwrite_main_grad = getattr( weights[0], "overwrite_main_grad", False ) if hasattr(weights[0], "__fsdp_param__"): - ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + ctx.main_grad_funcs = [weight.get_main_grad for weight in weights] else: ctx.main_grad_funcs = [ - lambda j=i: weights[j].main_grad for i in range(num_gemms) + lambda j=i: weights[j].main_grad for i in range(len(weights)) ] ctx.device = device ctx.dgrad_out = dgrad_out @@ -514,6 +640,9 @@ def forward( skip_fp8_weight_update, save_original_input, debug, + single_grouped_weight, + single_grouped_bias, + grouped_gemm_backend, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -525,8 +654,10 @@ def forward( save_original_input = False num_gemms = len(m_splits) - weights = weights_and_biases[:num_gemms] - biases = weights_and_biases[num_gemms:] + num_weight_args = 1 if single_grouped_weight else num_gemms + num_bias_args = 1 if single_grouped_bias else num_gemms + weights = weights_and_biases[:num_weight_args] + biases = weights_and_biases[num_weight_args : num_weight_args + num_bias_args] device = inp.device weight_requires_grad = weights[0].requires_grad @@ -590,7 +721,8 @@ def forward( f"weight tensor (shape={tuple(weights[0].size())})" ) - if _GroupedLinear._is_grouped_tensor_path_supported( + use_grouped_tensor_path = _GroupedLinear._is_grouped_tensor_path_supported( + grouped_gemm_backend=grouped_gemm_backend, fp8=fp8, fp8_calibration=fp8_calibration, debug=debug, @@ -600,7 +732,13 @@ def forward( activation_dtype=activation_dtype, input_quantizers=input_quantizers, output_quantizers=output_quantizers, - ): + ) + if grouped_gemm_backend == "grouped_tensor" and not use_grouped_tensor_path: + raise RuntimeError( + "The grouped_tensor backend is not supported by the active device, cuBLASLt " + "version, quantization recipe, or GroupedLinear feature configuration." + ) + if use_grouped_tensor_path: return _GroupedLinear._forward_grouped_tensor( ctx, inp=inp, @@ -620,6 +758,8 @@ def forward( weight_workspaces=weight_workspaces, cache_weight=cache_weight, skip_fp8_weight_update=skip_fp8_weight_update, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, weights=weights, biases=biases, out=out, @@ -834,12 +974,20 @@ def _backward_grouped_tensor( saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms grouped_x = saved_tensors[0] - weights = saved_tensors[1 : 1 + N] - split_sizes = saved_tensors[1 + N] - base_split_offsets = saved_tensors[2 + N] - - origin_weights = [None] * N - main_grads = [None] * N + if ctx.single_grouped_weight: + weights_for_gemm = saved_tensors[1] + weight_tensors = [weights_for_gemm] + split_sizes = saved_tensors[2] + base_split_offsets = saved_tensors[3] + else: + weight_tensors = saved_tensors[1 : 1 + N] + weights_for_gemm = weight_tensors + split_sizes = saved_tensors[1 + N] + base_split_offsets = saved_tensors[2 + N] + + num_weight_args = 1 if ctx.single_grouped_weight else N + origin_weights = [None] * num_weight_args + main_grads = [None] * num_weight_args if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: origin_weight_refs = ctx.origin_weight_refs ctx.origin_weight_refs = None @@ -891,11 +1039,16 @@ def _backward_grouped_tensor( dtype=ctx.activation_dtype, ) - grad_biases = [None] * N if ctx.use_bias: if dbias_packed is None: dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N) - grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + if ctx.single_grouped_bias: + grad_bias_args = [dbias_packed.to(dtype=ctx.activation_dtype)] + else: + grad_bias_args = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + else: + num_bias_args = 1 if ctx.single_grouped_bias else N + grad_bias_args = [None] * num_bias_args dgrad = None if ctx.requires_dgrad: @@ -904,7 +1057,7 @@ def _backward_grouped_tensor( recipe = ctx.fp8_recipe if hasattr(recipe, "fp8_gemm_dgrad"): dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - for weight in weights: + for weight in weight_tensors: if isinstance(weight, QuantizedTensorStorage): weight.update_usage(columnwise_usage=True) dgrad = _GroupedLinear._validate_or_alloc_output( @@ -923,7 +1076,7 @@ def _backward_grouped_tensor( dtype=ctx.activation_dtype, ) general_grouped_gemm_for_grouped_tensor( - weights, + weights_for_gemm, grouped_dy, grouped_dgrad, layout="NN", @@ -944,16 +1097,42 @@ def _backward_grouped_tensor( if hasattr(recipe, "fp8_gemm_wgrad"): wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator if ctx.fuse_wgrad_accumulation: - wgrad_list = main_grads + if ctx.single_grouped_weight: + main_grad = main_grads[0] + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=N, + tensor_shape=(ctx.weights_shape_0, ctx.weights_shape_1), + rowwise_data=main_grad.view(-1), + dtype=main_grad.dtype, + ) + wgrad_output = grouped_wgrad + wgrad_list = [main_grad] + else: + wgrad_output = main_grads + wgrad_list = main_grads else: - wgrad_packed = torch.empty( - N, - ctx.weights_shape_0, - ctx.weights_shape_1, - dtype=ctx.activation_dtype, - device=ctx.device, - ) - wgrad_list = [wgrad_packed[i] for i in range(N)] + if ctx.single_grouped_weight: + grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=N, + shapes=[(ctx.weights_shape_0, ctx.weights_shape_1)] * N, + quantizer=None, + device=ctx.device, + dtype=ctx.activation_dtype, + ) + wgrad_output = grouped_wgrad + wgrad_list = [ + grouped_wgrad.rowwise_data.view(N, ctx.weights_shape_0, ctx.weights_shape_1) + ] + else: + wgrad_packed = torch.empty( + N, + ctx.weights_shape_0, + ctx.weights_shape_1, + dtype=ctx.activation_dtype, + device=ctx.device, + ) + wgrad_output = [wgrad_packed[i] for i in range(N)] + wgrad_list = wgrad_output accumulate = ( accumulate_wgrad_into_param_main_grad @@ -973,9 +1152,9 @@ def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights): return None, [None] * N, None if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): - ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad) + ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], grouped_gemm_wgrad) else: - grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list) + grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_output) def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: @@ -1003,10 +1182,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) ] else: - wgrad_list = [None] * N - - if not ctx.use_bias: - grad_biases = [None] * N + wgrad_list = [None] * num_weight_args if ctx.reduce_and_update_bwd_fp8_tensors: FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) @@ -1017,7 +1193,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): None, # out None, # dgrad_out *wgrad_list, - *grad_biases, + *grad_bias_args, ) @staticmethod @@ -1379,6 +1555,13 @@ class GroupedLinear(TransformerEngineBaseModule): EXPERIMENTAL and subject to change. Gated by the ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var is not set this argument is forced to ``False`` with a warning. + grouped_gemm_backend : {"legacy", "grouped_tensor"}, default = None + Selects the per-expert legacy implementation or the native GroupedTensor + grouped GEMM implementation. The native backend requires CUDA ``m_splits``. + ``None`` preserves the deprecated + ``NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM`` environment-variable + selection for compatibility. New callers should select the backend + explicitly. Notes ----- @@ -1412,6 +1595,7 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, + grouped_gemm_backend: Optional[str] = None, ) -> None: super().__init__(name) @@ -1427,11 +1611,40 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + if grouped_gemm_backend is None: + grouped_gemm_backend_env = os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM") + if grouped_gemm_backend_env is not None: + warnings.warn( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM is deprecated and will be " + "removed in a future release. Pass grouped_gemm_backend='grouped_tensor' " + "or grouped_gemm_backend='legacy' to GroupedLinear instead.", + FutureWarning, + stacklevel=2, + ) + else: + grouped_gemm_backend_env = "0" + grouped_gemm_backend = ( + "grouped_tensor" if bool(int(grouped_gemm_backend_env)) else "legacy" + ) + if grouped_gemm_backend not in ("legacy", "grouped_tensor"): + raise ValueError( + "grouped_gemm_backend must be 'legacy' or 'grouped_tensor', " + f"got {grouped_gemm_backend!r}." + ) + self.grouped_gemm_backend = grouped_gemm_backend single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags( single_grouped_weight, single_grouped_bias ) self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias + if self.use_bias and self.single_grouped_weight and not self.single_grouped_bias: + warnings.warn( + "GroupedLinear has single_grouped_weight=True and bias=True, but " + "single_grouped_bias=False. This requires packing the per-GEMM biases on every " + "forward; enable single_grouped_bias to keep both parameters grouped.", + UserWarning, + stacklevel=2, + ) if ub_overlap_rs or ub_overlap_ag: raise ValueError("GroupedLinear doesn't support Userbuffer overlap.") self.init_method = init_method @@ -1846,21 +2059,29 @@ def forward( is_grad_enabled = torch.is_grad_enabled() num_gemms = self.num_gemms - if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = ( - FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor - ) - else: - skip_fp8_weight_update = None - if skip_fp8_weight_update is not None: - is_first_microbatch = False - # Make sure splits are in expected format + if ( + self.single_grouped_weight or self.single_grouped_bias + ) and self.grouped_gemm_backend != "grouped_tensor": + raise RuntimeError( + "single_grouped_weight and single_grouped_bias require " + "grouped_gemm_backend='grouped_tensor'; the legacy backend only supports " + "discrete parameters." + ) if not isinstance(m_splits, torch.Tensor): + if self.grouped_gemm_backend == "grouped_tensor": + raise TypeError( + "grouped_gemm_backend='grouped_tensor' requires m_splits to be a CUDA tensor." + ) # Convert list of ints to tensor for backward compatibility m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") elif m_splits.dtype != torch.int64: m_splits = m_splits.to(dtype=torch.int64) + if self.grouped_gemm_backend == "grouped_tensor" and m_splits.device.type != "cuda": + raise ValueError( + "grouped_gemm_backend='grouped_tensor' requires CUDA m_splits; CPU m_splits " + "would force the legacy per-expert path." + ) if m_splits.size() != (num_gemms,): raise ValueError( f"Shape of splits tensor ({tuple(m_splits.size())}) " @@ -1884,6 +2105,7 @@ def forward( try: weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() + use_grouped_bias = self.use_bias and self.single_grouped_bias and not self.return_bias quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() @@ -1909,11 +2131,14 @@ def forward( autograd_ctx = [None] cache_weight = is_first_microbatch is not None - weight_workspaces = ( - [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] - if cache_weight - else [None] * num_gemms - ) + if self.single_grouped_weight: + weight_workspaces = [self._fp8_workspaces.get("weight")] if cache_weight else [None] + else: + weight_workspaces = ( + [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] + if cache_weight + else [None] * num_gemms + ) non_tensor_args = ( self.apply_bias, @@ -1937,6 +2162,9 @@ def forward( skip_fp8_weight_update, self.save_original_input, debug, + self.single_grouped_weight, + use_grouped_bias, + self.grouped_gemm_backend, ) out, new_workspaces = linear_fn( *autograd_ctx, @@ -1954,7 +2182,8 @@ def forward( if ws is not None: if isinstance(ws, torch.Tensor): ws = ws.detach() - self._fp8_workspaces[f"weight{i}"] = ws + key = "weight" if self.single_grouped_weight else f"weight{i}" + self._fp8_workspaces[key] = ws finally: self.end_forward() @@ -1974,31 +2203,31 @@ def backward_dw(self): return with get_nvtx_range_context("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() - wgrad_list = tensor_list[2] + wgrad_output = tensor_list[2] weight_params = self._get_weight_tensors() if not self.fuse_wgrad_accumulation: - for i in range(self.num_gemms): - weight_params[i].grad = wgrad_list[i].to(weight_params[i].dtype) + if self.single_grouped_weight: + weight_params[0].grad = wgrad_output.rowwise_data.view( + self.num_gemms, self.out_features, self.in_features + ).to(weight_params[0].dtype) + else: + for i in range(self.num_gemms): + weight_params[i].grad = wgrad_output[i].to(weight_params[i].dtype) has_grad_biases = [ grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_ ] if self.use_bias and any(has_grad_biases): - grouped_bias = getattr(self, "bias", None) - if grouped_bias is not None: - if not all(has_grad_biases): - raise RuntimeError("Expected all grouped bias gradients to be present.") - gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype) - if grouped_bias.grad is None: - grouped_bias.grad = gstack - else: - grouped_bias.grad.add_(gstack) - else: - bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] - for i in range(self.num_gemms): - if has_grad_biases[i] and bias_params[i].grad is None: - bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) + if self.grouped_gemm_backend == "grouped_tensor": + raise RuntimeError( + "Delayed wgrad unexpectedly produced bias gradients; the grouped-tensor " + "path computes them during the main backward." + ) + bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + for i in range(self.num_gemms): + if has_grad_biases[i] and bias_params[i].grad is None: + bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) del grad_biases_ - del wgrad_list + del wgrad_output del tensor_list self._trigger_wgrad_accumulation_and_reduce_hooks() @@ -2041,10 +2270,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage """Get the weight tensors of the module.""" grouped_weight = getattr(self, "weight", None) if grouped_weight is not None: - weight_tensors = grouped_weight.quantized_tensors - if weight_tensors is None: - # TODO(ksivaman): Remove this after GEMM integration. - weight_tensors = grouped_weight.split_into_quantized_tensors() + weight_tensors = [grouped_weight] else: weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): @@ -2059,14 +2285,23 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage return weight_tensors def _get_bias_tensors(self) -> List[torch.Tensor]: - """Per-GEMM bias tensors (views into grouped storage when ``single_grouped_bias``).""" + """Get grouped compute bias or per-GEMM biases for ``return_bias=True``.""" grouped_bias = getattr(self, "bias", None) - if grouped_bias is not None: - parts = grouped_bias.quantized_tensors - if parts is None: - parts = grouped_bias.split_into_quantized_tensors() - return [p.reshape(-1) for p in parts] - return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + if grouped_bias is None: + return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + + # When bias is enabled, apply_bias and return_bias select mutually exclusive paths: + # apply_bias=True passes the grouped parent to GEMM, while return_bias=True leaves bias + # addition to the caller and preserves the public list-of-per-GEMM-biases contract. + if not self.return_bias: + # Here apply_bias=True, so the grouped-tensor GEMM consumes the parent directly. + return [grouped_bias] + + # Here apply_bias=False; expose views only for return, not for grouped GEMM compute. + parts = grouped_bias.quantized_tensors + if parts is None: + parts = grouped_bias.split_into_quantized_tensors() + return [part.reshape(-1) for part in parts] def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" From 79e428c64a694fd9ff547242358f26f4d21fed84 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Mon, 20 Jul 2026 17:09:32 -0700 Subject: [PATCH 2/9] proper support for grouped bias Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_grouped_linear.py | 73 +++++++++++++++++++ .../pytorch/module/grouped_linear.py | 39 +++++----- 2 files changed, 95 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 3ebda27cfc..8d74415d01 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -2143,6 +2143,79 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk grouped_linear.backward_dw() +def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): + """return_bias preserves the grouped parent and accumulates dbias into it. + + This mirrors how MCore applies a returned MoE bias:: + + x -> GroupedLinear (bias not applied) -> output + + + grouped bias [2, 128] + | + +-> repeat_interleave([256, 256]) + -> per-token bias [512, 128] + | + * routing probabilities + | + +-> loss.backward() -> grouped bias.grad + + Since the loss sums every output feature, each feature of expert ``i`` receives + ``sum(probs_for_expert_i)``. The identity assertion also ensures that TE returns the + registered grouped parent rather than copied or split bias tensors. + """ + _require_native_grouped_tensor_gemm() + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + dtype = torch.bfloat16 + num_gemms = 2 + in_features = 128 + out_features = 128 + m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + total_tokens = 512 + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=True, + return_bias=True, + params_dtype=dtype, + device="cuda", + single_grouped_weight=False, + single_grouped_bias=True, + grouped_gemm_backend="grouped_tensor", + ) + + x = torch.randn( + total_tokens, + in_features, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + probs = torch.rand(total_tokens, dtype=dtype, device="cuda") + output, returned_bias = grouped_linear(x, m_splits) + + assert returned_bias is grouped_linear.bias + assert returned_bias.shape == (num_gemms, out_features) + + bias_per_token = torch.repeat_interleave( + returned_bias, + m_splits, + dim=0, + output_size=total_tokens, + ) + (output + bias_per_token * probs.unsqueeze(-1)).sum().backward() + + # For token t and output feature j, the externally applied bias contributes + # bias[expert(t), j] * probs[t] to the summed loss. Therefore every bias feature for + # expert e receives sum(probs[t]) over that expert's token range. m_splits places the + # first 256 tokens on expert 0 and the remaining 256 tokens on expert 1. + expected_dbias = torch.stack( + (probs[:256].sum(), probs[256:].sum()), + ).unsqueeze(-1).expand(num_gemms, out_features) + torch.testing.assert_close(grouped_linear.bias.grad.float(), expected_dbias.float()) + + @pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"]) @pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"]) def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 4702ca84f2..fdabfaf996 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1531,7 +1531,9 @@ class GroupedLinear(TransformerEngineBaseModule): when set to ``True``, this module will not apply the additive bias itself, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when - the bias addition can be fused to subsequent operations. + the bias addition can be fused to subsequent operations. A single grouped + bias is returned as its packed ``GroupedTensor`` parameter; discrete biases + are returned as a list of per-GEMM tensors. params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters @@ -2105,7 +2107,7 @@ def forward( try: weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() - use_grouped_bias = self.use_bias and self.single_grouped_bias and not self.return_bias + use_grouped_bias = self.use_bias and self.single_grouped_bias quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() @@ -2189,6 +2191,8 @@ def forward( self.end_forward() if self.return_bias: + if use_grouped_bias: + return out, bias_tensors[0] return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors] return out @@ -2285,23 +2289,24 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage return weight_tensors def _get_bias_tensors(self) -> List[torch.Tensor]: - """Get grouped compute bias or per-GEMM biases for ``return_bias=True``.""" + """Get bias parameters in their registered grouped or per-GEMM layout. + + A single grouped bias remains one packed GroupedTensor. When return_bias=True, + an upper-level framework such as MCore must apply that packed bias accordingly; + Discrete bias parameters retain the existing list-of-per-GEMM contract. + + Example with 2 experts, 128 output features: + + single grouped bias: + GroupedTensor shape = [2, 128] -> [grouped_bias] + + discrete biases: + bias0 [128] + bias1 [128] -> [bias0, bias1] + """ grouped_bias = getattr(self, "bias", None) - if grouped_bias is None: - return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] - - # When bias is enabled, apply_bias and return_bias select mutually exclusive paths: - # apply_bias=True passes the grouped parent to GEMM, while return_bias=True leaves bias - # addition to the caller and preserves the public list-of-per-GEMM-biases contract. - if not self.return_bias: - # Here apply_bias=True, so the grouped-tensor GEMM consumes the parent directly. + if grouped_bias is not None: return [grouped_bias] - - # Here apply_bias=False; expose views only for return, not for grouped GEMM compute. - parts = grouped_bias.quantized_tensors - if parts is None: - parts = grouped_bias.split_into_quantized_tensors() - return [part.reshape(-1) for part in parts] + return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" From 879c0830142d404ed6e3dce6ebf1fde06dca41ca Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:10:35 +0000 Subject: [PATCH 3/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_grouped_linear.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 8d74415d01..1d18e74e6b 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -2210,9 +2210,13 @@ def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): # bias[expert(t), j] * probs[t] to the summed loss. Therefore every bias feature for # expert e receives sum(probs[t]) over that expert's token range. m_splits places the # first 256 tokens on expert 0 and the remaining 256 tokens on expert 1. - expected_dbias = torch.stack( - (probs[:256].sum(), probs[256:].sum()), - ).unsqueeze(-1).expand(num_gemms, out_features) + expected_dbias = ( + torch.stack( + (probs[:256].sum(), probs[256:].sum()), + ) + .unsqueeze(-1) + .expand(num_gemms, out_features) + ) torch.testing.assert_close(grouped_linear.bias.grad.float(), expected_dbias.float()) From e11c0a7e96cade92ac7a8025c20c2fd0ca3634af Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 23 Jul 2026 00:52:04 -0700 Subject: [PATCH 4/9] fixes and improvements Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_grouped_linear.py | 2 +- tests/pytorch/test_grouped_mlp.py | 20 +++++++ .../cast/mxfp8/group_quantize_mxfp8.cuh | 19 +++++-- .../pytorch/module/grouped_linear.py | 15 +++-- .../pytorch/ops/basic/grouped_linear.py | 56 ++++++++++--------- .../pytorch/ops/fused/grouped_mlp.py | 3 +- 6 files changed, 76 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 1d18e74e6b..feda63a56d 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1923,7 +1923,7 @@ def test_grouped_parameter_layout_matches_cpu_m_splits( m_splits=m_splits, ) - tolerances = dict(rtol=1e-5, atol=1e-5) + tolerances = dict(rtol=1e-2, atol=5e-3) for name in ("output", "dgrad", "wgrad", "dbias"): if reference[name] is None: diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 76550db5f8..8d03f5652b 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -232,6 +232,26 @@ def make_reference_and_test_tensors( class TestGroupedLinearOp: """Tests for advanced features with grouped linear basic op""" + def test_single_grouped_bias_uses_registered_packed_storage(self, monkeypatch) -> None: + """The grouped bias compute view must alias the registered trainable parent.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + op = te.ops.GroupedLinear( + 2, + 128, + 128, + bias=True, + device="cuda", + dtype=torch.bfloat16, + single_grouped_bias=True, + ) + + bias_packed = op._get_packed_bias_tensor(torch.bfloat16) + + assert bias_packed.shape == (op.num_groups, op.out_features) + assert bias_packed.untyped_storage().data_ptr() == op.bias.rowwise_data.data_ptr() + assert op.bias.requires_grad + assert dict(op.named_parameters())["bias"] is op.bias + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 14832573d7..016936c622 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -572,9 +572,18 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); const size_t tensor_base = current_block.tensor_base; - const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1) - ? static_cast(offsets_ptr[tensor_id]) - : tensor_base; + size_t tensor_base_for_scales = tensor_base; + size_t tensor_rows_for_scales = rows; + if constexpr (WITH_GEMM_SWIZZLED_SCALES && + SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + // The payload is one tall tensor, but GEMM scales are swizzled independently per member. + // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly. + tensor_rows_for_scales = first_logical_dim / num_tensors; + tensor_base_for_scales = tensor_id * tensor_rows_for_scales * cols; + } else if constexpr (WITH_GEMM_SWIZZLED_SCALES && + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + tensor_base_for_scales = static_cast(offsets_ptr[tensor_id]); + } const size_t block_id_Y = current_block.block_id_Y; const size_t block_id_X = current_block.block_id_X; const size_t block_offset_Y = current_block.block_offset_Y; @@ -680,8 +689,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel process_colwise_stage( buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, - scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr, - sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); + scale_stride_colwise, tensor_base_for_scales, tensor_rows_for_scales, cols, sIn_ptr, + sActIn_ptr, sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); } if constexpr (ROWWISE_SCALING) { diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index fdabfaf996..52f947a2ef 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -408,15 +408,20 @@ def _prepare_bias_for_grouped_tensor_gemm( bias = biases[0] if not isinstance(bias, GroupedTensorStorage): raise TypeError("single_grouped_bias requires a GroupedTensor parameter.") - if bias.rowwise_data.dtype == dtype: - return bias + bias_data = bias.rowwise_data + if bias_data.dtype != dtype: + bias_data = bias_data.to(dtype=dtype) + + # The parameter exposes a packed [num_gemms, out_features] tensor, but its grouped + # members are 1D vectors. The grouped bias-add kernel consumes those same bytes as + # num_gemms row matrices with shape [1, out_features]. return GroupedTensorStorage( - shape=bias.logical_shape, + shape=(num_gemms, out_features), dtype=dtype, num_tensors=num_gemms, - shapes=bias.tensor_shapes, + shapes=[(1, out_features)] * num_gemms, quantizer=None, - data=bias.rowwise_data.to(dtype=dtype), + data=bias_data.reshape(-1), ) @staticmethod diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index af3f8b5930..23fbe90092 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -294,17 +294,28 @@ def backward_dw(self) -> None: w.grad = grad_weights[group_idx].to(w.dtype) self._trigger_wgrad_accumulation_and_reduce_hooks() - def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: - """Retrieve per-group bias tensors in the given dtype.""" + def _get_discrete_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: + """Retrieve discrete per-group bias parameters in the given dtype.""" if self.single_grouped_bias: - bias_parts = self.bias.quantized_tensors - if bias_parts is None: - bias_parts = self.bias.split_into_quantized_tensors() - return [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts] + raise RuntimeError( + "Discrete bias tensors were requested for a single grouped bias parameter." + ) return [ maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(self.num_groups) ] + def _get_packed_bias_tensor(self, dtype: torch.dtype) -> torch.Tensor: + """Return all per-group biases as one dense tensor with shape [num_groups, out_features]. + + A single grouped bias is already stored in this layout, so return a view of the + registered parent parameter instead of splitting it into members and stacking it again. + Discrete biases require a stack because they are independent parameters. + """ + if self.single_grouped_bias: + bias_data = self.bias.rowwise_data.view(self.num_groups, self.out_features) + return bias_data if bias_data.dtype == dtype else bias_data.to(dtype=dtype) + return torch.stack(self._get_discrete_bias_tensors(dtype), dim=0) + def num_quantizers(self, mode: str) -> int: if mode == "forward": return 2 * self.num_groups @@ -945,16 +956,7 @@ def _get_grouped_bias_for_gemm( return None num_groups = self.num_groups - if self.single_grouped_bias: - # Already a contiguous (num_groups * out_features) buffer. - bias_data = self.bias.rowwise_data - if bias_data.dtype != dtype: - bias_data = bias_data.to(dtype=dtype) - else: - bias_list = [ - maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) - ] - bias_data = torch.stack(bias_list, dim=0).contiguous() + bias_data = self._get_packed_bias_tensor(dtype) return GroupedTensorStorage( shape=(num_groups, self.out_features), @@ -1033,6 +1035,13 @@ def fuser_forward( dtype=dtype, single_grouped_weight=self.single_grouped_weight, ) + if (self.single_grouped_weight or self.single_grouped_bias) and not use_grouped_tensor_path: + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor GroupedLinear path, " + "which is unavailable for the current device, dtype, or quantization recipe. " + "Disable single_grouped_weight/single_grouped_bias or use a supported grouped-" + "tensor configuration." + ) if use_grouped_tensor_path: out, tensors_to_save = self._fuser_forward_grouped_tensor( @@ -1178,16 +1187,11 @@ def _fuser_forward_split_quantize( # Need CPU split sizes for split_quantize / general_grouped_gemm. split_sizes_int = [int(s) for s in split_sizes.tolist()] - # Extract params - if self.single_grouped_weight: - weights = self.weight.quantized_tensors - if weights is None: - weights = self.weight.split_into_quantized_tensors() - else: - weights = self._forward_weight_list() # materialized when distributed + # Single grouped parameters are rejected before entering this legacy path. + weights = self._forward_weight_list() # materialized when distributed bs = None if has_bias: - bs = self._get_bias_tensors(dtype) + bs = self._get_discrete_bias_tensors(dtype) ws = self._get_discrete_weights_for_gemm( weights, @@ -1476,7 +1480,7 @@ def _fuser_backward_split_quantize( offsets = torch.zeros(num_groups + 1, dtype=torch.int64, device=device) offsets[1:] = split_sizes.cumsum(0) if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype)) + bias_packed = self._get_packed_bias_tensor(ctx.dtype) scales_f32 = scales.to(dtype=torch.float32) dbias_packed, grad_scales = compute_grouped_dbias_dscales( dy_2d, @@ -1696,7 +1700,7 @@ def _fuser_backward_grouped_tensor( grad_scales: Optional[torch.Tensor] = None if has_bias: if self._scale_bias: - bias_packed = torch.stack(self._get_bias_tensors(dtype)) + bias_packed = self._get_packed_bias_tensor(dtype) scales_f32 = scales.to(dtype=torch.float32) dbias_packed, grad_scales = compute_grouped_dbias_dscales( dy_2d, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 83954a9b3d..61e4fadad6 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1930,8 +1930,7 @@ def fuser_backward( fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc2_bias_grad_packed: Optional[torch.Tensor] = None if scale_bias: - fc2_biases = fc2_op._get_bias_tensors(dtype) - bias_packed = torch.stack(fc2_biases) + bias_packed = fc2_op._get_packed_bias_tensor(dtype) fc2_dbias_packed_result, grad_scales = compute_grouped_dbias_dscales( fc2_dy, scales_f32, From fe3f47f1e9c73f6e5c840036aa7f69e0ade7bec8 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 23 Jul 2026 02:12:12 -0700 Subject: [PATCH 5/9] fix unit test failure Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_grouped_linear.py | 11 ++++++- .../pytorch/module/grouped_linear.py | 33 ++++++++++++------- 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index feda63a56d..2cbb8aef0f 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1635,6 +1635,7 @@ def test_single_grouped_weight_mxfp8_workspace_cache(monkeypatch): with autocast(enabled=True, recipe=mxfp8_recipe): grouped_linear(x, m_splits, is_first_microbatch=True) workspace = grouped_linear._fp8_workspaces["weight"] + assert isinstance(workspace, GroupedTensor) pointers = ( workspace.rowwise_data.data_ptr(), workspace.columnwise_data.data_ptr(), @@ -2217,7 +2218,15 @@ def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): .unsqueeze(-1) .expand(num_gemms, out_features) ) - torch.testing.assert_close(grouped_linear.bias.grad.float(), expected_dbias.float()) + assert grouped_linear.bias.grad is not None + # repeat_interleave backward and the standalone sums above use different BF16 reduction + # orders, so compare the independently derived result with BF16-appropriate tolerance. + torch.testing.assert_close( + grouped_linear.bias.grad.float(), + expected_dbias.float(), + rtol=1e-2, + atol=5e-3, + ) @pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"]) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 52f947a2ef..a95634b3ae 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -301,12 +301,19 @@ def _prepare_weights_for_grouped_tensor_gemm( source = weight.rowwise_data.view(weight.logical_shape) update_workspace = is_first_microbatch is None or is_first_microbatch if workspace is None: + if cache_weight: + # Match quantize_weight(): persistent workspaces must be Tensor subclasses + # so autograd can save them without decomposing their storage metadata. + saved_internal = weight_quantizer.internal + weight_quantizer.internal = False grouped_weight = tex.group_quantize( source, weight_quantizer, num_gemms, None, ) + if cache_weight: + weight_quantizer.internal = saved_internal elif skip_fp8_weight_update is not None: grouped_weight = tex.group_quantize( source, @@ -738,12 +745,23 @@ def forward( input_quantizers=input_quantizers, output_quantizers=output_quantizers, ) - if grouped_gemm_backend == "grouped_tensor" and not use_grouped_tensor_path: + if ( + grouped_gemm_backend == "grouped_tensor" + and not use_grouped_tensor_path + and (single_grouped_weight or single_grouped_bias) + ): raise RuntimeError( - "The grouped_tensor backend is not supported by the active device, cuBLASLt " - "version, quantization recipe, or GroupedLinear feature configuration." + "The active device, cuBLASLt version, quantization recipe, or GroupedLinear " + "feature configuration does not support the grouped_tensor backend. Legacy " + "fallback is unavailable for single grouped parameters; use discrete parameters " + "or a supported grouped_tensor configuration." ) if use_grouped_tensor_path: + if m_splits.device.type != "cuda": + raise ValueError( + "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 " + "tensor, or select grouped_gemm_backend='legacy'." + ) return _GroupedLinear._forward_grouped_tensor( ctx, inp=inp, @@ -2076,19 +2094,10 @@ def forward( "discrete parameters." ) if not isinstance(m_splits, torch.Tensor): - if self.grouped_gemm_backend == "grouped_tensor": - raise TypeError( - "grouped_gemm_backend='grouped_tensor' requires m_splits to be a CUDA tensor." - ) # Convert list of ints to tensor for backward compatibility m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") elif m_splits.dtype != torch.int64: m_splits = m_splits.to(dtype=torch.int64) - if self.grouped_gemm_backend == "grouped_tensor" and m_splits.device.type != "cuda": - raise ValueError( - "grouped_gemm_backend='grouped_tensor' requires CUDA m_splits; CPU m_splits " - "would force the legacy per-expert path." - ) if m_splits.size() != (num_gemms,): raise ValueError( f"Shape of splits tensor ({tuple(m_splits.size())}) " From 4e989f14fbf38e60ee9eecdf99be59e8d9246083 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 23 Jul 2026 02:58:09 -0700 Subject: [PATCH 6/9] fix group linear UT Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_grouped_linear.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 2cbb8aef0f..c166b21000 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1528,7 +1528,7 @@ def _reset_fp8_state(monkeypatch): @pytest.mark.parametrize( "m_splits,exception", - [([256, 256], TypeError), (torch.tensor([256, 256]), ValueError)], + [([256, 256], ValueError), (torch.tensor([256, 256]), ValueError)], ids=["python-list", "cpu-tensor"], ) def test_single_grouped_weight_rejects_host_m_splits(monkeypatch, m_splits, exception): @@ -2205,7 +2205,8 @@ def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): dim=0, output_size=total_tokens, ) - (output + bias_per_token * probs.unsqueeze(-1)).sum().backward() + biased_output = (output + bias_per_token * probs.reshape(-1, 1)).to(output.dtype) + biased_output.sum().backward() # For token t and output feature j, the externally applied bias contributes # bias[expert(t), j] * probs[t] to the summed loss. Therefore every bias feature for @@ -2213,18 +2214,18 @@ def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch): # first 256 tokens on expert 0 and the remaining 256 tokens on expert 1. expected_dbias = ( torch.stack( - (probs[:256].sum(), probs[256:].sum()), + (probs[:256].float().sum(), probs[256:].float().sum()), ) .unsqueeze(-1) .expand(num_gemms, out_features) ) assert grouped_linear.bias.grad is not None - # repeat_interleave backward and the standalone sums above use different BF16 reduction - # orders, so compare the independently derived result with BF16-appropriate tolerance. + # Megatron applies the packed bias in BF16. Compare its gradient against an independently + # accumulated FP32 reference with a tolerance appropriate for a 256-element BF16 reduction. torch.testing.assert_close( grouped_linear.bias.grad.float(), - expected_dbias.float(), - rtol=1e-2, + expected_dbias, + rtol=5e-2, atol=5e-3, ) From 1169d6e1af1eda3d55bb91029bc8c4587f836293 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 23 Jul 2026 02:59:15 -0700 Subject: [PATCH 7/9] fix group mlp UT by disabling wrong fallback Signed-off-by: Zhongbo Zhu --- tests/pytorch/test_grouped_mlp.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 8d03f5652b..b85406b461 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -23,6 +23,7 @@ OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, ) +from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch import ( QuantizedTensor, Float8CurrentScalingQuantizer, @@ -129,6 +130,28 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +def grouped_tensor_path_supported( + *, + quantized_compute: bool, + quantization: Optional[str], + dtype: torch.dtype, + single_grouped_weight: bool, +) -> bool: + """Mirror GroupedLinear's native grouped-tensor backend selection.""" + compute_capability = get_device_compute_capability() + if not (9, 0) <= compute_capability <= (11, 0): + return False + if not quantized_compute: + return dtype in (torch.bfloat16, torch.float16) + if quantization == "fp8_current_scaling": + return compute_capability >= (10, 0) or tex.get_cublasLt_version() >= 130500 + if quantization == "mxfp8": + return compute_capability >= (10, 0) + if quantization == "nvfp4_rht": + return compute_capability >= (10, 0) and not single_grouped_weight + return False + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], @@ -312,6 +335,14 @@ def test_grouped_linear( if single_grouped_bias and not bias: pytest.skip("single_grouped_bias requires bias=True") + if (single_grouped_weight or single_grouped_bias) and not grouped_tensor_path_supported( + quantized_compute=quantized_compute, + quantization=quantization, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + ): + # Single grouped parameters intentionally have no split-quantize fallback. + pytest.skip("Single grouped parameters require the native grouped-tensor path") if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"): pytest.skip( "single_grouped_weight does not support FP8 delayed scaling " From 66dfa2c60876fecdcc7aaaa08040e2d9efe2d2a0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 10:27:53 +0000 Subject: [PATCH 8/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 016936c622..ffe0d608ac 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -574,8 +574,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t tensor_base = current_block.tensor_base; size_t tensor_base_for_scales = tensor_base; size_t tensor_rows_for_scales = rows; - if constexpr (WITH_GEMM_SWIZZLED_SCALES && - SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + if constexpr (WITH_GEMM_SWIZZLED_SCALES && SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { // The payload is one tall tensor, but GEMM scales are swizzled independently per member. // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly. tensor_rows_for_scales = first_logical_dim / num_tensors; From 3d37fd3b18af7227ab87bb08f291bd1a4d62c6fc Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu Date: Thu, 23 Jul 2026 03:55:16 -0700 Subject: [PATCH 9/9] lint Signed-off-by: Zhongbo Zhu --- .../common/cast/mxfp8/group_quantize_mxfp8.cuh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index ffe0d608ac..bc00c5824f 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -579,8 +579,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly. tensor_rows_for_scales = first_logical_dim / num_tensors; tensor_base_for_scales = tensor_id * tensor_rows_for_scales * cols; - } else if constexpr (WITH_GEMM_SWIZZLED_SCALES && - SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + } + if constexpr (WITH_GEMM_SWIZZLED_SCALES && + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { tensor_base_for_scales = static_cast(offsets_ptr[tensor_id]); } const size_t block_id_Y = current_block.block_id_Y;