diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 10d0dbd227..549f7de17d 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -97,6 +97,23 @@ def nvfp4_row_scaled(): return nvfp4_recipe +def nvfp4_row_scaled_quantized_backward(): + # Same row-scaled activation recipe as nvfp4_row_scaled(), but with + # backward_override=None so the backward runs in NVFP4 instead of falling back + # to high precision. + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override=None, + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + def nvfp4_4over6(): nvfp4_recipe = recipe.NVFP4BlockScaling( disable_rht=True, @@ -377,6 +394,84 @@ def test_grouped_linear_accuracy( torch.testing.assert_close(o, o_ref, rtol=0, atol=0) +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize("num_gemms", [1, 3]) +@pytest.mark.parametrize("bs", [2]) +@pytest.mark.parametrize("bias", all_boolean) +def test_grouped_linear_row_scaled_quantized_backward(dtype, num_gemms, bs, bias, model="126m"): + """Row-scaled NVFP4 GroupedLinear with quantized (non-fallback) NVFP4 backward. + + With ``backward_override=None`` the wgrad is computed in NVFP4: the row-scaled + activation becomes operand A of the ``NT`` grouped GEMM, which this PR routes + through the per-expert dense ``general_gemm`` loop. GroupedLinear must then + match a stack of independent dense ``Linear`` layers bit-for-bit, since both + execute the exact same per-expert quantize + GEMM kernels. + """ + recipe_row_scaled = nvfp4_row_scaled_quantized_backward() + config = model_configs[model] + if dtype not in get_nvfp4_inp_supported_dtypes(recipe_row_scaled, dtype): + pytest.skip(f"Input dtype {dtype} not supported for row-scaled NVFP4.") + + grouped_linear = ( + GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + device="cuda", + ) + .cuda() + .eval() + ) + sequential_linear = torch.nn.ModuleList( + [ + Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + device="cuda", + ).eval() + for _ in range(num_gemms) + ] + ) + + # Share weights/biases so the two paths are numerically comparable. + with torch.no_grad(): + for i in range(num_gemms): + sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) + if bias: + sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone()) + + outputs_ref = _test_grouped_linear_accuracy( + sequential_linear, + num_gemms, + bs, + dtype, + config, + recipe_row_scaled, + fp8=True, + fuse_wgrad_accumulation=False, + ) + outputs = _test_grouped_linear_accuracy( + grouped_linear, + num_gemms, + bs, + dtype, + config, + recipe_row_scaled, + fp8=True, + fuse_wgrad_accumulation=False, + ) + + # GroupedLinear is a per-expert loop over the same dense kernels, so the + # forward output, dgrad, and (row-scaled) wgrad must match bit-for-bit. + for o, o_ref in zip(outputs, outputs_ref): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + @pytest.mark.skipif( torch.cuda.get_device_capability() != (9, 0), reason="Only enable CUTLASS grouped gemm on Hopper", diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..016e005f9c 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -99,18 +99,13 @@ __launch_bounds__(BLOCK_SIZE) #endif } +// This is a general amax reduction, so it is not restricted to SM 10.0+. template -__global__ void -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -__launch_bounds__(BLOCK_SIZE) -#endif +__global__ void __launch_bounds__(BLOCK_SIZE) compute_columnwise_amax_kernel(const int num_rows, const int num_cols, const IType *__restrict__ input, float *__restrict__ output_columnwise_amax, const float *__restrict__ noop) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) - NVTE_DEVICE_ERROR("SM 10.0+ is required."); -#else if (noop != nullptr && noop[0] == 1.0f) { return; } @@ -128,7 +123,6 @@ __launch_bounds__(BLOCK_SIZE) if (threadIdx.x == 0) { output_columnwise_amax[col_idx] = col_amax; } -#endif } template diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index d46c31317f..6f47a5298c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -339,9 +339,12 @@ def general_grouped_gemm( else: bias_dtype = TE_DType[torch.bfloat16] - if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): - raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support row-scaled A.") - if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): + # TODO: Row-scaled NVFP4 grouped GEMM is handled by looping over the per-expert + # dense general_gemm, fprop and wgrad without a dedicated grouped kernel. + # The per-expert loop is not CUDA-graph-safe and will be fixed in a future work + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A) or any( + _is_nvfp4_row_scaled_tensor(tensor) for tensor in B + ): assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." if single_output: assert ( diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..49e77d1204 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1001,8 +1001,9 @@ std::tuple, std::vector, bool> bulk_alloc const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); - NVTE_CHECK(!columnwise_usage, - "Row-scaled NVFP4 bulk allocation does not support columnwise usage."); + // Columnwise (transpose) output is supported for row-scaled NVFP4: the + // columnwise amax is sized per-column below so the per-expert dense cast can + // emit the row-scaled transpose consumed by the wgrad GEMM. } const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto fp4_dtype = quantizer_cpp_list[0]->dtype; @@ -1146,7 +1147,10 @@ std::tuple, std::vector, bool> bulk_alloc dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); alignments.insert(alignments.end(), num_tensors, 16); for (size_t i = 0; i < num_tensors; ++i) { - shapes.emplace_back(amax_shape(columnwise_data_shapes[i])); + // columnwise_data_shapes[i] is the transposed shape, so its leading dim is + // the original last dim (number of columns). For row-scaled NVFP4 this + // yields a per-column amax vector; otherwise it stays a scalar {1}. + shapes.emplace_back(amax_shape(columnwise_data_shapes[i], row_scaled_nvfp4)); } dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); alignments.insert(alignments.end(), num_tensors, 16); @@ -1206,7 +1210,7 @@ std::tuple, std::vector, bool> bulk_alloc } if (columnwise_usage) { tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(amax_columnwise_list[i])); } tensor_cpp_list.emplace_back(std::move(tensor_wrapper));