diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d195eb2f78..28f7360641 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -755,6 +755,13 @@ def test_grouped_mlp( maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) if dtype == torch.bfloat16 and not is_bf16_available(): pytest.skip("BF16 requires SM 8.0+") + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) if single_grouped_weight and quantization != "mxfp8": pytest.skip("single_grouped_weight is only supported for MXFP8 quantization") if single_grouped_bias and not bias: @@ -1026,8 +1033,10 @@ def _make_module(): or ( quantization == "nvfp4_rht" and dtype == torch.bfloat16 - and activation == "scaled_srelu" - and glu_interleave_size is None + and ( + (not activation_is_glu and glu_interleave_size is None) + or (activation_is_glu and glu_interleave_size == 32) + ) ) ) if expected_grouped_mlp_fusion: @@ -1266,6 +1275,8 @@ def test_grouped_mlp_single_weight_numerics( ) -> None: """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP.""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0": + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -1584,6 +1595,8 @@ def test_grouped_mlp_overwrite_main_grad( that read ``.grad`` don't see stale bytes from the cached dummy). """ + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight: + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -1715,6 +1728,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( ) -> None: """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight: + pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if dtype not in (torch.bfloat16, torch.float16): diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 95f02c64f0..8248a63680 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -290,6 +290,31 @@ py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float l py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer, float limit, float alpha, float glu_linear_offset); + +/* Scaled activation */ +py::object scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, int64_t glu_interleave_size); + +py::object scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size); + +py::object scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer); + +py::tuple scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + int64_t glu_interleave_size, bool compute_scale_grad); + +py::tuple scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, float limit, + float alpha, float glu_linear_offset, int64_t glu_interleave_size, + bool compute_scale_grad); + +py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + bool compute_scale_grad); + /*************************************************************************************************** * LayerNorm **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 58a8f84f85..544ff92c1b 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -342,5 +342,150 @@ py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py:: glu_linear_offset); } +/* Scaled activation helpers (activation + per-row scale via nvte_scaled_*). */ + +template +at::Tensor scaled_activation_compute(const at::Tensor& input, const at::Tensor& act_scales, + int shape_divisor, Args&&... args) { + init_extension(); + NVTE_CHECK(input.dim() >= 1, "scaled activation input must have at least 1 dimension"); + NVTE_CHECK(shape_divisor > 0 && input.size(-1) % shape_divisor == 0, + "scaled activation input width is not compatible with activation"); + + auto input_tensor = input.contiguous(); + auto scales_tensor = act_scales.contiguous().reshape({-1}); + const int64_t rows = input_tensor.numel() / input_tensor.size(-1); + NVTE_CHECK(scales_tensor.numel() == rows, "scaled activation expects one scale per input row"); + + std::vector output_sizes(input_tensor.sizes().begin(), input_tensor.sizes().end()); + output_sizes.back() /= shape_divisor; + auto output = at::empty(output_sizes, input_tensor.options()); + + const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor); + const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor); + const TensorWrapper& output_nvte = makeTransformerEngineTensor(output); + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + act_func(input_nvte.data(), scales_nvte.data(), output_nvte.data(), std::forward(args)..., + stream); + }); + return output; +} + +template +std::tuple scaled_dactivation_compute(const at::Tensor& grad, + const at::Tensor& input, + const at::Tensor& act_scales, + bool compute_scale_grad, + Args&&... args) { + init_extension(); + NVTE_CHECK(input.dim() >= 1 && grad.dim() >= 1, + "scaled dactivation input and grad must have at least 1 dimension"); + + auto grad_tensor = grad.contiguous(); + auto input_tensor = input.contiguous(); + auto scales_tensor = act_scales.contiguous(); + const int64_t rows = input_tensor.numel() / input_tensor.size(-1); + NVTE_CHECK(scales_tensor.numel() == rows, "scaled dactivation expects one scale per input row"); + + auto scales_flat = scales_tensor.reshape({-1}); + auto grad_input = at::empty_like(input_tensor); + auto grad_scales = compute_scale_grad ? at::empty_like(scales_tensor) : at::Tensor(); + auto grad_scales_flat = compute_scale_grad ? grad_scales.reshape({-1}) : at::Tensor(); + + const TensorWrapper& grad_nvte = makeTransformerEngineTensor(grad_tensor); + const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor); + const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_flat); + const TensorWrapper& grad_input_nvte = makeTransformerEngineTensor(grad_input); + std::optional grad_scales_nvte; + if (compute_scale_grad) { + grad_scales_nvte.emplace(makeTransformerEngineTensor(grad_scales_flat)); + } + + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + dact_func(grad_nvte.data(), input_nvte.data(), scales_nvte.data(), grad_input_nvte.data(), + compute_scale_grad ? grad_scales_nvte->data() : nullptr, std::forward(args)..., + stream); + }); + return {grad_input, grad_scales}; +} + +py::object maybe_quantize(const at::Tensor& tensor, py::handle quantizer) { + if (quantizer.is_none()) { + return py::cast(tensor); + } + auto quantizer_cpp = convert_quantizer(quantizer); + const TensorWrapper& tensor_nvte = makeTransformerEngineTensor(tensor); + const auto shape_te = tensor_nvte.shape(); + const std::vector shape(shape_te.data, shape_te.data + shape_te.ndim); + auto fake_dtype = GetTransformerEngineDType(tensor.scalar_type()); + auto [out_nvte, out_py] = quantizer_cpp->create_tensor(shape, fake_dtype); + quantizer_cpp->quantize(tensor_nvte, out_nvte); + return out_py; +} + +template +py::object scaled_activation_helper(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, int shape_divisor, Args&&... args) { + auto output = scaled_activation_compute(input, act_scales, shape_divisor, + std::forward(args)...); + return maybe_quantize(output, quantizer); +} + +template +py::tuple scaled_dactivation_helper(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + bool compute_scale_grad, Args&&... args) { + auto [grad_input, grad_scales] = scaled_dactivation_compute( + grad, input, act_scales, compute_scale_grad, std::forward(args)...); + return py::make_tuple(maybe_quantize(grad_input, quantizer), + compute_scale_grad ? py::cast(grad_scales) : py::none()); +} + +py::object scaled_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, int64_t glu_interleave_size) { + return scaled_activation_helper(input, act_scales, quantizer, + /*shape_divisor=*/2, glu_interleave_size); +} + +py::object scaled_clamped_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, float limit, float alpha, + float glu_linear_offset, int64_t glu_interleave_size) { + return scaled_activation_helper( + input, act_scales, quantizer, /*shape_divisor=*/2, limit, alpha, glu_linear_offset, + glu_interleave_size); +} + +py::object scaled_srelu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer) { + return scaled_activation_helper(input, act_scales, quantizer, + /*shape_divisor=*/1); +} + +py::tuple scaled_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + int64_t glu_interleave_size, bool compute_scale_grad) { + return scaled_dactivation_helper(grad, input, act_scales, quantizer, + compute_scale_grad, glu_interleave_size); +} + +py::tuple scaled_clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, float limit, + float alpha, float glu_linear_offset, int64_t glu_interleave_size, + bool compute_scale_grad) { + return scaled_dactivation_helper( + grad, input, act_scales, quantizer, compute_scale_grad, limit, alpha, glu_linear_offset, + glu_interleave_size); +} + +py::tuple scaled_dsrelu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + bool compute_scale_grad) { + return scaled_dactivation_helper(grad, input, act_scales, quantizer, + compute_scale_grad); +} + } // namespace pytorch } // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c6bf7eb516..e406dc3446 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -285,6 +285,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f); + /* Scaled activation */ + m.def("scaled_swiglu", transformer_engine::pytorch::scaled_swiglu, "Scaled SwiGLU activation", + py::arg("input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("glu_interleave_size") = 0); + m.def("scaled_clamped_swiglu", transformer_engine::pytorch::scaled_clamped_swiglu, + "Scaled clamped SwiGLU activation", py::arg("input"), py::arg("act_scales"), + py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, + py::arg("glu_linear_offset") = 1.0f, py::arg("glu_interleave_size") = 0); + m.def("scaled_srelu", transformer_engine::pytorch::scaled_srelu, "Scaled SReLU activation", + py::arg("input"), py::arg("act_scales"), py::arg("quantizer")); + m.def("scaled_dswiglu", transformer_engine::pytorch::scaled_dswiglu, "Scaled SwiGLU backward", + py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true); + m.def("scaled_clamped_dswiglu", transformer_engine::pytorch::scaled_clamped_dswiglu, + "Scaled clamped SwiGLU backward", py::arg("grad"), py::arg("fwd_input"), + py::arg("act_scales"), py::arg("quantizer"), py::arg("limit") = 7.0f, + py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f, + py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true); + m.def("scaled_dsrelu", transformer_engine::pytorch::scaled_dsrelu, "Scaled SReLU backward", + py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), + py::arg("compute_scale_grad") = true); /* DBias + DAct fusions*/ m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer")); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 12497d0cb0..9860d48237 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -561,7 +561,7 @@ def _make_grouped_tensor( *, num_gemms: int, split_sizes: torch.Tensor, - base_split_offsets: torch.Tensor, + tensor_offsets: torch.Tensor, last_dim: int, dtype: torch.dtype, ) -> GroupedTensorStorage: @@ -573,7 +573,7 @@ def _make_grouped_tensor( quantizer=None, data=data.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * last_dim, + tensor_offsets=tensor_offsets, ) @staticmethod @@ -704,8 +704,18 @@ def _forward_grouped_tensor( out_features = weights[0].size(0) weight_requires_grad = weights[0].requires_grad - split_sizes = m_splits.to(device=device) - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_sizes, ( + base_split_offsets, + input_tensor_offsets, + output_tensor_offsets, + ) = tex.splits_to_offsets_multi( + m_splits, + device, + strides=[1, in_features, out_features], + include_leading_zero=[True, True, True], + dtypes=[torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) inp_view = inp.reshape(-1, in_features) x = cast_if_needed(inp_view, activation_dtype) @@ -716,13 +726,19 @@ def _forward_grouped_tensor( columnwise=is_grad_enabled and weight_requires_grad, ) input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes) + grouped_x = tex.group_quantize( + x, + input_quantizer, + num_gemms, + split_sizes, + tensor_offsets=input_tensor_offsets, + ) else: grouped_x = _GroupedLinear._make_grouped_tensor( x, num_gemms=num_gemms, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=input_tensor_offsets, last_dim=in_features, dtype=activation_dtype, ) @@ -751,7 +767,7 @@ def _forward_grouped_tensor( out, num_gemms=num_gemms, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=output_tensor_offsets, last_dim=out_features, dtype=activation_dtype, ) @@ -796,6 +812,8 @@ def _forward_grouped_tensor( *weights_to_save, split_sizes, base_split_offsets, + input_tensor_offsets, + output_tensor_offsets, ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects @@ -1217,6 +1235,8 @@ def _backward_grouped_tensor( weights = saved_tensors[1 : 1 + N] split_sizes = saved_tensors[1 + N] base_split_offsets = saved_tensors[2 + N] + input_tensor_offsets = saved_tensors[3 + N] + output_tensor_offsets = saved_tensors[4 + N] origin_weights = [None] * N main_grads = [None] * N @@ -1253,6 +1273,7 @@ def _backward_grouped_tensor( grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( @@ -1260,13 +1281,14 @@ def _backward_grouped_tensor( grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = _GroupedLinear._make_grouped_tensor( dy_2d, num_gemms=N, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=output_tensor_offsets, last_dim=ctx.weights_shape_0, dtype=ctx.activation_dtype, ) @@ -1298,7 +1320,7 @@ def _backward_grouped_tensor( dgrad, num_gemms=N, split_sizes=split_sizes, - base_split_offsets=base_split_offsets, + tensor_offsets=input_tensor_offsets, last_dim=ctx.weights_shape_1, dtype=ctx.activation_dtype, ) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index f4beffe90c..8137051322 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -348,18 +348,8 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: return tex.dsrelu(*args, **kwargs) -class ScaledSReLU(BasicOperation): - r"""Squared ReLU with per-row post-scaling. - - If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied - with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. - - Parameters - ---------- - activation_recompute_in_mlp : bool, default = ``False`` - Enable fused grouped MLP kernels to recompute activation outputs - during backward when supported instead of saving them. - """ +class _ScaledUnary(BasicOperation, metaclass=abc.ABCMeta): + """Unary activation with per-row scales (fused grouped MLP middle op).""" num_extra_inputs: int = 1 @@ -367,6 +357,25 @@ def __init__(self, *, activation_recompute_in_mlp: bool = False) -> None: super().__init__() self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp + @abc.abstractmethod + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + """Apply the scaled unary activation.""" + + @abc.abstractmethod + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Apply the scaled unary activation backward pass.""" + def op_forward(self, *args, **kwargs) -> None: raise RuntimeError( f"{self.__class__.__name__} operation has " @@ -410,7 +419,7 @@ def fuser_forward( x = maybe_dequantize(input_.contiguous(), dtype) scales = maybe_dequantize(extra_input, dtype) - y = tex.srelu(x, None) * scales.unsqueeze(-1) + y = self._scaled_unary_forward(x, scales) ctx = basic_op_ctxs[0] if ctx.requires_grad: @@ -448,21 +457,57 @@ def fuser_backward( scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype) - grad_input = None - if ctx.input_requires_grad: - grad_srelu_out = grad_output * scales.unsqueeze(-1) - grad_input = tex.dsrelu(grad_srelu_out, x, None) - - grad_extra_input = None - if ctx.extra_input_requires_grad: - srelu_out = tex.srelu(x, None) - grad_extra_input = torch.linalg.vecdot(srelu_out, grad_output) + grad_input, grad_extra_input = self._scaled_unary_backward( + grad_output, + x, + scales, + compute_scale_grad=ctx.extra_input_requires_grad, + ) + if not ctx.input_requires_grad: + grad_input = None clear_tensor_data(ctx.saved_tensors[0]) return grad_input, [()], [(grad_extra_input,)] +class ScaledSReLU(_ScaledUnary): + r"""Squared ReLU with per-row post-scaling. + + If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied + with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. + + Parameters + ---------- + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. + """ + + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + return tex.scaled_srelu(input_, scales, None) + + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return tex.scaled_dsrelu( + grad_output, + input_, + scales, + None, + compute_scale_grad, + ) + + class SReGLU(_ActivationOperation): r"""Squared Rectified Gated Linear Unit diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea..c94d78fecd 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1136,14 +1136,17 @@ def fuser_forward_save_ctx( # temporary workspaces freshly created in each forward pass. if is_cpu_offload_enabled(): saved = tensors_to_save[0] - offset = 4 if self._scale_bias else 3 + # Metadata prefix: + # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, (scales?)] + offset = 6 if self._scale_bias else 5 if use_grouped_tensor_path: - # Layout: [split_sizes, base_split_offsets, split_points, (scales?), grouped_x, *weights] + # Layout: [..., grouped_x, *weights] grouped_x = saved[offset] if grouped_x is not None: mark_activation_offload(grouped_x) else: - # Layout: [split_sizes, None, None, (scales?), *xs, *ws] + # Layout: [..., *xs, *ws] live_xs = [t for t in saved[offset : offset + self.num_groups] if t is not None] if live_xs: mark_activation_offload(*live_xs) @@ -1169,9 +1172,9 @@ def fuser_forward_save_ctx( ctx.weight_quantizers = weight_quantizers ctx.grad_output_quantizers = grad_output_quantizers ctx.grad_input_quantizers = None - # ``split_sizes`` and ``base_split_offsets`` are routed through - # ``save_for_backward`` (see ``_fuser_forward_split_quantize`` and - # ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). + # ``split_sizes``, offset metadata, and related tensors are routed + # through ``save_for_backward`` (see ``_fuser_forward_split_quantize`` + # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). if torch.is_autocast_enabled(): ctx.dtype = torch.get_autocast_dtype("cuda") else: @@ -1285,12 +1288,12 @@ def _fuser_forward_split_quantize( # Build the tuple of tensors to save for backward. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if scale_bias), *xs, *ws] - # ``base_split_offsets`` and ``split_points`` are unused on the - # split-quantize backward path but are included as ``None`` so the - # saved-tensor layout matches the graph-safe - # ``_fuser_forward_grouped_tensor`` path (and the fused MLP forward). - saved: list[Optional[torch.Tensor]] = [split_sizes, None, None] + # Offset metadata slots are unused on the split-quantize backward path + # but are included as ``None`` so the saved-tensor layout matches the + # graph-safe ``_fuser_forward_grouped_tensor`` path. + saved: list[Optional[torch.Tensor]] = [split_sizes, None, None, None, None] if self._scale_bias: saved.append(scales) saved.extend(xs) @@ -1312,28 +1315,36 @@ def _fuser_forward_grouped_tensor( device: torch.device, out_buffer: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: - """Graph-safe GroupedTensor forward path (pure compute). - Returns ``(output, tensors_to_save)``. ``split_sizes``, - ``base_split_offsets`` and ``split_points`` are returned so that - ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them. - """ + """Build graph-safe grouped input storage and run grouped GEMM.""" num_groups = self.num_groups - has_bias = self.has_bias - - base_split_offsets = tex.splits_to_offsets(split_sizes, 1) - split_points = base_split_offsets[1:].to(dtype=torch.int) - - # Flatten to 2D so the first dim is the total token count. + split_sizes, grouped_tensor_offsets = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[1, 1, self.in_features, self.out_features], + include_leading_zero=[False, True, True, True], + dtypes=[torch.int32, torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) + split_points = grouped_tensor_offsets[0] + base_split_offsets = grouped_tensor_offsets[1] + input_tensor_offsets = grouped_tensor_offsets[2] + output_tensor_offsets = grouped_tensor_offsets[3] original_shape = list(input_.size()) + # Flatten to 2D so the first dim is the total token count. x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) total_tokens = x.size(0) - - # Build the input GroupedTensor. + # Build the input GroupedTensorStorage for input. if with_quantized_compute: input_quantizer = input_quantizers[0] input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) input_quantizer.optimize_for_gemm = True - grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes) + grouped_x = tex.group_quantize( + x, + input_quantizer, + num_groups, + split_sizes, + tensor_offsets=input_tensor_offsets, + ) else: # No quantize: wrap the contiguous high-precision buffer. grouped_x = GroupedTensorStorage( @@ -1343,8 +1354,9 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=x.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.in_features, + tensor_offsets=input_tensor_offsets, ) + has_bias = self.has_bias if is_cpu_offload_enabled() and grouped_x is not None: start_offload(grouped_x) @@ -1379,7 +1391,7 @@ def _fuser_forward_grouped_tensor( quantizer=None, data=out.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + tensor_offsets=output_tensor_offsets, ) # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias @@ -1414,14 +1426,23 @@ def _fuser_forward_grouped_tensor( # Build the tuple of tensors to save for backward. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if _scale_bias), grouped_x, *weights] + # ``output_tensor_offsets`` matches the linear output row layout and is + # reused as ``grad_output`` offsets in backward. if grouped_x is not None: # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data # in backward pass) if with_quantized_compute and grouped_x.columnwise_data is not None: grouped_x.rowwise_data = None grouped_x.scale_inv = None - saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points] + saved: list[Optional[torch.Tensor]] = [ + split_sizes, + base_split_offsets, + split_points, + input_tensor_offsets, + output_tensor_offsets, + ] if self._scale_bias: saved.append(scales) saved.append(grouped_x) @@ -1471,13 +1492,13 @@ def _fuser_backward_split_quantize( # Saved tensors from forward pass. Layout: # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, # (scales if _scale_bias), *xs, *ws] - # ``base_split_offsets`` and ``split_points`` are unused on this path - # but are present so the saved-tensor layout matches the graph-safe - # path (and the fused MLP forward). + # Offset metadata beyond ``split_sizes`` is unused on this path but is + # present so the saved-tensor layout matches the graph-safe path. saved_tensors = ctx.saved_tensors split_sizes = saved_tensors[0] - saved_tensors = saved_tensors[3:] + saved_tensors = saved_tensors[5:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1653,24 +1674,24 @@ def _fuser_backward_grouped_tensor( Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: + """Graph-safe GroupedTensor backward path.""" num_groups = self.num_groups has_bias = self.has_bias weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup() device = weights[0].device dtype = ctx.dtype - with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) - # Saved tensors from forward pass - # Layout: [split_sizes, base_split_offsets, split_points, - # (scales if _scale_bias), grouped_x, *weights] - # ``split_points`` is unused on this path but is present so the - # saved-tensor layout matches the fused MLP forward (which needs it - # for the cuDNN grouped GEMM kernel). + # Saved tensors from forward pass. Layout: + # [split_sizes, base_split_offsets, split_points, + # input_tensor_offsets, output_tensor_offsets, + # (scales if _scale_bias), grouped_x, *weights] saved_tensors = ctx.saved_tensors split_sizes = saved_tensors[0] base_split_offsets = saved_tensors[1] - saved_tensors = saved_tensors[3:] + input_tensor_offsets = saved_tensors[3] + output_tensor_offsets = saved_tensors[4] + saved_tensors = saved_tensors[5:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1684,6 +1705,7 @@ def _fuser_backward_grouped_tensor( # to figure out total tokens. dy_2d = grad_output.reshape(-1, self.out_features) total_tokens = dy_2d.size(0) + grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] # Build the grad_output GroupedTensor. # Optionally get dbias is fusion available with bgrad_group_quantize @@ -1691,10 +1713,10 @@ def _fuser_backward_grouped_tensor( if with_quantized_compute: grad_output_quantizer = ctx.grad_output_quantizers[0] grad_output_quantizer.set_usage( - rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad + rowwise=ctx.input_requires_grad, + columnwise=ctx.weight_requires_grad, ) grad_output_quantizer.optimize_for_gemm = True - # FP8 block scaling computes dbias in the rowwise (dgrad) pass, so only fuse # when dgrad is required. fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or ( @@ -1702,15 +1724,22 @@ def _fuser_backward_grouped_tensor( ) if has_bias and not self._scale_bias and fuse_bgrad: grouped_dy, dbias_packed = tex.bgrad_group_quantize( - dy_2d, grad_output_quantizer, num_groups, split_sizes + dy_2d, + grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( - dy_2d, grad_output_quantizer, num_groups, split_sizes + dy_2d, + grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=output_tensor_offsets, ) else: dy_2d = maybe_dequantize(dy_2d, dtype) - # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm grouped_dy = GroupedTensorStorage( shape=(total_tokens, self.out_features), dtype=dtype, @@ -1718,7 +1747,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=dy_2d.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.out_features, + tensor_offsets=output_tensor_offsets, ) # Bias Grads compute if not already computed in bgrad_group_quantize @@ -1745,7 +1774,6 @@ def _fuser_backward_grouped_tensor( # ---- dgrad GEMM ---------------------------------------------------- grad_input = None if ctx.input_requires_grad: - grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] grad_input = validate_or_alloc_output( getattr(ctx, "dgrad_out", None), grad_input_shape, dtype, device ) @@ -1756,7 +1784,7 @@ def _fuser_backward_grouped_tensor( quantizer=None, data=grad_input.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * self.in_features, + tensor_offsets=input_tensor_offsets, ) general_grouped_gemm_for_grouped_tensor( dist_dgrad_weights if is_dist_weight else ws, diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 02f330ede3..fb663c0480 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -387,14 +387,21 @@ def __init__( self.glu_interleave_size: Optional[int] = glu_interleave_size self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: + def _scaled_glu_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: raise NotImplementedError - def _glu_backward( + def _scaled_glu_backward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, - ) -> torch.Tensor: + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: raise NotImplementedError def op_forward(self, *args, **kwargs) -> None: @@ -442,22 +449,7 @@ def fuser_forward( # Make sure inputs are in correct dtype input_ = maybe_dequantize(input_, dtype) scales = maybe_dequantize(extra_input, dtype) - - # Remove gate interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - swiglu_out = self._glu_forward(swiglu_in) - out = swiglu_out * scales.unsqueeze(-1) + out = self._scaled_glu_forward(input_, scales) # Save state for backward pass ctx = basic_op_ctxs[0] @@ -469,7 +461,7 @@ def fuser_forward( ctx.dtype = dtype ctx.save_for_backward( input_, - scales if ctx.input_requires_grad else None, + scales if ctx.input_requires_grad or ctx.extra_input_requires_grad else None, ) return out, [()] @@ -498,41 +490,14 @@ def fuser_backward( scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output, ctx.dtype) - # Remove gate interleaving if needed - swiglu_in = input_ - if self.glu_interleave_size is not None: - shape = swiglu_in.size() - swiglu_in = swiglu_in.reshape( - -1, - shape[-1] // (2 * self.glu_interleave_size), - 2, - self.glu_interleave_size, - ) - swiglu_in = swiglu_in.transpose(1, 2).contiguous() - swiglu_in = swiglu_in.view(shape) - - # Compute input grad - grad_input = None - if ctx.input_requires_grad: - grad_swiglu_out = grad_output * scales.unsqueeze(-1) - grad_swiglu_in = self._glu_backward(grad_swiglu_out, swiglu_in) - grad_input = grad_swiglu_in - if self.glu_interleave_size is not None: - shape = grad_input.size() - grad_input = grad_input.reshape( - -1, - 2, - shape[-1] // (2 * self.glu_interleave_size), - self.glu_interleave_size, - ) - grad_input = grad_input.transpose(1, 2).contiguous() - grad_input = grad_input.view(shape) - - # Compute scales grad by recomputing GLU - grad_extra_input = None - if ctx.extra_input_requires_grad: - swiglu_out = self._glu_forward(swiglu_in) - grad_extra_input = torch.linalg.vecdot(swiglu_out, grad_output) + grad_input, grad_extra_input = self._scaled_glu_backward( + grad_output, + input_, + scales, + compute_scale_grad=ctx.extra_input_requires_grad, + ) + if not ctx.input_requires_grad: + grad_input = None # Clear input tensor if possible clear_tensor_data(ctx.saved_tensors[0]) # input_ @@ -558,15 +523,34 @@ class ScaledSwiGLU(_ScaledGLU): """ - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: - return tex.swiglu(swiglu_in, None) - - def _glu_backward( + def _scaled_glu_forward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, ) -> torch.Tensor: - return tex.dswiglu(grad_swiglu_out, swiglu_in, None) + return tex.scaled_swiglu( + input_, + scales, + None, + int(self.glu_interleave_size or 0), + ) + + def _scaled_glu_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return tex.scaled_dswiglu( + grad_output, + input_, + scales, + None, + int(self.glu_interleave_size or 0), + compute_scale_grad, + ) class ScaledClampedQGeGLU(_ScaledGLU): @@ -614,16 +598,39 @@ def __init__( glu_linear_offset=glu_linear_offset, ) - def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: - return self._clamped._tex_clamped_swiglu_forward(swiglu_in, None) - - def _glu_backward( + def _scaled_glu_forward( self, - grad_swiglu_out: torch.Tensor, - swiglu_in: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, ) -> torch.Tensor: - return self._clamped._tex_clamped_dswiglu( - grad_swiglu_out, - swiglu_in, + clamped = self._clamped + return tex.scaled_clamped_swiglu( + input_, + scales, + None, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(self.glu_interleave_size or 0), + ) + + def _scaled_glu_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + clamped = self._clamped + return tex.scaled_clamped_dswiglu( + grad_output, + input_, + scales, None, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(self.glu_interleave_size or 0), + compute_scale_grad, ) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 0113833647..db2c4e7674 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1091,21 +1091,40 @@ def fuser_forward( # shared-expert path never consumes it. base_split_offsets = split_sizes fc1_x_tensor_offsets = None + fc1_out_tensor_offsets = None fc2_x_tensor_offsets = None fc2_out_tensor_offsets = None else: + # Bulk-allocate every grouped-tensor offset the forward and backward + # passes need, so the backward can reuse them from the context + # instead of recomputing offsets per GEMM. split_sizes, ( split_points, base_split_offsets, fc1_x_tensor_offsets, + fc1_out_tensor_offsets, fc2_x_tensor_offsets, fc2_out_tensor_offsets, ) = tex.splits_to_offsets_multi( split_sizes, device, - strides=[1, 1, fc1_weight_shape[1], fc2_weight_shape[1], fc2_weight_shape[0]], - include_leading_zero=[False, True, True, True, True], - dtypes=[torch.int32, torch.int64, torch.int64, torch.int64, torch.int64], + strides=[ + 1, + 1, + fc1_weight_shape[1], + fc1_weight_shape[0], + fc2_weight_shape[1], + fc2_weight_shape[0], + ], + include_leading_zero=[False, True, True, True, True, True], + dtypes=[ + torch.int32, + torch.int64, + torch.int64, + torch.int64, + torch.int64, + torch.int64, + ], bulk_allocate=True, ) @@ -1743,6 +1762,10 @@ def fuser_forward( split_sizes, base_split_offsets, split_points, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, saved_fc1_x, *fc1_weight_tensors, activation_in, @@ -1796,12 +1819,22 @@ def fuser_backward( # Saved tensors from the joint forward. # Layout: [split_sizes, base_split_offsets, split_points, + # fc1_x_tensor_offsets, fc1_out_tensor_offsets, + # fc2_x_tensor_offsets, fc2_out_tensor_offsets, # grouped_fc1_x, *fc1_weights, # activation_in, scales, # grouped_fc2_x, *fc2_weights] saved_tensors = fc1_ctx.saved_tensors - split_sizes, base_split_offsets, split_points = saved_tensors[:3] - saved_tensors = saved_tensors[3:] + ( + split_sizes, + base_split_offsets, + split_points, + fc1_x_tensor_offsets, + fc1_out_tensor_offsets, + fc2_x_tensor_offsets, + fc2_out_tensor_offsets, + ) = saved_tensors[:7] + saved_tensors = saved_tensors[7:] grouped_fc1_x, saved_tensors = saved_tensors[0], saved_tensors[1:] if fc1_op.single_grouped_weight: grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -1868,6 +1901,7 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, + tensor_offsets=fc2_out_tensor_offsets, ) else: grouped_fc2_dy = _group_quantize_for_grouped_mlp( @@ -1875,9 +1909,7 @@ def fuser_backward( fc2_grad_output_quantizer, num_groups, split_sizes, - tensor_offsets=( - None if num_groups == 1 else base_split_offsets * fc2_weight_shape[0] - ), + tensor_offsets=fc2_out_tensor_offsets, ) use_nvfp4 = ( @@ -2150,7 +2182,7 @@ def fuser_backward( fc2_input_quantizer, num_groups, split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[1], + tensor_offsets=fc2_x_tensor_offsets, ) else: sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") @@ -2172,7 +2204,7 @@ def fuser_backward( scale_inv=None, columnwise_scale_inv=fc2_x_col_scale.reshape(-1), first_dims=split_sizes, - tensor_offsets=base_split_offsets * fc2_weight_shape[1], + tensor_offsets=fc2_x_tensor_offsets, with_gemm_swizzled_scales=True, ) @@ -2215,9 +2247,7 @@ def fuser_backward( fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)] # FC1 grad output for dgrad and wgrad GEMMs - fc1_dy_tensor_offsets = ( - None if num_groups == 1 else base_split_offsets * fc1_weight_shape[0] - ) + fc1_dy_tensor_offsets = fc1_out_tensor_offsets fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] if use_nvfp4: fc1_grad_output_quantizer.set_usage( @@ -2304,7 +2334,6 @@ def fuser_backward( ) elif use_nvfp4: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) - fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] grouped_grad_input = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[1]), dtype=dtype,