diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index f5dc7da4dd..9ee5827bd1 100644 --- a/docs/_static/css/diagram-colors.css +++ b/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; } diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 54981c9086..87a81e8c65 100644 --- a/docs/api/pytorch.rst +++ b/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 ------------------------------------- diff --git a/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py new file mode 100644 index 0000000000..b511fc41dc --- /dev/null +++ b/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") 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..301b655b3f --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -0,0 +1,364 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fine-grained-quantization-recipes: +.. _heterogeneous-quantization-recipes: + +Fine-grained quantization recipes +================================= + +Standard TE recipes quantize the whole model the same way. That is often too +coarse: one sensitive layer may need BF16 while the rest runs in MXFP8, or a +gradient GEMM may tolerate a cheaper format than the forward pass. Fine-grained +recipes lift this restriction: you write a small factory function that picks a +quantizer for each slot TE asks about, and pass it via +:class:`~transformer_engine.common.recipe.CustomRecipe` to the usual +:class:`~transformer_engine.pytorch.autocast`. +"Fine-grained" refers to the granularity of that choice (per module, tensor +role, and GEMM direction), not to the block size of the scaling factors. + +.. warning:: + + Fine-grained recipes are currently available only in the PyTorch API of + TE. + +.. warning:: + + 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. + + +Example: mixing MXFP8, NVFP4, and BF16 +-------------------------------------- + +The `runnable example `__ +makes the following assignments: + +.. raw:: html + :file: img/fine_grained_assignments.svg + +*Figure 1. Precision assignments per module and GEMM used throughout this +guide.* + +A minimal factory implementing these assignments, plugged into the standard +TE autocast path: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + 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, + ) + + + def quantizer_factory(role): + if role is not None and role.name == "demo.fc1": + if role.tensor_type == "input": + # wgrad keeps the original BF16 input + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # fprop in MXFP8, dgrad in NVFP4 + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # dgrad in NVFP4, wgrad keeps 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 == "demo.fc2": + return te.IdentityQuantizer() # whole module stays in BF16 + return mxfp8_factory(role) # every other TE module in MXFP8 + + + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + +The complete, runnable version is available +`on GitHub `__ +(requires Blackwell or later); run it from the repository root after +installing TE: + +.. code-block:: bash + + python docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py + +CustomRecipe and quantizer factory +---------------------------------- + +:class:`~transformer_engine.common.recipe.CustomRecipe` is used like any +other TE recipe (``DelayedScaling``, ``MXFP8BlockScaling``, ...), but carries +no quantization logic of its own: TE asks your ``qfactory`` for a quantizer +whenever a module needs one. + +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. + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + # QuantizerRole describes the slot being configured (fields below): + # + # @dataclasses.dataclass(frozen=True) + # class QuantizerRole: + # module_type: str = "" + # tensor_type: str = "" + # name: str = "" + + + def quantizer_factory(role: Optional[te.QuantizerRole]): + # construct a fresh quantizer on every call + ... + # Boundary slots may pass role=None or a role with empty fields, so + # always end with a default that covers every remaining role. + return mxfp8_factory(role) + + + # The factory plugs into the standard TE autocast path: + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + + **Module type** + + The kind of TE module that owns the slot, filled in by TE itself: + + * ``"linear"`` — ``Linear``, ``LayerNormLinear``, ``fc1``/``fc2`` in + ``LayerNormMLP``, ``qkv``/``proj`` in ``MultiheadAttention``; + * ``"grouped_linear"`` — ``GroupedLinear``; + * ``"dpa"`` — ``DotProductAttention``. + + **Tensor type** + + Which tensor of that module the quantizer will process, also filled in + by TE. For ``"linear"`` and ``"grouped_linear"``: + + * ``"input"`` — the activation (fprop, wgrad); + * ``"weight"`` — (fprop, dgrad); + * ``"grad_output"`` — the incoming gradient (dgrad, wgrad). + + For ``"dpa"``: + + * ``"qkv"`` — the query/key/value tensor; + * ``"s"`` — the softmax output; + * ``"do"`` — the output gradient; + * ``"dp"`` — the gradient of ``"s"``. + + **Name** + + The identity of one concrete module instance, supplied by the caller: + ``te.Linear(..., name="decoder.39.fc2")``. Composite TE modules may + append suffixes such as ``.fc1``, ``.fc2``, and ``.proj``. + + The role vocabulary is experimental and may grow between releases — + one more reason to end the factory with a total default. Treat the role + strings as selectors, not a fixed enumeration. Prefer a module-level + function for the factory itself, so that launchers and checkpointing + setups can import or pickle it. + + TE provides factories for its native quantizers in + ``transformer_engine.pytorch.custom_recipes.quantizer_factories`` + (``mxfp8_factory``, ``nvfp4_factory``, ...). They can be used as + defaults or to construct ``HybridQuantizer`` children. Additional + specialized recipes are available in + ``transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo``. + + The factory is not limited to TE-native quantizers: it may return your + own :class:`~transformer_engine.pytorch.Quantizer` subclass, and custom + quantizers can also serve as ``HybridQuantizer`` children. The GEMMs + still need to receive representations in formats they support. + +HybridQuantizer +--------------- + +During training, each tensor of a ``Linear`` or ``GroupedLinear`` layer feeds +two different GEMMs: its rowwise representation feeds one, its columnwise +representation the other (the exact operand layout is described in the +:doc:`Introduction <../introduction/introduction>`). Since those two GEMMs may +want different formats, the tensor needs a quantizer per direction: +:class:`~transformer_engine.pytorch.HybridQuantizer` composes a rowwise and a +columnwise quantizer, and its output, +:class:`~transformer_engine.pytorch.HybridQuantizedTensor`, composes the +corresponding representations. + +.. tabs:: + + .. tab:: PyTorch + + The following is pseudocode illustrating the composition: + + .. code-block:: text + + quantizer = te.HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), + columnwise_quantizer=NVFP4Quantizer(), + columnwise_source="original", # or "rowwise_dequantized" + ) + + # Quantization yields a HybridQuantizedTensor whose rowwise + # representation is MXFP8 and columnwise representation is NVFP4; + # each GEMM consumes the representation it needs. + qtensor = quantizer(tensor) + +.. raw:: html + :file: img/hybrid_quantizer.svg + +*Figure 2. HybridQuantizer composes a rowwise and a columnwise quantizer; each +representation of the result feeds a different GEMM.* + +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. + +.. raw:: html + :file: img/hybrid_columnwise_source.svg + +*Figure 3. 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. +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. + +IdentityQuantizer +----------------- + +: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``: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + # whole slot in high precision (e.g. a module kept in BF16) + quantizer = te.IdentityQuantizer() + + # one direction in high precision, the other quantized + quantizer = te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + +Note that in the second example the columnwise direction is high precision but +holds the value reconstructed from MXFP8, not the original input — see +`Choosing the columnwise source`_ above. + +Example: one format per GEMM +---------------------------- + +A natural way to design a recipe is to pick one format for each GEMM. To +translate that into quantizers, look at what each GEMM consumes — both of its +operands must be in that GEMM's format: + +* **fprop** consumes ``input.rowwise`` and ``weight.rowwise``; +* **dgrad** consumes ``grad_output.rowwise`` and ``weight.columnwise``; +* **wgrad** consumes ``input.columnwise`` and ``grad_output.columnwise``. + +Reading the same table per tensor gives the ``HybridQuantizer`` for each role. +For the example assignments (fprop in MXFP8, dgrad in NVFP4, wgrad in BF16): + +.. code-block:: text + + input = HybridQuantizer(rowwise=MXFP8, columnwise=BF16) # fprop | wgrad + weight = HybridQuantizer(rowwise=MXFP8, columnwise=NVFP4) # fprop | dgrad + grad_output = HybridQuantizer(rowwise=NVFP4, columnwise=BF16) # dgrad | wgrad + +.. raw:: html + :file: img/fine_grained_linear_mapping.svg + +*Figure 4. Each GEMM consumes one representation of each of its two operand +tensors; giving both operands the same format sets that GEMM's precision.* + +If two directions use the same quantizer configuration, a plain quantizer may +replace the corresponding hybrid; one factory may return both plain and hybrid +quantizers. +The two operands of each GEMM still need a combination supported by that GEMM +backend. TE may reject incompatible quantizer pairs or unsupported layouts. + +.. note:: + + On supported hardware these recipes run TE's regular quantized kernels: the + tensors are quantized on the GPU and the GEMMs execute in the selected + low-precision formats. TE does not fall back to fake quantization + (quantize-dequantize followed by a high-precision GEMM). + + +Validating and optimizing a recipe +---------------------------------- + +The factory API can express more recipes than TE has kernels for, so any +assignment lands in one of three buckets: + +* **Fast** — quantization hits TE's fused kernels and every GEMM runs a + native low-precision implementation. +* **Correct but potentially unoptimized** — the recipe executes, but some + selected paths may not have fused or optimized implementations in the + current TE release. For example, ``HybridQuantizer`` may produce its rowwise + and columnwise representations in separate kernel launches; future releases + may fuse this work. +* **Rejected** — the two operands of some GEMM end up in a combination of + formats or layouts that no GEMM backend supports, and TE raises an error. + This can happen with plain and hybrid quantizers alike. + +Before adopting a recipe for a real workload, check that: + +* it executes at all on the target GPU, software version, and modules; +* it runs on optimized kernels rather than fallback paths; +* accuracy and convergence hold on the target model and distributed setup; +* throughput and memory actually improve on the target workload. + +The unoptimized paths are still useful: accuracy and convergence experiments can run +on them before dedicated kernels exist, so the precision of each GEMM can be +treated as an accuracy/performance trade-off to explore. + +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/img/fine_grained_assignments.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg new file mode 100644 index 0000000000..b9aa48333d --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg @@ -0,0 +1,95 @@ + + + Precision assignment by tensor role and module + Each tensor role provides a rowwise and a columnwise representation, consumed by fprop, dgrad, and wgrad GEMMs. demo.fc1: input is MXFP8 rowwise and original BF16 columnwise; weight is MXFP8 rowwise and NVFP4 columnwise; grad_output is NVFP4 rowwise and original BF16 columnwise. demo.fc2 keeps every tensor in BF16. Other TE modules use MXFP8 everywhere. + + + + + Precision assignment by tensor role and module + + demo.fc1 + demo.fc2 + Other TE modules + + + + input + rowwise (fprop) + + MXFP8 + + BF16 + + MXFP8 + + columnwise (wgrad) + + BF16 (original) + + BF16 + + MXFP8 + + + + + weight + rowwise (fprop) + + MXFP8 + + BF16 + + MXFP8 + + columnwise (dgrad) + + NVFP4 + + BF16 + + MXFP8 + + + + + grad_output + rowwise (dgrad) + + NVFP4 + + BF16 + + MXFP8 + + columnwise (wgrad) + + BF16 (original) + + BF16 + + MXFP8 + + + + + MXFP8 + + NVFP4 + + BF16 (high precision) + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg new file mode 100644 index 0000000000..0e17ffdfa4 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg @@ -0,0 +1,59 @@ + + + Per-GEMM formats and the operands each GEMM consumes + Three GEMM cards. Fprop consumes input.rowwise and weight.rowwise, both MXFP8. Dgrad consumes weight.columnwise and grad_output.rowwise, both NVFP4. Wgrad consumes input.columnwise and grad_output.columnwise, both BF16. + + + + + + + fprop + format: MXFP8 + + input.rowwise + MXFP8 + × + + weight.rowwise + MXFP8 + + + + + dgrad + format: NVFP4 + + grad_output.rowwise + NVFP4 + × + + weight.columnwise + NVFP4 + + + + + wgrad + format: BF16 + + input.columnwise + BF16 + × + + grad_output.columnwise + BF16 + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg new file mode 100644 index 0000000000..ebc077eb04 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg @@ -0,0 +1,74 @@ + + + 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/fine_grained_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg new file mode 100644 index 0000000000..6e542306c7 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg @@ -0,0 +1,58 @@ + + + HybridQuantizer data flow + A high-precision tensor enters a HybridQuantizer whose rowwise child is an MXFP8 quantizer and columnwise child is an NVFP4 quantizer. The result is a HybridQuantizedTensor with an MXFP8 rowwise representation and an NVFP4 columnwise representation, each consumed by a different GEMM. + + + + + + + + + + tensor + high precision (BF16) + + + + + + + HybridQuantizer + + rowwise_quantizer + MXFP8Quantizer + + columnwise_quantizer + NVFP4Quantizer + + + + + + + HybridQuantizedTensor + + rowwise + MXFP8 + + columnwise + NVFP4 + + + + + GEMM 1 + GEMM 2 + diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst index 48d17db8d5..557a4b09e8 100644 --- a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst @@ -180,6 +180,37 @@ Blackwell and later (SM >= 10.0) – the recipe is emulated with MXFP8. Note tha ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Blockwise scaling uses + :class:`~transformer_engine.pytorch.Float8BlockQuantizer`. Each block of + the tensor gets its own power-of-two scale: ``block_scaling_dim=1`` + scales 1x128 blocks, ``block_scaling_dim=2`` (the default) scales + 128x128 blocks. This recipe is not available in TE/JAX. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=1, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + Developer Notes --------------- diff --git a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst index cac3792194..2436a07566 100644 --- a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst +++ b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst @@ -164,6 +164,57 @@ Here's how to use FP8 Current Scaling recipe in PyTorch and JAX: ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Current scaling uses + :class:`~transformer_engine.pytorch.Float8CurrentScalingQuantizer`. It + needs no external state: at each call it computes the amax of the input + tensor, derives the scale from it, and then quantizes. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8CurrentScalingQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + device="cuda", + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + Current scaling uses ``CurrentScaleQuantizer``. At each call it computes + the amax of the input tensor, derives the scale from it, and then + quantizes. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- @@ -177,4 +228,4 @@ On Blackwell and later, rowwise and columnwise tensors share the same memory lay so all-gather of columnwise tensors is directly supported. For Hopper and Ada, all-gather of transposed FP8 tensors is not supported. -The rowwise tensor is gathered first, then transposed to columnwise format. \ No newline at end of file +The rowwise tensor is gathered first, then transposed to columnwise format. diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst index d39787f6f5..99a379eed1 100644 --- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst +++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst @@ -160,4 +160,69 @@ However, amax reduction works slightly differently in different frameworks. Supported devices ----------------- -Ada and later (SM 8.9+) \ No newline at end of file +Ada and later (SM 8.9+) + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Delayed scaling uses + :class:`~transformer_engine.pytorch.Float8Quantizer`. It does not + compute the scaling factor from the current tensor: one-element + ``scale`` and ``amax`` buffers are supplied at construction. + Quantization applies the given scale and records the tensor's amax into + the ``amax`` buffer. + + During training both buffers are views into the recipe state: ``scale`` + into its per-quantizer scale vector, ``amax`` into the current row of + its ``(amax_history_len, num_quantizers)`` amax history. At the end of + each step the recipe state computes a new scale from the history (its + max or most recent entry, per ``amax_compute_algo``), rolls the history + by one slot, and zeroes the current row — all in place, so the views + held by the quantizer stay valid for the whole training run. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8Quantizer( + scale=torch.ones(1, device="cuda"), + amax=torch.zeros(1, device="cuda"), + fp8_dtype=te.DType.kFloat8E4M3, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + Delayed scaling uses ``DelayedScaleQuantizer``. The ``scale`` and the + ``amax_history`` (1024 entries by default) are fields of the quantizer + itself, carried through JAX transformations as its pytree state. Each + ``quantize()`` call applies the current ``scale``, then updates the + state: the tensor's amax is written into the history, a new scale is + computed from the history (max or most-recent entry, per + ``amax_compute_algo``), and the history is rolled by one slot. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.DELAYED_TENSOR_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() 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/docs/features/low_precision_training/introduction/introduction.rst b/docs/features/low_precision_training/introduction/introduction.rst index fba7796ece..2255308b04 100644 --- a/docs/features/low_precision_training/introduction/introduction.rst +++ b/docs/features/low_precision_training/introduction/introduction.rst @@ -283,3 +283,71 @@ so GEMM with tensors ``A`` and ``B`` returns ``B * A^T``. :file: img/fp8_linear_flow.svg *Figure 4: Forward pass of a Linear layer with low precision data flow.* + +Quantizers +---------- + +Every recipe implements its quantization logic in a **quantizer** — an object +that converts a high-precision tensor into a quantized one. TE modules create +and use quantizers internally according to the active recipe, but a quantizer +can also be used directly: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + qtensor = quantizer(tensor) # quantize + roundtrip = qtensor.dequantize() # back to high precision + + The main parts of the interface are: + + * ``quantize(tensor)`` — quantizes a high-precision tensor and returns a + ``QuantizedTensor``; calling the quantizer (``quantizer(tensor)``) is + a shorthand; + * ``update_quantized(src, dst)`` — quantizes ``src`` in place into an + already-allocated quantized tensor ``dst``; + * ``make_empty(shape)`` — allocates an uninitialized quantized tensor to + be filled later; + * ``rowwise_usage`` / ``columnwise_usage`` — flags selecting which of + the two GEMM-oriented representations the produced tensor holds; + * the returned ``QuantizedTensor`` supports ``dequantize()`` back to + high precision. + + .. tab:: JAX + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + + The main parts of the interface are: + + * ``quantize(x, is_rowwise=..., is_colwise=...)`` — quantizes a tensor + and returns a ``ScaledTensor`` holding the requested representations + (the default comes from the quantizer's ``q_layout``); + * the returned ``ScaledTensor`` supports ``dequantize()`` back to high + precision; + * quantizers are registered pytrees, so they can be passed through JAX + transformations. + +Each recipe section ends with a short description of that recipe's quantizer. diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index 1fbcc43af9..1827d42cb6 100644 --- a/docs/features/low_precision_training/mxfp8/mxfp8.rst +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -152,6 +152,57 @@ SM 10.0, SM 10.3 ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + MXFP8 uses :class:`~transformer_engine.pytorch.MXFP8Quantizer`. Every + 32-element block shares one power-of-two (E8M0) scale, computed from the + block's amax at quantization time; no external state is needed. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + MXFP8 uses ``BlockScaleQuantizer`` — the JAX quantizer for block-based + scaling, selected by ``ScalingMode.MXFP8_1D_SCALING``. Instead of one + scale per tensor, the tensor is split along the quantization axis into + 32-element blocks and each block gets its own power-of-two (E8M0) scale, + computed from that block's amax at quantization time. Because the scale + is derived from the current data, no external state (scale buffers or + amax history) is needed. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- @@ -210,4 +261,4 @@ All-gather of columnwise tensors All-gather of columnwise tensors is supported and necessary because: - columnwise quantized tensors cannot be computed from rowwise quantized ones, -- gathering high-precision tensors is avoided in most cases for performance reasons. \ No newline at end of file +- gathering high-precision tensors is avoided in most cases for performance reasons. diff --git a/docs/features/low_precision_training/nvfp4/nvfp4.rst b/docs/features/low_precision_training/nvfp4/nvfp4.rst index 900789b0d3..26d798651d 100644 --- a/docs/features/low_precision_training/nvfp4/nvfp4.rst +++ b/docs/features/low_precision_training/nvfp4/nvfp4.rst @@ -250,6 +250,60 @@ Supported devices ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + NVFP4 uses :class:`~transformer_engine.pytorch.NVFP4Quantizer`. It + implements the two-level scaling described above: an FP8 (E4M3) scale + per 16-element block plus one FP32 scale per tensor. Further keyword + options select the recipe variations from this page (random Hadamard + transforms, stochastic rounding, 2D weight scaling); they are internal + knobs and may change without notice. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + NVFP4 uses its own ``NVFP4Quantizer``, with the same two-level scaling. + ``ScalingMode.NVFP4_1D_SCALING`` selects per-block scaling only, + ``ScalingMode.NVFP4_2D_SCALING`` adds 2D weight scaling. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.NVFP4_1D_SCALING, + q_dtype=jnp.float4_e2m1fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- diff --git a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst index 2c21799dd6..afa5d16c86 100644 --- a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst +++ b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst @@ -143,6 +143,61 @@ Transformer Engine chooses the best possible fusion internally taking the recipe *Figure 3: Three scenarios of producing quantized tensors in rowwise and columnwise usages.* +**Usages in the quantizer API** + +The usages are visible directly in the quantizer API: + +.. tabs:: + + .. tab:: PyTorch + + At quantization time, the quantizer's ``rowwise_usage`` and + ``columnwise_usage`` flags select which representations ``quantize()`` + produces; when both are set, the representations are computed together + in one fused kernel (scenario 1 above). + + After quantization, ``update_usage()`` on the quantized tensor removes a + representation or, when supported by the format, generates a missing one. + Passing ``rowwise_usage=False`` after the forward pass frees the rowwise + data while keeping the columnwise data for backward. Some formats also + support ``columnwise_usage=True`` to create the columnwise representation + from the data already present (e.g. by a transpose on Hopper — scenario 3 + above); unsupported requests raise an error. Arguments left as ``None`` + preserve the current state. + + .. code-block:: python + + quantizer = te.MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + + qtensor = quantizer(tensor) # both representations, one fused kernel + + qtensor.update_usage(rowwise_usage=False) # drop rowwise, keep columnwise + + .. tab:: JAX + + The usages are selected when the tensor is quantized: the quantizer's + ``q_layout`` (``QuantizeLayout.ROWWISE``, ``COLWISE``, or + ``ROWWISE_COLWISE``) sets the default, and ``quantize()`` accepts + ``is_rowwise``/``is_colwise`` overrides. Requesting both usages returns + a ``ScaledTensor2x`` holding the two representations. There is no + in-place ``update_usage()``: JAX arrays are immutable, so a + representation is not added or dropped later — unneeded ones are simply + not requested and get dropped by XLA's dead-code elimination. + + .. code-block:: python + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE_COLWISE, + ) + + qtensor = quantizer.quantize(x) # ScaledTensor2x, both representations + rowwise_only = quantizer.quantize(x, is_rowwise=True, is_colwise=False) Memory usage @@ -470,4 +525,3 @@ Actual behavior depends on the recipe and module configuration. *Figure 5: All-gather of quantized tensors for input and gradient tensors. This is one possible scenario — actual behavior varies depending on the recipe and module configuration.* - diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 5d5ce1f6cf..128e8280bb 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -639,13 +639,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. @@ -663,7 +668,8 @@ class CustomRecipe(Recipe): See ``transformer_engine.pytorch.quantization.QuantizerRole`` and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` - for full documentation. + for API details. See :ref:`heterogeneous-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 8df2ec8b4b..26d0798b92 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 8310afc653..ec171564fe 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):