From 0097e7890ec27212900d4f29a56222cbd29121f6 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 10 Aug 2026 14:17:54 +0200 Subject: [PATCH 1/3] Fine-grained recipe docs Signed-off-by: Evgeny --- docs/api/pytorch.rst | 18 ++ .../fine_grained_quantization.rst | 270 ++++++++++++++++++ ...torch_fine_grained_quantization_example.py | 152 ++++++++++ .../features/low_precision_training/index.rst | 1 + transformer_engine/common/recipe/__init__.py | 12 +- .../pytorch/tensor/hybrid_tensor.py | 4 + .../pytorch/tensor/identity_tensor.py | 19 ++ 7 files changed, 473 insertions(+), 3 deletions(-) create mode 100644 docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst create mode 100644 docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py 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..fdc204457b --- /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(MXFP8-dequantized 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): From 2f4f64fec21e864d4cc46e7c3664722797bdc740 Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Mon, 10 Aug 2026 16:02:25 +0200 Subject: [PATCH 2/3] Update docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Evgeny Tsykunov --- .../pytorch_fine_grained_quantization_example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index fdc204457b..a1d3e3b377 100644 --- 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 @@ -9,7 +9,7 @@ * fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` * dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` -* wgrad: ``input.col(MXFP8-dequantized BF16) x grad_output.col(original BF16)`` +* 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. From c3b03986e332d064df26a2539d3292bc8746de17 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 17 Aug 2026 14:31:25 +0200 Subject: [PATCH 3/3] resolve comments Signed-off-by: Evgeny --- docs/api/pytorch.rst | 4 +- .../heterogeneous_quantization.rst | 121 ++++++++++ ...rch_heterogeneous_quantization_example.py} | 122 +++++----- .../heterogeneous_quantization.rst} | 208 ++++++++---------- .../img/heterogeneous_linear_mapping.svg | 106 +++++++++ .../img/hybrid_columnwise_source.svg | 79 +++++++ .../features/low_precision_training/index.rst | 2 +- docs/index.rst | 1 + transformer_engine/common/recipe/__init__.py | 2 +- 9 files changed, 456 insertions(+), 189 deletions(-) create mode 100644 docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst rename docs/{features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py => examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py} (63%) rename docs/features/low_precision_training/{fine_grained_quantization/fine_grained_quantization.rst => heterogeneous_quantization/heterogeneous_quantization.rst} (52%) create mode 100644 docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg create mode 100644 docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 963cdadb1f..60556d6cac 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -109,8 +109,8 @@ Communication-computation overlap :members: FP8, NONE -Fine-grained quantization recipes ---------------------------------- +Heterogeneous quantization recipes +---------------------------------- .. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="") diff --git a/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst b/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst new file mode 100644 index 0000000000..e4ec7b2be1 --- /dev/null +++ b/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst @@ -0,0 +1,121 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fine-grained-quantization-tutorial: +.. _heterogeneous-quantization-tutorial: + +Building a heterogeneous quantization recipe +============================================== + +This tutorial demonstrates factory composition, module and name targeting, +the general three-format GEMM mapping, and high-precision directions. See +:doc:`Heterogeneous quantization recipes +<../../features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization>` +for the API concepts and direction mapping. + +Constructing a factory +---------------------- + +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 a TE-native factory by keeping one named ``linear`` +module in high precision, using NVFP4 for every ``grouped_linear`` role, and +retaining MXFP8 as the global fallback: + +.. tabs:: + + .. tab:: PyTorch + + .. 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) + +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. + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: pytorch_heterogeneous_quantization_example.py + :language: python + :start-after: # START_HETEROGENEOUS_QUANTIZATION_EXAMPLE + :end-before: # END_HETEROGENEOUS_QUANTIZATION_EXAMPLE + +Run it from the repository root after installing TE: + +.. code-block:: bash + + python docs/examples/heterogeneous_quantization/pytorch_heterogeneous_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. + +The TE-native base factories in +``transformer_engine/pytorch/custom_recipes/quantizer_factories.py`` construct +standard TE quantizers. They can be used directly as factory fallbacks or +composed as children of a ``HybridQuantizer``. + +More specialized 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 remain illustrative examples rather than official, +broadly validated defaults; read each factory's rationale before adapting it. + +Validating and optimizing a recipe +---------------------------------- + +A factory can describe assignments beyond current optimized kernel coverage. +Before adopting an assignment for a workload: + +* Confirm that the module, GEMM layout and shape, software version, and GPU can + execute it. +* Check whether the selected module path has an appropriate optimized kernel + and integration. +* Validate accuracy and convergence on the target model and distributed setup. +* Benchmark throughput and memory 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. diff --git a/docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py similarity index 63% rename from docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py rename to docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py index a1d3e3b377..c2834d2624 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py +++ b/docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py @@ -3,7 +3,7 @@ # # See LICENSE for license information. -"""Runnable fine-grained quantization recipe example. +"""Runnable heterogeneous quantization recipe example. The factory assigns one precision to each ``demo.fc1`` Linear GEMM: @@ -16,14 +16,38 @@ Run from the Transformer Engine repository root:: - python docs/features/low_precision_training/fine_grained_quantization/\ - pytorch_fine_grained_quantization_example.py + python docs/examples/heterogeneous_quantization/\ + pytorch_heterogeneous_quantization_example.py """ -# START_FINE_GRAINED_QUANTIZATION_EXAMPLE - from __future__ import annotations +import torch +import transformer_engine.pytorch as te + + +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)) + + +require_supported_hardware() + +# START_HETEROGENEOUS_QUANTIZATION_EXAMPLE + from typing import Optional import torch @@ -78,75 +102,29 @@ def quantizer_factory(role: Optional[te.QuantizerRole]): 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() +linear_options = {"bias": False, "params_dtype": torch.bfloat16, "device": "cuda"} +model = torch.nn.Sequential( + te.Linear(128, 256, name=THREE_FORMAT_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 256, name=HIGH_PRECISION_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 128, name="demo.output", **linear_options), +) +inputs = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) +recipe = CustomRecipe(qfactory=quantizer_factory) - 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) +with te.autocast(enabled=True, recipe=recipe): + outputs = model(inputs) - 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") +loss = outputs.float().square().mean() +loss.backward() +# END_HETEROGENEOUS_QUANTIZATION_EXAMPLE -if __name__ == "__main__": - main() +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) -# END_FINE_GRAINED_QUANTIZATION_EXAMPLE +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") diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst similarity index 52% rename from docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst rename to docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst index 8b28371c6a..e7087e422b 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst @@ -4,14 +4,16 @@ See LICENSE for license information. .. _fine-grained-quantization-recipes: +.. _heterogeneous-quantization-recipes: -Fine-grained quantization recipes -================================= +Heterogeneous 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 +Transformer Engine (TE) supports heterogeneous quantization recipes that select +quantizers by module or operation type, tensor role, module or operation +instance name, and rowwise/columnwise direction or GEMM type. Heterogeneous recipes +provide role-aware mixed-precision and quantization configuration at module, +tensor-role, and GEMM-direction granularity. 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 @@ -19,53 +21,56 @@ factory can compose TE-native quantizers with :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:: +Mixing formats at a glance +-------------------------- - 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. +A single ``CustomRecipe`` can mix quantization formats and high precision +across modules and GEMM directions. The accompanying +:doc:`tutorial <../../../examples/heterogeneous_quantization/heterogeneous_quantization>` +makes the following assignments: -.. warning:: +.. list-table:: + :header-rows: 1 + :widths: 34 22 22 22 - 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. + * - Module assignment + - Fprop + - Dgrad + - Wgrad + * - ``demo.fc1`` + - MXFP8 + - NVFP4 + - BF16 + * - ``demo.fc2`` + - BF16 + - BF16 + - BF16 + * - Other TE modules + - MXFP8 + - MXFP8 + - MXFP8 -Configuration readiness ------------------------ +Once the factory defines these assignments, the recipe uses the standard TE +autocast path: -Treat fine-grained construction as an experimental recipe exploration surface. -Keep these readiness levels distinct: +.. tabs:: -.. list-table:: - :header-rows: 1 - :widths: 25 75 + .. tab:: PyTorch - * - 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. + .. code-block:: python + + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + +The complete factory appears in the tutorial. The same machinery can also: + +* assign formats by module or operation type; +* override a named module instance; +* choose fprop, dgrad, and wgrad formats independently; and +* keep selected slots or directions in high precision. Factory contract ---------------- @@ -99,43 +104,14 @@ The role vocabulary includes: 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")``. +See the +:doc:`tutorial <../../../examples/heterogeneous_quantization/heterogeneous_quantization>` +for factory construction rules and complete examples. + Linear GEMM direction mapping ----------------------------- @@ -168,13 +144,19 @@ for wgrad, map to tensor quantizers as: weight = Hybrid(rowwise=F, columnwise=D) grad_output = Hybrid(rowwise=D, columnwise=W) +.. raw:: html + :file: img/heterogeneous_linear_mapping.svg + +*Figure 1. Fine-grained tensor representations provide matching operand +formats for each Linear GEMM.* + 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). +One factory may return both plain and hybrid quantizers (see the +:doc:`tutorial <../../../examples/heterogeneous_quantization/heterogeneous_quantization>`). Combining rowwise and columnwise quantizers ------------------------------------------- @@ -201,6 +183,12 @@ source for the columnwise representation: * - ``"rowwise_dequantized"`` - Dequantized rowwise representation. +.. raw:: html + :file: img/hybrid_columnwise_source.svg + +*Figure 2. The columnwise representation can be derived from the original +high-precision tensor or from the 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. @@ -215,52 +203,46 @@ Keeping directions in high precision 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 +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python - return te.HybridQuantizer( - rowwise_quantizer=mxfp8_factory(role), - columnwise_quantizer=te.IdentityQuantizer(), - columnwise_source="rowwise_dequantized", - ) + 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. +Tutorial +-------- -.. literalinclude:: pytorch_fine_grained_quantization_example.py - :language: python - :start-after: # START_FINE_GRAINED_QUANTIZATION_EXAMPLE - :end-before: # END_FINE_GRAINED_QUANTIZATION_EXAMPLE +See :doc:`Building a heterogeneous quantization recipe +<../../../examples/heterogeneous_quantization/heterogeneous_quantization>` for +factory composition, a runnable example, recipe starting points, and workload +validation guidance. -Run it from the repository root after installing TE: +Support status +-------------- -.. code-block:: bash +.. note:: - python docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py + 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. -Recipe starting points ----------------------- +.. warning:: -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. + Fine-grained recipes and their construction APIs are experimental. API, + validation, and kernel coverage may change without notice. This guide does + not define a supported recipe or an expected accuracy/performance ordering. API reference ------------- diff --git a/docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg b/docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg new file mode 100644 index 0000000000..0910629db0 --- /dev/null +++ b/docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg @@ -0,0 +1,106 @@ + + + Heterogeneous precision mapping for Linear GEMMs + Input, weight, and output-gradient tensors each provide rowwise and columnwise representations. Forward uses format F, input-gradient uses format D, and weight-gradient uses format W. + + + + + + + + Tensor representations and Linear GEMMs + Quantizer outputs by tensor role + + + + Input + + rowwise + F + + columnwise + W + + + + + Weight + + rowwise + F + + columnwise + D + + + + + Grad output + + rowwise + D + + columnwise + W + + + TE selects the required direction for each GEMM + + + FPROP + + weight.rowwise + format F + × + + input.rowwise + format F + + + Output + + + + DGRAD + + weight.columnwise + format D + × + + grad_output.rowwise + format D + + + Input grad. + + + + WGRAD + + input.columnwise + format W + × + + grad_output.columnwise + format W + + + Weight grad. + + diff --git a/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg new file mode 100644 index 0000000000..d7aa328fe7 --- /dev/null +++ b/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg @@ -0,0 +1,79 @@ + + + Hybrid quantizer columnwise source choices + With original provenance, both quantizers consume the original high-precision tensor. With rowwise-dequantized provenance, the columnwise quantizer consumes the dequantized rowwise representation. + + + + + + + + Choosing the columnwise source + + + + columnwise_source="original" + + + High-precision tensor + + + + same original source + + + Rowwise quantizer + + Columnwise quantizer + + + + + Rowwise + representation + + Columnwise + representation + + + + + columnwise_source="rowwise_dequantized" + + + High-precision tensor + + + + Rowwise quantizer + + + + Rowwise + representation + + + + Dequantize + + + + Columnwise quantizer + + + Columnwise + representation + + diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst index b9649c00a4..8e963c9d1b 100644 --- a/docs/features/low_precision_training/index.rst +++ b/docs/features/low_precision_training/index.rst @@ -15,5 +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 + heterogeneous_quantization/heterogeneous_quantization.rst speedups.rst diff --git a/docs/index.rst b/docs/index.rst index fcd15a7a11..414c4d87b3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -60,6 +60,7 @@ Transformer Engine documentation examples/onnx/onnx_export.ipynb examples/te_jax_integration.rst examples/op_fuser/op_fuser.rst + examples/heterogeneous_quantization/heterogeneous_quantization.rst examples/gemm_profiling/gemm_profiling.rst .. toctree:: diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 36826541d3..bdbaf224bd 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -664,7 +664,7 @@ class CustomRecipe(Recipe): See ``transformer_engine.pytorch.quantization.QuantizerRole`` and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` - for API details. See :ref:`fine-grained-quantization-recipes` for + for API details. See :ref:`heterogeneous-quantization-recipes` for construction rules and direction mapping. backward_override : {None, 'high_precision', 'dequantized'}, default = None