diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 39d3e79e62..1dc714faf3 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -70,6 +70,7 @@ NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_ PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py" +NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp_without_cutedsl_fusion.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py without CuTe DSL fusion" if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d195eb2f78..245bc355cd 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -24,6 +24,12 @@ OUTPUT_BUFFER_KEY, GRAD_INPUT_BUFFER_KEY, ) +from transformer_engine.pytorch.ops.fused.backward_activation_grouped_linear import ( + BackwardScaledActivationGroupedLinear, +) +from transformer_engine.pytorch.ops.fused.forward_activation_grouped_linear import ( + ForwardScaledActivationGroupedLinear, +) from transformer_engine.pytorch import ( QuantizedTensor, Float8CurrentScalingQuantizer, @@ -73,6 +79,10 @@ if nvfp4_available: _grouped_mlp_quantization_list.append("nvfp4_rht") +# Quantization recipes supported by ScaledActivation + GroupedLinear fusion +if fp8_available: + _grouped_mlp_quantization_list.append("fp8_current_scaling") + @pytest.fixture(autouse=True, scope="function") def _reset_rng_states_per_test(): @@ -755,6 +765,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,18 +1043,21 @@ 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) + ) ) ) + forward_ops = module._module_groups[0]._forward_ops + backward_ops = module._module_groups[0]._backward_ops + full_grouped_mlp_fusion = False if expected_grouped_mlp_fusion: if activation_is_glu: fused_cls = te.ops.fused.GroupedMLP_CuTeGEMMGLU else: fused_cls = te.ops.fused.GroupedMLP_CuTeGEMMUnary if fused_cls.is_supported(): - forward_ops = module._module_groups[0]._forward_ops - backward_ops = module._module_groups[0]._backward_ops assert len(forward_ops) == 1 assert len(backward_ops) == 1 assert isinstance( @@ -1045,6 +1065,24 @@ def _make_module(): fused_cls, ) assert backward_ops[0][0] is forward_ops[0][0] + full_grouped_mlp_fusion = True + + # When the full FC1 + activation + FC2 fusion is unavailable, verify + # that ScaledActivation + GroupedLinear fusions cover both boundaries + # whenever grouped quantized compute is supported. + act_grouped_linear_fusion_expected = ( + not full_grouped_mlp_fusion + and te.ops.fused.act_grouped_linear_fusion_supported(fc2, module[1], recipe) + and te.ops.fused.act_grouped_linear_fusion_supported(fc1, module[1], recipe) + ) + assert ( + any(isinstance(op, ForwardScaledActivationGroupedLinear) for op, _ in forward_ops) + == act_grouped_linear_fusion_expected + ) + assert ( + any(isinstance(op, BackwardScaledActivationGroupedLinear) for op, _ in backward_ops) + == act_grouped_linear_fusion_expected + ) # Loose tols for sanity checking tols = {"rtol": 0.125, "atol": 0.25} @@ -1266,6 +1304,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 +1624,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 +1757,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..9b51df7e09 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -290,6 +290,68 @@ 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); + +/* Scaled activation + grouped quantize */ +py::object grouped_scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, + int64_t glu_interleave_size); + +py::object grouped_scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, float limit, + float alpha, float glu_linear_offset, + int64_t glu_interleave_size); + +py::object grouped_scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets); + +py::tuple grouped_scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets, + int64_t glu_interleave_size, bool compute_scale_grad); + +py::tuple grouped_scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, float limit, + float alpha, float glu_linear_offset, + int64_t glu_interleave_size, bool compute_scale_grad); + +py::tuple grouped_scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input, + const at::Tensor &act_scales, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets, 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..7c486a522e 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -342,5 +342,256 @@ 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_*). + * + * Grouped variants reuse the dense compute helpers, then optionally apply + * group_quantize. Keep the nvte_scaled_* launch path in one place. + */ + +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; +} + +py::object maybe_group_quantize(const at::Tensor& tensor, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets) { + if (quantizer.is_none()) { + return py::cast(tensor); + } + return group_quantize(tensor, quantizer, num_tensors, first_dims, std::nullopt, tensor_offsets, + std::nullopt); +} + +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()); +} + +template +py::object grouped_scaled_activation_helper(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, + int shape_divisor, Args&&... args) { + NVTE_CHECK(input.dim() == 2, "grouped scaled activation input must be 2D"); + auto output = scaled_activation_compute(input, act_scales, shape_divisor, + std::forward(args)...); + return maybe_group_quantize(output, quantizer, num_tensors, first_dims, tensor_offsets); +} + +template +py::tuple grouped_scaled_dactivation_helper(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, + bool compute_scale_grad, Args&&... args) { + NVTE_CHECK(input.dim() == 2 && grad.dim() == 2, + "grouped scaled dactivation input and grad must be 2D"); + auto [grad_input, grad_scales] = scaled_dactivation_compute( + grad, input, act_scales, compute_scale_grad, std::forward(args)...); + // Return both the (optionally) grouped-quantized grad input for the next + // grouped GEMM and the dense high-precision grad input so callers can reuse + // it (e.g. bias gradient) without a lossy dequantize. + return py::make_tuple( + maybe_group_quantize(grad_input, quantizer, num_tensors, first_dims, tensor_offsets), + py::cast(grad_input), 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); +} + +py::object grouped_scaled_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, + int64_t glu_interleave_size) { + return grouped_scaled_activation_helper( + input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, /*shape_divisor=*/2, + glu_interleave_size); +} + +py::object grouped_scaled_clamped_swiglu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, float limit, + float alpha, float glu_linear_offset, + int64_t glu_interleave_size) { + return grouped_scaled_activation_helper( + input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, /*shape_divisor=*/2, + limit, alpha, glu_linear_offset, glu_interleave_size); +} + +py::object grouped_scaled_srelu(const at::Tensor& input, const at::Tensor& act_scales, + py::handle quantizer, const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets) { + return grouped_scaled_activation_helper( + input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, + /*shape_divisor=*/1); +} + +py::tuple grouped_scaled_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets, + int64_t glu_interleave_size, bool compute_scale_grad) { + return grouped_scaled_dactivation_helper( + grad, input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, + compute_scale_grad, glu_interleave_size); +} + +py::tuple grouped_scaled_clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + const size_t num_tensors, + std::optional first_dims, + std::optional tensor_offsets, float limit, + float alpha, float glu_linear_offset, + int64_t glu_interleave_size, bool compute_scale_grad) { + return grouped_scaled_dactivation_helper( + grad, input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, + compute_scale_grad, limit, alpha, glu_linear_offset, glu_interleave_size); +} + +py::tuple grouped_scaled_dsrelu(const at::Tensor& grad, const at::Tensor& input, + const at::Tensor& act_scales, py::handle quantizer, + const size_t num_tensors, std::optional first_dims, + std::optional tensor_offsets, bool compute_scale_grad) { + return grouped_scaled_dactivation_helper(grad, input, act_scales, quantizer, + num_tensors, first_dims, + tensor_offsets, 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..386125368a 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -285,6 +285,57 @@ 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); + /* Scaled activation + grouped quantize */ + m.def("grouped_scaled_swiglu", transformer_engine::pytorch::grouped_scaled_swiglu, + "Scaled SwiGLU + grouped quantize", py::arg("input"), py::arg("act_scales"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none(), py::arg("glu_interleave_size") = 0); + m.def("grouped_scaled_clamped_swiglu", transformer_engine::pytorch::grouped_scaled_clamped_swiglu, + "Scaled clamped SwiGLU + grouped quantize", py::arg("input"), py::arg("act_scales"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none(), 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("grouped_scaled_srelu", transformer_engine::pytorch::grouped_scaled_srelu, + "Scaled SReLU + grouped quantize", py::arg("input"), py::arg("act_scales"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none()); + m.def("grouped_scaled_dswiglu", transformer_engine::pytorch::grouped_scaled_dswiglu, + "Scaled SwiGLU backward + optional grouped quantize", py::arg("grad"), py::arg("fwd_input"), + py::arg("act_scales"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none(), py::arg("glu_interleave_size") = 0, + py::arg("compute_scale_grad") = true); + m.def("grouped_scaled_clamped_dswiglu", + transformer_engine::pytorch::grouped_scaled_clamped_dswiglu, + "Scaled clamped SwiGLU backward + optional grouped quantize", py::arg("grad"), + py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"), py::arg("num_tensors"), + py::arg("first_dims"), py::arg("tensor_offsets") = py::none(), 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("grouped_scaled_dsrelu", transformer_engine::pytorch::grouped_scaled_dsrelu, + "Scaled SReLU backward + optional grouped quantize", py::arg("grad"), py::arg("fwd_input"), + py::arg("act_scales"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"), + py::arg("tensor_offsets") = py::none(), 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 0b51ca6a1e..fd1210c4ec 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -195,7 +195,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: @@ -207,7 +207,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 @@ -338,8 +338,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) @@ -350,13 +360,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, ) @@ -385,7 +401,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, ) @@ -430,6 +446,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 @@ -840,6 +858,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 @@ -876,6 +896,7 @@ def _backward_grouped_tensor( grad_output_quantizer, N, split_sizes, + tensor_offsets=output_tensor_offsets, ) else: grouped_dy = tex.group_quantize( @@ -883,13 +904,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, ) @@ -921,7 +943,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..bc8dbed6c9 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1064,7 +1064,7 @@ def fuser_forward( ) if use_grouped_tensor_path: - out, tensors_to_save = self._fuser_forward_grouped_tensor( + out, tensors_to_save = self._fuser_forward_graph_safe( input_=input_, split_sizes=split_sizes, scales=scales, @@ -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,19 +1288,20 @@ 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 (and fused forwards + # that share this GroupedLinear context contract). + saved: list[Optional[torch.Tensor]] = [split_sizes, None, None, None, None] if self._scale_bias: saved.append(scales) saved.extend(xs) saved.extend(ws) return out, tuple(saved) - def _fuser_forward_grouped_tensor( + def _fuser_forward_graph_safe( self, *, input_: torch.Tensor, @@ -1312,30 +1316,32 @@ 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, + ) + input_tensor_offsets = grouped_tensor_offsets[2] original_shape = list(input_.size()) + total_tokens = math.prod(original_shape[:-1]) x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) - total_tokens = x.size(0) - - # Build the input GroupedTensor. 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( shape=(total_tokens, self.in_features), dtype=dtype, @@ -1343,8 +1349,61 @@ 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, + ) + return self._fuser_forward_grouped_tensor( + grouped_input=grouped_x, + split_sizes=split_sizes, + scales=scales, + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + 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], + out_buffer=out_buffer, + out_shape=original_shape[:-1] + [self.out_features], + ) + + def _fuser_forward_grouped_tensor( + self, + *, + grouped_input: GroupedTensorStorage, + split_sizes: torch.Tensor, + scales: Optional[torch.Tensor], + with_quantized_compute: bool, + input_quantizers: list[Optional[Quantizer]], + weight_quantizers: list[Optional[Quantizer]], + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + device: torch.device, + split_points: torch.Tensor, + base_split_offsets: torch.Tensor, + input_tensor_offsets: torch.Tensor, + output_tensor_offsets: torch.Tensor, + out_buffer: Optional[torch.Tensor] = None, + out_shape: list[int], + ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: + """Run grouped GEMM with a pre-built grouped input.""" + num_groups = self.num_groups + has_bias = self.has_bias + total_tokens, in_features = grouped_input.logical_shape + expected_quantizer = input_quantizers[0] if with_quantized_compute else None + if grouped_input.quantizer is not expected_quantizer or in_features != self.in_features: + raise ValueError( + "GroupedLinear received an incompatible grouped input " + f"(quantizer={grouped_input.quantizer}, " + f"logical_shape={grouped_input.logical_shape}; " + f"expected quantizer={expected_quantizer}, " + f"in_features={self.in_features})" ) + grouped_x = grouped_input if is_cpu_offload_enabled() and grouped_x is not None: start_offload(grouped_x) @@ -1370,7 +1429,6 @@ def _fuser_forward_grouped_tensor( ) # Allocate output buffer and wrap as a GroupedTensor view. - out_shape = original_shape[:-1] + [self.out_features] out = validate_or_alloc_output(out_buffer, out_shape, dtype, device) grouped_out = GroupedTensorStorage( shape=(total_tokens, self.out_features), @@ -1379,7 +1437,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 +1472,24 @@ 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 (including fused + # activation + grouped linear 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) @@ -1445,7 +1513,7 @@ def fuser_backward( ctx = basic_op_ctxs[0] # Dispatch to the path used in forward (saved as ``ctx.use_grouped_tensor_path``). if getattr(ctx, "use_grouped_tensor_path", False): - return self._fuser_backward_grouped_tensor( + return self._fuser_backward_graph_safe( ctx=ctx, grad_output=grad_output, ) @@ -1471,13 +1539,14 @@ 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 (and + # fused forwards that share this GroupedLinear context contract). 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:] @@ -1640,10 +1709,7 @@ def _fuser_backward_split_quantize( grad_extra = (None, grad_scales) if self._scale_bias else (None,) return grad_input, [grad_params], [grad_extra] - # ================================================================== - # Graph-safe backward: counterpart of `_fuser_forward_grouped_tensor`. - # ================================================================== - def _fuser_backward_grouped_tensor( + def _fuser_backward_graph_safe( self, *, ctx: OperationContext, @@ -1653,48 +1719,24 @@ def _fuser_backward_grouped_tensor( Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: + """Build graph-safe grouped grad-output storage and run grouped GEMMs.""" 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 = ctx.saved_tensors - split_sizes = saved_tensors[0] - base_split_offsets = saved_tensors[1] - saved_tensors = saved_tensors[3:] - scales = None - if self._scale_bias: - scales, saved_tensors = saved_tensors[0], saved_tensors[1:] - grouped_x, saved_tensors = saved_tensors[0], saved_tensors[1:] - if self.single_grouped_weight: - ws, saved_tensors = saved_tensors[0], saved_tensors[1:] - else: - ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] - - # Flatten grad_output to 2D (total_tokens, out_features) - # to figure out total tokens. + split_sizes = ctx.saved_tensors[0] + output_tensor_offsets = ctx.saved_tensors[4] dy_2d = grad_output.reshape(-1, self.out_features) total_tokens = dy_2d.size(0) - # Build the grad_output GroupedTensor. - # Optionally get dbias is fusion available with bgrad_group_quantize dbias_packed = None 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 +1744,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,8 +1767,76 @@ 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, + ) + + return self._fuser_backward_grouped_tensor( + ctx=ctx, + grad_output=grad_output, + grouped_grad_output=grouped_dy, + dbias_packed=dbias_packed, + ) + + def _fuser_backward_grouped_tensor( + self, + *, + ctx: OperationContext, + grad_output: torch.Tensor, + grouped_grad_output: GroupedTensorStorage, + dbias_packed: Optional[torch.Tensor] = None, + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + 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, + # input_tensor_offsets, output_tensor_offsets, + # (scales if _scale_bias), grouped_x, *weights] + # ``split_points`` / ``output_tensor_offsets`` are unused on this path + # but are present so the saved-tensor layout matches the fused MLP / + # activation-fusion forwards that share this GroupedLinear context + # contract. + saved_tensors = ctx.saved_tensors + split_sizes = saved_tensors[0] + base_split_offsets = saved_tensors[1] + input_tensor_offsets = saved_tensors[3] + saved_tensors = saved_tensors[5:] + scales = None + if self._scale_bias: + scales, saved_tensors = saved_tensors[0], saved_tensors[1:] + grouped_x, saved_tensors = saved_tensors[0], saved_tensors[1:] + if self.single_grouped_weight: + ws, saved_tensors = saved_tensors[0], saved_tensors[1:] + else: + ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] + + # Keep the dense high-precision grad for bias-gradient computation while + # optionally using a pre-quantized grouped view for the grouped GEMMs. + 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] + + expected_quantizer = ctx.grad_output_quantizers[0] if with_quantized_compute else None + if grouped_grad_output.quantizer is not expected_quantizer or tuple( + grouped_grad_output.logical_shape + ) != (total_tokens, self.out_features): + raise ValueError( + "GroupedLinear received an incompatible grouped grad_output " + f"(quantizer={grouped_grad_output.quantizer}, " + f"logical_shape={grouped_grad_output.logical_shape}; " + f"expected quantizer={expected_quantizer}, " + f"logical_shape={(total_tokens, self.out_features)})" ) + grouped_dy = grouped_grad_output # Bias Grads compute if not already computed in bgrad_group_quantize final_bias_grads: Optional[torch.Tensor] = None @@ -1745,7 +1862,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 +1872,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/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index dc9dcd6dc3..85191a9807 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -6,9 +6,14 @@ from ..fuser import register_backward_fusion, register_forward_fusion from .backward_activation_bias import BackwardActivationBias +from .backward_activation_grouped_linear import BackwardScaledActivationGroupedLinear from .backward_add_rmsnorm import BackwardAddRMSNorm from .backward_linear_add import BackwardLinearAdd from .backward_linear_scale import BackwardLinearScale +from .forward_activation_grouped_linear import ( + ForwardScaledActivationGroupedLinear, + act_grouped_linear_fusion_supported, +) from .forward_linear_bias_activation import ForwardLinearBiasActivation from .forward_linear_bias_add import ForwardLinearBiasAdd from .forward_linear_scale_add import ForwardLinearScaleAdd @@ -21,6 +26,7 @@ register_forward_fusion(ForwardLinearBiasAdd.fuse_forward_ops) register_forward_fusion(ForwardLinearBiasActivation.fuse_forward_ops) register_forward_fusion(ForwardLinearScaleAdd.fuse_forward_ops) +register_forward_fusion(ForwardScaledActivationGroupedLinear.fuse_forward_ops) # Register backward fusions register_backward_fusion(UserbuffersBackwardLinear.fuse_backward_ops) @@ -28,6 +34,7 @@ register_backward_fusion(BackwardLinearScale.fuse_backward_ops) register_backward_fusion(BackwardActivationBias.fuse_backward_ops) register_backward_fusion(BackwardAddRMSNorm.fuse_backward_ops) +register_backward_fusion(BackwardScaledActivationGroupedLinear.fuse_backward_ops) # Import experimental fusions # Note: Registration logic is non-trivial, so submodule handles it internally. diff --git a/transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py b/transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py new file mode 100644 index 0000000000..0cb17fe0d1 --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/backward_activation_grouped_linear.py @@ -0,0 +1,181 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused scaled activation + grouped linear backward.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Optional + +import torch + +import transformer_engine_torch as tex +from ...quantization import Recipe +from ...tensor import Quantizer +from ...utils import clear_tensor_data +from .._common import maybe_dequantize +from ..basic import GroupedLinear, ScaledClampedQGeGLU, ScaledSReLU, ScaledSwiGLU +from ..op import FusedOperation, FusibleOperation, OperationContext +from .forward_activation_grouped_linear import ( + _SCALED_ACTIVATION_TYPES, + _ScaledActivation, + act_grouped_linear_fusion_supported, +) + + +def _grouped_scaled_dactivation( + activation: _ScaledActivation, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + quantizer: Quantizer, + num_groups: int, + split_sizes: torch.Tensor, + tensor_offsets: torch.Tensor, + compute_scale_grad: bool, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Dispatch a grouped scaled activation backward pass.""" + dy = grad_output.reshape(-1, grad_output.size(-1)) + x = input_.reshape(-1, input_.size(-1)) + s = scales.reshape(-1) + if isinstance(activation, ScaledSwiGLU): + return tex.grouped_scaled_dswiglu( + dy, + x, + s, + quantizer, + num_groups, + split_sizes, + tensor_offsets, + int(activation.glu_interleave_size or 0), + compute_scale_grad, + ) + if isinstance(activation, ScaledClampedQGeGLU): + clamped = activation._clamped + return tex.grouped_scaled_clamped_dswiglu( + dy, + x, + s, + quantizer, + num_groups, + split_sizes, + tensor_offsets, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(activation.glu_interleave_size or 0), + compute_scale_grad, + ) + if isinstance(activation, ScaledSReLU): + return tex.grouped_scaled_dsrelu( + dy, + x, + s, + quantizer, + num_groups, + split_sizes, + tensor_offsets, + compute_scale_grad, + ) + raise TypeError(f"Unsupported scaled activation type ({type(activation).__name__})") + + +class BackwardScaledActivationGroupedLinear(FusedOperation): + """Scaled activation backward + grouped quantize + grouped linear backward.""" + + def __init__(self, *, linear: GroupedLinear, activation: _ScaledActivation) -> None: + super().__init__((linear, activation)) + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + ) -> tuple[ + Optional[torch.Tensor], + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + linear = self.basic_ops[0] + activation = self.basic_ops[1] + linear_ctx, activation_ctx = basic_op_ctxs + + if not linear_ctx.requires_grad: + _, _, act_grad_extra_inputs = activation.fuser_backward( + [activation_ctx], + grad_output, + basic_op_grad_extra_outputs=[basic_op_grad_extra_outputs[1]], + ) + return None, [(), ()], [(), act_grad_extra_inputs[0]] + + del basic_op_grad_extra_outputs + input_, scales = activation_ctx.saved_tensors + input_ = maybe_dequantize(input_, activation_ctx.dtype) + scales = maybe_dequantize(scales, activation_ctx.dtype) + grad_output = maybe_dequantize(grad_output, activation_ctx.dtype) + + split_sizes = linear_ctx.saved_tensors[0] + activation_input_tensor_offsets = linear_ctx.saved_tensors[4] + grad_output_quantizer = linear_ctx.grad_output_quantizers[0] + grad_output_quantizer.set_usage( + rowwise=linear_ctx.input_requires_grad, + columnwise=linear_ctx.weight_requires_grad, + ) + grad_output_quantizer.optimize_for_gemm = True + grouped_dy, dense_dy, grad_scales = _grouped_scaled_dactivation( + activation, + grad_output, + input_, + scales, + quantizer=grad_output_quantizer, + num_groups=linear.num_groups, + split_sizes=split_sizes, + tensor_offsets=activation_input_tensor_offsets, + compute_scale_grad=activation_ctx.extra_input_requires_grad, + ) + + grad_input, grad_params, grad_extra_inputs = linear._fuser_backward_grouped_tensor( + ctx=linear_ctx, + grad_output=dense_dy, + grouped_grad_output=grouped_dy, + ) + + clear_tensor_data(activation_ctx.saved_tensors[0]) + return ( + grad_input, + [grad_params[0], ()], + [grad_extra_inputs[0], (grad_scales,)], + ) + + @staticmethod + def fuse_backward_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Fuse each supported GroupedLinear + ScaledActivation pair.""" + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + if ( + idx + 1 < len(ops) + and isinstance(ops[idx], GroupedLinear) + and isinstance(ops[idx + 1], _SCALED_ACTIVATION_TYPES) + and act_grouped_linear_fusion_supported(ops[idx], ops[idx + 1], recipe) + ): + out.append( + BackwardScaledActivationGroupedLinear( + linear=ops[idx], + activation=ops[idx + 1], + ) + ) + idx += 2 + else: + out.append(ops[idx]) + idx += 1 + return out diff --git a/transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py b/transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py new file mode 100644 index 0000000000..4f0b0d627c --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/forward_activation_grouped_linear.py @@ -0,0 +1,229 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused scaled activation + grouped linear forward.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any, Optional + +import torch + +import transformer_engine_torch as tex +from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...quantization import Recipe +from ...tensor import Quantizer +from .._common import maybe_dequantize +from ..basic import GroupedLinear, ScaledClampedQGeGLU, ScaledSReLU, ScaledSwiGLU +from ..basic.activation import _ScaledUnary +from ..basic.swiglu import _ScaledGLU +from ..op import FusedOperation, FusibleOperation, OperationContext + + +_ScaledActivation = _ScaledGLU | _ScaledUnary +_SCALED_ACTIVATION_TYPES = (_ScaledGLU, _ScaledUnary) + + +def _grouped_scaled_activation( + activation: _ScaledActivation, + input_: torch.Tensor, + scales: torch.Tensor, + quantizer: Quantizer, + num_groups: int, + split_sizes: torch.Tensor, + tensor_offsets: torch.Tensor, +) -> torch.Tensor: + """Dispatch a grouped scaled activation.""" + x = input_.reshape(-1, input_.size(-1)) + s = scales.reshape(-1) + if isinstance(activation, ScaledSwiGLU): + return tex.grouped_scaled_swiglu( + x, + s, + quantizer, + num_groups, + split_sizes, + tensor_offsets, + int(activation.glu_interleave_size or 0), + ) + if isinstance(activation, ScaledClampedQGeGLU): + clamped = activation._clamped + return tex.grouped_scaled_clamped_swiglu( + x, + s, + quantizer, + num_groups, + split_sizes, + tensor_offsets, + clamped.limit, + clamped.alpha, + clamped.glu_linear_offset, + int(activation.glu_interleave_size or 0), + ) + if isinstance(activation, ScaledSReLU): + return tex.grouped_scaled_srelu( + x, + s, + quantizer, + num_groups, + split_sizes, + tensor_offsets, + ) + raise TypeError(f"Unsupported scaled activation type ({type(activation).__name__})") + + +def act_grouped_linear_fusion_supported( + linear: GroupedLinear, + activation: _ScaledActivation, + recipe: Optional[Recipe], +) -> bool: + """Whether ScaledActivation + GroupedLinear can use grouped quantized compute.""" + if recipe is None or activation.activation_recompute_in_mlp: + return False + input_quantizers = [ + linear.get_quantizer("forward", 2 * group_idx) for group_idx in range(linear.num_groups) + ] + weight = linear.weight if linear.single_grouped_weight else linear.weight0 + dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else weight.dtype + return linear._is_graph_safe_path_supported( + with_quantized_compute=True, + input_quantizers=input_quantizers, + dtype=dtype, + single_grouped_weight=linear.single_grouped_weight, + ) + + +class ForwardScaledActivationGroupedLinear(FusedOperation): + """Scaled activation + grouped quantize + grouped linear forward.""" + + def __init__(self, *, activation: _ScaledActivation, linear: GroupedLinear) -> None: + super().__init__((activation, linear)) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + activation = self.basic_ops[0] + linear = self.basic_ops[1] + activation_ctx, linear_ctx = basic_op_ctxs + if basic_op_kwargs[0] or basic_op_kwargs[1]: + raise ValueError("Scaled activation and GroupedLinear do not expect keyword arguments") + + weight = linear.weight if linear.single_grouped_weight else linear.weight0 + device = weight.device + dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else weight.dtype + input_ = maybe_dequantize(input_, dtype) + scales = maybe_dequantize(basic_op_extra_inputs[0][0], dtype) + + split_sizes = basic_op_extra_inputs[1][0] + if int(split_sizes.numel()) != linear.num_groups: + raise ValueError( + f"Expected {linear.num_groups} splits, but got {int(split_sizes.numel())}." + ) + split_sizes = split_sizes.to(device=device, dtype=torch.int64) + linear_scales = basic_op_extra_inputs[1][1] if linear._scale_bias else None + split_sizes, grouped_tensor_offsets = tex.splits_to_offsets_multi( + split_sizes, + device, + strides=[1, 1, linear.in_features, linear.out_features], + include_leading_zero=[False, True, True, True], + dtypes=[torch.int32, torch.int64, torch.int64, torch.int64], + bulk_allocate=True, + ) + + input_quantizers = [ + linear.get_quantizer("forward", 2 * group_idx) for group_idx in range(linear.num_groups) + ] + weight_quantizers = [ + linear.get_quantizer("forward", 2 * group_idx + 1) + for group_idx in range(linear.num_groups) + ] + input_quantizer = input_quantizers[0] + weight_requires_grad = linear_ctx.requires_grad and weight.requires_grad + input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + input_quantizer.optimize_for_gemm = True + + grouped_x = _grouped_scaled_activation( + activation, + input_, + scales, + input_quantizer, + linear.num_groups, + split_sizes, + grouped_tensor_offsets[2], + ) + + if activation_ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(input_) + activation_ctx.input_requires_grad = True + activation_ctx.extra_input_requires_grad = basic_op_extra_inputs[0][0].requires_grad + activation_ctx.dtype = dtype + activation_ctx.save_for_backward(input_, scales) + + out, tensors_to_save = linear._fuser_forward_grouped_tensor( + grouped_input=grouped_x, + split_sizes=split_sizes, + scales=linear_scales, + with_quantized_compute=True, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + dtype=dtype, + input_requires_grad=linear_ctx.requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + 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], + out_shape=list(input_.size())[:-1] + [linear.out_features], + ) + linear.fuser_forward_save_ctx( + basic_op_ctxs=[linear_ctx], + input_=input_, + tensors_to_save=[tensors_to_save], + requires_grad=[linear_ctx.requires_grad], + basic_op_extra_inputs=[basic_op_extra_inputs[1]], + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + basic_op_kwargs=[basic_op_kwargs[1]], + use_grouped_tensor_path=True, + ) + return out, [(), ()] + + @staticmethod + def fuse_forward_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Fuse each supported ScaledActivation + GroupedLinear pair.""" + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + if ( + idx + 1 < len(ops) + and isinstance(ops[idx], _SCALED_ACTIVATION_TYPES) + and isinstance(ops[idx + 1], GroupedLinear) + and act_grouped_linear_fusion_supported(ops[idx + 1], ops[idx], recipe) + ): + out.append( + ForwardScaledActivationGroupedLinear( + activation=ops[idx], + linear=ops[idx + 1], + ) + ) + idx += 2 + else: + out.append(ops[idx]) + idx += 1 + return out diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 97ea587bf6..d0beab2dcb 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, ) @@ -1738,6 +1757,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, @@ -1791,12 +1814,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:] @@ -1863,6 +1896,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( @@ -1870,9 +1904,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 = ( @@ -2145,7 +2177,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") @@ -2167,7 +2199,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, ) @@ -2210,9 +2242,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( @@ -2299,7 +2329,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,