diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 63e69beafc..e2d5a3f513 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -28,6 +28,8 @@ def check_nvfp4_gemm_versus_reference( x_columnwise: bool = False, w_columnwise: bool = False, row_scaled_nvfp4: bool = False, + err_corrected_nvfp4: bool = False, + use_bias: bool = False, use_4over6: bool = False, nvfp4_e4m3_max: int = 448, nvfp4_4over6_err_mode: str = "MAE", @@ -47,6 +49,7 @@ def check_nvfp4_gemm_versus_reference( w_shape = (K, N) if w_columnwise else (N, K) x = torch.randn(x_shape, dtype=x_dtype, device=device) w = torch.randn(w_shape, dtype=w_dtype, device=device) + bias = torch.randn((N,), dtype=torch.bfloat16, device=device) if use_bias else None # Setup out tensor if accumulate is True if accumulate: @@ -64,6 +67,7 @@ def check_nvfp4_gemm_versus_reference( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=row_scaled_nvfp4, + err_corrected_nvfp4=err_corrected_nvfp4, nvfp4_use_4over6=use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, @@ -134,6 +138,7 @@ def check_nvfp4_gemm_versus_reference( eps=0.0, quant_tile_shape=(1, 16), row_scaled_nvfp4=row_scaled_nvfp4, + err_corrected_nvfp4=err_corrected_nvfp4, nvfp4_use_4over6=use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, @@ -166,7 +171,7 @@ def check_nvfp4_gemm_versus_reference( qx=qx_data, qw=qw_data, m_params=None, # MMParams not used in reference - out_dtype=out_dtype, + out_dtype=torch.float32 if use_bias else out_dtype, sx=sx_trimmed, sw=sw_trimmed, bias=None, # No bias for this test @@ -176,6 +181,8 @@ def check_nvfp4_gemm_versus_reference( qresult_x=x_nvfp4_ref, qresult_w=w_nvfp4_ref, ) + if bias is not None: + y_ref = (y_ref + bias.to(torch.float32)).to(out_dtype) # Native TE GEMM using tex.generic_gemm (cuBLAS GEMM) # Allocate cuBLAS workspace @@ -184,7 +191,6 @@ def check_nvfp4_gemm_versus_reference( transa = True if not w_columnwise else False transb = False if not x_columnwise else True out_quantizer = None - bias = None bias_dtype = TE_DType[torch.bfloat16] use_gelu = False gelu_input = None @@ -204,6 +210,7 @@ def check_nvfp4_gemm_versus_reference( accumulate=accumulate, layout=layout, out=out.clone() if accumulate else None, + bias=bias, )[0] else: # Native cuBLAS GEMM @@ -489,6 +496,76 @@ def test_nvfp4_gemm_versus_reference( ) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("M, K, N", [(128, 128, 128), (112, 288, 96)]) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +@pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"]) +def test_error_corrected_row_scaled_gemm_versus_reference( + M: int, + K: int, + N: int, + out_dtype: torch.dtype, + use_bias: bool, +) -> None: + """The FP32 sum of primary and residual GEMMs matches the blockwise reference.""" + check_nvfp4_gemm_versus_reference( + x_dtype=torch.bfloat16, + w_dtype=torch.bfloat16, + out_dtype=out_dtype, + M=M, + K=K, + N=N, + accumulate=False, + row_scaled_nvfp4=True, + err_corrected_nvfp4=True, + use_bias=use_bias, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_error_corrected_row_scaled_gemm_improves_accuracy() -> None: + """The residual product reduces MSE relative to the primary row-scaled GEMM.""" + torch.manual_seed(41) + torch.cuda.manual_seed(41) + device = "cuda" + x = torch.randn((128, 256), dtype=torch.bfloat16, device=device) + w = torch.randn((128, 256), dtype=torch.bfloat16, device=device) + + def _activation_quantizer(*, corrected: bool) -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + err_corrected_nvfp4=corrected, + ) + + weight_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_q = weight_quantizer(w) + x_primary = _activation_quantizer(corrected=False)(x) + x_corrected = _activation_quantizer(corrected=True)(x) + + y_primary = general_gemm(w_q, x_primary, out_dtype=torch.float32, layout="TN")[0] + y_corrected = general_gemm(w_q, x_corrected, out_dtype=torch.float32, layout="TN")[0] + y_ref = x.to(torch.float32) @ w.to(torch.float32).T + + primary_mse = torch.mean((y_primary - y_ref) ** 2) + corrected_mse = torch.mean((y_corrected - y_ref) ** 2) + assert corrected_mse < primary_mse + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "m_splits, k, n", diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 85953687a7..0aaadf73d6 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -108,6 +108,7 @@ def check_quantization_nvfp4_versus_reference( use_cpp_allocator: bool, with_2d_quantization: bool, row_scaled_nvfp4: bool = False, + err_corrected_nvfp4: bool = False, use_4over6: bool = False, nvfp4_e4m3_max: int = 448, nvfp4_4over6_err_mode: str = "MAE", @@ -140,6 +141,7 @@ def check_quantization_nvfp4_versus_reference( with_post_rht_amax=False, with_2d_quantization=with_2d_quantization, row_scaled_nvfp4=row_scaled_nvfp4, + err_corrected_nvfp4=err_corrected_nvfp4, nvfp4_use_4over6=use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, @@ -176,6 +178,12 @@ def check_quantization_nvfp4_versus_reference( sx_t = x_nvfp4_sut._columnwise_scale_inv qx_amax = x_nvfp4_sut._amax_rowwise qx_amax_t = x_nvfp4_sut._amax_columnwise + qx_err = ( + x_nvfp4_sut._rowwise_data_err.view(dtype=torch.uint8) + if x_nvfp4_sut._rowwise_data_err is not None + else None + ) + sx_err = x_nvfp4_sut._rowwise_scale_inv_err # Reference quantization quant_tile_shape = (1, 16) if not with_2d_quantization else (16, 16) @@ -187,6 +195,7 @@ def check_quantization_nvfp4_versus_reference( eps=0.0, quant_tile_shape=quant_tile_shape, row_scaled_nvfp4=row_scaled_nvfp4, + err_corrected_nvfp4=err_corrected_nvfp4, nvfp4_use_4over6=use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, @@ -211,6 +220,14 @@ def check_quantization_nvfp4_versus_reference( ) ref_amax = x_nvfp4_ref.global_amax_row ref_amax_t = x_nvfp4_ref.global_amax_col + qx_err_ref = ( + unpack_fp4(x_nvfp4_ref.data_err.view(dtype=torch.uint8)) + if x_nvfp4_ref.data_err is not None + else None + ) + sx_err_ref = ( + x_nvfp4_ref.scale_err.view(dtype=torch.uint8) if x_nvfp4_ref.scale_err is not None else None + ) qx = unpack_fp4(qx) qx_t = unpack_fp4(qx_t) if qx_t is not None else None @@ -234,6 +251,47 @@ def check_quantization_nvfp4_versus_reference( torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) + if err_corrected_nvfp4: + assert qx_err is not None and qx_err_ref is not None + assert sx_err is not None and sx_err_ref is not None + torch.testing.assert_close(unpack_fp4(qx_err), qx_err_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close( + sx_err[: sx_err_ref.shape[0], : sx_err_ref.shape[1]], + sx_err_ref, + atol=0.0, + rtol=0.0, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + pytest.param(128, 128, id="aligned-tuned-1d"), + pytest.param(112, 272, id="unaligned-generic-vector"), + ], +) +@pytest.mark.parametrize( + "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] +) +def test_error_corrected_row_scaled_quantization_versus_reference( + M: int, + N: int, + use_cpp_allocator: bool, +) -> None: + """Primary and residual FP4 bytes/scales exactly match the BF16 reference.""" + check_quantization_nvfp4_versus_reference( + x_dtype=torch.bfloat16, + M=M, + N=N, + return_transpose=False, + swizzled_scale=False, + use_cpp_allocator=use_cpp_allocator, + with_2d_quantization=False, + row_scaled_nvfp4=True, + err_corrected_nvfp4=True, + ) + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 3e6fdb816b..2440022251 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -1623,6 +1623,104 @@ def test_nvfp4_recipe_state_role_dispatch_backward(): assert q.stochastic_rounding == nvfp4_recipe.fp4_quant_bwd_grad.stochastic_rounding +def test_nvfp4_error_correction_is_forward_activation_only(): + """Recipe opt-in is attached to forward activations, never weights or backward tensors.""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState + + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + row_scaled_activation=True, + err_corrected_activation=True, + ) + forward = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _nvfp4_role("input"), + _nvfp4_role("weight"), + _nvfp4_role("output"), + ], + ).make_quantizers() + assert [q.err_corrected_nvfp4 for q in forward] == [True, False, True] + assert [q.row_scaled_nvfp4 for q in forward] == [True, False, True] + + backward = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="backward", + num_quantizers=2, + roles=[_nvfp4_role("grad_output"), _nvfp4_role("grad_input")], + ).make_quantizers() + assert all(not q.err_corrected_nvfp4 for q in backward) + assert all(not q.row_scaled_nvfp4 for q in backward) + + +def test_nvfp4_error_correction_requires_row_scaled_activation(): + with pytest.raises(ValueError, match="row_scaled_activation=True"): + recipe.NVFP4BlockScaling( + row_scaled_activation=False, + err_corrected_activation=True, + ) + with pytest.raises(ValueError, match="does not support activation 4over6"): + recipe.NVFP4BlockScaling( + disable_rht=True, + row_scaled_activation=True, + err_corrected_activation=True, + nvfp4_4over6="activations", + ) + + +@pytest.mark.parametrize("bias", [False, True]) +def test_nvfp4_error_correction_is_batch_invariant(bias): + """Shared rows produce identical outputs when the surrounding batch changes.""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + torch.manual_seed(0) + batch_sizes = (1, 2, 4) + sequence_length = 16 + in_features = 128 + out_features = 128 + inp = torch.randn( + max(batch_sizes), + sequence_length, + in_features, + device="cuda", + dtype=torch.bfloat16, + ) + model = Linear( + in_features, + out_features, + bias=bias, + params_dtype=torch.bfloat16, + device="cuda", + ).eval() + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + row_scaled_activation=True, + err_corrected_activation=True, + ) + + outputs = {} + with torch.inference_mode(): + for batch_size in batch_sizes: + with autocast(enabled=True, recipe=nvfp4_recipe): + outputs[batch_size] = model(inp[:batch_size]).clone() + + reference = outputs[max(batch_sizes)] + for batch_size in batch_sizes: + assert torch.isfinite(outputs[batch_size]).all() + torch.testing.assert_close( + outputs[batch_size], + reference[:batch_size], + rtol=0, + atol=0, + ) + + def test_nvfp4_recipe_state_positional_fallback_matches_explicit_roles(): """``roles=None`` matches explicit ``[input, weight, output]`` slot-for-slot.""" if not torch.cuda.is_available() or not te.is_nvfp4_available(): diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 033d464bcf..4933769087 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -103,6 +103,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const auto [rows, cols] = input_tensor->flat_2d_dims(); auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; + const bool err_corrected_nvfp4 = output_tensor->err_corrected_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, "Non-4over6 NVFP4 quantization requires E4M3 max 448."); @@ -115,6 +116,10 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, "Row-scaled NVFP4 quantization does not produce columnwise output."); nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream); } + NVTE_CHECK(!err_corrected_nvfp4 || row_scaled_nvfp4, + "Error-corrected NVFP4 quantization requires row scaling."); + NVTE_CHECK(!err_corrected_nvfp4 || !nvfp4_use_4over6, + "Error-corrected NVFP4 quantization does not support 4over6."); // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and // both-directions keep their existing routing. Columnwise-only 1D and non-bf16 fall back to // quantize_transpose_vector_blockwise_fp4. @@ -155,6 +160,9 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, /*row_scaled_nvfp4=*/row_scaled_nvfp4, + /*err_corrected_nvfp4=*/err_corrected_nvfp4, + /*scale_inv_err=*/output_tensor->scale_inv_err, + /*output_err=*/output_tensor->data_err, /*noop_tensor=*/noop_tensor->data, /*stream=*/stream); } @@ -315,6 +323,9 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, /*row_scaled_nvfp4=*/false, + /*err_corrected_nvfp4=*/false, + /*scale_inv_err=*/output_tensor->scale_inv_err, + /*output_err=*/output_tensor->data_err, /*noop_tensor=*/noop_tensor->data, /*stream=*/stream); } diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 3a34c76de8..26bd1491a1 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -237,13 +237,15 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / template + bool ROW_SCALED_NVFP4, bool ERR_CORRECTED_NVFP4> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, const __grid_constant__ CUtensorMap tensor_map_output_t, + const __grid_constant__ CUtensorMap tensor_map_output_err, nvfp4_scale_t *const scales_ptr, - nvfp4_scale_t *const scales_t_ptr, const float *noop, + nvfp4_scale_t *const scales_t_ptr, + nvfp4_scale_t *const scales_err_ptr, const float *noop, const float *const amax_rowwise_ptr, const float *const amax_colwise_ptr, const size_t rows, const size_t cols, const size_t scale_stride, @@ -327,6 +329,7 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; constexpr size_t out_mem_colwise_data = buff_size_aligned_out; + constexpr size_t out_mem_error_data = ERR_CORRECTED_NVFP4 ? buff_size_aligned_out : 0; constexpr size_t out_mem_rowwise_scales = 0; extern __shared__ char dynamic_shmem[]; @@ -340,11 +343,14 @@ __global__ void __launch_bounds__(THREADS_NUM) IType *in_sh = reinterpret_cast(dshmem); fp4e2m1x2 *out_data_sh = reinterpret_cast(dshmem + in_mem); fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + fp4e2m1x2 *out_err_data_sh = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_error_data); nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_error_data + + out_mem_rowwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -637,28 +643,24 @@ __global__ void __launch_bounds__(THREADS_NUM) } } + const size_t scales_offset_Y = + scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; + const size_t scales_offset_X = scales_offset_X_rowwise; + const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + const bool rowwise_scale_is_within_bounds_Y = + (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; + float S_enc_rowwise_block = S_enc_rowwise; + float S_dec_rowwise_block = S_dec_rowwise; + nvfp4_scale_t S_dec_b_fp8; float block_scale_inverse; if constexpr (ROW_SCALED_NVFP4) { // 2. Compute E4M3 scaling factor - const size_t scales_offset_Y = - scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; - const float S_enc_rowwise_block = + S_enc_rowwise_block = scales_offset_Y < rows ? compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[scales_offset_Y]) : 1.0f; - const float S_dec_rowwise_block = 1.0f / S_enc_rowwise_block; - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); - - // Check boundaries - const size_t scales_offset_X = scales_offset_X_rowwise; - const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; - - const bool rowwise_scale_is_within_bounds_Y = - (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; - if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { - scales_ptr[scale_idx_global] = S_dec_b_fp8; - } + S_dec_rowwise_block = 1.0f / S_enc_rowwise_block; + S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); // Compute "correct" per-block encoding scaling factor constexpr float float_max = detail::TypeExtrema::max; @@ -667,29 +669,20 @@ __global__ void __launch_bounds__(THREADS_NUM) float_max); // S_enc_b_fp8 } else { // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise); - - // Check boundaries - const size_t scales_offset_Y = - scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; - const size_t scales_offset_X = scales_offset_X_rowwise; - const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; - - const bool rowwise_scale_is_within_bounds_Y = - (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; - if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { - scales_ptr[scale_idx_global] = S_dec_b_fp8; - } + S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise); // Compute "correct" per-block encoding scaling factor constexpr float float_max = detail::TypeExtrema::max; block_scale_inverse = fminf(1.0f / (static_cast(S_dec_b_fp8) * S_dec_rowwise), float_max); // S_enc_b_fp8 } + if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { + scales_ptr[scale_idx_global] = S_dec_b_fp8; + } const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; -// 3. Scale elements + // 3. Scale elements + IType residual[SCALE_DIM]; #pragma unroll for (int w = 0; w < WAVES; ++w) { Vec out; @@ -713,12 +706,68 @@ __global__ void __launch_bounds__(THREADS_NUM) out.data.elt[e] = ptx::mul_cvt_fp32_to_fp4_4x( in01, in23, block_scale_inverse_2x, rbits); } + if constexpr (ERR_CORRECTED_NVFP4) { + const float4 q = static_cast(out.data.elt[e]); + const float q_values[4] = {q.x, q.y, q.z, q.w}; + const IType *input_values = reinterpret_cast(&in_IType[w]); + constexpr float fp8_max = detail::TypeExtrema::max; + constexpr float fp4_max = detail::TypeExtrema::max; + constexpr float global_decode_scale_multiplier = 1.0f / (fp4_max * fp8_max); + const float primary_global_decode_scale = + scales_offset_Y < rows + ? amax_rowwise_ptr[scales_offset_Y] * global_decode_scale_multiplier + : 1.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + float primary_dequantized_fp32 = q_values[j] * static_cast(S_dec_b_fp8); + primary_dequantized_fp32 *= primary_global_decode_scale; + const IType primary_dequantized = static_cast(primary_dequantized_fp32); + residual[w * PACK_SIZE + 4 * e + j] = + static_cast(static_cast(input_values[4 * e + j]) - + static_cast(primary_dequantized)); + } + } } const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; out.store_to(&out_data_sh[shmem_offset_rowwise]); } + + if constexpr (ERR_CORRECTED_NVFP4) { + IType2 error_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int e = 0; e < SCALE_DIM; e += 2) { + const IType2 error_2x = {residual[e], residual[e + 1]}; + ptx::abs_max_2x(error_amax_2x, error_amax_2x, error_2x); + } + const float error_amax = + static_cast(__hmax(__habs(error_amax_2x.x), __habs(error_amax_2x.y))); + const nvfp4_scale_t S_dec_err_fp8 = + compute_decoding_scaling_factor(error_amax, S_enc_rowwise_block); + constexpr float float_max = detail::TypeExtrema::max; + const float error_scale_inverse = + fminf(1.0f / (static_cast(S_dec_err_fp8) * S_dec_rowwise_block), float_max); + const float2 error_scale_inverse_2x{error_scale_inverse, error_scale_inverse}; + if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { + scales_err_ptr[scale_idx_global] = S_dec_err_fp8; + } +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + Vec out_err; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 4; ++e) { + const uint64_t error_values = + *reinterpret_cast(&residual[w * PACK_SIZE + 4 * e]); + out_err.data.elt[e] = + ptx::mul_cvt_bf16_to_fp4_4x(error_values, error_scale_inverse_2x, 0); + } + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; + out_err.store_to(&out_err_data_sh[shmem_offset_rowwise]); + } + } } } @@ -742,6 +791,12 @@ __global__ void __launch_bounds__(THREADS_NUM) reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, reinterpret_cast(&out_data_sh[buff_offset_out])); + if constexpr (ERR_CORRECTED_NVFP4) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_err), global_offset_X, + global_offset_Y, reinterpret_cast(&out_err_data_sh[buff_offset_out])); + } + if constexpr (RETURN_TRANSPOSE) { ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output_t), global_offset_X_t, @@ -1347,6 +1402,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; const bool row_scaled_nvfp4 = output->row_scaled_nvfp4; + const bool err_corrected_nvfp4 = output->err_corrected_nvfp4; NVTE_CHECK(!row_scaled_nvfp4 || !use_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); @@ -1390,6 +1446,15 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, NVTE_CHECK(!with_gemm_swizzled_scales || use_2d_quantization, "In-kernel GEMM-swizzled NVFP4 scales are only supported on the 2D quantization path; " "other paths must emit scales in compact format."); + NVTE_CHECK(!err_corrected_nvfp4 || row_scaled_nvfp4, + "Error-corrected NVFP4 quantization requires row scaling."); + NVTE_CHECK(!err_corrected_nvfp4 || input.dtype() == DType::kBFloat16, + "Error-corrected NVFP4 quantization requires BF16 input."); + NVTE_CHECK( + !err_corrected_nvfp4 || (output->has_data_err() && output->scale_inv_err.dptr != nullptr), + "Error-corrected NVFP4 data and scales must be allocated."); + NVTE_CHECK(!err_corrected_nvfp4 || !with_gemm_swizzled_scales, + "Error-corrected NVFP4 output must have scales in compact format."); if (return_transpose) { NVTE_CHECK(output->has_columnwise_data(), "NVFP4 transposed output tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), @@ -1421,6 +1486,8 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); nvfp4_scale_t *const scales_transpose_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + nvfp4_scale_t *const scales_err_ptr = + reinterpret_cast(output->scale_inv_err.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); @@ -1443,6 +1510,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_output{}; alignas(64) CUtensorMap tensor_map_output_transpose{}; + alignas(64) CUtensorMap tensor_map_output_err{}; create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, sizeof(IType) * 8); @@ -1451,6 +1519,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, create_2D_tensor_map(tensor_map_output, output->data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, 4); } + if (err_corrected_nvfp4) { + create_2D_tensor_map(tensor_map_output_err, output->data_err, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols, 0, 4); + } if (return_transpose) { create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); @@ -1467,41 +1539,56 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, constexpr size_t out_data_mem = buff_size_aligned_out; constexpr size_t out_data_transpose_mem = buff_size_aligned_out; + const size_t out_data_error_mem = err_corrected_nvfp4 ? buff_size_aligned_out : 0; constexpr size_t out_scales_transpose_mem = buff_size_scales; - constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; + const size_t out_mem = out_data_mem + out_data_transpose_mem + out_data_error_mem; - constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + const size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_rowwise, RETURN_ROWWISE, { - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - // The 1D kernel always produces rowwise output (no RETURN_ROWWISE); the dispatch only - // routes columnwise-only requests here when use_2d_quantization is true. - auto kernel = quantize_transpose_nvfp4_kernel; - - if constexpr (use_2d_quantization) { - if (with_gemm_swizzled_scales) { - kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, - RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + TRANSFORMER_ENGINE_SWITCH_CONDITION(err_corrected_nvfp4, ERR_CORRECTED_NVFP4, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_rowwise, RETURN_ROWWISE, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + if constexpr (use_2d_quantization) { + if (with_gemm_swizzled_scales) { + auto kernel = quantize_transpose_nvfp4_2D_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, + scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, + cols, scale_stride, scale_stride_transpose, rng_state); + } else { + auto kernel = quantize_transpose_nvfp4_2D_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, + scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, + cols, scale_stride, scale_stride_transpose, rng_state); + } } else { - kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, - RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; + auto kernel = + quantize_transpose_nvfp4_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, + tensor_map_output_err, scales_ptr, scales_transpose_ptr, scales_err_ptr, + noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, scale_stride, + scale_stride_transpose, rng_state); } - } - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - kernel<<>>( - tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, - scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, - scale_stride, scale_stride_transpose, rng_state); + }); }); }); });); diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 8f37229fd5..279155633d 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -139,6 +139,7 @@ using IType2 = typename ptx::FPx2; using IType3D = IType[BUFFS_NUM_IN][BUFF_IN_DIM_Y][BUFF_IN_DIM_X]; using IType2x3D = IType2[BUFFS_NUM_IN][BUFF_IN_DIM_Y][BUFF_IN_DIM_X / 2]; using OType2x3D = fp4e2m1x2[BUFFS_NUM_OUT][BUFF_OUT_DIM_Y][BUFF_OUT_DIM_X]; +using OTypeErr2x3D = fp4e2m1x2[BUFFS_NUM_OUT][BUFF_OUT_DIM_Y][BUFF_OUT_DIM_X]; using OType2xt3D = fp4e2m1x2[BUFFS_NUM_OUT_TR][BUFF_OUT_TR_DIM_Y][BUFF_OUT_TR_DIM_X]; using ScalesType2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_Y][SCALES_PER_CHUNK_X]; using ScalesTypeTr2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_X][SCALES_PER_CHUNK_Y]; @@ -261,17 +262,21 @@ __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_pt } } -template +template __device__ __forceinline__ void rowwise_scaling( const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_ptr, - nvfp4_scale_t *__restrict__ sSFrowwise_ptr, const float S_enc_rowwise, const int stage_Y, + nvfp4_scale_t *__restrict__ sSFrowwise_ptr, fp4e2m1x2 *__restrict__ sOutErr_ptr, + nvfp4_scale_t *__restrict__ sSFerr_ptr, const float S_enc_rowwise, const int stage_Y, const int stage_X, const int buff_in, const int buff_out, const float *amax_rowwise_ptr, const size_t row_offset, const size_t rows, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn = *reinterpret_cast(sIn_ptr); auto &sOut = *reinterpret_cast(sOut_ptr); + auto &sOutErr = *reinterpret_cast(sOutErr_ptr); auto &sSFrowwise = *reinterpret_cast(sSFrowwise_ptr); + auto &sSFerr = *reinterpret_cast(sSFerr_ptr); const int thread_lane = threadIdx.x % THREADS_PER_WARP; const int bank_group = thread_lane / THREADS_PER_BANK; @@ -314,10 +319,11 @@ __device__ __forceinline__ void rowwise_scaling( const float block_amax = get_amax_of_pair(thread_amax_2x); nvfp4_scale_t S_dec_b_fp8; + float S_enc_rowwise_block = S_enc_rowwise; + const size_t row_idx = row_offset + stage_Y * TILE_DIM_Y + it_offset_Y_rowwise; scaling_coeff_type SFcoefficient; if constexpr (ROW_SCALED_NVFP4) { - const size_t row_idx = row_offset + stage_Y * TILE_DIM_Y + it_offset_Y_rowwise; - const float S_enc_rowwise_block = + S_enc_rowwise_block = row_idx < rows ? core::compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[row_idx]) : 1.0f; S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); @@ -336,7 +342,8 @@ __device__ __forceinline__ void rowwise_scaling( sSFrowwise[scales_offset_Y][scales_offset_X] = S_dec_b_fp8; } -// Scale elements + // Scale elements + IType residual[ELTS_PER_THREAD]; #pragma unroll for (int w = 0; w < WAVES; ++w) { const uint64_t elts03 = *reinterpret_cast(&rIn[w][0]); @@ -356,19 +363,72 @@ __device__ __forceinline__ void rowwise_scaling( const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % ELTS_PER_THREAD; const int swizzled_idx = (swizzled_group_idx + thread_offset_X_rowwise) / 2; ptx::st_shared_b32(&sOut[buff_out][it_offset_Y_rowwise][swizzled_idx], out_x8); + + if constexpr (ERR_CORRECTED_NVFP4) { + const auto *out_4x = reinterpret_cast(&out_x8); + const float4 q0 = static_cast(out_4x[0]); + const float4 q1 = static_cast(out_4x[1]); + const float q_values[8] = {q0.x, q0.y, q0.z, q0.w, q1.x, q1.y, q1.z, q1.w}; + constexpr float fp8_max = detail::TypeExtrema::max; + constexpr float fp4_max = detail::TypeExtrema::max; + constexpr float global_decode_scale_multiplier = 1.0f / (fp4_max * fp8_max); + const float primary_global_decode_scale = + row_idx < rows ? amax_rowwise_ptr[row_idx] * global_decode_scale_multiplier : 1.0f; +#pragma unroll + for (int e = 0; e < PACK_SIZE; ++e) { + float primary_dequantized_fp32 = q_values[e] * static_cast(S_dec_b_fp8); + primary_dequantized_fp32 *= primary_global_decode_scale; + const IType primary_dequantized = static_cast(primary_dequantized_fp32); + const IType input_value = reinterpret_cast(&rIn[w])[e]; + residual[w * PACK_SIZE + e] = static_cast(static_cast(input_value) - + static_cast(primary_dequantized)); + } + } + } + + if constexpr (ERR_CORRECTED_NVFP4) { + IType2 error_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int e = 0; e < ELTS_PER_THREAD; e += 2) { + const IType2 error_2x = {residual[e], residual[e + 1]}; + ptx::abs_max_2x(error_amax_2x, error_amax_2x, error_2x); + } + const float error_amax = get_amax_of_pair(error_amax_2x); + const nvfp4_scale_t S_dec_err_fp8 = + compute_decoding_scaling_factor(error_amax, S_enc_rowwise_block); + const scaling_coeff_type error_coefficient = + compute_nvfp4_scaling_coefficient(S_dec_err_fp8, S_enc_rowwise_block); + if (SF_storing_thread) { + const int scales_offset_Y = stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE; + const int scales_offset_X = stage_rowwise_scales_offset_X; + sSFerr[scales_offset_Y][scales_offset_X] = S_dec_err_fp8; + } +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const uint64_t error03 = *reinterpret_cast(&residual[w * PACK_SIZE]); + const uint64_t error47 = *reinterpret_cast(&residual[w * PACK_SIZE + 4]); + const uint32_t error_out_x8 = + ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest(error03, error47, + error_coefficient); + const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % ELTS_PER_THREAD; + const int swizzled_idx = (swizzled_group_idx + thread_offset_X_rowwise) / 2; + ptx::st_shared_b32(&sOutErr[buff_out][it_offset_Y_rowwise][swizzled_idx], error_out_x8); + } } } } template + bool ROW_SCALED_NVFP4, bool ERR_CORRECTED_NVFP4> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D_kernel( const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, - const __grid_constant__ CUtensorMap tensor_map_output_t, nvfp4_scale_t *const scales_ptr, - nvfp4_scale_t *const scales_t_ptr, const float *noop, const float *const amax_rowwise_ptr, - const float *const amax_colwise_ptr, const size_t rows, const size_t cols, - const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { + const __grid_constant__ CUtensorMap tensor_map_output_t, + const __grid_constant__ CUtensorMap tensor_map_output_err, nvfp4_scale_t *const scales_ptr, + nvfp4_scale_t *const scales_t_ptr, nvfp4_scale_t *const scales_err_ptr, const float *noop, + const float *const amax_rowwise_ptr, const float *const amax_colwise_ptr, const size_t rows, + const size_t cols, const size_t scale_stride, const size_t scale_stride_t, + const size_t *rng_state) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -400,8 +460,10 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D constexpr int out_mem_rowwise_data = buff_size_aligned_out; constexpr int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + constexpr int out_mem_error_data = ERR_CORRECTED_NVFP4 ? buff_size_aligned_out : 0; constexpr int out_mem_rowwise_scales = DIVUP_TO_MULTIPLE( TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + constexpr int out_mem_error_scales = ERR_CORRECTED_NVFP4 ? out_mem_rowwise_scales : 0; // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned extern __shared__ unsigned char dynamic_shmem[]; @@ -410,17 +472,25 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D IType *sIn_ptr = reinterpret_cast(dshmem); fp4e2m1x2 *sOut_ptr = reinterpret_cast(dshmem + in_mem); fp4e2m1x2 *sOut_tr_ptr = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + fp4e2m1x2 *sOutErr_ptr = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); auto &sIn = *reinterpret_cast(sIn_ptr); auto &sOut = *reinterpret_cast(sOut_ptr); auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); + auto &sOutErr = *reinterpret_cast(sOutErr_ptr); nvfp4_scale_t *sSFrowwise_ptr = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_error_data); + nvfp4_scale_t *sSFerr_ptr = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_error_data + + out_mem_rowwise_scales); nvfp4_scale_t *sSFcolwise_ptr = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_error_data + + out_mem_rowwise_scales + out_mem_error_scales); auto &sSFrowwise = *reinterpret_cast(sSFrowwise_ptr); + auto &sSFerr = *reinterpret_cast(sSFerr_ptr); auto &sSFcolwise = *reinterpret_cast(sSFcolwise_ptr); constexpr int shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -582,9 +652,11 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D ptx::cp_async_bulk_wait_group_read(); // NVFP4 Quantization - rowwise_scaling( - sIn_ptr, sOut_ptr, sSFrowwise_ptr, S_enc_rowwise, stage_Y, stage_X, buff_in, buff_out, - amax_rowwise_ptr, block_offset_Y, rows, rng, random_uint4, rnd_idx); + rowwise_scaling(sIn_ptr, sOut_ptr, sSFrowwise_ptr, sOutErr_ptr, + sSFerr_ptr, S_enc_rowwise, stage_Y, stage_X, buff_in, + buff_out, amax_rowwise_ptr, block_offset_Y, rows, rng, + random_uint4, rnd_idx); if constexpr (RETURN_TRANSPOSE) { colwise_scaling( @@ -608,6 +680,12 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, reinterpret_cast(&sOut[buff_out])); + if constexpr (ERR_CORRECTED_NVFP4) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_err), global_offset_X, + global_offset_Y, reinterpret_cast(&sOutErr[buff_out])); + } + if constexpr (RETURN_TRANSPOSE) { ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output_t), global_offset_X_tr, @@ -642,6 +720,20 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D } } + if constexpr (ERR_CORRECTED_NVFP4) { + using ScalesVec = Vec; + const int count = min(SCALES_PER_CHUNK_X, chunk_cols / SCALE_DIM); + for (size_t row = threadIdx.x; row < TunableConfig::CHUNK_DIM_Y; row += THREADS_NUM) { + const size_t row_global = scales_block_offset_Y_rowwise + row; + if (row_global < rows) { + ScalesVec &scales_vec = *reinterpret_cast(sSFerr[row]); + const size_t scale_idx_global = + row_global * scale_stride + scales_block_offset_X_rowwise; + scales_vec.store_to_elts(&scales_err_ptr[scale_idx_global], 0, count); + } + } + } + // Colwise if constexpr (RETURN_TRANSPOSE) { using ScalesVec = Vec; @@ -692,6 +784,7 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, const bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; const bool row_scaled_nvfp4 = output->row_scaled_nvfp4; + const bool err_corrected_nvfp4 = output->err_corrected_nvfp4; // If transposed output is allocated, return the transposed data // Otherwise, it's not necesary to return the transposed data. @@ -710,6 +803,11 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, "Row-scaled NVFP4 quantization requires rowwise amax."); NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), "Row-scaled NVFP4 quantization does not produce columnwise output."); + NVTE_CHECK(!err_corrected_nvfp4 || row_scaled_nvfp4, + "Error-corrected NVFP4 quantization requires row scaling."); + NVTE_CHECK( + !err_corrected_nvfp4 || (output->has_data_err() && output->scale_inv_err.dptr != nullptr), + "Error-corrected NVFP4 data and scales must be allocated."); if (return_transpose) { NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), @@ -737,6 +835,8 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); nvfp4_scale_t *const scales_transpose_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + nvfp4_scale_t *const scales_err_ptr = + reinterpret_cast(output->scale_inv_err.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); @@ -757,12 +857,17 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_output{}; alignas(64) CUtensorMap tensor_map_output_transpose{}; + alignas(64) CUtensorMap tensor_map_output_err{}; create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, sizeof(IType) * 8); create_2D_tensor_map(tensor_map_output, output->data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, 4); + if (err_corrected_nvfp4) { + create_2D_tensor_map(tensor_map_output_err, output->data_err, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols, 0, 4); + } if (return_transpose) { create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); @@ -786,13 +891,15 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, const int out_data_mem = buff_size_aligned_out; const int out_data_transpose_mem = return_transpose ? buff_size_aligned_out_t : 0; + const int out_data_error_mem = err_corrected_nvfp4 ? buff_size_aligned_out : 0; const int out_scales_mem = buff_size_scales; + const int out_scales_error_mem = err_corrected_nvfp4 ? buff_size_scales : 0; const int out_scales_transpose_mem = return_transpose ? buff_size_scales_transpose : 0; - const int out_mem = out_data_mem + out_data_transpose_mem; + const int out_mem = out_data_mem + out_data_transpose_mem + out_data_error_mem; - const int dshmem_size = - in_mem + out_mem + out_scales_transpose_mem + out_scales_mem + TMA_SHMEM_ALIGNMENT; + const int dshmem_size = in_mem + out_mem + out_scales_transpose_mem + out_scales_mem + + out_scales_error_mem + TMA_SHMEM_ALIGNMENT; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, @@ -800,18 +907,21 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, use_fast_math, USE_FAST_MATH, TRANSFORMER_ENGINE_SWITCH_CONDITION( row_scaled_nvfp4, ROW_SCALED_NVFP4, - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - auto kernel = - quantize_transpose_nvfp4_tuned_1D_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - dshmem_size); - kernel<<>>( - tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, - scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, - scale_stride, scale_stride_transpose, rng_state); - });););); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + err_corrected_nvfp4, ERR_CORRECTED_NVFP4, + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + auto kernel = quantize_transpose_nvfp4_tuned_1D_kernel< + USE_STOCHASTIC_ROUNDING, USE_FAST_MATH, RETURN_TRANSPOSE, ROW_SCALED_NVFP4, + ERR_CORRECTED_NVFP4>; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, + tensor_map_output_err, scales_ptr, scales_transpose_ptr, scales_err_ptr, + noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, scale_stride, + scale_stride_transpose, rng_state); + }););););); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..07ef7bd15a 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -289,6 +289,8 @@ struct Tensor { SimpleTensor scale; SimpleTensor scale_inv; SimpleTensor columnwise_scale_inv; + SimpleTensor data_err; + SimpleTensor scale_inv_err; NVTEScalingMode scaling_mode; NVTETensor nvte_tensor; @@ -302,6 +304,8 @@ struct Tensor { * Only meaningful for NVFP4 tensors. */ bool row_scaled_nvfp4 = false; + /*! \brief Whether rowwise NVFP4 error-correction data is present. */ + bool err_corrected_nvfp4 = false; /*! \brief Global E4M3 scale bound used by NVFP4. * * Standard NVFP4 uses 448. Some 4over6 tensors use 256 to leave room for @@ -320,7 +324,10 @@ struct Tensor { sizeof(NVTEBasicTensor), // kNVTEColumnwiseAmax sizeof(uint8_t), // kNVTEWithGEMMSwizzledScales sizeof(uint8_t), // kNVTERowScaledNVFP4 - sizeof(int) // kNVTENVFP4E4M3Max + sizeof(int), // kNVTENVFP4E4M3Max + sizeof(NVTEBasicTensor), // kNVTERowwiseDataErr + sizeof(NVTEBasicTensor), // kNVTERowwiseScaleInvErr + sizeof(uint8_t) // kNVTEErrCorrectedNVFP4 }; Tensor() : scaling_mode{NVTE_DELAYED_TENSOR_SCALING}, nvte_tensor{0} {} @@ -334,9 +341,12 @@ struct Tensor { scale.clear(); scale_inv.clear(); columnwise_scale_inv.clear(); + data_err.clear(); + scale_inv_err.clear(); scaling_mode = NVTE_DELAYED_TENSOR_SCALING; with_gemm_swizzled_scales = false; row_scaled_nvfp4 = false; + err_corrected_nvfp4 = false; nvfp4_e4m3_max = 448; } @@ -361,6 +371,9 @@ struct Tensor { */ bool has_columnwise_data() const { return columnwise_data.has_data(); } + /*! Whether the tensor has rowwise NVFP4 residual data. */ + bool has_data_err() const { return data_err.has_data(); } + /*! Datatype of tensor elements. */ DType dtype() const { if (!has_data() && has_columnwise_data()) { diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index aa0405e177..3e538802af 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -90,6 +90,9 @@ enum NVTETensorParam { * Standard NVFP4 uses 448; 4over6 may use 256 for map-to-4 headroom. */ kNVTENVFP4E4M3Max = 9, + kNVTERowwiseDataErr = 10, /*!< Error-correction FP4 data in rowwise layout */ + kNVTERowwiseScaleInvErr = 11, /*!< Block scales for error-correction rowwise data */ + kNVTEErrCorrectedNVFP4 = 12, /*!< Whether an NVFP4 tensor carries error correction */ kNVTENumTensorParams }; @@ -864,6 +867,17 @@ class TensorWrapper { return set_parameter(kNVTEColumnwiseAmax, dptr, type, shape); } + template + TensorWrapper &set_rowwise_data_err(void *dptr, DType type, const ShapeType &shape) noexcept { + return set_parameter(kNVTERowwiseDataErr, dptr, type, shape); + } + + template + TensorWrapper &set_rowwise_scale_inv_err(void *dptr, DType type, + const ShapeType &shape) noexcept { + return set_parameter(kNVTERowwiseScaleInvErr, dptr, type, shape); + } + void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales) { const auto val = static_cast(with_gemm_swizzled_scales); nvte_set_tensor_param_v2(tensor_, kNVTEWithGEMMSwizzledScales, &val, sizeof(val)); @@ -879,6 +893,11 @@ class TensorWrapper { nvte_set_tensor_param_v2(tensor_, kNVTENVFP4E4M3Max, &val, sizeof(val)); } + void set_err_corrected_nvfp4(bool err_corrected_nvfp4) { + const auto val = static_cast(err_corrected_nvfp4); + nvte_set_tensor_param_v2(tensor_, kNVTEErrCorrectedNVFP4, &val, sizeof(val)); + } + // Parameter getters NVTEBasicTensor get_parameter(const NVTETensorParam param) const noexcept { @@ -909,6 +928,14 @@ class TensorWrapper { return get_parameter(kNVTEColumnwiseAmax); } + NVTEBasicTensor get_rowwise_data_err() const noexcept { + return get_parameter(kNVTERowwiseDataErr); + } + + NVTEBasicTensor get_rowwise_scale_inv_err() const noexcept { + return get_parameter(kNVTERowwiseScaleInvErr); + } + bool get_with_gemm_swizzled_scales() const { uint8_t val = 0; nvte_get_tensor_param_v2(tensor_, kNVTEWithGEMMSwizzledScales, &val, sizeof(val), nullptr); @@ -927,6 +954,12 @@ class TensorWrapper { return val; } + bool get_err_corrected_nvfp4() const { + uint8_t val = 0; + nvte_get_tensor_param_v2(tensor_, kNVTEErrCorrectedNVFP4, &val, sizeof(val), nullptr); + return static_cast(val); + } + /*! \brief Get an underlying NVTETensor. * * \return NVTETensor held by this TensorWrapper. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8a03f2f51a..2eb7cff7db 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -524,6 +524,10 @@ class NVFP4BlockScaling(Recipe): If set to `True`, forward activation quantizers emit row-scaled NVFP4 tensors. In this mode, rowwise ``amax`` metadata is stored as a vector with one FP32 value per tensor row. + err_corrected_activation : bool, default = False + If set to `True`, forward row-scaled activation quantizers also + emit a second NVFP4 representation of the BF16 quantization + residual. Forward GEMMs add the residual contribution in FP32. nvfp4_4over6 : {'none', 'weights', 'activations', 'all'}, default = 'none' Enable 4over6 adaptive NVFP4 block scaling for selected tensor scopes. For each selected FP4 block, quantization compares @@ -551,6 +555,7 @@ class NVFP4BlockScaling(Recipe): ) disable_2d_quantization: bool = os.getenv("NVTE_NVFP4_DISABLE_2D_QUANTIZATION", "0") == "1" row_scaled_activation: bool = os.getenv("NVTE_NVFP4_ROW_SCALED_ACTIVATION", "0") == "1" + err_corrected_activation: bool = os.getenv("NVTE_NVFP4_ERR_CORRECTED_ACTIVATION", "0") == "1" nvfp4_4over6: str = os.getenv("NVTE_NVFP4_4OVER6", "none") nvfp4_4over6_e4m3_use_256: str = os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") nvfp4_4over6_err_mode: str = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").upper() @@ -578,6 +583,14 @@ def __post_init__(self) -> None: assert ( self.nvfp4_4over6_err_mode in _NVFP4_4OVER6_ERR_MODES ), "NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'." + if self.err_corrected_activation and not self.row_scaled_activation: + raise ValueError( + "NVFP4 error-corrected activation quantization requires row_scaled_activation=True." + ) + if self.err_corrected_activation and self.nvfp4_4over6 in ("activations", "all"): + raise ValueError( + "NVFP4 error-corrected activation quantization does not support activation 4over6." + ) # Quantization params # Note: RHT is currently only applied to column-wise usage so that @@ -607,6 +620,7 @@ def _make_repr(self) -> str: f"fp8_mha={self.fp8_mha}, " f"backward_override={self.backward_override}, " f"row_scaled_activation={self.row_scaled_activation}, " + f"err_corrected_activation={self.err_corrected_activation}, " f"nvfp4_4over6={self.nvfp4_4over6}, " f"nvfp4_4over6_e4m3_use_256={self.nvfp4_4over6_e4m3_use_256}, " f"nvfp4_4over6_err_mode={self.nvfp4_4over6_err_mode}, " diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index b3179d38fd..1d7f7ed86a 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -847,6 +847,19 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256, "Unsupported NVFP4 E4M3 max (got ", t.nvfp4_e4m3_max, ")"); break; + case kNVTERowwiseDataErr: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.data_err = *basic_tensor; + break; + } + case kNVTERowwiseScaleInvErr: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.scale_inv_err = *basic_tensor; + break; + } + case kNVTEErrCorrectedNVFP4: + t.err_corrected_nvfp4 = static_cast(*reinterpret_cast(buf)); + break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } @@ -933,6 +946,19 @@ void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, vo case kNVTENVFP4E4M3Max: std::memcpy(buf, &t->nvfp4_e4m3_max, attr_size); break; + case kNVTERowwiseDataErr: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->data_err); + break; + } + case kNVTERowwiseScaleInvErr: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->scale_inv_err); + break; + } + case kNVTEErrCorrectedNVFP4: + *reinterpret_cast(buf) = static_cast(t->err_corrected_nvfp4); + break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h index c462b30147..e593e55c32 100644 --- a/transformer_engine/common/transpose/cast_transpose.h +++ b/transformer_engine/common/transpose/cast_transpose.h @@ -68,6 +68,7 @@ void quantize_transpose_vector_blockwise_fp4( const bool return_identity, const bool return_transpose, const bool pow2_scale, const bool swizzled_scale, const bool use_stochastic_rounding, const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4, + const bool err_corrected_nvfp4, SimpleTensor &scale_inv_err, SimpleTensor &output_err, const SimpleTensor &noop_tensor, cudaStream_t stream); } // namespace transformer_engine::detail diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index d5f2fa9a2c..ba98b8e0ac 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -316,14 +316,15 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x(const float2 in01, template + bool kApplyStochasticRounding, bool kIs2DBlockScaling, bool kRowScaledNVFP4, + bool kErrCorrectedNVFP4> __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel( const IType* const input, const float* global_amax, OType* const output_c, OType* const output_t, ScaleType* const tile_scales_inv_c, ScaleType* const tile_scales_inv_t, - const size_t row_length, const size_t num_rows, const size_t scale_stride_x, - const size_t scale_stride_y, const size_t scale_t_stride_x, const size_t scale_t_stride_y, - const size_t kScaleBlockDim, const float epsilon, const size_t* rng_state, - const float* noop_ptr) { + OType* const output_err, ScaleType* const tile_scales_inv_err, const size_t row_length, + const size_t num_rows, const size_t scale_stride_x, const size_t scale_stride_y, + const size_t scale_t_stride_x, const size_t scale_t_stride_y, const size_t kScaleBlockDim, + const float epsilon, const size_t* rng_state, const float* noop_ptr) { constexpr int kNVecContainer = kNVecOut / kNFP4PerContainer; using SMemVec = Vec; using OVec = Vec; @@ -437,6 +438,10 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo : 0); // For not aligned case OType* output_g = &output_c[(r_g * row_length + c_g) / kNFP4PerContainer]; // Output address in global memory + OType* output_err_g = nullptr; + if constexpr (kErrCorrectedNVFP4) { + output_err_g = &output_err[(r_g * row_length + c_g) / kNFP4PerContainer]; + } // Each kNumThreadsStore threads form a warp process one row, we need to find the lane id of // the first thread to do the reduction. const unsigned src_lane = @@ -539,6 +544,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } // Step 2.6: Quantize OVec output_vec; + IType residual[kNVecOut]; #pragma unroll for (int i = 0; i < kNVecOut / kNVecSMem; i += 2) { // Pack two elements into __nv_bfloat162 @@ -554,18 +560,82 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo output_vec.data.elt[i] = reinterpret_cast<__nv_fp4x2_storage_t*>(&out_4x)[0]; output_vec.data.elt[i + 1] = reinterpret_cast<__nv_fp4x2_storage_t*>(&out_4x)[1]; + + if constexpr (kErrCorrectedNVFP4) { + const float4 q4 = static_cast(out_4x); + const int residual_offset = i * kNVecSMem; + const float q_values[4] = {q4.x, q4.y, q4.z, q4.w}; + constexpr float fp8_max = TypeExtrema::max; + constexpr float fp4_max = TypeExtrema::max; + constexpr float global_decode_scale_multiplier = 1.0f / (fp4_max * fp8_max); + const float primary_global_decode_scale = + row_idx < num_rows ? global_amax[row_idx] * global_decode_scale_multiplier : 1.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int smem_vec_idx = i + j / kNVecSMem; + const int smem_elt_idx = j % kNVecSMem; + float primary_dequantized_fp32 = q_values[j] * static_cast(scale_inv); + primary_dequantized_fp32 *= primary_global_decode_scale; + const IType primary_dequantized = static_cast(primary_dequantized_fp32); + const float error = static_cast(smem_vec[smem_vec_idx].data.elt[smem_elt_idx]) - + static_cast(primary_dequantized); + residual[residual_offset + j] = static_cast(error); + } + } + } + + OVec output_err_vec; + if constexpr (kErrCorrectedNVFP4) { + float error_amax = 0.0f; +#pragma unroll + for (int i = 0; i < kNVecOut; ++i) { + error_amax = fmaxf(error_amax, fabsf(static_cast(residual[i]))); + } + const ScaleType scale_inv_err_value = + ComputeDecodeScaleFP4(error_amax, row_global_encode_scale_multiplier); + const float encode_scale_err = + ComputeEncodeScaleFP4(scale_inv_err_value, row_global_decode_scale); + if (write_scale_inv) { + const size_t scale_row_idx = block_idx_y * kTileDim + r_s; + const size_t scale_col_idx = block_idx_x * (kNumThreadsStore / kNumThreadsReduce) + + (threadIdx.x % kNumThreadsStore) / kNumThreadsReduce; + tile_scales_inv_err[scale_row_idx * scale_stride_y + scale_col_idx * scale_stride_x] = + scale_inv_err_value; + } +#pragma unroll + for (int i = 0; i < kNVecOut / kNVecSMem; i += 2) { + const float2 error01 = + make_float2(static_cast(residual[2 * i]) * encode_scale_err, + static_cast(residual[2 * i + 1]) * encode_scale_err); + const float2 error23 = + make_float2(static_cast(residual[2 * i + 2]) * encode_scale_err, + static_cast(residual[2 * i + 3]) * encode_scale_err); + const __nv_fp4x4_e2m1 error_4x = cvt_fp32_to_fp4_4x(error01, error23, 0); + output_err_vec.data.elt[i] = reinterpret_cast(&error_4x)[0]; + output_err_vec.data.elt[i + 1] = + reinterpret_cast(&error_4x)[1]; + } } // Step 2.7: Store output_c if constexpr (kAligned) { output_vec.store_to(output_g); + if constexpr (kErrCorrectedNVFP4) { + output_err_vec.store_to(output_err_g); + } } else { if (r_g < num_rows) { output_vec.store_to_elts(output_g, 0, num_ele); + if constexpr (kErrCorrectedNVFP4) { + output_err_vec.store_to_elts(output_err_g, 0, num_ele); + } } } // Step 2.8: Update output address, row index of shared memory (and row index of global memory // for not aligned case) output_g += stride_g / kNFP4PerContainer; + if constexpr (kErrCorrectedNVFP4) { + output_err_g += stride_g / kNFP4PerContainer; + } r_s += r_stride; if constexpr (!kAligned) { r_g += r_stride; @@ -778,6 +848,7 @@ void quantize_transpose_vector_blockwise_fp4( const bool return_identity, const bool return_transpose, const bool pow2_scale, const bool swizzled_scale, const bool use_stochastic_rounding, const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4, + const bool err_corrected_nvfp4, SimpleTensor& scale_inv_err, SimpleTensor& output_err, const SimpleTensor& noop_tensor, cudaStream_t stream) { NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4); #if CUDA_VERSION >= 12080 @@ -793,6 +864,24 @@ void quantize_transpose_vector_blockwise_fp4( "Row-scaled NVFP4 quantization only supports rowwise quantization."); NVTE_CHECK(!row_scaled_nvfp4 || !use_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!err_corrected_nvfp4 || row_scaled_nvfp4, + "Error-corrected NVFP4 quantization requires row scaling."); + NVTE_CHECK(!err_corrected_nvfp4 || input.dtype == DType::kBFloat16, + "Error-corrected NVFP4 quantization requires BF16 input."); + NVTE_CHECK(!err_corrected_nvfp4 || (!use_stochastic_rounding && !return_transpose), + "Error-corrected NVFP4 quantization only supports deterministic rowwise output."); + NVTE_CHECK(!err_corrected_nvfp4 || !swizzled_scale, + "Error-corrected NVFP4 quantization requires compact block scales."); + NVTE_CHECK(!err_corrected_nvfp4 || (output_err.dptr != nullptr && scale_inv_err.dptr != nullptr), + "Error-corrected NVFP4 output data and scales must be allocated."); + NVTE_CHECK(!err_corrected_nvfp4 || output_err.dtype == output.dtype, + "Error-corrected NVFP4 primary and residual data types must match."); + NVTE_CHECK(!err_corrected_nvfp4 || scale_inv_err.dtype == scale_inv.dtype, + "Error-corrected NVFP4 primary and residual scale types must match."); + NVTE_CHECK(!err_corrected_nvfp4 || output_err.shape == output.shape, + "Error-corrected NVFP4 primary and residual data shapes must match."); + NVTE_CHECK(!err_corrected_nvfp4 || scale_inv_err.shape == scale_inv.shape, + "Error-corrected NVFP4 primary and residual scale shapes must match."); const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u; size_t num_elements = row_length; @@ -875,38 +964,45 @@ void quantize_transpose_vector_blockwise_fp4( TRANSFORMER_ENGINE_SWITCH_CONDITION( row_scaled_nvfp4, kRowScaledNVFP4, - size_t smem_bytes = kSMemSize * sizeof(InputType); - auto kernel = block_scaled_1d_cast_transpose_kernel< - kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned, - float, InputType, OutputType, ScaleType, kSwizzledScale, - kApplyStochasticRounding, kIs2DBlockScaling, - kRowScaledNVFP4>; - if (smem_bytes >= 48 * 1024) { - cudaError_t err = cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_bytes); - NVTE_CHECK(err == cudaSuccess, - "Failed to set dynamic shared memory size."); - } kernel<<>>( - reinterpret_cast(input.dptr), - reinterpret_cast(global_amax.dptr), - reinterpret_cast(output.dptr), - reinterpret_cast(output_t.dptr), - reinterpret_cast(scale_inv.dptr), - reinterpret_cast(scale_inv_t.dptr), - row_length, num_rows, scale_stride_x, scale_stride_y, - scale_t_stride_x, scale_t_stride_y, kScaleBlockDim, - epsilon, rng_state, - noop_ptr);) // kRowScaledNVFP4 - ) // kIs2DBlockScaling - ) // kApplyStochasticRounding - ) // kSwizzledScale - ) // kAligned - ) // kReturnTranspose - ) // kReturnIdentity - ) // OutputType - ) // InputType + TRANSFORMER_ENGINE_SWITCH_CONDITION( + err_corrected_nvfp4, kErrCorrectedNVFP4, + + size_t smem_bytes = kSMemSize * sizeof(InputType); + auto kernel = block_scaled_1d_cast_transpose_kernel< + kReturnIdentity, kReturnTranspose, kPow2Scale, + kAligned, float, InputType, OutputType, ScaleType, + kSwizzledScale, kApplyStochasticRounding, + kIs2DBlockScaling, kRowScaledNVFP4, + kErrCorrectedNVFP4>; + if (smem_bytes >= 48 * 1024) { + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes); + NVTE_CHECK(err == cudaSuccess, + "Failed to set dynamic shared memory size."); + } kernel<<>>( + reinterpret_cast(input.dptr), + reinterpret_cast(global_amax.dptr), + reinterpret_cast(output.dptr), + reinterpret_cast(output_t.dptr), + reinterpret_cast(scale_inv.dptr), + reinterpret_cast(scale_inv_t.dptr), + reinterpret_cast(output_err.dptr), + reinterpret_cast(scale_inv_err.dptr), + row_length, num_rows, scale_stride_x, scale_stride_y, + scale_t_stride_x, scale_t_stride_y, kScaleBlockDim, + epsilon, rng_state, + noop_ptr);) // kErrCorrectedNVFP4 + ) // kRowScaledNVFP4 + ) // kIs2DBlockScaling + ) // kApplyStochasticRounding + ) // kSwizzledScale + ) // kAligned + ) // kReturnTranspose + ) // kReturnIdentity + ) // OutputType + ) // InputType NVTE_CHECK_CUDA(cudaGetLastError()); #else diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3b066d50b..777d2f6ee2 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -80,24 +80,50 @@ def _nvfp4_row_scaled_gemm_inputs( B: NVFP4TensorStorage, *, transa: bool, -) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor]: - """Return GEMM aliases and FP32 output scales for row-scaled NVFP4.""" +) -> Tuple[ + NVFP4TensorStorage, + NVFP4TensorStorage, + Optional[NVFP4TensorStorage], + torch.Tensor, +]: + """Return GEMM aliases and FP32 output scales for row-scaled NVFP4. + + The aliases replace the global amax metadata with one so both the primary + and residual GEMMs produce unscaled FP32 partials. The caller applies the + shared row/weight global scale after summing those partials. + """ A_metadata = A.get_metadata() weight_amax = A._amax_rowwise if transa else A._amax_columnwise - assert weight_amax is not None and weight_amax.numel() == 1 + if weight_amax is None or weight_amax.numel() != 1: + raise RuntimeError("Row-scaled NVFP4 GEMM requires scalar weight amax metadata.") A_metadata["amax_rowwise" if transa else "amax_columnwise"] = weight_amax.new_ones(1) A_metadata["row_scaled_nvfp4"] = False B_metadata = B.get_metadata() rhs_rowwise_amax = B._amax_rowwise - assert rhs_rowwise_amax is not None + if rhs_rowwise_amax is None: + raise RuntimeError("Row-scaled NVFP4 B is missing rowwise amax metadata.") B_metadata["amax_rowwise"] = rhs_rowwise_amax.new_ones(1) B_metadata["row_scaled_nvfp4"] = False - - assert rhs_rowwise_amax.dtype == torch.float32 and weight_amax.dtype == torch.float32 + B_metadata["rowwise_data_err"] = None + B_metadata["rowwise_scale_inv_err"] = None + B_metadata["err_corrected_nvfp4"] = False + + error_B = None + if B._err_corrected_nvfp4: + if B._rowwise_data_err is None or B._rowwise_scale_inv_err is None: + raise RuntimeError("Error-corrected NVFP4 B is missing residual data or scales.") + error_metadata = dict(B_metadata) + error_metadata["rowwise_data"] = B._rowwise_data_err + error_metadata["rowwise_scale_inv"] = B._rowwise_scale_inv_err + error_B = NVFP4TensorStorage(**error_metadata) + + if rhs_rowwise_amax.dtype != torch.float32 or weight_amax.dtype != torch.float32: + raise RuntimeError("Row-scaled NVFP4 GEMM requires FP32 global amax metadata.") return ( NVFP4TensorStorage(**A_metadata), NVFP4TensorStorage(**B_metadata), + error_B, (rhs_rowwise_amax * weight_amax).view(-1, 1), ) @@ -236,7 +262,9 @@ def general_gemm( ), "Row-scaled NVFP4 GEMM currently requires NVFP4 A." # cuBLAS folds NVFP4 global amax values into GEMM alpha. Keep the row-scaled # recipe's global scales out of alpha and apply them in FP32 below. - gemm_A, gemm_B, rowwise_global_scales = _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa) + gemm_A, gemm_B, error_B, rowwise_global_scales = _nvfp4_row_scaled_gemm_inputs( + A, B, transa=transa + ) requested_out, requested_out_dtype = out, out_dtype fp32_out = ( @@ -252,6 +280,18 @@ def general_gemm( gemm_args[6] = TE_DType[torch.float32] # out_dtype gemm_args[7] = None # bias out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **kwargs) + + if error_B is not None: + # Accumulate the residual GEMM before applying the shared global + # row/weight scale. Bias and the final dtype conversion remain a + # single epilogue below. + error_gemm_args = list(gemm_args) + error_gemm_args[2] = error_B + error_gemm_args[4] = out + error_gemm_args[14] = True + error_gemm_kwargs = dict(kwargs) + error_gemm_kwargs["beta"] = 1.0 + out, _, _, _ = tex.generic_gemm(*error_gemm_args, **error_gemm_kwargs) out_2d = out.reshape(-1, out.shape[-1]) assert rowwise_global_scales.dtype == torch.float32 and out.dtype == torch.float32 diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index aa0e0c87fe..c9c7080eb9 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -354,6 +354,8 @@ class NVFP4Quantizer : public Quantizer { int nvfp4_e4m3_max; // Whether tensors emitted by this quantizer use row-scaled NVFP4 metadata. bool row_scaled_nvfp4; + // Whether row-scaled tensors also carry a quantized BF16 residual. + bool err_corrected_nvfp4; int rht_matrix_random_sign_mask_t; at::Tensor rht_matrix; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8d77a9e349..aae3baf5af 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -80,7 +80,8 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob // Post-quantize swizzle for quantizers whose kernel does not bake // the GEMM-swizzled scale layout in directly - if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales()) { + if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales() && + !output_cpp.get_err_corrected_nvfp4()) { inplace_swizzle_scale_for_gemm(output_py); } @@ -984,6 +985,7 @@ std::tuple, std::vector, bool> bulk_alloc const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage; const bool row_scaled_nvfp4 = quantizer_cpp_list[0]->row_scaled_nvfp4; + const bool err_corrected_nvfp4 = quantizer_cpp_list[0]->err_corrected_nvfp4; const bool nvfp4_use_4over6 = quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max; @@ -993,6 +995,8 @@ std::tuple, std::vector, bool> bulk_alloc NVTE_CHECK(!columnwise_usage, "Row-scaled NVFP4 bulk allocation does not support columnwise usage."); } + NVTE_CHECK(!err_corrected_nvfp4, + "Error-corrected NVFP4 bulk/grouped quantization is not supported."); const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto fp4_dtype = quantizer_cpp_list[0]->dtype; diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4308464c52..04100a0b5b 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1884,6 +1884,7 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize NVTE_ERROR("Unsupported NVFP4 4over6 error mode: ", nvfp4_4over6_err_mode); } this->row_scaled_nvfp4 = quantizer.attr("row_scaled_nvfp4").cast(); + this->err_corrected_nvfp4 = quantizer.attr("err_corrected_nvfp4").cast(); // Get amax reduction group if needed for NVFP4 AG const bool with_amax_reduction = quantizer.attr("with_amax_reduction").cast(); @@ -1908,10 +1909,15 @@ void NVFP4Quantizer::set_quantization_params(TensorWrapper* tensor) const { auto columnwise_data = tensor->get_columnwise_data(); columnwise_data.dtype = static_cast(this->dtype); + auto rowwise_data_err = tensor->get_rowwise_data_err(); + rowwise_data_err.dtype = static_cast(this->dtype); + tensor->set_rowwise_data(rowwise_data.data_ptr, static_cast(rowwise_data.dtype), rowwise_data.shape); tensor->set_columnwise_data(columnwise_data.data_ptr, static_cast(columnwise_data.dtype), columnwise_data.shape); + tensor->set_rowwise_data_err(rowwise_data_err.data_ptr, + static_cast(rowwise_data_err.dtype), rowwise_data_err.shape); } bool NVFP4Quantizer::is_eligible_for_rht_cast_fusion(const std::vector& shape, @@ -1968,6 +1974,7 @@ std::pair NVFP4Quantizer::create_tensor( "NVFP4 requires tensor dims that are divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool err_corrected_nvfp4 = this->err_corrected_nvfp4; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -1975,11 +1982,18 @@ std::pair NVFP4Quantizer::create_tensor( NVTE_CHECK(!columnwise_usage, "Row-scaled NVFP4 quantization does not support columnwise usage."); } + if (err_corrected_nvfp4) { + NVTE_CHECK(row_scaled_nvfp4, "Error-corrected NVFP4 quantization requires row-scaled NVFP4."); + NVTE_CHECK(dtype == DType::kBFloat16, + "Error-corrected NVFP4 quantization requires BF16 input."); + NVTE_CHECK(!nvfp4_use_4over6, "Error-corrected NVFP4 quantization does not support 4over6."); + } const auto rowwise_scale_inv_shape = get_scale_shape(shape, false); const auto columnwise_scale_inv_shape = get_scale_shape(shape, true); // Allocate tensors at::Tensor rowwise_data_tensor, rowwise_scale_inv_tensor, amax_rowwise; + at::Tensor rowwise_data_err_tensor, rowwise_scale_inv_err_tensor; at::Tensor columnwise_data_tensor, columnwise_scale_inv_tensor, amax_columnwise; const auto bit8_tensor_opts = at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); @@ -1994,6 +2008,10 @@ std::pair NVFP4Quantizer::create_tensor( // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed amax_rowwise = at::empty({amax_rows}, bit32_tensor_opts); + if (err_corrected_nvfp4) { + rowwise_data_err_tensor = at::empty(convert_shape_for_fp4(shape_int64), bit8_tensor_opts); + rowwise_scale_inv_err_tensor = at::empty(scale_inv_shape_int64, bit8_tensor_opts); + } } if (columnwise_usage) { const std::vector scale_inv_shape_int64(columnwise_scale_inv_shape.begin(), @@ -2021,6 +2039,8 @@ std::pair NVFP4Quantizer::create_tensor( auto columnwise_scale_inv_py = py_cast(columnwise_scale_inv_tensor, columnwise_usage); auto amax_rowwise_py = py_cast(amax_rowwise, rowwise_usage); auto amax_columnwise_py = py_cast(amax_columnwise, columnwise_usage); + auto rowwise_data_err_py = py_cast(rowwise_data_err_tensor, err_corrected_nvfp4); + auto rowwise_scale_inv_err_py = py_cast(rowwise_scale_inv_err_tensor, err_corrected_nvfp4); // Construct Python NVFP4 tensor py::object out_py; @@ -2037,6 +2057,9 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); + kwargs["rowwise_data_err"] = rowwise_data_err_py; + kwargs["rowwise_scale_inv_err"] = rowwise_scale_inv_err_py; + kwargs["err_corrected_nvfp4"] = py::cast(err_corrected_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); kwargs["fake_dtype"] = GetATenDType(dtype); @@ -2069,6 +2092,9 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["device"] = py::cast(device); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); + kwargs["rowwise_data_err"] = rowwise_data_err_py; + kwargs["rowwise_scale_inv_err"] = rowwise_scale_inv_err_py; + kwargs["err_corrected_nvfp4"] = py::cast(err_corrected_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); py::tuple args(0); @@ -2089,6 +2115,11 @@ std::pair NVFP4Quantizer::create_tensor( out_cpp.set_rowwise_scale_inv(rowwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, rowwise_scale_inv_shape); out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); + if (err_corrected_nvfp4) { + out_cpp.set_rowwise_data_err(rowwise_data_err_tensor.data_ptr(), DType::kFloat4E2M1, shape); + out_cpp.set_rowwise_scale_inv_err(rowwise_scale_inv_err_tensor.data_ptr(), DType::kFloat8E4M3, + getTensorShape(rowwise_scale_inv_err_tensor)); + } } if (columnwise_usage) { // enforce 2D shape to avoid [S, B, H] shape and B and be 1 @@ -2104,6 +2135,7 @@ std::pair NVFP4Quantizer::create_tensor( } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); + out_cpp.set_err_corrected_nvfp4(err_corrected_nvfp4); out_cpp.set_nvfp4_e4m3_max(nvfp4_e4m3_max); this->set_quantization_params(&out_cpp); @@ -2118,6 +2150,9 @@ std::pair NVFP4Quantizer::create_grouped_tenso const size_t logical_last_dim) const { using namespace pybind11::literals; + NVTE_CHECK(!this->err_corrected_nvfp4, + "Error-corrected NVFP4 grouped quantization is not supported."); + const auto tensor_offsets = resolve_grouped_tensor_offsets(num_tensors, first_dims, last_dims, precomputed_tensor_offsets, logical_first_dim, logical_last_dim); @@ -2268,6 +2303,8 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( }; auto rowwise_data = get_tensor("_rowwise_data"); auto rowwise_scale_inv = get_tensor("_rowwise_scale_inv"); + auto rowwise_data_err = get_tensor("_rowwise_data_err"); + auto rowwise_scale_inv_err = get_tensor("_rowwise_scale_inv_err"); auto columnwise_data = get_tensor("_columnwise_data"); auto columnwise_scale_inv = get_tensor("_columnwise_scale_inv"); auto amax_rowwise = get_tensor("_amax_rowwise"); @@ -2295,6 +2332,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const bool with_gemm_swizzled_scales = nvfp4_emits_gemm_swizzled_scales(*this, shape); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool err_corrected_nvfp4 = this->err_corrected_nvfp4; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -2302,7 +2340,12 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( NVTE_CHECK(!columnwise_usage, "Row-scaled NVFP4 quantization does not support columnwise usage."); } + if (err_corrected_nvfp4) { + NVTE_CHECK(row_scaled_nvfp4, "Error-corrected NVFP4 quantization requires row-scaled NVFP4."); + NVTE_CHECK(!nvfp4_use_4over6, "Error-corrected NVFP4 quantization does not support 4over6."); + } tensor.attr("_row_scaled_nvfp4") = row_scaled_nvfp4; + tensor.attr("_err_corrected_nvfp4") = err_corrected_nvfp4; tensor.attr("_with_gemm_swizzled_scales") = with_gemm_swizzled_scales; tensor.attr("_nvfp4_use_4over6") = py::cast(nvfp4_use_4over6); tensor.attr("_nvfp4_e4m3_max") = py::cast(nvfp4_e4m3_max); @@ -2323,6 +2366,27 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( rowwise_scale_inv = at::empty(scale_inv_shape_int64, opts); tensor.attr("_rowwise_scale_inv") = *rowwise_scale_inv; } + if (err_corrected_nvfp4) { + if (!rowwise_data_err) { + const std::vector shape_int64(shape.begin(), shape.end()); + const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + rowwise_data_err = at::empty(convert_shape_for_fp4(shape_int64), opts); + tensor.attr("_rowwise_data_err") = *rowwise_data_err; + } + if (!rowwise_scale_inv_err) { + const auto scale_inv_shape = get_scale_shape(shape, false); + const std::vector scale_inv_shape_int64(scale_inv_shape.begin(), + scale_inv_shape.end()); + const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + rowwise_scale_inv_err = at::empty(scale_inv_shape_int64, opts); + tensor.attr("_rowwise_scale_inv_err") = *rowwise_scale_inv_err; + } + } else { + rowwise_data_err.reset(); + rowwise_scale_inv_err.reset(); + tensor.attr("_rowwise_data_err") = py::none(); + tensor.attr("_rowwise_scale_inv_err") = py::none(); + } const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; if (!amax_rowwise || amax_rowwise->numel() != amax_rows) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); @@ -2344,6 +2408,10 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( amax_rowwise.reset(); tensor.attr("_amax_rowwise") = py::none(); } + rowwise_data_err.reset(); + rowwise_scale_inv_err.reset(); + tensor.attr("_rowwise_data_err") = py::none(); + tensor.attr("_rowwise_scale_inv_err") = py::none(); } // Coerce column-wise data @@ -2395,6 +2463,11 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*rowwise_scale_inv)); out_cpp.set_amax(amax_rowwise->data_ptr(), DType::kFloat32, getTensorShape(*amax_rowwise)); + if (err_corrected_nvfp4) { + out_cpp.set_rowwise_data_err(rowwise_data_err->data_ptr(), DType::kFloat4E2M1, shape); + out_cpp.set_rowwise_scale_inv_err(rowwise_scale_inv_err->data_ptr(), DType::kFloat8E4M3, + getTensorShape(*rowwise_scale_inv_err)); + } } if (columnwise_usage) { // enforce 2D shape to avoid [S, B, H] shape and B and be 1 @@ -2410,6 +2483,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); + out_cpp.set_err_corrected_nvfp4(err_corrected_nvfp4); out_cpp.set_nvfp4_e4m3_max(nvfp4_e4m3_max); this->set_quantization_params(&out_cpp); @@ -2554,6 +2628,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou const auto [rows, cols] = get_2d_dims(input.shape()); const bool row_scaled_nvfp4 = out.get_row_scaled_nvfp4(); + const bool err_corrected_nvfp4 = out.get_err_corrected_nvfp4(); if (row_scaled_nvfp4) { NVTE_CHECK(!this->with_rht, "Row-scaled NVFP4 quantization does not support RHT."); NVTE_CHECK(!this->with_2d_quantization, @@ -2564,6 +2639,17 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou "Row-scaled NVFP4 quantization does not support amax reduction."); NVTE_CHECK(cols % 16 == 0, "Row-scaled NVFP4 quantization requires last dim divisible by 16."); } + if (err_corrected_nvfp4) { + NVTE_CHECK(row_scaled_nvfp4, "Error-corrected NVFP4 quantization requires row-scaled NVFP4."); + NVTE_CHECK(input.dtype() == DType::kBFloat16, + "Error-corrected NVFP4 quantization requires BF16 input."); + NVTE_CHECK(this->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "Error-corrected NVFP4 quantization does not support 4over6."); + NVTE_CHECK(out.get_rowwise_data_err().data_ptr != nullptr, + "Error-corrected NVFP4 output data must be allocated."); + NVTE_CHECK(out.get_rowwise_scale_inv_err().data_ptr != nullptr, + "Error-corrected NVFP4 output scales must be allocated."); + } // Restriction for the RHT cast fusion kernel because we are using MMA hardware for computing RHT const bool eligible_for_rht_cast_fusion = diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index ddb85808a5..60ec3971d2 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -135,6 +135,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) const bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); const bool with_gemm_swizzled_scales = tensor.attr("_with_gemm_swizzled_scales").cast(); const bool row_scaled_nvfp4 = tensor.attr("_row_scaled_nvfp4").cast(); + const bool err_corrected_nvfp4 = tensor.attr("_err_corrected_nvfp4").cast(); const int nvfp4_e4m3_max = tensor.attr("_nvfp4_e4m3_max").cast(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for NVFP4 Tensor."); @@ -148,6 +149,14 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) convert_shape_back_from_fp4(getTensorShape(data), false)); ret.set_rowwise_scale_inv(scale_inv.data_ptr(), DType::kFloat8E4M3, getTensorShape(scale_inv)); ret.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); + if (err_corrected_nvfp4) { + const auto &data_err = tensor.attr("_rowwise_data_err").cast(); + const auto &scale_inv_err = tensor.attr("_rowwise_scale_inv_err").cast(); + ret.set_rowwise_data_err(data_err.data_ptr(), dtype, + convert_shape_back_from_fp4(getTensorShape(data_err), false)); + ret.set_rowwise_scale_inv_err(scale_inv_err.data_ptr(), DType::kFloat8E4M3, + getTensorShape(scale_inv_err)); + } } // Column-scaled data @@ -166,6 +175,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) // Scale layout ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); ret.set_row_scaled_nvfp4(row_scaled_nvfp4); + ret.set_err_corrected_nvfp4(err_corrected_nvfp4); ret.set_nvfp4_e4m3_max(nvfp4_e4m3_max); // Quantizer state diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py index d09b95ace3..9db2d7d778 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py @@ -219,6 +219,8 @@ class NVFP4TensorRef(QuantizedTensorStorage): scale: Optional[torch.Tensor] = None data_t: Optional[torch.Tensor] = None scale_t: Optional[torch.Tensor] = None + data_err: Optional[torch.Tensor] = None + scale_err: Optional[torch.Tensor] = None global_amax_row: Optional[torch.Tensor] = None global_amax_col: Optional[torch.Tensor] = None nvfp4_use_4over6: bool = False @@ -239,11 +241,13 @@ def prepare_for_saving( self, ) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]: """Prepare the quantization result for saving for backward""" - tensors = [self.data, self.data_t, self.scale, self.scale_t] + tensors = [self.data, self.data_t, self.scale, self.scale_t, self.data_err, self.scale_err] self.data = None self.data_t = None self.scale = None self.scale_t = None + self.data_err = None + self.scale_err = None return tensors, self def restore_from_saved( @@ -254,7 +258,9 @@ def restore_from_saved( self.data_t = tensors[1] self.scale = tensors[2] self.scale_t = tensors[3] - return tensors[4:] + self.data_err = tensors[4] + self.scale_err = tensors[5] + return tensors[6:] # Compatibility @property @@ -352,6 +358,7 @@ def __init__( eps: float = 0.0, quant_tile_shape: Tuple[int, int] = (1, 16), row_scaled_nvfp4: bool = False, + err_corrected_nvfp4: bool = False, nvfp4_use_4over6: bool = False, nvfp4_e4m3_max: int = 448, nvfp4_4over6_err_mode: str = "MAE", @@ -367,6 +374,10 @@ def __init__( raise ValueError( "Row-scaled NVFP4 reference quantization does not support columnwise usage." ) + if err_corrected_nvfp4 and not row_scaled_nvfp4: + raise ValueError("Error-corrected NVFP4 reference requires row scaling.") + if err_corrected_nvfp4 and nvfp4_use_4over6: + raise ValueError("Error-corrected NVFP4 reference does not support 4over6.") if nvfp4_use_4over6: if nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError(f"Unsupported NVFP4 4over6 error mode: {nvfp4_4over6_err_mode}.") @@ -382,6 +393,7 @@ def __init__( self.eps = eps self.quant_tile_shape = quant_tile_shape self.row_scaled_nvfp4 = row_scaled_nvfp4 + self.err_corrected_nvfp4 = err_corrected_nvfp4 self.nvfp4_use_4over6 = nvfp4_use_4over6 self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 if self.nvfp4_e4m3_max not in (448, 256): @@ -833,6 +845,8 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ Optional[torch.Tensor], torch.Tensor, torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], ]: """ Python implementation of microblock FP4 quantization. @@ -844,8 +858,10 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ Returns ------- - Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor] - (qx, sx, qx_t, sx_t, global_amax) where: + Tuple containing quantized primary/error data, block scales, and global amax values. + + ``(qx, sx, qx_t, sx_t, global_amax_row, global_amax_col, + qx_err, sx_err)`` where: - qx: quantized data in row-major order (if rowwise_usage), None otherwise - sx: scale tensor for qx (if rowwise_usage), None otherwise - qx_t: quantized data in column-major order (if columnwise_usage), None otherwise @@ -920,9 +936,41 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ qx = self._rm_pad_tensor(qx, (M, N // 2)) + qx_err = None + sx_err = None + if self.err_corrected_nvfp4: + if tensor.dtype != torch.bfloat16: + raise ValueError("Error-corrected NVFP4 reference requires BF16 input.") + qx_dequantized = cast_from_fp4x2(qx, torch.float32) + block_scales = sx.to(torch.float32).repeat_interleave( + self.quant_tile_shape[1], dim=1 + ) + global_decode_scale = global_amax_row.view(M, 1) / (6.0 * 448.0) + primary_dequantized = (qx_dequantized * block_scales * global_decode_scale).to( + torch.bfloat16 + ) + error = (row_input - primary_dequantized).to(torch.bfloat16) + error_padded = self._pad_tensor( + error, + row_divisor=self.quant_tile_shape[0], + col_divisor=self.quant_tile_shape[1], + ) + qx_err, sx_err = self._quantize_blockwise_reference( + error_padded, + global_amax_row, + self.quant_tile_shape[1], + self.quant_tile_shape[0], + pow_2_scales=False, + row_scaled_nvfp4=True, + eps=self.eps, + ) + qx_err = self._rm_pad_tensor(qx_err, (M, N // 2)) + else: qx = None sx = None + qx_err = None + sx_err = None if self.columnwise_usage: x_t = col_input @@ -951,7 +999,7 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ qx_t = None sx_t = None - return qx, sx, qx_t, sx_t, global_amax_row, global_amax_col + return qx, sx, qx_t, sx_t, global_amax_row, global_amax_col, qx_err, sx_err def quantize( self, @@ -970,13 +1018,17 @@ def quantize( if tensor.ndim > 2: tensor = tensor.view(-1, tensor.shape[-1]) - qx, sx, qx_t, sx_t, global_amax_row, global_amax_col = self._quantize(tensor) + qx, sx, qx_t, sx_t, global_amax_row, global_amax_col, qx_err, sx_err = self._quantize( + tensor + ) return NVFP4TensorRef( data=qx, scale=sx, data_t=qx_t, scale_t=sx_t, + data_err=qx_err, + scale_err=sx_err, global_amax_row=global_amax_row, global_amax_col=global_amax_col, nvfp4_use_4over6=self.nvfp4_use_4over6, @@ -1019,13 +1071,15 @@ def update_quantized( if src.ndim > 2: src = src.view(-1, src.shape[-1]) - qx, sx, qx_t, sx_t, global_amax_row, global_amax_col = self._quantize(src) + qx, sx, qx_t, sx_t, global_amax_row, global_amax_col, qx_err, sx_err = self._quantize(src) # Update the destination with new data dst.data = qx dst.scale = sx dst.data_t = qx_t dst.scale_t = sx_t + dst.data_err = qx_err + dst.scale_err = sx_err dst.global_amax_row = global_amax_row dst.global_amax_col = global_amax_col dst.nvfp4_use_4over6 = self.nvfp4_use_4over6 @@ -1089,6 +1143,13 @@ def qgemm( high_precision_x = cast_from_fp4x2(qx, out_dtype) high_precision_w = cast_from_fp4x2(qw, out_dtype) + high_precision_x_err = None + sx_err = None + if qresult_x is not None and qresult_x.data_err is not None: + if qresult_x.scale_err is None: + raise ValueError("qresult_x has error data but no error block scales") + high_precision_x_err = cast_from_fp4x2(qresult_x.data_err, out_dtype) + sx_err = qresult_x.scale_err.to(torch.float32) if self.pow_2_scales: @@ -1198,6 +1259,10 @@ def qgemm( raise ValueError( f"sw shape mismatch: expected ({N}, {K // block_length}), got {sw.shape}" ) + if sx_err is not None and sx_err.shape != (M, K // block_length): + raise ValueError( + f"sx_err shape mismatch: expected ({M}, {K // block_length}), got {sx_err.shape}" + ) y = torch.zeros(M, N, dtype=torch.float32, device=qx.device) @@ -1219,6 +1284,11 @@ def qgemm( y += torch.outer(sx_block, sw_block) * high_precision_gemm_ref( qx_block, qw_block, torch.float32, is_b_transposed=True ) + if high_precision_x_err is not None: + qx_err_block = high_precision_x_err[:, k_start:k_end].clone().contiguous() + y += torch.outer(sx_err[:, k], sw_block) * high_precision_gemm_ref( + qx_err_block, qw_block, torch.float32, is_b_transposed=True + ) if not self.pow_2_scales and K > 0: # only apply global scale for NVFP4 and non-empty cases diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f1bffb122f..7722790d73 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1714,6 +1714,11 @@ def _make(tensor_type: str) -> NVFP4Quantizer: and tensor_type != "weight" and self.recipe.row_scaled_activation ), + err_corrected_nvfp4=( + self.mode == "forward" + and tensor_type != "weight" + and self.recipe.err_corrected_activation + ), nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.recipe.nvfp4_4over6_err_mode, diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ccf06ac166..06821b53d2 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -129,6 +129,8 @@ class NVFP4Quantizer(Quantizer): """Whether emitted NVFP4 tensors store one FP32 amax per row.""" row_scaled_nvfp4: bool + """Whether emitted row-scaled tensors include an NVFP4 residual.""" + err_corrected_nvfp4: bool """Whether to use NVFP4 4over6 map-to-4/map-to-6 block selection.""" nvfp4_use_4over6: bool """Global E4M3 scale bound used by emitted NVFP4 tensors.""" @@ -151,6 +153,7 @@ def __init__( with_2d_quantization: bool = False, stochastic_rounding: bool = False, row_scaled_nvfp4: bool = False, + err_corrected_nvfp4: bool = False, nvfp4_use_4over6: bool = False, nvfp4_e4m3_max: int = 448, nvfp4_4over6_err_mode: str = "MAE", @@ -165,6 +168,11 @@ def __init__( self.with_2d_quantization = with_2d_quantization self.stochastic_rounding = stochastic_rounding self.row_scaled_nvfp4 = row_scaled_nvfp4 + self.err_corrected_nvfp4 = err_corrected_nvfp4 + if self.err_corrected_nvfp4 and not self.row_scaled_nvfp4: + raise ValueError("Error-corrected NVFP4 requires row_scaled_nvfp4=True.") + if self.err_corrected_nvfp4 and nvfp4_use_4over6: + raise ValueError("Error-corrected NVFP4 does not support 4over6 quantization.") self.nvfp4_use_4over6 = nvfp4_use_4over6 self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 if self.nvfp4_e4m3_max not in (448, 256): @@ -240,6 +248,7 @@ def copy(self) -> NVFP4Quantizer: with_2d_quantization=self.with_2d_quantization, stochastic_rounding=self.stochastic_rounding, row_scaled_nvfp4=self.row_scaled_nvfp4, + err_corrected_nvfp4=self.err_corrected_nvfp4, nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, @@ -400,6 +409,9 @@ def __new__( quantizer: Quantizer, with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, + rowwise_data_err: Optional[torch.Tensor] = None, + rowwise_scale_inv_err: Optional[torch.Tensor] = None, + err_corrected_nvfp4: bool = False, nvfp4_use_4over6: bool = False, nvfp4_e4m3_max: int = 448, **kwargs, @@ -417,6 +429,9 @@ def __new__( with_gemm_swizzled_scales, *args, row_scaled_nvfp4=row_scaled_nvfp4, + rowwise_data_err=rowwise_data_err, + rowwise_scale_inv_err=rowwise_scale_inv_err, + err_corrected_nvfp4=err_corrected_nvfp4, nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, **kwargs, @@ -487,11 +502,18 @@ def clone(self) -> NVFP4Tensor: columnwise_data = None if self._columnwise_data is not None: columnwise_data = self._columnwise_data.detach().clone() + rowwise_data_err = None + rowwise_scale_inv_err = None + if self._rowwise_data_err is not None: + rowwise_data_err = self._rowwise_data_err.detach().clone() + rowwise_scale_inv_err = self._rowwise_scale_inv_err.detach().clone() return _IdentityFunc.apply( self, { "rowwise_data": rowwise_data, "columnwise_data": columnwise_data, + "rowwise_data_err": rowwise_data_err, + "rowwise_scale_inv_err": rowwise_scale_inv_err, }, ) @@ -732,6 +754,8 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): qt._columnwise_scale_inv, qt._amax_rowwise, qt._amax_columnwise, + qt._rowwise_data_err, + qt._rowwise_scale_inv_err, ): if t is not None and t.is_cuda: t.record_stream(stream) @@ -751,6 +775,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dst._amax_rowwise.copy_(src._amax_rowwise.detach()) if dst._amax_columnwise is not None and src._amax_columnwise is not None: dst._amax_columnwise.copy_(src._amax_columnwise.detach()) + if dst._rowwise_data_err is not None and src._rowwise_data_err is not None: + dst._rowwise_data_err.copy_(src._rowwise_data_err.detach()) + dst._rowwise_scale_inv_err.copy_(src._rowwise_scale_inv_err.detach()) return dst # NVFP4 dequantize not supported. Add manual support for needed funcs. @@ -767,8 +794,19 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): tensor._rowwise_scale_inv, *args[1:], **kwargs ) amax_rowwise = torch.zeros_like(tensor._amax_rowwise, *args[1:], **kwargs) + rowwise_data_err = ( + data_init_func(tensor._rowwise_data_err, *args[1:], **kwargs) + if tensor._rowwise_data_err is not None + else None + ) + rowwise_scale_inv_err = ( + scale_inv_init_func(tensor._rowwise_scale_inv_err, *args[1:], **kwargs) + if tensor._rowwise_scale_inv_err is not None + else None + ) else: rowwise_data, rowwise_scale_inv, amax_rowwise = None, None, None + rowwise_data_err, rowwise_scale_inv_err = None, None if tensor._columnwise_data is not None: columnwise_data = data_init_func(tensor._columnwise_data, *args[1:], **kwargs) @@ -798,6 +836,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, row_scaled_nvfp4=tensor._row_scaled_nvfp4, + rowwise_data_err=rowwise_data_err, + rowwise_scale_inv_err=rowwise_scale_inv_err, + err_corrected_nvfp4=tensor._err_corrected_nvfp4, nvfp4_use_4over6=tensor._nvfp4_use_4over6, nvfp4_e4m3_max=tensor._nvfp4_e4m3_max, ) @@ -824,6 +865,9 @@ def __reduce_ex__(self, protocol: int) -> tuple: self._row_scaled_nvfp4, self._nvfp4_use_4over6, self._nvfp4_e4m3_max, + self._rowwise_data_err, + self._rowwise_scale_inv_err, + self._err_corrected_nvfp4, ), ) @@ -917,6 +961,9 @@ def _set_data(self, tensor: torch.Tensor) -> None: self._amax_columnwise = tensor._amax_columnwise self._with_gemm_swizzled_scales = tensor._with_gemm_swizzled_scales self._row_scaled_nvfp4 = tensor._row_scaled_nvfp4 + self._rowwise_data_err = tensor._rowwise_data_err + self._rowwise_scale_inv_err = tensor._rowwise_scale_inv_err + self._err_corrected_nvfp4 = tensor._err_corrected_nvfp4 self._nvfp4_use_4over6 = tensor._nvfp4_use_4over6 self._nvfp4_e4m3_max = tensor._nvfp4_e4m3_max return @@ -975,6 +1022,9 @@ def _make_nvfp4_tensor_in_reduce_ex( row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, nvfp4_e4m3_max: int = 448, + rowwise_data_err: Optional[torch.Tensor] = None, + rowwise_scale_inv_err: Optional[torch.Tensor] = None, + err_corrected_nvfp4: bool = False, ) -> NVFP4Tensor: """Reconstruct an ``NVFP4Tensor`` from its ``__reduce_ex__`` payload.""" # Infer device from whichever inner buffer is populated so the wrapper @@ -1000,6 +1050,9 @@ def _make_nvfp4_tensor_in_reduce_ex( with_gemm_swizzled_scales=with_gemm_swizzled_scales, device=device, row_scaled_nvfp4=row_scaled_nvfp4, + rowwise_data_err=rowwise_data_err, + rowwise_scale_inv_err=rowwise_scale_inv_err, + err_corrected_nvfp4=err_corrected_nvfp4, nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, ) @@ -1086,6 +1139,13 @@ def forward( with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, row_scaled_nvfp4=tensor._row_scaled_nvfp4, + rowwise_data_err=( + tensor._rowwise_data_err.view(list(shape[:-1]) + [shape[-1] // 2]) + if tensor._rowwise_data_err is not None + else None + ), + rowwise_scale_inv_err=tensor._rowwise_scale_inv_err, + err_corrected_nvfp4=tensor._err_corrected_nvfp4, nvfp4_use_4over6=tensor._nvfp4_use_4over6, nvfp4_e4m3_max=tensor._nvfp4_e4m3_max, ) @@ -1132,6 +1192,13 @@ def backward( with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, device=grad.device, row_scaled_nvfp4=grad._row_scaled_nvfp4, + rowwise_data_err=( + grad._rowwise_data_err.view(list(ctx.shape[:-1]) + [ctx.shape[-1] // 2]) + if grad._rowwise_data_err is not None + else None + ), + rowwise_scale_inv_err=grad._rowwise_scale_inv_err, + err_corrected_nvfp4=grad._err_corrected_nvfp4, nvfp4_use_4over6=grad._nvfp4_use_4over6, nvfp4_e4m3_max=grad._nvfp4_e4m3_max, ) @@ -1220,6 +1287,13 @@ def forward( with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, row_scaled_nvfp4=tensor._row_scaled_nvfp4, + rowwise_data_err=( + tensor._rowwise_data_err.reshape(list(shape[:-1]) + [shape[-1] // 2]) + if tensor._rowwise_data_err is not None + else None + ), + rowwise_scale_inv_err=tensor._rowwise_scale_inv_err, + err_corrected_nvfp4=tensor._err_corrected_nvfp4, nvfp4_use_4over6=tensor._nvfp4_use_4over6, nvfp4_e4m3_max=tensor._nvfp4_e4m3_max, ) @@ -1266,6 +1340,13 @@ def backward( with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, device=grad.device, row_scaled_nvfp4=grad._row_scaled_nvfp4, + rowwise_data_err=( + grad._rowwise_data_err.reshape(list(ctx.shape[:-1]) + [ctx.shape[-1] // 2]) + if grad._rowwise_data_err is not None + else None + ), + rowwise_scale_inv_err=grad._rowwise_scale_inv_err, + err_corrected_nvfp4=grad._err_corrected_nvfp4, nvfp4_use_4over6=grad._nvfp4_use_4over6, nvfp4_e4m3_max=grad._nvfp4_e4m3_max, ) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 09f040ba67..8bee93869f 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -86,6 +86,10 @@ class NVFP4TensorStorage(QuantizedTensorStorage): _columnwise_data: Optional[torch.Tensor] # Block scaling factors for row-scaled FP4 data _rowwise_scale_inv: torch.Tensor + # Packed FP4 representation of the BF16 primary-quantization residual + _rowwise_data_err: Optional[torch.Tensor] + # Block scaling factors for the error-correction FP4 data + _rowwise_scale_inv_err: Optional[torch.Tensor] # Block scaling factors for column-scaled FP4 data _columnwise_scale_inv: torch.Tensor # Input absolute maximum value (used to compute tensor scale for @@ -104,6 +108,8 @@ class NVFP4TensorStorage(QuantizedTensorStorage): _with_gemm_swizzled_scales: bool # Whether this NVFP4 tensor uses row-scaled amax metadata _row_scaled_nvfp4: bool + # Whether error-correction data and scales are populated + _err_corrected_nvfp4: bool # Whether this NVFP4 tensor uses 4over6 map-to-4/map-to-6 block selection _nvfp4_use_4over6: bool # Global E4M3 scale bound used by this NVFP4 tensor @@ -123,6 +129,9 @@ def __new__( *args, fake_dtype: Optional[torch.dtype] = None, row_scaled_nvfp4: bool = False, + rowwise_data_err: Optional[torch.Tensor] = None, + rowwise_scale_inv_err: Optional[torch.Tensor] = None, + err_corrected_nvfp4: bool = False, nvfp4_use_4over6: bool = False, nvfp4_e4m3_max: int = 448, **kwargs, @@ -143,6 +152,14 @@ def __new__( instance._amax_columnwise = amax_columnwise instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales instance._row_scaled_nvfp4 = row_scaled_nvfp4 + instance._rowwise_data_err = rowwise_data_err + instance._rowwise_scale_inv_err = rowwise_scale_inv_err + instance._err_corrected_nvfp4 = err_corrected_nvfp4 + if err_corrected_nvfp4: + if not row_scaled_nvfp4: + raise ValueError("Error-corrected NVFP4 requires row-scaled metadata.") + if rowwise_data_err is None or rowwise_scale_inv_err is None: + raise ValueError("Error-corrected NVFP4 requires error data and block scales.") instance._nvfp4_use_4over6 = nvfp4_use_4over6 instance._nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 @@ -157,6 +174,8 @@ def clear(self): self._columnwise_scale_inv, self._amax_rowwise, self._amax_columnwise, + self._rowwise_data_err, + self._rowwise_scale_inv_err, ): if t is not None: t.data = _empty_tensor() @@ -171,6 +190,8 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: raise RuntimeError("Scale layout mismatch in copy_from_storage") if self._row_scaled_nvfp4 != src._row_scaled_nvfp4: raise RuntimeError("Rowwise amax scaling mode mismatch in copy_from_storage") + if self._err_corrected_nvfp4 != src._err_corrected_nvfp4: + raise RuntimeError("NVFP4 error-correction mode mismatch in copy_from_storage") if self._nvfp4_use_4over6 != src._nvfp4_use_4over6: raise RuntimeError("NVFP4 4over6 mode mismatch in copy_from_storage") if self._nvfp4_e4m3_max != src._nvfp4_e4m3_max: @@ -186,6 +207,8 @@ def _copy_optional(dst: Optional[torch.Tensor], src_tensor: Optional[torch.Tenso _copy_optional(self._columnwise_scale_inv, src._columnwise_scale_inv) _copy_optional(self._amax_rowwise, src._amax_rowwise) _copy_optional(self._amax_columnwise, src._amax_columnwise) + _copy_optional(self._rowwise_data_err, src._rowwise_data_err) + _copy_optional(self._rowwise_scale_inv_err, src._rowwise_scale_inv_err) def get_metadata(self) -> Dict[str, Any]: """Get this tensor's metadata.""" @@ -200,6 +223,9 @@ def get_metadata(self) -> Dict[str, Any]: "quantizer": self._quantizer, "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, "row_scaled_nvfp4": self._row_scaled_nvfp4, + "rowwise_data_err": self._rowwise_data_err, + "rowwise_scale_inv_err": self._rowwise_scale_inv_err, + "err_corrected_nvfp4": self._err_corrected_nvfp4, "nvfp4_use_4over6": self._nvfp4_use_4over6, "nvfp4_e4m3_max": self._nvfp4_e4m3_max, "fake_dtype": self._dtype, @@ -214,6 +240,8 @@ def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], NVFP4TensorS self._columnwise_scale_inv, self._amax_rowwise, self._amax_columnwise, + self._rowwise_data_err, + self._rowwise_scale_inv_err, ] self._rowwise_data = None self._columnwise_data = None @@ -221,6 +249,8 @@ def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], NVFP4TensorS self._columnwise_scale_inv = None self._amax_rowwise = None self._amax_columnwise = None + self._rowwise_data_err = None + self._rowwise_scale_inv_err = None return tensors, self def restore_from_saved( @@ -233,10 +263,14 @@ def restore_from_saved( self._columnwise_scale_inv = tensors[3] self._amax_rowwise = tensors[4] self._amax_columnwise = tensors[5] - return tensors[6:] + self._rowwise_data_err = tensors[6] + self._rowwise_scale_inv_err = tensors[7] + return tensors[8:] def get_data_tensors(self): """Get this Tensor's data.""" + if self._rowwise_data_err is not None: + return self._rowwise_data, self._columnwise_data, self._rowwise_data_err return self._rowwise_data, self._columnwise_data def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: @@ -335,6 +369,13 @@ def view(self, shape: torch.Size): fp4_dtype=self._fp4_dtype, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, row_scaled_nvfp4=self._row_scaled_nvfp4, + rowwise_data_err=( + self._rowwise_data_err.view(byte_shape) + if self._rowwise_data_err is not None + else None + ), + rowwise_scale_inv_err=self._rowwise_scale_inv_err, + err_corrected_nvfp4=self._err_corrected_nvfp4, nvfp4_use_4over6=self._nvfp4_use_4over6, nvfp4_e4m3_max=self._nvfp4_e4m3_max, fake_dtype=self._dtype,