[PyTorch] Document heterogeneous quantization recipes

docs/api/pytorch.rst

@@ -109,6 +109,12 @@ Communication-computation overlap
   :members: FP8, NONE
 
 
+Heterogeneous 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
 -------------------------------------
 

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
+
+         <div style="background: #f0f4f8; border-left: 3px solid #5c7cfa; padding: 6px 12px; font-size: 13px; color: #495057; margin-bottom: 0; border-radius: 4px 4px 0 0;">
+            Requires SM100 (Blackwell) or later
+         </div>
+
+      .. 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.

docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Runnable heterogeneous 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/examples/heterogeneous_quantization/\
+        pytorch_heterogeneous_quantization_example.py
+"""
+
+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
+
+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)
+
+
+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)
+
+with te.autocast(enabled=True, recipe=recipe):
+    outputs = model(inputs)
+
+loss = outputs.float().square().mean()
+loss.backward()
+
+# END_HETEROGENEOUS_QUANTIZATION_EXAMPLE
+
+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")