diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5fac0a89a6..963cdadb1f 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -109,6 +109,12 @@ Communication-computation overlap :members: FP8, NONE +Fine-grained quantization recipes +--------------------------------- + +.. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="") + + Quantized tensors ----------------- @@ -126,6 +132,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensorStorage(*, rowwise_storage, columnwise_storage, quantizer, fake_dtype=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensorStorage(*, hp_data, fake_dtype=None, quantizer=None) + .. autoapiclass:: transformer_engine.pytorch.Float8Tensor(shape, dtype, data, fp8_scale_inv, fp8_dtype, requires_grad=False, data_transpose=None, quantizer=None) .. autoapiclass:: transformer_engine.pytorch.MXFP8Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer) @@ -134,6 +144,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensor(shape, dtype, *, rowwise_storage, columnwise_storage, quantizer, requires_grad=False, device=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensor(shape, dtype, *, hp_data, quantizer=None, requires_grad=False, device=None) + Quantizers ---------- @@ -150,6 +164,10 @@ Quantizers .. autoapiclass:: transformer_engine.pytorch.NVFP4Quantizer(fp4_dtype, *, rowwise=True, columnwise=True, **kwargs) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizer(*, rowwise_quantizer, columnwise_quantizer, columnwise_source="original") + +.. autoapiclass:: transformer_engine.pytorch.IdentityQuantizer(*, dtype=None, rowwise=True, columnwise=True) + Tensor saving and restoring functions ------------------------------------- diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst new file mode 100644 index 0000000000..8b28371c6a --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -0,0 +1,270 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fine-grained-quantization-recipes: + +Fine-grained quantization recipes +================================= + +Transformer Engine (TE) can select quantizers by module or operation type, +tensor role, module or operation instance name, and rowwise or columnwise +direction. This enables **per-GEMM granularity of the precision format and/or +quantization logic**. A +:class:`~transformer_engine.common.recipe.CustomRecipe` supplies a quantizer +factory to the standard :class:`~transformer_engine.pytorch.autocast` path. The +factory can compose TE-native quantizers with +:class:`~transformer_engine.pytorch.HybridQuantizer` and +:class:`~transformer_engine.pytorch.IdentityQuantizer`. + +This guide covers PyTorch, TE-native quantizers, and static recipe construction. +It does not define a supported recipe or an expected accuracy/performance +ordering. Validate every configuration on the target model, hardware, and +distributed setup. + +.. note:: + + With TE-native low-precision quantizers on supported hardware and kernel + paths, recipes use TE's native GPU quantization and low-precision GEMM + implementations. No fake quantization or high-precision GEMM emulation is + involved on these paths. + +.. warning:: + + Fine-grained recipes and their construction APIs are experimental. API, + validation, and kernel coverage may change without notice. A configuration + that can be expressed by the API is not necessarily executable for every + module, GEMM shape, software version, or GPU. An executable configuration + is not necessarily optimized or validated for accuracy and convergence on + a particular workload. + +Configuration readiness +----------------------- + +Treat fine-grained construction as an experimental recipe exploration surface. +Keep these readiness levels distinct: + +.. list-table:: + :header-rows: 1 + :widths: 25 75 + + * - Level + - Meaning + * - Expressible + - A factory can describe the role and direction assignment. + * - Executable + - The current module, GEMM backend, layout, software, and GPU accept it. + * - Optimized + - The selected path has an appropriate optimized kernel and integration. + * - Workload-validated + - Accuracy, convergence, throughput, and memory have been measured on the + target workload. + +Fine-grained assignments provide a way to explore an accuracy/performance +slider by varying precision and quantization logic per GEMM. The recipes that +can be realized efficiently are constrained by available kernels and +integrations. Accuracy and convergence experiments can run on functionally +executable, non-optimized paths before dedicated kernels are available. + +Factory contract +---------------- + +Each TE module defines an ordered role list for the forward and backward +quantizer slots it needs. When module recipe state is initialized or rebuilt, +a ``CustomRecipe`` calls ``qfactory(role)`` once for every slot in that list. +It does not call the factory on every unchanged forward. + +The role vocabulary includes: + +.. list-table:: + :header-rows: 1 + :widths: 25 35 40 + + * - Field + - Examples + - Meaning + * - ``module_type`` + - ``"linear"``, ``"grouped_linear"``, ``"dpa"`` + - TE-defined module or operation type, populated by the TE module. + * - ``tensor_type`` + - ``"input"``, ``"weight"``, ``"grad_output"`` + - TE-defined slot in that module's vocabulary, populated by the TE module. + * - ``name`` + - ``"decoder.39.qkv"``, ``"decoder.39.fc2"`` + - Caller or framework-provided instance identity. Composite TE modules + may append suffixes for nested operations. + +``module_type`` and ``tensor_type`` are TE-defined selectors populated by the +module. The caller or framework supplies the root ``name``; composite TE +modules may extend it with suffixes such as ``.fc1``, ``.fc2``, and ``.proj``. + +A robust factory follows four construction rules: + +* Return a quantizer for every role. Return ``IdentityQuantizer`` for an + intentional high-precision slot; do not return ``None``. +* Constructing a fresh quantizer for every call is recommended. + ``HybridQuantizer`` owns and configures its rowwise and columnwise children. +* A module-level function is the most portable factory definition, especially + when a launcher or checkpointing setup needs to import or pickle it. +* Treat role strings as selectors, not a fixed enumeration. Preserve a base + factory fallback for roles the factory does not recognize. + +For example, compose TE-native factories by keeping one named ``linear`` module +in high precision, using NVFP4 for every ``grouped_linear`` role, and retaining +MXFP8 as the global fallback: + +.. code-block:: python + + from typing import Optional + + import transformer_engine.pytorch as te + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, + ) + + def my_factory(role: Optional[te.QuantizerRole]): + if role is not None: + if role.module_type == "linear" and role.name == "decoder.39.fc2": + return te.IdentityQuantizer() + if role.module_type == "grouped_linear": + return nvfp4_factory(role) + return mxfp8_factory(role) + +The training framework or caller must pass semantic names to TE modules for +name-based selection, for example +``te.Linear(..., name="decoder.39.fc2")``. + +Linear GEMM direction mapping +----------------------------- + +``Linear`` and ``GroupedLinear`` training consume rowwise and columnwise +representations as follows: + +.. list-table:: + :header-rows: 1 + :widths: 20 40 40 + + * - GEMM + - First operand + - Second operand + * - Forward (fprop) + - ``weight.rowwise`` + - ``input.rowwise`` + * - Input gradient (dgrad) + - ``weight.columnwise`` + - ``grad_output.rowwise`` + * - Weight gradient (wgrad) + - ``input.columnwise`` + - ``grad_output.columnwise`` + +Therefore three per-GEMM formats, ``F`` for fprop, ``D`` for dgrad, and ``W`` +for wgrad, map to tensor quantizers as: + +.. code-block:: text + + input = Hybrid(rowwise=F, columnwise=W) + weight = Hybrid(rowwise=F, columnwise=D) + grad_output = Hybrid(rowwise=D, columnwise=W) + +If two directions use the same quantizer configuration, a plain quantizer may +replace the corresponding hybrid. The two operands of each GEMM still need a +combination supported by that GEMM backend. TE may reject incompatible +quantizer pairs or unsupported layouts. + +One factory may return both plain and hybrid quantizers (see the runnable +example below). + +Combining rowwise and columnwise quantizers +------------------------------------------- + +:class:`~transformer_engine.pytorch.HybridQuantizer` composes a rowwise and a +columnwise quantizer. Its output, +:class:`~transformer_engine.pytorch.HybridQuantizedTensor`, composes the +corresponding representations. + +Choosing the columnwise source +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``columnwise_source`` is a separate numerical recipe choice that controls the +source for the columnwise representation: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Value + - Columnwise source + * - ``"original"`` + - The original high-precision tensor. + * - ``"rowwise_dequantized"`` + - Dequantized rowwise representation. + +For forward inputs and weights, ``"rowwise_dequantized"`` derives the backward +representation from the value consumed in the forward direction. This +can improve forward/backward numerical consistency and may affect convergence. +It does not recover information discarded by rowwise quantization. +``"original"`` instead derives both representations from the original tensor. +Choose the provenance as part of the numerical recipe. + +Keeping directions in high precision +------------------------------------ + +:class:`~transformer_engine.pytorch.IdentityQuantizer` stores its input in the +held compute dtype, typically BF16, FP16, or FP32. It can keep a complete slot +in high precision or act as one child of a ``HybridQuantizer``: + +.. code-block:: python + + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + +In this example, the rowwise direction uses MXFP8. The columnwise direction is +held in high precision, but its value is reconstructed from MXFP8. Use +``columnwise_source="original"`` when the high-precision direction should +retain the original input value instead. + +Runnable example +---------------- + +The following synthetic example demonstrates base-factory composition, the +general three-format mapping, high-precision directions, and module/name +targeting. It uses only TE-native quantizers. MXFP8 and NVFP4 execution requires +supported hardware and software. + +.. literalinclude:: pytorch_fine_grained_quantization_example.py + :language: python + :start-after: # START_FINE_GRAINED_QUANTIZATION_EXAMPLE + :end-before: # END_FINE_GRAINED_QUANTIZATION_EXAMPLE + +Run it from the repository root after installing TE: + +.. code-block:: bash + + python docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py + +Recipe starting points +---------------------- + +The runnable example above is deliberately synthetic: it demonstrates the +expressiveness of the API, not a recommended training recipe. More realistic +starting points are available in +``transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py``. Some +zoo factories encode externally described recipe structures or have specific +motivating evidence. They are still illustrative examples rather than +official, broadly validated defaults. Read each factory's rationale and +validate accuracy, convergence, and performance on the target workload. +Realizing the intended performance may require dedicated kernel enablement for +the selected operand formats, layouts, or module path; functional execution +does not imply that an optimized kernel path exists. + +API reference +------------- + +See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``, +``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types. +See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``. diff --git a/docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py new file mode 100644 index 0000000000..a1d3e3b377 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Runnable fine-grained quantization recipe example. + +The factory assigns one precision to each ``demo.fc1`` Linear GEMM: + +* fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` +* dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` +* wgrad: ``input.col(original BF16) x grad_output.col(original BF16)`` + +``demo.fc2`` runs every GEMM in high precision. ``demo.output`` is not +special-cased and therefore exercises the MXFP8 base-factory fallback. + +Run from the Transformer Engine repository root:: + + python docs/features/low_precision_training/fine_grained_quantization/\ + pytorch_fine_grained_quantization_example.py +""" + +# START_FINE_GRAINED_QUANTIZATION_EXAMPLE + +from __future__ import annotations + +from typing import Optional + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, +) + + +THREE_FORMAT_MODULE = "demo.fc1" +HIGH_PRECISION_MODULE = "demo.fc2" +BASE_FACTORY = mxfp8_factory + + +def quantizer_factory(role: Optional[te.QuantizerRole]): + """Return a fresh quantizer for every role, including ``None``. + + ``BASE_FACTORY`` makes the factory total: unknown roles, future role values, + and untargeted modules all retain valid MXFP8 behavior. + """ + + if role is not None and role.name == THREE_FORMAT_MODULE: + # Constructing fresh child quantizers for every call is recommended. + if role.tensor_type == "input": + # Wgrad retains the original BF16 input. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # Dgrad uses NVFP4 quantized from the dequantized MXFP8 fprop weight. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # Dgrad uses NVFP4 while wgrad retains the original BF16 gradient. + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + + if role is not None and role.name == HIGH_PRECISION_MODULE: + return te.IdentityQuantizer() + + return BASE_FACTORY(role) + + +def require_supported_hardware() -> None: + """Fail early with TE's reason when either required format is unavailable.""" + + if not torch.cuda.is_available(): + raise SystemExit("This example requires a CUDA-capable NVIDIA GPU.") + + failures = [] + for name, check in ( + ("MXFP8", te.is_mxfp8_available), + ("NVFP4", te.is_nvfp4_available), + ): + available, reason = check(return_reason=True) + if not available: + failures.append(f"{name}: {reason}") + if failures: + raise SystemExit("Required formats are unavailable: " + "; ".join(failures)) + + +def build_model() -> torch.nn.Module: + """Build aligned TE Linear layers with stable semantic names.""" + + common = { + "bias": False, + "params_dtype": torch.bfloat16, + "device": "cuda", + } + return torch.nn.Sequential( + te.Linear(128, 256, name=THREE_FORMAT_MODULE, **common), + torch.nn.GELU(), + te.Linear(256, 256, name=HIGH_PRECISION_MODULE, **common), + torch.nn.GELU(), + te.Linear(256, 128, name="demo.output", **common), + ) + + +def main() -> None: + """Run one training step through the custom recipe.""" + + require_supported_hardware() + torch.manual_seed(1234) + torch.cuda.manual_seed_all(1234) + + model = build_model() + recipe = CustomRecipe(qfactory=quantizer_factory) + inputs = torch.randn( + 64, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + with te.autocast(enabled=True, recipe=recipe): + outputs = model(inputs) + + loss = outputs.float().square().mean() + # Backward uses the quantizers selected and saved during the forward pass. + loss.backward() + + gradients = [inputs.grad, *(parameter.grad for parameter in model.parameters())] + assert all(gradient is not None for gradient in gradients) + assert all(torch.isfinite(gradient).all() for gradient in gradients) + + print(f"GPU: {torch.cuda.get_device_name()}") + print(f"TE Linear names: {[model[index].name for index in (0, 2, 4)]}") + print(f"loss: {loss.item():.6f}; forward and backward completed") + + +if __name__ == "__main__": + main() + +# END_FINE_GRAINED_QUANTIZATION_EXAMPLE diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst index 0a798f1364..b9649c00a4 100644 --- a/docs/features/low_precision_training/index.rst +++ b/docs/features/low_precision_training/index.rst @@ -15,4 +15,5 @@ Low precision training fp8_blockwise_scaling/fp8_blockwise_scaling.rst mxfp8/mxfp8.rst nvfp4/nvfp4.rst + fine_grained_quantization/fine_grained_quantization.rst speedups.rst diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index a89ddba917..36826541d3 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -635,13 +635,18 @@ class CustomRecipe(Recipe): ---------- qfactory : Callable Factory callable that returns a quantizer instance *or* a - ``QuantizerRequest`` subclass for a given ``QuantizerRole``. + ``QuantizerRequest`` subclass for a given optional ``QuantizerRole``. The callable is invoked as:: qfactory( - role: QuantizerRole, + role: Optional[QuantizerRole], ) -> Union[Quantizer, QuantizerRequest] + Boundary slots may provide ``None`` or a role with empty fields. The + factory must return a valid object for every call. Return an + ``IdentityQuantizer`` for an intentional high-precision slot instead + of returning ``None``. + ``QuantizerRole`` is a frozen dataclass with the following fields: - ``module_type`` (str): module type (empty string when not set), e.g. @@ -659,7 +664,8 @@ class CustomRecipe(Recipe): See ``transformer_engine.pytorch.quantization.QuantizerRole`` and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` - for full documentation. + for API details. See :ref:`fine-grained-quantization-recipes` for + construction rules and direction mapping. backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index dc65c9894b..4b9564ee23 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -19,6 +19,10 @@ class HybridQuantizer(Quantizer): """Quantizer that composes rowwise and columnwise representations. + .. warning:: + **EXPERIMENTAL**: ``HybridQuantizer`` is under active development and + its API is subject to change without notice. + When both representations are requested, applies ``rowwise_quantizer`` to produce the rowwise representation and ``columnwise_quantizer`` to produce the columnwise representation. The results are wrapped in a diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 9fb980a755..e21aafc3d9 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -26,6 +26,10 @@ class IdentityQuantizer(Quantizer): """Quantizer that produces a high-precision passthrough representation. + .. warning:: + **EXPERIMENTAL**: ``IdentityQuantizer`` is under active development and + its API is subject to change without notice. + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) holding the tensor directly, without a low-precision encoding. ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it @@ -174,6 +178,21 @@ class IdentityTensor(IdentityTensorStorage, QuantizedTensor): Presents as a standard tensor of its nominal dtype; internally it just holds data directly in that dtype, without a low-precision encoding. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Logical tensor datatype. + hp_data : torch.Tensor + Held high-precision data. + quantizer : IdentityQuantizer, optional + Quantizer that produced the tensor. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + device : torch.device, optional + Device containing the tensor. """ def __repr__(self, *, tensor_contents=None):