diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index e1e7866a7..848e08253 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -317,7 +317,13 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims); std::vector multi_tensor_quantize(const std::vector &tensor_list, - std::vector quantizer_list); + std::vector quantizer_list, + // ROCm adds a fused multi-quantize + // path for GroupedLinear expert weights. The cached + // workspace can be passed directly as `outputs`, so the + // kernel writes in place instead of allocating separate + // storage. + const py::object &outputs = py::none()); std::vector split_quantize(const at::Tensor &tensor, const std::vector &split_sections, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 56086b11c..38c226470 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -491,12 +491,23 @@ void multi_tensor_quantize_impl(const std::vector &input_list, } // namespace std::vector multi_tensor_quantize(const std::vector &tensor_list, - std::vector quantizer_list) { + std::vector quantizer_list, + const py::object &outputs) { // Check number of tensors const size_t num_tensors = tensor_list.size(); NVTE_CHECK(quantizer_list.size() == num_tensors, "Expected ", num_tensors, " quantizers, but got ", quantizer_list.size()); + const bool use_provided_outputs = !outputs.is_none(); + py::sequence outputs_seq; + if (use_provided_outputs) { + NVTE_CHECK(py::isinstance(outputs) || py::isinstance(outputs), + "multi_tensor_quantize: outputs must be None, a list, or a tuple."); + outputs_seq = py::reinterpret_borrow(outputs); + NVTE_CHECK(static_cast(outputs_seq.size()) == num_tensors, "multi_tensor_quantize: ", + "len(outputs) is ", outputs_seq.size(), " but expected ", num_tensors, "."); + } + // Convert quantizers to C++ objects std::vector> quantizer_cpp_list; for (size_t i = 0; i < num_tensors; i++) { @@ -516,9 +527,17 @@ std::vector multi_tensor_quantize(const std::vector &ten const auto input_shape = input_cpp.shape(); const auto input_dtype = GetTransformerEngineDType(input_py.scalar_type()); - // Construct output tensor - std::vector output_shape(input_shape.data, input_shape.data + input_shape.ndim); - auto [output_cpp, output_py] = quantizer_cpp_list[i]->create_tensor(output_shape, input_dtype); + TensorWrapper output_cpp; + py::object output_py; + if (use_provided_outputs) { + py::object output_obj = outputs_seq[static_cast(i)]; + std::tie(output_cpp, output_py) = + quantizer_cpp_list[i]->convert_and_update_tensor(std::move(output_obj)); + } else { + std::vector output_shape(input_shape.data, input_shape.data + input_shape.ndim); + std::tie(output_cpp, output_py) = + quantizer_cpp_list[i]->create_tensor(output_shape, input_dtype); + } output_cpp_list.emplace_back(std::move(output_cpp)); output_py_list.emplace_back(std::move(output_py)); } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index cc90ff396..77d71d6a1 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -301,7 +301,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("rmsnorm_bwd_add", &transformer_engine::pytorch::rmsnorm_bwd_add, "Fused backward of RMSNorm + add"); m.def("multi_tensor_quantize", &transformer_engine::pytorch::multi_tensor_quantize, - "Multi-tensor quantize", py::arg("tensor_list"), py::arg("quantizer_list")); + "Multi-tensor quantize", py::arg("tensor_list"), py::arg("quantizer_list"), + py::arg("outputs") = py::none()); m.def("split_quantize", &transformer_engine::pytorch::split_quantize, "Split and multi-tensor quantize", py::arg("tensor"), py::arg("split_sections"), py::arg("quantizer_list"), py::arg("disable_bulk_allocation") = false diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index cea02e5b5..144540389 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -801,6 +801,127 @@ def quantize_weight( return out, None +def quantize_multi_weight( + *, + tensors: List[torch.Tensor], + quantizers: List[Quantizer], + workspaces: Optional[List[Optional[QuantizedTensorStorage]]] = None, + update_workspace: bool = True, + skip_update_flag: Optional[torch.Tensor] = None, + workspace_dtype: Optional[torch.dtype] = None, + cache: bool = False, +) -> Tuple[List[QuantizedTensorStorage], List[Optional[QuantizedTensorStorage]]]: + """Quantize a group of weights, optionally reusing cached workspaces. + + Group analogue of :func:`quantize_weight`. For delayed-scaling FP8 the whole + group is cast and transposed with a single fused ``multi_cast_transpose`` + kernel, and on ROCm an MXFP8 group is quantized with a single fused + multi-quantize kernel, instead of one quantize call per tensor. When fusion + is not applicable (other recipes, FP8 rowwise-only usage, already-quantized + weights, or a CUDA-graph skip flag) the call falls back to + :func:`quantize_weight` for each tensor. + + Parameters + ---------- + tensors: list of torch.Tensor + Weight tensors to quantize. + quantizers: list of Quantizer + Quantizers for casting the weights. + workspaces: list of QuantizedTensorStorage, optional + Previously cached workspaces (from the module's ``_fp8_workspaces``). + ``None`` entries indicate a cache miss. + update_workspace: bool, default = True + Whether to update existing workspaces with fresh values. + skip_update_flag: torch.Tensor, optional + GPU flag to conditionally skip the update. + workspace_dtype: torch.dtype, optional + High-precision dtype for debug quantization workspaces. + cache: bool, default = False + If ``True``, brand-new workspaces are returned so the caller can store + them. + + Returns + ------- + (weights, new_workspaces) + *weights*: quantized weights ready for GEMM. + *new_workspaces*: per-tensor entries that are non-``None`` only when a + workspace should be (re)stored in ``_fp8_workspaces`` (i.e. ``cache`` is + ``True`` and the fused/per-tensor path produced a workspace to cache). + """ + num_tensors = len(tensors) + if workspaces is None: + workspaces = [None] * num_tensors + + # Fused path eligibility. Two recipes have a fused multi-tensor quantize kernel: + # * delayed-scaling FP8 (Float8Quantizer): cast + transpose, so it requires + # columnwise usage. + # * MXFP8 (ROCm only): fused multi-quantize kernel that handles whichever + # rowwise/columnwise buffers are allocated. + # Both require high-precision (not already quantized) weights and no CUDA-graph + # skip flag (the fused kernels have no device-side noop and would not preserve + # cached buffer pointers). + fused_fp8 = all( + isinstance(quantizer, Float8Quantizer) and quantizer.columnwise_usage + for quantizer in quantizers + ) + fused_mxfp8 = IS_HIP_EXTENSION and all( + isinstance(quantizer, MXFP8Quantizer) for quantizer in quantizers + ) + can_fuse = ( + num_tensors > 0 + and skip_update_flag is None + and (fused_fp8 or fused_mxfp8) + and not any(isinstance(tensor, QuantizedTensorStorage) for tensor in tensors) + ) + + if can_fuse: + # Validate cached workspaces; treat any invalid/missing entry as a cache miss. + cache_miss = False + for i, (workspace, quantizer) in enumerate(zip(workspaces, quantizers)): + if workspace is not None and not _is_weight_workspace_valid(workspace, quantizer): + workspaces[i] = None + if workspaces[i] is None: + cache_miss = True + + new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * num_tensors + if not cache_miss and not update_workspace: + # All workspaces valid and no refresh requested. + return list(workspaces), new_workspaces + + # Force internal=False so cached workspaces survive prepare_for_saving. + if cache: + saved_internal = [quantizer.internal for quantizer in quantizers] + for quantizer in quantizers: + quantizer.internal = False + if cache_miss: + # Single fused kernel allocates all outputs. + weights = tex.multi_tensor_quantize(list(tensors), quantizers) + else: + # In-place refresh of cached workspaces. + weights = tex.multi_tensor_quantize(list(tensors), quantizers, outputs=list(workspaces)) + if cache: + for quantizer, internal in zip(quantizers, saved_internal): + quantizer.internal = internal + new_workspaces = list(weights) + return list(weights), new_workspaces + + # Fallback: quantize each weight individually. + weights = [] + new_workspaces = [None] * num_tensors + for i in range(num_tensors): + weight, new_workspaces[i] = quantize_weight( + tensor=tensors[i], + quantizer=quantizers[i], + workspace=workspaces[i], + update_workspace=update_workspace, + skip_update_flag=skip_update_flag, + workspace_dtype=workspace_dtype, + cache=cache, + ) + weights.append(weight) + return weights, new_workspaces + + class TransformerEngineBaseModule(torch.nn.Module, ABC): """Base TE module.""" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f3fd19581..d134cdcf1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -20,6 +20,7 @@ from .base import ( get_dummy_wgrad, quantize_weight, + quantize_multi_weight, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -198,20 +199,20 @@ def forward( weights_fp8: list new_workspaces = [None] * num_gemms if fp8 or debug: - weights_fp8 = [] + # FP8 cast to workspace buffer. For delayed-scaling FP8 the whole group is + # cast and transposed with a single fused multi_cast_transpose kernel, and on + # ROCm an MXFP8 group is quantized with a single fused multi-quantize kernel; + # other cases fall back to a per-weight quantize inside quantize_multi_weight. update_ws = is_first_microbatch is None or is_first_microbatch - for i in range(num_gemms): - weight_fp8, new_workspaces[i] = quantize_weight( - tensor=weights[i], - quantizer=weight_quantizers[i], - workspace=weight_workspaces[i] if weight_workspaces else None, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - weights_fp8.append(weight_fp8) - + weights_fp8, new_workspaces = quantize_multi_weight( + tensors=list(weights), + quantizers=weight_quantizers, + workspaces=weight_workspaces if weight_workspaces else [None] * num_gemms, + update_workspace=update_ws, + skip_update_flag=skip_fp8_weight_update, + workspace_dtype=activation_dtype, + cache=cache_weight, + ) else: weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]