[PyTorch] Docs: fine-grained quantization recipes (rework of NVIDIA#3336)

docs/_static/css/diagram-colors.css

@@ -279,3 +279,18 @@ html[data-theme="dark"] .subtitle,
 html[data-theme="dark"] .memory-label { fill: #e0e0e0; }
 html[data-theme="dark"] .connector { stroke: #bdbdbd; }
 
+
+/* fine_grained_quantization diagrams */
+html[data-theme="dark"] .fmt-mxfp8 { fill: #10375c; stroke: #64b5f6; }
+html[data-theme="dark"] .fmt-nvfp4 { fill: #5c3a10; stroke: #ffb74d; }
+html[data-theme="dark"] .fmt-bf16 { fill: #1e4620; stroke: #81c784; }
+html[data-theme="dark"] .fmt-mxfp8-text { fill: #90caf9; }
+html[data-theme="dark"] .fmt-nvfp4-text { fill: #ffcc80; }
+html[data-theme="dark"] .fmt-bf16-text { fill: #a5d6a7; }
+html[data-theme="dark"] .source { fill: #3a2f5c; stroke: #b39ddb; }
+html[data-theme="dark"] .quantizer { fill: #1e4620; stroke: #81c784; }
+html[data-theme="dark"] .representation { fill: #10375c; stroke: #64b5f6; }
+html[data-theme="dark"] .dequantize { fill: #5c3a10; stroke: #ffb74d; }
+html[data-theme="dark"] .rowlabel,
+html[data-theme="dark"] .legend,
+html[data-theme="dark"] .op { fill: #e0e0e0; }

docs/api/pytorch.rst

@@ -112,6 +112,12 @@ Communication-computation overlap
   :members: FP8, NONE
 
 
+Fine-grained quantization recipes
+---------------------------------
+
+.. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="")
+
+
 Quantized tensors
 -----------------
 
@@ -129,6 +135,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)
@@ -137,6 +147,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
 ----------
 
@@ -153,6 +167,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/fine_grained_quantization/pytorch_fine_grained_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 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/examples/fine_grained_quantization/\
+        pytorch_fine_grained_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_FINE_GRAINED_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_FINE_GRAINED_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")