diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 10d0dbd227..2f22c0143f 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -14,6 +14,8 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, Float8Quantizer, Fp8Padding, Fp8Unpadding, @@ -30,6 +32,10 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) +from transformer_engine.pytorch.module.grouped_linear import ( + _GroupedLinear, + is_module_grouped_tensor_path_supported, +) from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, get_align_size_for_quantization, @@ -1154,46 +1160,124 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_b torch.testing.assert_close(o, o_ref, **tols) +@pytest.mark.parametrize("use_bias_scale", [False, True]) +def test_grouped_gemm_grouped_tensor_zero_work_bias(use_bias_scale) -> None: + """A grouped bias operation is a no-op when every group has zero rows. + + Zero-sized CUDA tensors may legally have a null data pointer. Exercise both bias entry points + so neither the ordinary nor scaled path mistakes that pointer for missing output storage. + This BF16 case runs on both Hopper and Blackwell when grouped cuBLASLt GEMM is available. + """ + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor GEMM is unavailable.") + + num_groups = 4 + in_features = 256 + out_features = 256 + m_sizes = [0] * num_groups + dtype = torch.bfloat16 + device = torch.device("cuda") + + weights = [ + torch.randn(out_features, in_features, dtype=dtype, device=device) + for _ in range(num_groups) + ] + biases = [torch.randn(1, out_features, dtype=dtype, device=device) for _ in range(num_groups)] + grouped_weights = _make_grouped_tensor_uniform( + num_groups, out_features, in_features, device, dtype + ) + grouped_input = _make_grouped_tensor_from_splits(m_sizes, in_features, device, dtype) + grouped_output = _make_grouped_tensor_from_splits(m_sizes, out_features, device, dtype) + grouped_bias = _make_grouped_tensor_uniform(num_groups, 1, out_features, device, dtype) + _pack_grouped_tensor(grouped_weights, weights) + _pack_grouped_tensor(grouped_bias, biases) + + bias_scale = torch.empty(0, dtype=torch.float32, device=device) if use_bias_scale else None + general_grouped_gemm_for_grouped_tensor( + grouped_weights, + grouped_input, + grouped_output, + layout="TN", + bias=grouped_bias, + bias_scale=bias_scale, + ) + torch.cuda.synchronize() + + assert grouped_output.rowwise_data.numel() == 0 + + @pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) @pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("quant_type", ["bf16", "mxfp8"]) +@pytest.mark.parametrize( + "quant_type", ["bf16", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"] +) def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -> None: """Grouped GEMM with all-zero split sizes (zero total work). For wgrad (NT layout) the output should be zero when not accumulating, or unchanged when accumulating with beta=1. """ - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") if not is_bf16_available(): pytest.skip("bfloat16 is required for grouped GEMM test.") - if quant_type == "mxfp8" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) z = 4 k, n = 256, 256 dtype = torch.bfloat16 device = torch.device("cuda") - use_mxfp8 = quant_type == "mxfp8" + + test_recipe = { + "bf16": None, + "fp8_current_scaling": recipe.Float8CurrentScaling(), + "mxfp8": recipe.MXFP8BlockScaling(), + "fp8_block_scaling": recipe.Float8BlockScaling(), + }[quant_type] + if not is_module_grouped_tensor_path_supported( + test_recipe, + dtype, + ): + pytest.skip(f"{quant_type} grouped-tensor GEMM is unavailable") transa = layout[0] == "T" transb = layout[1] == "T" zero_first_dims = torch.zeros(z, dtype=torch.int64, device=device) + def _make_quantizer(fp8_dtype, rowwise, columnwise): + if quant_type == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=device) + quantizer.set_usage(rowwise=rowwise, columnwise=columnwise) + elif quant_type == "mxfp8": + quantizer = MXFP8Quantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + ) + elif quant_type == "fp8_block_scaling": + quantizer = Float8BlockQuantizer( + fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + else: + raise ValueError(f"Unsupported quantized zero-work test type {quant_type}") + quantizer.optimize_for_gemm = True + return quantizer + def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): """Create a GroupedTensor with non-zero logical_shape but zero first_dims.""" buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device) - if use_mxfp8: + if test_recipe is not None: if is_a: rowwise, columnwise = transa, not transa else: rowwise, columnwise = not transb, transb - quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=rowwise, - columnwise=columnwise, - ) - quantizer.optimize_for_gemm = True + fp8_dtype = tex.DType.kFloat8E4M3 + quantizer = _make_quantizer(fp8_dtype, rowwise, columnwise) return tex.group_quantize(buf, quantizer, z, zero_first_dims) return GroupedTensor.make_grouped_tensor( num_tensors=z, @@ -1208,12 +1292,18 @@ def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): if layout in ("TN", "NN"): weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] - if use_mxfp8: - grouped_A = _make_grouped_tensor_quantized_mxfp8( - weight_tensors, + if test_recipe is not None: + grouped_weight = torch.cat(weight_tensors, dim=0) + weight_quantizer = _make_quantizer( + tex.DType.kFloat8E4M3, rowwise=transa, columnwise=not transa, - device=device, + ) + grouped_A = tex.group_quantize( + grouped_weight, + weight_quantizer, + z, + torch.full((z,), n, dtype=torch.int64, device=device), ) else: grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype) @@ -1506,16 +1596,530 @@ def test_fp8_grouped_gemm(shape, accumulate): @pytest.fixture(autouse=True) def _reset_fp8_state(monkeypatch): - monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "0") + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) yield FP8GlobalStateManager.reset() monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) +@pytest.mark.parametrize( + "m_splits,exception", + [([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): + """A single parent parameter must never fall back to host-split per-expert GEMMs.""" + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("Native GroupedTensor GEMM is unavailable on this system.") + + 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, + use_grouped_tensor=True, + ) + x = torch.randn(512, 64, dtype=torch.bfloat16, device="cuda") + with pytest.raises(exception, match="requires.*CUDA"): + grouped_linear(x, m_splits) + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + pytest.param(None, id="bf16"), + pytest.param( + recipe.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + ], +) +def test_single_grouped_weight_matches_discrete_grouped_tensor_path(monkeypatch, fp8_recipe): + """Match single and discrete weights while both use CUDA m_splits and grouped GEMM.""" + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + torch.bfloat16, + ): + pytest.skip("Recipe is not supported with a single grouped weight on this system.") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + + num_gemms = 3 + in_features = 256 + out_features = 256 + 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", + use_grouped_tensor=True, + ) + single = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + 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) + + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + 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.""" + mxfp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported( + mxfp8_recipe, + torch.bfloat16, + ): + pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + 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"] + assert isinstance(workspace, GroupedTensor) + 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) +@pytest.mark.parametrize("fp8_recipe", [recipe.MXFP8BlockScaling()], ids=recipe_id) +def test_single_grouped_weight_with_disabled_weight_preswizzle(monkeypatch, fp8_recipe): + """Grouped weight preparation preserves a disabled preswizzle decision.""" + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + with quantized_model_init(enabled=True, recipe=fp8_recipe): + grouped_linear = GroupedLinear( + 2, + 256, + 256, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + single_grouped_weight=True, + use_grouped_tensor=True, + ) + + weight_quantizer = grouped_linear.weight.quantizer + assert weight_quantizer is not None + preswizzle = grouped_linear._enable_weight_preswizzle( + weight_quantizer, + grouped_linear.weight, + ) + assert preswizzle is False + weight_quantizer.optimize_for_gemm = preswizzle + grouped_weight, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( + (grouped_linear.weight,), + [weight_quantizer], + [None], + num_gemms=grouped_linear.num_gemms, + single_grouped_weight=True, + with_quantized_compute=True, + columnwise_usage=True, + activation_dtype=torch.bfloat16, + is_first_microbatch=True, + skip_fp8_weight_update=None, + cache_weight=True, + ) + + assert weight_quantizer.optimize_for_gemm is False + assert len(new_workspaces) == 1 + assert new_workspaces[0] is None + assert grouped_weight is grouped_linear.weight + assert hasattr(grouped_weight, "_with_gemm_swizzled_scales") + assert grouped_weight._with_gemm_swizzled_scales is False + + +@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.""" + mxfp8_recipe = recipe.MXFP8BlockScaling() + if not is_module_grouped_tensor_path_supported( + mxfp8_recipe, + torch.bfloat16, + ): + pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.") + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + FP8GlobalStateManager.reset() + 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, + use_grouped_tensor=True, + ) + + 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( + *, + use_grouped_tensor, + 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, + use_grouped_tensor=use_grouped_tensor, + ) + + 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 use_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 use_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.Float8CurrentScaling(), + marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8), + id="fp8-current-scaling", + ), + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + id="mxfp8", + ), + pytest.param( + recipe.Float8BlockScaling(), + marks=pytest.mark.skipif( + not _fp8_block_scaling_available, + reason=_reason_for_no_fp8_block_scaling, + ), + id="fp8-block-scaling", + ), + ], +) +@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.""" + if not is_module_grouped_tensor_path_supported( + fp8_recipe, + torch.bfloat16, + ): + pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.") + FP8GlobalStateManager.reset() + + torch.manual_seed(1234) + # MXFP8 and the grouped FP8 recipes require aligned expert problems. Use the same + # 256-aligned shapes for every recipe so all precision modes exercise one 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( + use_grouped_tensor=False, + 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( + use_grouped_tensor=True, + 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-2, atol=5e-3) + + 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, @@ -1609,33 +2213,16 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( fp8_recipe, bias, fp8_model_params, delay_wgrad_compute, monkeypatch ): use_fp8 = fp8_recipe is not None - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." - ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor - # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. - is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() - if is_block_scaling and not (9, 0) <= device_capability < (10, 0): - pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") - if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): - pytest.skip( - "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500: - pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + 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.") if fp8_model_params and not use_fp8: pytest.skip("fp8_model_params requires FP8") - dtype = torch.bfloat16 num_gemms = 3 in_features = 128 out_features = 128 @@ -1692,8 +2279,11 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy( def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monkeypatch): - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("GroupedTensor grouped GEMM path requires SM100+") + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") @@ -1722,6 +2312,102 @@ 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. + """ + if not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") + 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, + use_grouped_tensor=True, + ) + + x = torch.randn( + total_tokens, + in_features, + dtype=dtype, + device="cuda", + requires_grad=True, + ) + + probs = torch.cat( + ( + torch.full((256,), 0.25, dtype=dtype, device="cuda"), + torch.full((256,), 0.5, 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, + ) + 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 + # 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].float().sum(), probs[256:].float().sum()), + ) + .unsqueeze(-1) + .expand(num_gemms, out_features) + ) + assert grouped_linear.bias.grad is not None + # 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, + rtol=5e-2, + atol=5e-3, + ) + + @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): @@ -1730,18 +2416,11 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc Checks that a supplied buffer is written in place (bit-for-bit vs internal allocation) and, on the fused path with a padded input, that only the valid rows are touched. """ - if use_fused_path: - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell" - " (SM10x and SM110)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + if use_fused_path and not is_module_grouped_tensor_path_supported( + None, + torch.bfloat16, + ): + pytest.skip("BF16 GroupedTensor path is unavailable on this system.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0") give_out = supply in ("out", "both") @@ -1819,19 +2498,16 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4) -def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): - """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct. - - Graph-safe grouped quantization currently requires RHT, so requesting NVFP4 with - ``disable_rht=True`` while the fused grouped-tensor path is enabled falls back to the - legacy path internally. We verify the output and gradients against a reference built from - per-GEMM ``te.Linear`` modules that share the same weights and use the same NVFP4 recipe; - the grouped GEMM should match the loop of single GEMMs. +def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(): + """Non-RHT NVFP4 falls back to split-quantize for discrete parameters. + + Graph-safe grouped quantization currently requires RHT. Discrete parameters have a valid + split-quantize fallback, so enabling the grouped-tensor path is a preference rather than a + hard requirement for this parameter layout. """ if torch.cuda.get_device_capability() < (10, 0): pytest.skip("NVFP4 GroupedTensor grouped GEMM path requires SM100+") - monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") FP8GlobalStateManager.reset() dtype = torch.bfloat16 @@ -1854,7 +2530,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): disable_stochastic_rounding=True, ) - # Grouped path: fused path enabled, but non-RHT NVFP4 falls back to legacy internally. grouped_linear = GroupedLinear( num_gemms, in_features, @@ -1862,6 +2537,7 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): bias=False, params_dtype=dtype, device="cuda", + use_grouped_tensor=True, ) with torch.no_grad(): for i in range(num_gemms): @@ -1872,7 +2548,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): y = grouped_linear(x, m_splits) y.backward(dy) - # Reference: one te.Linear per GEMM sharing the same weights and NVFP4 recipe. ref_linears = torch.nn.ModuleList( [ Linear(in_features, out_features, bias=False, params_dtype=dtype, device="cuda") @@ -1890,7 +2565,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): ) y_ref.backward(dy) - # cuBLAS grouped GEMM should match the loop of single GEMMs bit-for-bit. tols = dict(rtol=0, atol=0) torch.testing.assert_close(y.float(), y_ref.float(), **tols) torch.testing.assert_close(x.grad.float(), x_ref.grad.float(), **tols) @@ -1931,33 +2605,16 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch): def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): """Fused GroupedTensor GEMM path should be CUDA graph capturable.""" use_fp8 = fp8_recipe is not None - device_capability = torch.cuda.get_device_capability() - if not (9, 0) <= device_capability <= (11, 0): - pytest.skip( - "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)." - ) - # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor - # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only. - is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling() - is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling() - if is_block_scaling and not (9, 0) <= device_capability < (10, 0): - pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).") - if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0): - pytest.skip( - "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)." - ) - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500: - pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + 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.") monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") FP8GlobalStateManager.reset() - dtype = torch.bfloat16 device = "cuda" num_gemms = 3 in_features = 128 diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index cadfd6c1f6..1698362584 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -22,6 +22,7 @@ from transformer_engine.pytorch.ops.basic.grouped_linear import ( OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, + is_op_fuser_grouped_tensor_path_supported, ) from transformer_engine.pytorch import ( QuantizedTensor, @@ -259,6 +260,26 @@ def test_meta_single_grouped_weight_with_delayed_wgrad(self, monkeypatch) -> Non assert grouped_weight.skip_backward_post_hook + 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) @@ -319,6 +340,16 @@ def test_grouped_linear( if single_grouped_bias and not bias: pytest.skip("single_grouped_bias requires bias=True") + recipe = make_recipe(quantization) + compute_recipe = recipe if quantized_compute else None + if ( + single_grouped_weight or single_grouped_bias + ) and not is_op_fuser_grouped_tensor_path_supported( + compute_recipe, + dtype, + ): + # 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 " @@ -374,7 +405,6 @@ def test_grouped_linear( y_ref.backward(dy_ref) # Construct fusible operation - recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): op = te.ops.GroupedLinear( group_size, @@ -505,22 +535,6 @@ def test_grouped_linear_cuda_graph_safe( "single_grouped_weight/single_grouped_bias requires" " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" ) - device_capability = torch.cuda.get_device_capability() - if device_capability < (9, 0): - pytest.skip( - "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)" - ) - # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path, - # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+). - requires_blackwell = quantization is not None and quantization != "fp8_current_scaling" - if requires_blackwell and device_capability < (10, 0): - pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") - # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+. - cublaslt_version = tex.get_cublasLt_version() - if device_capability < (10, 0) and cublaslt_version < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if cublaslt_version < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") if quantization is None and quantized_weight: pytest.skip("quantized_weight requires a quantization recipe") if ( @@ -538,6 +552,13 @@ def test_grouped_linear_cuda_graph_safe( "only discrete weights (single_grouped_weight=False) are supported." ) + recipe = make_recipe(quantization) + if not is_op_fuser_grouped_tensor_path_supported( + recipe, + dtype, + ): + pytest.skip("Configuration falls back to the non-CUDA-graph-safe path") + single_grouped_bias = bias and single_grouped_weight # Split sizes (statically pinned for graph capture) @@ -550,7 +571,6 @@ def test_grouped_linear_cuda_graph_safe( in_shape = (num_active_tokens + token_padding, in_features) out_shape = (in_shape[0], out_features) - recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): op = te.ops.GroupedLinear( group_size, diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 4dd52ee2bd..ac8c493290 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -577,7 +577,7 @@ def test_quantize_grouped_mxfp8(self, shape_case: str, output_dbias: bool) -> No if output_dbias: expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) - assert torch.allclose(dbias, expected_dbias) + torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3) @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @@ -805,6 +805,152 @@ def _run(inp): assert torch.equal(static_output.rowwise_data, expected.rowwise_data) assert torch.equal(static_output.scale_inv, expected.scale_inv) + @pytest.mark.parametrize( + "quantization", + [ + pytest.param( + "fp8_current_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + ), + pytest.param( + "fp8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_grouped_available, + reason=reason_for_no_fp8_block_scaling_grouped, + ), + ), + ], + ) + def test_group_quantize_reuses_destination_and_honors_noop_fp8(self, quantization: str) -> None: + """Check pointer-stable weight-cache updates and no-op CUDA graph replays.""" + num_tensors = 2 + shape = (num_tensors * 256, 256) + if quantization == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + force_pow_2_scales=False, + amax_epsilon=0.0, + ) + quantizer.set_usage(rowwise=True, columnwise=True) + else: + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=False, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + output = tex.group_quantize(source, quantizer, num_tensors, None) + output_buffers = tuple( + buffer + for buffer in ( + output.rowwise_data, + output.columnwise_data, + output.scale_inv, + output.columnwise_scale_inv, + output.amax, + output.columnwise_amax, + output.scale, + ) + if buffer is not None + ) + assert output_buffers + pointers = tuple(buffer.data_ptr() for buffer in output_buffers) + + # Re-quantization must update the existing allocations because a captured GEMM keeps + # their raw addresses rather than looking them up again through the Python object. + 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) + assert updated is output + assert pointers == tuple(buffer.data_ptr() for buffer in output_buffers) + reference_buffers = tuple( + buffer + for buffer in ( + reference.rowwise_data, + reference.columnwise_data, + reference.scale_inv, + reference.columnwise_scale_inv, + reference.amax, + reference.columnwise_amax, + reference.scale, + ) + if buffer is not None + ) + assert len(output_buffers) == len(reference_buffers) + for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True): + assert torch.equal(output_buffer, reference_buffer) + + # A nonzero device flag must preserve every cached payload and metadata buffer. + old_buffers = tuple(buffer.clone() for buffer in output_buffers) + noop = torch.ones(1, dtype=torch.float32, device="cuda") + tex.group_quantize( + updated_source * 2, + quantizer, + num_tensors, + None, + noop_flag=noop, + output=output, + ) + for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True): + assert torch.equal(output_buffer, old_buffer) + + # Capture one pointer-stable update. Replays change only source contents and the device + # no-op value, matching TE's CUDA graph microbatch weight-cache protocol. + graph_source = updated_source.clone() + 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, + ) + + graph_source.copy_(source - 1) + noop.zero_() + graph.replay() + torch.cuda.synchronize() + reference = tex.group_quantize(graph_source, quantizer, num_tensors, None) + reference_buffers = tuple( + buffer + for buffer in ( + reference.rowwise_data, + reference.columnwise_data, + reference.scale_inv, + reference.columnwise_scale_inv, + reference.amax, + reference.columnwise_amax, + reference.scale, + ) + if buffer is not None + ) + assert len(output_buffers) == len(reference_buffers) + for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True): + assert torch.equal(output_buffer, reference_buffer) + + old_buffers = tuple(buffer.clone() for buffer in output_buffers) + graph_source.copy_(source + 2) + noop.fill_(1) + graph.replay() + torch.cuda.synchronize() + for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True): + assert torch.equal(output_buffer, old_buffer) + @pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"]) @pytest.mark.parametrize( "shape_case", @@ -1121,7 +1267,7 @@ def test_quantize_grouped_fp8_blockwise( if output_dbias: expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) - assert torch.allclose(dbias, expected_dbias) + torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3) @pytest.mark.parametrize( "shape", @@ -1204,6 +1350,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/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index dff94cef9d..8c3ae7c47b 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -37,6 +37,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data +from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported from utils import ModelConfig, recipe_id, skip_unsupported_backward_override # Only run FP8 tests on supported devices. @@ -86,11 +87,11 @@ def is_fp8_supported(config: ModelConfig): def nvfp4_vanilla(): - 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() - return nvfp4_recipe + return recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) def nvfp4_row_scaled(): @@ -123,7 +124,7 @@ def nvfp4_4over6(): if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: - fp8_recipes.append(nvfp4_vanilla()) # TODO: fix check for this + fp8_recipes.append(nvfp4_vanilla()) fp8_recipes.append(nvfp4_4over6()) if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) @@ -563,6 +564,7 @@ def test_sanity_linear_with_zero_tokens( out = te_linear(inp_hidden_states) loss = out.sum() loss.backward() + torch.cuda.synchronize() assert out.shape == (num_tokens, ffn_hidden_size) @@ -603,17 +605,34 @@ def test_sanity_grouped_linear( if fp8_recipe is not None: fp8_recipe = copy.deepcopy(fp8_recipe) fp8_recipe.backward_override = backward_override + if single_param and not is_module_grouped_tensor_path_supported( + fp8_recipe, + dtype, + ): + pytest.skip("Single grouped parameters require the native grouped-tensor path") + if single_param: + # Single grouped parameters intentionally have no split-quantize fallback, so this + # test must satisfy the native grouped kernels' shape contract. MCore pads each + # expert's token count to 256; TE requires at least 128-row alignment. Weight K must + # be 64-aligned. + tokens_per_nonempty_expert = bs * config.max_seqlen_q + if tokens_per_nonempty_expert % 128 != 0: + pytest.skip("Single grouped parameters require each nonempty m_split to be 128-aligned") + k_alignment = 64 + if config.hidden_size % k_alignment != 0: + pytest.skip(f"Single grouped parameters require GEMM K to be {k_alignment}-aligned") if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") if fp8_recipe.nvfp4(): - if not getattr(fp8_recipe, "row_scaled_activation", False): - pytest.skip("NVFP4 not supported for grouped linear") - if single_param: - pytest.skip("Row-scaled NVFP4 does not support GroupedTensor grouped linear") - if dtype == torch.float16: - pytest.skip("FP16 output for NVFP4 not supported") + if dtype != torch.bfloat16: + pytest.skip("NVFP4 GroupedLinear requires BF16") + if single_param and not fp8_model_params: + pytest.skip( + "NVFP4 single grouped BF16 primary weights require unsupported non-RHT " + "grouped weight quantization; enable quantized model initialization" + ) use_fp8 = fp8_recipe is not None with quantized_model_init(enabled=use_fp8 and fp8_model_params, recipe=fp8_recipe): @@ -625,6 +644,7 @@ def test_sanity_grouped_linear( params_dtype=dtype, single_grouped_weight=single_param, single_grouped_bias=single_param, + use_grouped_tensor=single_param, ).cuda() # Verify grouped linear exposes a single grouped weight parameter(and bias when applicable). @@ -644,11 +664,25 @@ def test_sanity_grouped_linear( m_splits[-1] = 0 elif empty_split == "middle": m_splits[num_gemms // 2] = 0 + if single_param: + m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cuda") + + if NVTE_TEST_NVINSPECT_ENABLED and single_param: + # DebugQuantizer operates on per-GEMM tensors, while single grouped parameters + # intentionally have no split-quantize fallback. + with pytest.raises( + RuntimeError, + match="TE debug features do not support single grouped parameters", + ): + with autocast(enabled=use_fp8, recipe=fp8_recipe): + te_grouped_linear(inp_hidden_states, m_splits) + return with autocast(enabled=use_fp8, recipe=fp8_recipe): out = te_grouped_linear(inp_hidden_states, m_splits) loss = out.sum() loss.backward() + torch.cuda.synchronize() assert out.shape == (num_tokens, ffn_hidden_size) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 14832573d7..bc00c5824f 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; + } + 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/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 3997e5249d..3663106d50 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1324,7 +1324,8 @@ __global__ void setup_grouped_gemm_kernel( char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta, TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_bits_per_elem, size_t b_bits_per_elem, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, - float *beta_ptr, bool use_per_group_alpha_beta, + float *beta_ptr, bool use_per_group_alpha_beta, bool a_is_discrete, bool c_is_discrete, + bool d_is_discrete, // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise, @@ -1339,12 +1340,9 @@ __global__ void setup_grouped_gemm_kernel( if (idx >= num_tensors) return; // Get dimensions for this tensor (from array or uniform value) - const bool has_a_multi_tensor = (a_base == nullptr); - const bool has_c_multi_tensor = (c_base == nullptr); - const bool has_d_multi_tensor = (d_base == nullptr); int64_t a_first = 0; int64_t a_last = 0; - if (!has_a_multi_tensor) { + if (!a_is_discrete) { a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first; a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last; } @@ -1354,24 +1352,25 @@ __global__ void setup_grouped_gemm_kernel( int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last; // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform) - int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx); + int64_t a_offset = a_is_discrete ? 0 : compute_grouped_tensor_offset(A_meta, idx); int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx); int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx); int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx); // Compute data pointers - A_ptrs[idx] = has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] - : (a_base + (a_offset * a_bits_per_elem) / 8); - B_ptrs[idx] = b_base + (b_offset * b_bits_per_elem) / 8; - C_ptrs[idx] = - has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size); - D_ptrs[idx] = - has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size); + A_ptrs[idx] = a_is_discrete + ? a_multi_tensor_args.data_ptrs[idx] + : (a_base == nullptr ? nullptr : a_base + (a_offset * a_bits_per_elem) / 8); + B_ptrs[idx] = b_base == nullptr ? nullptr : b_base + (b_offset * b_bits_per_elem) / 8; + C_ptrs[idx] = c_is_discrete ? c_multi_tensor_args.data_ptrs[idx] + : (c_base == nullptr ? nullptr : c_base + c_offset * c_elem_size); + D_ptrs[idx] = d_is_discrete ? d_multi_tensor_args.data_ptrs[idx] + : (d_base == nullptr ? nullptr : d_base + d_offset * d_elem_size); // Compute storage dimensions for cuBLAS matrix layouts from logical dims. // Rowwise and MXFP8 columnwise storage use logical row-major layout, viewed as // column-major rows=last, cols=first. Transposed columnwise storage reverses this. - if (has_a_multi_tensor) { + if (a_is_discrete) { a_rows[idx] = a_multi_tensor_args.rows[idx]; a_cols[idx] = a_multi_tensor_args.cols[idx]; } else if (a_storage_transposed) { @@ -1388,7 +1387,7 @@ __global__ void setup_grouped_gemm_kernel( b_rows[idx] = static_cast(b_last); b_cols[idx] = static_cast(b_first); } - if (has_d_multi_tensor) { + if (d_is_discrete) { d_rows[idx] = d_multi_tensor_args.rows[idx]; d_cols[idx] = d_multi_tensor_args.cols[idx]; } else { @@ -1403,7 +1402,7 @@ __global__ void setup_grouped_gemm_kernel( if (use_per_group_alpha_beta) { float a_amax_val = 0.0f; bool has_a_amax = false; - if (has_a_multi_tensor) { + if (a_is_discrete) { auto *a_amax_p = static_cast(a_multi_tensor_args.amax_ptrs[idx]); if (a_amax_p != nullptr) { a_amax_val = *a_amax_p; @@ -1471,10 +1470,12 @@ __global__ void setup_grouped_gemm_kernel( } }; - if (a_scale_base) { + if (a_is_discrete) { + a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; + } else if (a_scale_base) { fill_scale_ptr(a_scale_inv_ptrs, a_scale_base, A_meta, a_rowwise, a_scaling_mode); } else { - a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; + a_scale_inv_ptrs[idx] = nullptr; } if (b_scale_base) { fill_scale_ptr(b_scale_inv_ptrs, b_scale_base, B_meta, b_rowwise, b_scaling_mode); @@ -1491,7 +1492,7 @@ inline void launch_grouped_gemm_setup( const transformer_engine::Tensor *beta_tensor, bool use_per_group_alpha_beta, size_t num_tensors, cudaStream_t stream, const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list, - const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype, + const NVTETensor *D_list, bool a_is_discrete, char *a_base, transformer_engine::DType c_dtype, transformer_engine::DType d_dtype) { // Use logical shape info from selection; storage transposes are tracked separately. TensorShapeInfo A_meta = A_sel.logical_tensor_shape; @@ -1499,31 +1500,31 @@ inline void launch_grouped_gemm_setup( TensorShapeInfo C_meta{}; TensorShapeInfo D_meta{}; - const bool has_d_multi_tensor = (D_list != nullptr); - const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor; + const bool d_is_discrete = (D_list != nullptr); + const bool c_is_discrete = (C_list != nullptr) || d_is_discrete; MultiTensorGroupGemmOutputArgs c_multi_tensor_args{}; MultiTensorGroupGemmOutputArgs d_multi_tensor_args{}; - if (has_d_multi_tensor) { + if (d_is_discrete) { d_multi_tensor_args = build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, "D"); } if (C_list != nullptr) { c_multi_tensor_args = build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, "C"); - } else if (has_d_multi_tensor) { + } else if (d_is_discrete) { c_multi_tensor_args = d_multi_tensor_args; } char *c_base = nullptr; char *d_base = nullptr; - if (!has_c_multi_tensor) { + if (!c_is_discrete) { NVTE_CHECK(C != nullptr && D != nullptr, "Grouped GEMM: C/D grouped tensors are required when no C list is provided"); C_meta = TensorShapeInfo::create_shape_info_for_C(C, D); c_base = static_cast(C->data.dptr); } - if (!has_d_multi_tensor) { + if (!d_is_discrete) { NVTE_CHECK(D != nullptr, "Grouped GEMM: D grouped tensor is required when no D list is provided"); D_meta = TensorShapeInfo::from_tensor(D); @@ -1544,8 +1545,8 @@ inline void launch_grouped_gemm_setup( const bool b_rowwise = B_sel.rowwise; // NVFP4 alpha needs A's amax from either A_sel.amax (grouped) or amax_ptrs (discrete). - const bool a_has_amax = (A_sel.amax != nullptr) || - (A_sel.dptr == nullptr && a_multi_tensor_args.amax_ptrs[0] != nullptr); + const bool a_has_amax = + a_is_discrete ? a_multi_tensor_args.amax_ptrs[0] != nullptr : A_sel.amax != nullptr; const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && a_has_amax && (B_sel.amax != nullptr); @@ -1554,11 +1555,12 @@ inline void launch_grouped_gemm_setup( ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs, A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_bits_per_elem, b_bits_per_elem, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr), - static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, - reinterpret_cast(A_sel.scale_inv), reinterpret_cast(B_sel.scale_inv), - a_rowwise, b_rowwise, A_sel.storage_transposed, B_sel.storage_transposed, A_sel.scaling_mode, - B_sel.scaling_mode, num_tensors, a_multi_tensor_args, c_multi_tensor_args, - d_multi_tensor_args, A_sel.amax ? static_cast(A_sel.amax) : nullptr, + static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, a_is_discrete, + c_is_discrete, d_is_discrete, reinterpret_cast(A_sel.scale_inv), + reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.storage_transposed, + B_sel.storage_transposed, A_sel.scaling_mode, B_sel.scaling_mode, num_tensors, + a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args, + A_sel.amax ? static_cast(A_sel.amax) : nullptr, B_sel.amax ? static_cast(B_sel.amax) : nullptr, needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr); @@ -1633,8 +1635,8 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, stream, - a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, - inputC->dtype(), outputD->dtype()); + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, + /*a_is_discrete=*/false, A_sel.dptr, inputC->dtype(), outputD->dtype()); // Compute average dimensions for heuristics // K dimension: if transa, K is A's last dim; if not, K is A's first dim @@ -1788,8 +1790,8 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, stream, - a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, - inputC->dtype(), outputD->dtype()); + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, + /*a_is_discrete=*/true, nullptr, inputC->dtype(), outputD->dtype()); GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1873,8 +1875,8 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr, alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, - stream, a_multi_tensor_args, C_list, D_list, A_sel.dptr, d_dtype, - d_dtype); + stream, a_multi_tensor_args, C_list, D_list, + /*a_is_discrete=*/false, A_sel.dptr, d_dtype, d_dtype); GroupedGemmConfig gemm_config; gemm_config.use_split_accumulator = config_.use_split_accumulator; @@ -1909,8 +1911,6 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, NVTE_CHECK(outputD->num_tensors >= 1, api_name, ": number of tensors must be at least 1"); NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name, ": output and bias must have the same number of tensors"); - NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); - NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name, ": output and bias must have matching dtypes"); NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name, @@ -1921,15 +1921,23 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name, ": output and bias last dims must match"); + const int num_tensors = static_cast(outputD->num_tensors); + NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, + " tensors, got ", num_tensors); + const int total_rows = static_cast(outputD->logical_shape.data[0]); + // A valid zero-sized CUDA allocation may have a null data pointer. + if (total_rows == 0) { + return; + } + + NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); + NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); + const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD); const DType dtype = outputD->dtype(); constexpr int kThreads = 128; - const int num_tensors = static_cast(outputD->num_tensors); - NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, - " tensors, got ", num_tensors); - const int total_rows = static_cast(outputD->logical_shape.data[0]); const int n = static_cast(outputD->get_common_last_dim()); const size_t elem_size = typeToSize(dtype); @@ -1997,8 +2005,6 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); const Tensor *scale_tensor = convertNVTETensorCheck(scale); - NVTE_CHECK(scale_tensor->data.dptr != nullptr, - "Grouped scaled bias add: scale tensor must not be null"); NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32, "Grouped scaled bias add: scale must be float32"); NVTE_CHECK(scale_tensor->data.shape.size() == 1, @@ -2007,6 +2013,10 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou const size_t total_rows = static_cast(outputD->logical_shape.data[0]); NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, "Grouped scaled bias add: scale size (", scale_tensor->data.shape[0], ") must equal total rows (", total_rows, ")"); + if (total_rows > 0) { + NVTE_CHECK(scale_tensor->data.dptr != nullptr, + "Grouped scaled bias add: scale tensor must not be null"); + } const float *scale_ptr = static_cast(scale_tensor->data.dptr); launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream); diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 95f02c64f0..ce015bd1b9 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..400735bb72 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 { @@ -380,6 +403,9 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const QuantizationConfigWrapper quant_config_cpp; quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales); quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon); + if (noop_flag_cpp.has_value()) { + quant_config_cpp.set_noop_tensor(noop_flag_cpp->data()); + } NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), quant_config_cpp, at::cuda::getCurrentCUDAStream()); 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/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index ddb85808a5..7a3b284aca 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -219,12 +219,23 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { } auto ret = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); + auto get_initialized_storage_shape = [](const at::Tensor &data) { + auto shape = getTensorShape(data); + if (data.numel() == 0) { + // PyTorch may use a null pointer for a valid zero-sized allocation. TE Common reserves + // {nullptr, {0}} for uninitialized storage, so use an orientation-neutral 2D empty shape + // to distinguish explicitly provided rowwise or columnwise storage from missing storage. + shape = {0, 0}; + } + return shape; + }; + // Rowwise data if (!tensor.attr("rowwise_data").is_none()) { const auto &data = tensor.attr("rowwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; - ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + ret.set_rowwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data)); } else if (quantizer_dtype != DType::kNumTypes) { ret.set_rowwise_data(nullptr, quantizer_dtype, std::vector{0}); } @@ -234,7 +245,7 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { const auto &data = tensor.attr("columnwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; - ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + ret.set_columnwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data)); } else if (quantizer_dtype != DType::kNumTypes) { ret.set_columnwise_data(nullptr, quantizer_dtype, std::vector{0}); } diff --git a/transformer_engine/pytorch/module/__init__.py b/transformer_engine/pytorch/module/__init__.py index 3cf15efc11..9e0227070e 100644 --- a/transformer_engine/pytorch/module/__init__.py +++ b/transformer_engine/pytorch/module/__init__.py @@ -1,14 +1,14 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Module level PyTorch APIs""" -from .layernorm_linear import LayerNormLinear -from .linear import Linear -from .grouped_linear import GroupedLinear -from .layernorm_mlp import LayerNormMLP -from .layernorm import LayerNorm -from .rmsnorm import RMSNorm -from .fp8_padding import Fp8Padding -from .fp8_unpadding import Fp8Unpadding -from .base import initialize_ub, destroy_ub, UserBufferQuantizationMode +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Module level PyTorch APIs""" +from .layernorm_linear import LayerNormLinear +from .linear import Linear +from .grouped_linear import GroupedLinear, is_module_grouped_tensor_path_supported +from .layernorm_mlp import LayerNormMLP +from .layernorm import LayerNorm +from .rmsnorm import RMSNorm +from .fp8_padding import Fp8Padding +from .fp8_unpadding import Fp8Unpadding +from .base import initialize_ub, destroy_ub, UserBufferQuantizationMode diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 0b51ca6a1e..23c21c89fb 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -69,7 +69,6 @@ Float8CurrentScalingQuantizer, Float8Quantizer, MXFP8Quantizer, - NVFP4Quantizer, ) from ..quantized_tensor import ( QuantizedTensorStorage, @@ -80,7 +79,73 @@ from ...debug.pytorch.debug_quantization import DebugQuantizer from ...debug.pytorch.debug_state import TEDebugState -__all__ = ["GroupedLinear"] +__all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"] + + +def is_module_grouped_tensor_path_supported( + recipe: Optional[Recipe], + dtype: torch.dtype, +) -> bool: + """Whether the module grouped-tensor path supports this recipe and dtype. + + The grouped-tensor path dispatches to ``general_grouped_gemm_for_grouped_tensor`` + and does not inspect split values because they may reside in a CUDA tensor. + Inspecting them on the host would add synchronization and break CUDA Graph safety. + + Supported Compute Capability (CC) and precisions: + + * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 + block scaling. + * Blackwell (CC 10.x and 11.0): BF16/FP16, FP8 per-tensor current scaling, + MXFP8, and NVFP4 with RHT. + * Custom recipes are unsupported because they may assign different + quantizers to input, weight, and grad-output roles. This predicate + currently supports only built-in recipes with known uniform layouts. + * FP8 delayed scaling is unsupported because the required grouped + quantization kernels are unavailable. + * FP8 block scaling is unsupported by this path on Blackwell because it + does not implement the legacy path's MXFP8-broadcast emulation. + * Grouped GEMM requires cuBLASLt 13.3+, with 13.4+ required on Hopper, + 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. + + Runtime-only restrictions such as debug mode, CPU offloading, calibration, + output quantization, and backend selection are checked separately by + ``GroupedLinear``. + """ + if dtype not in (torch.bfloat16, torch.float16): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if recipe is None: + return True + if recipe.custom(): + return False + if recipe.backward_override is not None: + return False + if recipe.float8_current_scaling(): + return device_capability >= (10, 0) or cublaslt_version >= 130500 + if recipe.float8_block_scaling(): + # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8. + return device_capability < (10, 0) and cublaslt_version >= 130600 + if recipe.mxfp8(): + return device_capability >= (10, 0) + if recipe.nvfp4(): + return ( + device_capability >= (10, 0) + and not recipe.disable_rht + and not recipe.row_scaled_activation + ) + return False class _GroupedLinear(torch.autograd.Function): @@ -98,97 +163,6 @@ def _maybe_dequantize( return tensor.dequantize(dtype=dtype) return cast_if_needed(tensor, dtype) - @staticmethod - def _is_grouped_tensor_path_supported( - *, - fp8: bool, - fp8_calibration: bool, - debug: bool, - cpu_offloading: bool, - backward_override: Optional[str], - save_original_input: bool, - activation_dtype: torch.dtype, - input_quantizers: List[Optional[Quantizer]], - output_quantizers: List[Optional[Quantizer]], - ) -> bool: - """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata. - - There are no checks whether split sizes are supported. Splits - may be in a CUDA tensor, so checking would hurt performance - and be incompatible with CUDA Graphs. - - Supported Compute Capability (CC) and precisions: - * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8 - block scaling (1D/2D, including power-of-2 scales). - * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8 - per-tensor current scaling. - FP8 delayed scaling is not supported because the corresponding grouped - quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and - SM110) raises instead of falling back: the fused path is Hopper-only and has - no MXFP8-broadcast emulation. Architectures outside the fused-path window - (e.g. SM120) fall back to the legacy path like every other recipe. - Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8 - per-tensor current scaling on Hopper); otherwise the legacy path is used. - Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization - currently requires RHT. - - 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"))): - return False - # 2. Filter out advanced features - if ( - debug - or cpu_offloading - or fp8_calibration - or backward_override is not None - or save_original_input - ): - return False - # 3. Filter by compute capability and cuBLAS version - device_capability = get_device_compute_capability() - if not (9, 0) <= device_capability <= (11, 0): - return False - cublaslt_version = tex.get_cublasLt_version() - if cublaslt_version < 130300: - return False - if device_capability < (10, 0) and cublaslt_version < 130400: - return False - # 4. Output quantization is not supported. - if any(q is not None for q in output_quantizers): - return False - # 5. Filter by quantization recipes. - if fp8: - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+. - if device_capability < (10, 0) and cublaslt_version < 130500: - return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM - # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast - # emulation. On Blackwell (SM100/SM110, the only other arch that reaches - # this branch) fail loudly rather than silently falling back to the - # unfused path the user explicitly opted out of. - if get_device_compute_capability() >= (10, 0): - raise RuntimeError( - "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" - " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" - " FP8 block-scaling path is Hopper-only. Unset" - " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" - " path (emulated via MXFP8 GEMM on Blackwell)." - ) - return True - # MXFP8 and NVFP4 require Blackwell+. - if not (10, 0) <= device_capability <= (11, 0): - return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) - return activation_dtype in (torch.bfloat16, torch.float16) - @staticmethod def _make_grouped_tensor( data: torch.Tensor, @@ -238,14 +212,109 @@ 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("Quantized grouped compute requires a weight quantizer.") + 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. + + 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: + 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, + 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: @@ -285,13 +354,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}.") @@ -305,6 +371,43 @@ 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.") + 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=(num_gemms, out_features), + dtype=dtype, + num_tensors=num_gemms, + shapes=[(1, out_features)] * num_gemms, + quantizer=None, + data=bias_data.reshape(-1), + ) + @staticmethod def _forward_grouped_tensor( ctx, @@ -326,6 +429,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, @@ -335,7 +440,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) @@ -366,6 +471,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, @@ -392,8 +499,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, @@ -424,7 +532,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, @@ -442,16 +552,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 @@ -517,19 +629,22 @@ def forward( skip_fp8_weight_update, save_original_input, debug, + single_grouped_weight, + single_grouped_bias, + use_grouped_tensor, ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override - else: - backward_override = None + recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + backward_override = recipe.backward_override if recipe is not None else None if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": 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 @@ -593,17 +708,53 @@ def forward( f"weight tensor (shape={tuple(weights[0].size())})" ) - if _GroupedLinear._is_grouped_tensor_path_supported( - fp8=fp8, - fp8_calibration=fp8_calibration, - debug=debug, - cpu_offloading=cpu_offloading, - backward_override=backward_override, - save_original_input=save_original_input, - activation_dtype=activation_dtype, - input_quantizers=input_quantizers, - output_quantizers=output_quantizers, + recipe_supports_grouped_tensor = is_module_grouped_tensor_path_supported( + recipe, + activation_dtype, + ) + grouped_tensor_features_supported = ( + not fp8_calibration + and not debug + and not cpu_offloading + and not save_original_input + and not any(q is not None for q in output_quantizers) + ) + if ( + use_grouped_tensor + and grouped_tensor_features_supported + and fp8 + and recipe.float8_block_scaling() + and (10, 0) <= get_device_compute_capability() <= (11, 0) ): + raise RuntimeError( + "use_grouped_tensor=True does not support the FP8 block-scaling recipe on " + "Blackwell GPUs: the native grouped FP8 block-scaling path is Hopper-only. " + "Set use_grouped_tensor=False, or unset " + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM if it enabled this path, to use " + "the MXFP8-emulated path on Blackwell." + ) + use_grouped_tensor_path = ( + use_grouped_tensor + and grouped_tensor_features_supported + and recipe_supports_grouped_tensor + ) + if ( + use_grouped_tensor + and not use_grouped_tensor_path + and (single_grouped_weight or single_grouped_bias) + ): + raise RuntimeError( + "Single grouped parameters require the native grouped-tensor path, but the active " + "device, cuBLASLt version, quantization recipe, or GroupedLinear feature " + "configuration does not support it. Disable single_grouped_weight and " + "single_grouped_bias to allow the split-quantize fallback." + ) + 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 set use_grouped_tensor=False." + ) return _GroupedLinear._forward_grouped_tensor( ctx, inp=inp, @@ -623,6 +774,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, @@ -837,12 +990,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 @@ -894,11 +1055,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: @@ -907,7 +1073,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( @@ -926,7 +1092,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", @@ -947,16 +1113,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 @@ -976,9 +1168,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: @@ -1006,10 +1198,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) @@ -1020,7 +1209,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): None, # out None, # dgrad_out *wgrad_list, - *grad_biases, + *grad_bias_args, ) @staticmethod @@ -1358,7 +1547,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 @@ -1382,6 +1573,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. + use_grouped_tensor : bool or None, default = None + Prefer the native GroupedTensor grouped GEMM path. Discrete parameters + fall back to split-quantize when the path is unsupported. Single grouped + parameters require the native path and raise instead of falling back. + 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. Notes ----- @@ -1415,6 +1613,7 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, name: Optional[str] = None, + use_grouped_tensor: Optional[bool] = None, ) -> None: super().__init__(name) @@ -1430,11 +1629,37 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + 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: + warnings.warn( + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM is deprecated and will be " + "removed in a future release. Pass use_grouped_tensor=True or " + "use_grouped_tensor=False to GroupedLinear instead.", + FutureWarning, + stacklevel=2, + ) + else: + use_grouped_tensor_env = "0" + use_grouped_tensor = bool(int(use_grouped_tensor_env)) + if not isinstance(use_grouped_tensor, bool): + raise TypeError( + f"use_grouped_tensor must be a bool or None, got {type(use_grouped_tensor)}." + ) + self.use_grouped_tensor = use_grouped_tensor 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 @@ -1583,7 +1808,7 @@ def make_grouped_weights(self, defer_init=False) -> None: if weight_quantizers and weight_quantizers[0] is not None else None ) - if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): + if recipe is not None and recipe.delayed(): self.set_tensor_parallel_attributes(defer_init=defer_init) return @@ -1871,16 +2096,13 @@ 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 not self.use_grouped_tensor: + raise RuntimeError( + "single_grouped_weight and single_grouped_bias require " + "use_grouped_tensor=True; the split-quantize path only supports discrete " + "parameters." + ) if not isinstance(m_splits, torch.Tensor): # Convert list of ints to tensor for backward compatibility m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") @@ -1909,6 +2131,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 quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() @@ -1916,6 +2139,13 @@ def forward( if self.no_debug_features_active(list(chain(*quantizers))): debug = False quantizers = self._get_quantizers() + if debug and (self.single_grouped_weight or self.single_grouped_bias): + raise RuntimeError( + "TE debug features do not support single grouped parameters. DebugQuantizer " + "uses the split-quantize path, which only supports discrete parameters. " + "Disable single_grouped_weight and single_grouped_bias, or disable TE debug " + "features for this GroupedLinear." + ) ( input_quantizers, @@ -1941,11 +2171,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, @@ -1969,6 +2202,9 @@ def forward( skip_fp8_weight_update, self.save_original_input, debug, + self.single_grouped_weight, + use_grouped_bias, + self.use_grouped_tensor, ) out, new_workspaces = linear_fn( *autograd_ctx, @@ -1986,12 +2222,15 @@ 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() 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 @@ -2006,31 +2245,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.use_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() @@ -2073,10 +2312,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): @@ -2091,13 +2327,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 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 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 [grouped_bias] return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] def _get_weight_quantizers(self) -> List[Quantizer]: diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 6def36ffc7..caab9c9c8d 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -24,7 +24,7 @@ from .bias import Bias from .constant_scale import ConstantScale from .dropout import Dropout -from .grouped_linear import GroupedLinear +from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported from .identity import Identity from .l2normalization import L2Normalization from .layer_norm import LayerNorm diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..7496b0d1e1 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -14,6 +14,7 @@ import torch import transformer_engine_torch as tex +from transformer_engine.common.recipe import Recipe from ...constants import DType from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker @@ -24,14 +25,12 @@ _2X_ACC_WGRAD, ) from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload -from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe +from ...quantization import FP8GlobalStateManager, QuantizerRole from ...quantized_tensor import QuantizedTensorStorage from ...tensor import ( Float8BlockQuantizer, - Float8CurrentScalingQuantizer, MXFP8Quantizer, MXFP8Tensor, - NVFP4Quantizer, Quantizer, ) from ...utils import ( @@ -72,6 +71,67 @@ GRAD_INPUT_BUFFER_KEY = "grad_input" +def is_op_fuser_grouped_tensor_path_supported( + recipe: Optional[Recipe], + dtype: torch.dtype, +) -> bool: + """Whether the op-fuser grouped-tensor path supports this recipe and dtype. + + * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, + which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common + library. + * MXFP8 and NVFP4 are supported on Blackwell GPUs with Compute Capability + (CC) 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped + quantization currently requires it. + * FP8 per-tensor current scaling uses grouped current-scaling quantization + through ``tex.group_quantize`` and cuBLASLt grouped GEMM with per-batch + scalar FP8 scaling. It is supported on Hopper and Blackwell, with + cuBLASLt 13.5+ required on Hopper. + * FP8 block scaling uses the grouped-tensor path only on Hopper with + cuBLASLt 13.6+. On other architectures or older cuBLAS versions it + falls back to the split-quantize path for discrete parameters. + * Custom recipes are unsupported because they may assign different + quantizers to input, weight, and grad-output roles. This predicate + currently supports only built-in recipes with known uniform layouts. + * Other quantization recipes, including FP8 delayed scaling, fall back to + split quantization because their grouped quantization kernels are missing. + * Unquantized BF16/FP16 compute is supported on Hopper and Blackwell. FP32 + is excluded because cuBLASLt grouped GEMM does not support it. + * Single grouped parameters have no split-quantize fallback, so callers + must reject them when this function returns ``False``. + """ + if dtype not in (torch.bfloat16, torch.float16): + return False + + device_capability = get_device_compute_capability() + if not (9, 0) <= device_capability <= (11, 0): + return False + cublaslt_version = tex.get_cublasLt_version() + if cublaslt_version < 130300: + return False + if device_capability < (10, 0) and cublaslt_version < 130400: + return False + + if recipe is None: + return True + if recipe.custom(): + return False + if recipe.float8_current_scaling(): + return device_capability >= (10, 0) or cublaslt_version >= 130500 + if recipe.float8_block_scaling(): + # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8. + return device_capability < (10, 0) and cublaslt_version >= 130600 + if recipe.mxfp8(): + return device_capability >= (10, 0) + if recipe.nvfp4(): + return ( + device_capability >= (10, 0) + and not recipe.disable_rht + and not recipe.row_scaled_activation + ) + return False + + class GroupedLinear(BasicOperation): r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` @@ -316,17 +376,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 @@ -791,68 +862,6 @@ def op_backward(self, *args, **kwargs): "It overrides `fuser_backward` instead of `op_backward`." ) - @staticmethod - def _is_graph_safe_path_supported( - *, - with_quantized_compute: bool, - input_quantizers: Sequence[Optional[Quantizer]], - dtype: torch.dtype, - single_grouped_weight: bool, - ) -> bool: - """Whether the graph-safe grouped-tensor flow can be used. - - * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, - which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common - library. This filter mirrors cuBLASLt grouped GEMM's architecture - requirement without duplicating its cuBLAS version checks. - * Quantized compute supports MXFP8 and NVFP4 on Blackwell GPUs with Compute Capability (CC) - 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped quantization currently - requires RHT. NVFP4 is additionally restricted to discrete weights: with - ``single_grouped_weight=True`` the weight quantizer is non-RHT and cannot use the - graph-safe grouped quantize kernel, so we fall back to the split-quantize flow. - * FP8 per-tensor current scaling is backed by grouped current-scaling quantization - (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling, - which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0). - * FP8 block scaling uses the grouped-tensor path on Hopper (CC 9.0) with cuBLAS 13.4+; - it is Hopper-only (no MXFP8-broadcast emulation), so elsewhere it falls back. - * Every other quantization recipe (fp8 delayed scaling, ...) falls back to the legacy flow - because the corresponding grouped quantization kernels are missing. - * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0) - -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it. - * 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. - """ - if not (9, 0) <= get_device_compute_capability() <= (11, 0): - return False - if with_quantized_compute: - # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM - # path; the compute-capability range was already checked above. On Hopper it - # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS. - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - if ( - get_device_compute_capability() < (10, 0) - and tex.get_cublasLt_version() < 130500 - ): - return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere - # fall back to the split-quantize (MXFP8-emulated) flow. - if get_device_compute_capability() >= (10, 0): - return False - return tex.get_cublasLt_version() >= 130400 - # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. - if not (10, 0) <= get_device_compute_capability() <= (11, 0): - return False - if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): - return True - # NVFP4 graph-safe grouped quantization requires RHT and only supports - # discrete weights; otherwise fall back to the split-quantize flow. - if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): - return not single_grouped_weight - return False - return dtype in (torch.bfloat16, torch.float16) - def _get_grouped_weight_for_gemm( self, weight_param: GroupedTensor, @@ -974,16 +983,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), @@ -1052,16 +1052,22 @@ def fuser_forward( out_buffer = basic_op_kwargs[0].get(OUTPUT_BUFFER_KEY) # Dispatch: graph-safe GroupedTensor flow whenever it can be used. - # See ``_is_graph_safe_path_supported`` for the gating rationale -- + # See ``is_op_fuser_grouped_tensor_path_supported`` for the gating rationale -- # in short it requires Hopper (SM90+) plus a supported dtype / # quantization recipe. Otherwise we fall back to the legacy # ``tex.split_quantize`` + ``general_grouped_gemm`` flow. - use_grouped_tensor_path = self._is_graph_safe_path_supported( - with_quantized_compute=with_quantized_compute, - input_quantizers=input_quantizers, - dtype=dtype, - single_grouped_weight=self.single_grouped_weight, + recipe = FP8GlobalStateManager.get_fp8_recipe() if with_quantized_compute else None + use_grouped_tensor_path = is_op_fuser_grouped_tensor_path_supported( + recipe, + dtype, ) + 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( @@ -1207,16 +1213,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, @@ -1505,7 +1506,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, @@ -1726,7 +1727,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 4a5e7d0cb8..fe4f1e859e 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1938,8 +1938,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,