diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 18f9d3f772..6adb641372 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2854,9 +2854,18 @@ def test_grouped_mlp_activation_recompute( if bias: torch.testing.assert_close(on["fc1_db"], off["fc1_db"], **tols) torch.testing.assert_close(on["fc2_db"], off["fc2_db"], **tols) - assert_close_rms( - on["fc2_dw"], off["fc2_dw"], rtol=self._RECOMPUTE_FC2_WGRAD_RMS_TOL[quantization] - ) + # The dSReLU kernel regenerates fc2_x bit-exactly from the saved pre-activation and + # the NVFP4 backward re-quantizes it with the forward amax, so FC2's weight gradient + # matches bitwise. Plain SReLU needs a cuDNN frontend whose Hadamard-fused kernel + # rounds the accumulator through C's dtype before the activation; older versions + # regenerate a tensor that differs by a few BF16 ulps, hence the tolerance. MXFP8 + # quantizes the regenerated tensor inside the dSReLU kernel with its own scales. + if quantization == "nvfp4_rht" and activation == "scaled_tanh_srelu": + torch.testing.assert_close(on["fc2_dw"], off["fc2_dw"], rtol=0, atol=0) + else: + assert_close_rms( + on["fc2_dw"], off["fc2_dw"], rtol=self._RECOMPUTE_FC2_WGRAD_RMS_TOL[quantization] + ) @pytest.mark.parametrize("quantization", ("mxfp8", "nvfp4_rht")) def test_grouped_mlp_caller_buffers( diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 7bcd8430ce..02bf677344 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -427,6 +427,18 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q const at::Tensor &columnwise_amax, std::optional tensor_offsets); +/*! @brief Per-group NVFP4 amax of a grouped BF16 tensor without quantizing it. + * + * Returns (pre-RHT amax, post-RHT amax), one float32 value per group, as the group + * Hadamard-amax kernel computes them for the RHT recipe. Lets a caller quantize row-wise + * now and column-wise later from the same amax, e.g. when the column-wise copy is + * regenerated in the backward pass. + */ +std::tuple nvfp4_group_hadamard_amax( + const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, + std::optional first_dims, std::optional last_dims, + std::optional tensor_offsets); + py::object group_dequantize(const py::handle &input, DType otype); py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8c9b6b1079..3191d1b7bd 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -691,6 +691,74 @@ py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle q return py::reinterpret_borrow(grouped_output_py); } +std::tuple nvfp4_group_hadamard_amax( + const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, + std::optional first_dims, std::optional last_dims, + std::optional tensor_offsets) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + NVTE_CHECK(!last_dims.has_value(), + "nvfp4_group_hadamard_amax: varying last dim is not supported with NVFP4."); + NVTE_CHECK(first_dims.has_value() && tensor_offsets.has_value(), + "nvfp4_group_hadamard_amax: first_dims and tensor_offsets are required."); + NVTE_CHECK(first_dims->is_cuda() && first_dims->scalar_type() == at::kLong && + first_dims->numel() == static_cast(num_tensors), + "nvfp4_group_hadamard_amax: first_dims must be an int64 CUDA tensor with one " + "entry per group."); + NVTE_CHECK(tensor_offsets->is_cuda() && tensor_offsets->scalar_type() == at::kLong && + tensor_offsets->numel() == static_cast(num_tensors + 1), + "nvfp4_group_hadamard_amax: tensor_offsets must be an int64 CUDA tensor with " + "num_tensors + 1 entries."); + + auto quantizer_cpp = convert_quantizer(quantizer); + NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()), + "nvfp4_group_hadamard_amax only supports NVFP4 quantizers."); + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + NVTE_CHECK(nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax, + "nvfp4_group_hadamard_amax expects the NVFP4 RHT recipe with post-RHT amax."); + + std::vector logical_shape; + for (const auto &d : tensor.sizes()) { + logical_shape.push_back(d); + } + const auto float_opts = at::TensorOptions().dtype(at::kFloat).device(tensor.device()); + auto rowwise_amax = at::zeros({static_cast(num_tensors)}, float_opts); + auto columnwise_amax = at::zeros({static_cast(num_tensors)}, float_opts); + if (logical_shape[0] == 0 || logical_shape[1] == 0) { + return {rowwise_amax, columnwise_amax}; + } + + const auto input_dtype = GetTransformerEngineDType(tensor.scalar_type()); + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_input_tensor.set_rowwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); + + // The amax launcher derives which amax to compute from the storage the output declares: + // row-wise data requests the pre-RHT amax and column-wise data the post-RHT (transposed) + // amax. It only ever writes the amax buffers, so alias the input as both and allocate no + // quantized storage. + auto grouped_amax_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_amax_tensor.set_rowwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); + grouped_amax_tensor.set_columnwise_data(tensor.data_ptr(), input_dtype, getTensorShape(tensor)); + grouped_amax_tensor.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(rowwise_amax)); + grouped_amax_tensor.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32, + getTensorShape(columnwise_amax)); + grouped_amax_tensor.set_first_dims(first_dims->data_ptr(), DType::kInt64, + getTensorShape(*first_dims)); + grouped_amax_tensor.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_amax_graph_safe( + grouped_input_tensor.data(), grouped_amax_tensor.data(), 0, + nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, at::cuda::getCurrentCUDAStream()); + }); + + return {rowwise_amax, columnwise_amax}; +} + py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 6fbed1436f..2c25aae11b 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -187,6 +187,10 @@ void bind_quantize_with_amax_extensions(py::module_ &m) { py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"), py::arg("tensor_offsets") = py::none()); + m.def("nvfp4_group_hadamard_amax", nvfp4_group_hadamard_amax, py::arg("tensor"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(), + "Per-group (pre-RHT, post-RHT) NVFP4 amax of a grouped tensor without quantizing it"); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d29ae1fd43..4fed03a2f4 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -126,6 +126,55 @@ def _cudnn_frontend_supports_grouped_gemm_srelu_tanh() -> bool: return False +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_supports_grouped_gemm_srelu_tanh_hadamard() -> bool: + """Whether the cuDNN FE grouped GEMM + Hadamard wrapper accepts tanh_clamp_scale. + + With it, the clamped SReLU uses the fused RHT-amax kernel like the unclamped one + instead of the plain SReLU kernel plus a separate amax pass. Detected by signature + for the same reason as _cudnn_frontend_supports_grouped_gemm_srelu_tanh. + """ + if not ( + _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() + and _cudnn_frontend_supports_grouped_gemm_srelu_tanh() + ): + return False + try: + from cudnn import ( # pylint: disable=import-outside-toplevel + grouped_gemm_glu_hadamard_wrapper_sm100, + ) + except ImportError: + return False + try: + signature = inspect.signature(grouped_gemm_glu_hadamard_wrapper_sm100) + return "tanh_clamp_scale" in signature.parameters + except (TypeError, ValueError): + return False + + +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_supports_grouped_gemm_hadamard_sign_mask() -> bool: + """Whether the cuDNN FE grouped GEMM + Hadamard wrapper accepts rht_sign_mask. + + The kernel's post-RHT amax must be computed for the same randomized Hadamard + transform the NVFP4 quantizer applies; without the mask it describes the plain + transform and the column-wise cast gets a wrong global scale. + """ + if not _cudnn_frontend_supports_grouped_gemm_srelu_hadamard(): + return False + try: + from cudnn import ( # pylint: disable=import-outside-toplevel + grouped_gemm_glu_hadamard_wrapper_sm100, + ) + except ImportError: + return False + try: + signature = inspect.signature(grouped_gemm_glu_hadamard_wrapper_sm100) + return "rht_sign_mask" in signature.parameters + except (TypeError, ValueError): + return False + + @functools.lru_cache(maxsize=None) def _cudnn_frontend_supports_grouped_gemm_situglu() -> bool: """Feature-detect complete cuDNN frontend grouped SiTU-GLU support.""" @@ -384,6 +433,41 @@ def _group_quantize_with_amax_for_grouped_mlp( ) +def _fc2_x_amax_reuse_enabled() -> bool: + """Whether the backward reuses the forward NVFP4 amax of a recomputed fc2 input. + + The grouped dSReLU kernel regenerates fc2_x bit-exactly from the saved pre-activation, + so the per-group amax computed in the forward is also the amax of the regenerated + tensor. Reusing it turns the backward re-quantization into a single cast kernel. + ``NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX=0`` disables the reuse. + """ + return int(os.environ.get("NVTE_FUSED_GROUPED_MLP_REUSE_FC2_X_AMAX", "1")) > 0 + + +def _nvfp4_group_hadamard_amax( + tensor: torch.Tensor, + quantizer: Quantizer, + num_groups: int, + split_sizes: Optional[torch.Tensor], + *, + tensor_offsets: Optional[torch.Tensor] = None, +) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Per-group (pre-RHT, post-RHT) NVFP4 amax of ``tensor`` without quantizing it. + + Returns ``None`` when the extension lacks the kernel binding or the grouped metadata + is incomplete; callers then quantize with amax computed on the fly. + """ + amax_fn = getattr(tex, "nvfp4_group_hadamard_amax", None) + if ( + amax_fn is None + or not isinstance(quantizer, NVFP4Quantizer) + or split_sizes is None + or tensor_offsets is None + ): + return None + return amax_fn(tensor, quantizer, num_groups, split_sizes, tensor_offsets=tensor_offsets) + + def _weight_quantizers(op: GroupedLinear, num_groups: int) -> list[Quantizer]: """Per-expert weight quantizers of ``op``. @@ -1437,7 +1521,14 @@ def fuser_forward( if unit_activation_scale and num_groups != 1: unit_activation_scale = False - activation_is_srelu = isinstance(activation_op, ScaledSReLU) + # The hadamard kernel is the GLU one; it applies the SReLU soft clamp only in + # cuDNN frontends whose wrapper accepts tanh_clamp_scale. Without that, + # ScaledTanhSReLU keeps the plain SReLU kernel and a separate RHT amax pass -- + # a performance limitation, not a correctness one. + activation_is_srelu = isinstance(activation_op, ScaledSReLU) or ( + isinstance(activation_op, ScaledTanhSReLU) + and _cudnn_frontend_supports_grouped_gemm_srelu_tanh_hadamard() + ) supports_single_group_runtime_offsets = ( _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) ) @@ -1658,6 +1749,13 @@ def fuser_forward( if self.grouped_gemm_act_hadamard_kernel() is None: # Kernel is not available pass + elif ( + fc2_input_quantizer.rht_matrix_random_sign_mask_t != 0 + and not _cudnn_frontend_supports_grouped_gemm_hadamard_sign_mask() + ): + # Kernel would compute the post-RHT amax for the plain Hadamard + # transform while the quantizer applies a randomized one + pass elif self._cudnn_act_func in ("swiglu", "situglu"): kernel_impl = "gemm_act_rht_amax" elif activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard(): @@ -1785,6 +1883,10 @@ def fuser_forward( elif kernel_impl == "gemm_act_rht_amax": fc1_activation_kwargs["d_dtype"] = torch.bfloat16 fc1_activation_kwargs["use_tmem_post_rht_amax"] = _use_tmem_post_rht_amax() + if _cudnn_frontend_supports_grouped_gemm_hadamard_sign_mask(): + fc1_activation_kwargs["rht_sign_mask"] = ( + fc2_input_quantizer.rht_matrix_random_sign_mask_t + ) elif kernel_impl == "gemm_act_rht_quant": fc1_activation_kwargs["d_dtype"] = torch.float4_e2m1fn_x2 fc1_activation_kwargs["rht_colwise_dtype"] = torch.float4_e2m1fn_x2 @@ -1936,36 +2038,83 @@ def fuser_forward( ): grouped_fc2_weight._with_gemm_swizzled_scales = False + # Decide up front whether the backward rebuilds fc2_x from the dSReLU kernel + # (which applies the tanh clamp when regenerating fc2_x). When it does, the + # forward only needs row-wise fc2_x for the fc2 GEMM; the column-wise copy + # for wgrad is produced in the backward from the recomputed tensor. + recompute_srelu_fc2_x = ( + isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)) + and bool(getattr(activation_op, "activation_recompute_in_mlp", False)) + and weight_requires_grad + and _grouped_gemm_dsrelu_backward_supported() + and _nvidia_cudnn_frontend_supports_wgrad() + ) + # Forward (pre-RHT, post-RHT) amax of fc2_x, kept for its column-wise re-quantization + # in the backward when the regenerated tensor is bit-identical to this one. + fc2_x_amaxes_for_recompute = None + reuse_fc2_x_amax = recompute_srelu_fc2_x and _fc2_x_amax_reuse_enabled() + # Unpack kernel outputs if use_nvfp4: - fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc2_input_quantizer.set_usage( + rowwise=True, + columnwise=weight_requires_grad and not recompute_srelu_fc2_x, + ) fc2_input_quantizer.optimize_for_gemm = True if kernel_impl == "gemm_act": - # Quantize to NVFP4 + # Quantize to NVFP4. When the backward rebuilds fc2_x, compute the amax + # here with the kernel the row-wise quantization needs anyway and keep + # it, so the backward re-quantization is only the column-wise cast. fc2_in = fc1_kernel_out["d_tensor"] fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() - grouped_fc2_x = _group_quantize_for_grouped_mlp( - fc2_in, - fc2_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - use_dense_single_group=use_dense_single_group, - ) + fc2_x_amaxes = None + if reuse_fc2_x_amax: + fc2_x_amaxes = _nvfp4_group_hadamard_amax( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + ) + if fc2_x_amaxes is not None: + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + *fc2_x_amaxes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) + fc2_x_amaxes_for_recompute = fc2_x_amaxes + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) elif kernel_impl == "gemm_act_rht_amax": - # Quantize to NVFP4 using precomputed amax + # Quantize to NVFP4 using the amax the fc1 kernel computed fc2_in = fc1_kernel_out["d_tensor"] fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + fc2_x_amaxes = ( + fc1_kernel_out["amax_tensor"].view(-1), + fc1_kernel_out["post_rht_amax_tensor"].view(-1), + ) grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( fc2_in, fc2_input_quantizer, num_groups, split_sizes, - fc1_kernel_out["amax_tensor"].view(-1), - fc1_kernel_out["post_rht_amax_tensor"].view(-1), + *fc2_x_amaxes, tensor_offsets=fc2_x_tensor_offsets, use_dense_single_group=use_dense_single_group, ) + if reuse_fc2_x_amax: + fc2_x_amaxes_for_recompute = fc2_x_amaxes elif kernel_impl == "gemm_act_rht_quant": # Unpack NVFP4 output fc2_in_row_data = fc1_kernel_out["d_tensor"] @@ -2266,22 +2415,12 @@ def fuser_forward( # Save state for backward pass if requires_grad: saved_fc1_x = grouped_fc1_x - mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) - activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() - # The dSReLU kernel applies the tanh clamp when regenerating fc2_x. - activation_is_srelu = isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)) - activation_recompute_in_mlp = bool( - getattr(activation_op, "activation_recompute_in_mlp", False) - ) - recompute_srelu_fc2_x = ( - activation_is_srelu - and activation_recompute_in_mlp - and weight_requires_grad - and _grouped_gemm_dsrelu_backward_supported() - and _nvidia_cudnn_frontend_supports_wgrad() - ) + # recompute_srelu_fc2_x was decided before the fc2 input quantization above. + # A recomputed fc2_x is row-wise only and is not saved, so only mark what + # actually goes to save_for_backward. saved_grouped_fc2_x = None if recompute_srelu_fc2_x else grouped_fc2_x + mark_grouped_tensor(saved_fc1_x, activation_in, scales, saved_grouped_fc2_x) # wgrad only needs column-wise data, so clear row-wise data for grouped_fc_x in (saved_fc1_x, saved_grouped_fc2_x): @@ -2346,6 +2485,7 @@ def fuser_forward( fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad fc2_ctx.recompute_input_from_dsrelu = recompute_srelu_fc2_x + fc2_ctx.fc2_x_amaxes_for_recompute = fc2_x_amaxes_for_recompute return fc2_out, [(), (), ()] @@ -2833,14 +2973,29 @@ def fuser_backward( fc2_input_quantizer = fc2_ctx.input_quantizers[0] fc2_input_quantizer.set_usage(rowwise=False, columnwise=True) fc2_input_quantizer.optimize_for_gemm = True - grouped_fc2_x = _group_quantize_for_grouped_mlp( - fc2_x_bf16, - fc2_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - use_dense_single_group=use_dense_single_group, - ) + fc2_x_amaxes = getattr(fc2_ctx, "fc2_x_amaxes_for_recompute", None) + fc2_ctx.fc2_x_amaxes_for_recompute = None + if fc2_x_amaxes is not None: + # The regenerated fc2_x is bit-identical to the forward one, so its + # forward amax applies and the Hadamard amax pass is skipped. + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_x_bf16, + fc2_input_quantizer, + num_groups, + split_sizes, + *fc2_x_amaxes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_x_bf16, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, + ) else: sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") if sfd_col_d_srelu_tensor is None: