From 6fe9542db8090a37f156497bb3bc9a198815edb0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 12:18:36 +0200 Subject: [PATCH 1/5] Rework heterogeneous quantization docs into mixed-format quantization; add per-recipe Quantizer sections Signed-off-by: Pawel Gadzinski --- docs/api/pytorch.rst | 2 +- .../heterogeneous_quantization.rst | 121 ------ ...orch_mixed_format_quantization_example.py} | 10 +- .../fp8_blockwise_scaling.rst | 31 ++ .../fp8_current_scaling.rst | 53 ++- .../fp8_delayed_scaling.rst | 67 +++- .../heterogeneous_quantization.rst | 252 ------------ .../img/heterogeneous_linear_mapping.svg | 106 ----- .../features/low_precision_training/index.rst | 2 +- .../introduction/introduction.rst | 68 ++++ .../img/hybrid_columnwise_source.svg | 38 +- .../img/hybrid_quantizer.svg | 60 +++ .../img/mixed_format_assignments.svg | 97 +++++ .../img/mixed_format_linear_mapping.svg | 60 +++ .../mixed_format_quantization.rst | 364 ++++++++++++++++++ .../low_precision_training/mxfp8/mxfp8.rst | 53 ++- .../low_precision_training/nvfp4/nvfp4.rst | 54 +++ .../performance_considerations.rst | 56 ++- docs/index.rst | 1 - 19 files changed, 985 insertions(+), 510 deletions(-) delete mode 100644 docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst rename docs/examples/{heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py => mixed_format_quantization/pytorch_mixed_format_quantization_example.py} (94%) delete mode 100644 docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst delete mode 100644 docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg rename docs/features/low_precision_training/{heterogeneous_quantization => mixed_format_quantization}/img/hybrid_columnwise_source.svg (74%) create mode 100644 docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg create mode 100644 docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg create mode 100644 docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg create mode 100644 docs/features/low_precision_training/mixed_format_quantization/mixed_format_quantization.rst diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 60556d6cac..09cd663ef4 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -109,7 +109,7 @@ Communication-computation overlap :members: FP8, NONE -Heterogeneous quantization recipes +Mixed-format 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 deleted file mode 100644 index e4ec7b2be1..0000000000 --- a/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst +++ /dev/null @@ -1,121 +0,0 @@ -.. - 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/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py b/docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py similarity index 94% rename from docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py rename to docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py index c2834d2624..0471a248b8 100644 --- a/docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py +++ b/docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py @@ -3,7 +3,7 @@ # # See LICENSE for license information. -"""Runnable heterogeneous quantization recipe example. +"""Runnable mixed-format quantization recipe example. The factory assigns one precision to each ``demo.fc1`` Linear GEMM: @@ -16,8 +16,8 @@ Run from the Transformer Engine repository root:: - python docs/examples/heterogeneous_quantization/\ - pytorch_heterogeneous_quantization_example.py + python docs/examples/mixed_format_quantization/\ + pytorch_mixed_format_quantization_example.py """ from __future__ import annotations @@ -46,7 +46,7 @@ def require_supported_hardware() -> None: require_supported_hardware() -# START_HETEROGENEOUS_QUANTIZATION_EXAMPLE +# START_MIXED_FORMAT_QUANTIZATION_EXAMPLE from typing import Optional @@ -119,7 +119,7 @@ def quantizer_factory(role: Optional[te.QuantizerRole]): loss = outputs.float().square().mean() loss.backward() -# END_HETEROGENEOUS_QUANTIZATION_EXAMPLE +# END_MIXED_FORMAT_QUANTIZATION_EXAMPLE gradients = [inputs.grad, *(parameter.grad for parameter in model.parameters())] assert all(gradient is not None for gradient in gradients) 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 a4830a3fd5..516af71e7a 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 9d05305eda..779d6a2944 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/heterogeneous_quantization/heterogeneous_quantization.rst b/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst deleted file mode 100644 index e7087e422b..0000000000 --- a/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst +++ /dev/null @@ -1,252 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -.. _fine-grained-quantization-recipes: -.. _heterogeneous-quantization-recipes: - -Heterogeneous quantization recipes -================================== - -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 -:class:`~transformer_engine.pytorch.HybridQuantizer` and -:class:`~transformer_engine.pytorch.IdentityQuantizer`. - -This guide covers PyTorch, TE-native quantizers, and static recipe construction. - -Mixing formats at a glance --------------------------- - -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: - -.. list-table:: - :header-rows: 1 - :widths: 34 22 22 22 - - * - Module assignment - - Fprop - - Dgrad - - Wgrad - * - ``demo.fc1`` - - MXFP8 - - NVFP4 - - BF16 - * - ``demo.fc2`` - - BF16 - - BF16 - - BF16 - * - Other TE modules - - MXFP8 - - MXFP8 - - MXFP8 - -Once the factory defines these assignments, the recipe uses the standard TE -autocast path: - -.. tabs:: - - .. tab:: PyTorch - - .. 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 ----------------- - -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``. - -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 ------------------------------ - -``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) - -.. 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 -:doc:`tutorial <../../../examples/heterogeneous_quantization/heterogeneous_quantization>`). - -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. - -.. 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. -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``: - -.. tabs:: - - .. tab:: PyTorch - - .. 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. - -Tutorial --------- - -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. - -Support status --------------- - -.. 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. This guide does - not define a supported recipe or an expected accuracy/performance ordering. - -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/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg b/docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg deleted file mode 100644 index 0910629db0..0000000000 --- a/docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg +++ /dev/null @@ -1,106 +0,0 @@ - - - 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/index.rst b/docs/features/low_precision_training/index.rst index 8e963c9d1b..8e45d1558d 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 - heterogeneous_quantization/heterogeneous_quantization.rst + mixed_format_quantization/mixed_format_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 760a63b0b1..7ffae16a06 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/heterogeneous_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg similarity index 74% rename from docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg rename to docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg index d7aa328fe7..adbe63d8c4 100644 --- a/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg +++ b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg @@ -52,28 +52,28 @@ columnwise_source="rowwise_dequantized" - - High-precision tensor - + + High-precision tensor + - - Rowwise quantizer - + + Rowwise quantizer + - - Rowwise - representation - + + Rowwise + representation + - - Dequantize - + + Dequantize + - - Columnwise quantizer - - - Columnwise - representation + + Columnwise quantizer + + + Columnwise + representation diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg new file mode 100644 index 0000000000..cbc146a2a7 --- /dev/null +++ b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg @@ -0,0 +1,60 @@ + + + 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/mixed_format_quantization/img/mixed_format_assignments.svg b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg new file mode 100644 index 0000000000..f8fdce0126 --- /dev/null +++ b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg @@ -0,0 +1,97 @@ + + + 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/mixed_format_quantization/img/mixed_format_linear_mapping.svg b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg new file mode 100644 index 0000000000..982dc3913e --- /dev/null +++ b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg @@ -0,0 +1,60 @@ + + + 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/mixed_format_quantization/mixed_format_quantization.rst b/docs/features/low_precision_training/mixed_format_quantization/mixed_format_quantization.rst new file mode 100644 index 0000000000..8c636ca388 --- /dev/null +++ b/docs/features/low_precision_training/mixed_format_quantization/mixed_format_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: +.. _mixed-format-quantization-recipes: + +Mixed-format 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. Mixed-format +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`. + +.. warning:: + + Mixed-format recipes are currently available only in the PyTorch API of + TE. + +.. warning:: + + Mixed-format 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/mixed_format_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/mixed_format_quantization/pytorch_mixed_format_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 ships ready-made base factories for its native quantizers in + ``transformer_engine.pytorch.custom_recipes.quantizer_factories`` + (``mxfp8_factory``, ``nvfp4_factory``, ...) — the examples on this page + use them as fallbacks and as ``HybridQuantizer`` children. More + specialized factories live in + ``transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo``; + they are illustrative examples rather than validated defaults. + + 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/mixed_format_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 slow** — the recipe executes, but parts of it take fallback + paths. For example, ``HybridQuantizer`` launches a separate kernel per + direction instead of producing both representations in one fused kernel, + and under tensor parallelism it currently gathers the high-precision tensor + instead of the quantized one. +* **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 slow 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/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index f8f8f48b0d..8c67542533 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 0415963a71..3cc7dc39b2 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 a495af56c1..9754ceb6c6 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/docs/index.rst b/docs/index.rst index 414c4d87b3..fcd15a7a11 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -60,7 +60,6 @@ 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:: From 6137be83042b3eb45c8ecc01bea14e68972894f0 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 12:48:23 +0200 Subject: [PATCH 2/5] Align mixed-format quantization diagrams with shared diagram-colors.css and dark mode Signed-off-by: Pawel Gadzinski --- docs/_static/css/diagram-colors.css | 15 +++ .../img/hybrid_columnwise_source.svg | 23 ++-- .../img/hybrid_quantizer.svg | 54 +++++---- .../img/mixed_format_assignments.svg | 108 +++++++++--------- .../img/mixed_format_linear_mapping.svg | 63 +++++----- 5 files changed, 134 insertions(+), 129 deletions(-) diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index f5dc7da4dd..b96487302f 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; } + +/* mixed_format_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/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg index adbe63d8c4..ebc077eb04 100644 --- a/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg +++ b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg @@ -4,34 +4,29 @@ 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 + Choosing the columnwise source - columnwise_source="original" + columnwise_source="original" High-precision tensor - same original source + same original source Rowwise quantizer @@ -42,15 +37,15 @@ Rowwise - representation + representation Columnwise - representation + representation - columnwise_source="rowwise_dequantized" + columnwise_source="rowwise_dequantized" High-precision tensor @@ -62,7 +57,7 @@ Rowwise - representation + representation @@ -74,6 +69,6 @@ Columnwise - representation + representation diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg index cbc146a2a7..6e542306c7 100644 --- a/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg +++ b/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg @@ -4,17 +4,15 @@ 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. @@ -22,39 +20,39 @@ - - tensor - high precision (BF16) + + tensor + high precision (BF16) - + HybridQuantizer - - rowwise_quantizer - MXFP8Quantizer - - columnwise_quantizer - NVFP4Quantizer + + rowwise_quantizer + MXFP8Quantizer + + columnwise_quantizer + NVFP4Quantizer - + HybridQuantizedTensor - - rowwise - MXFP8 - - columnwise - NVFP4 + + rowwise + MXFP8 + + columnwise + NVFP4 - GEMM 1 - GEMM 2 + GEMM 1 + GEMM 2 diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg index f8fdce0126..b9aa48333d 100644 --- a/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg +++ b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg @@ -4,94 +4,92 @@ 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 + Precision assignment by tensor role and module demo.fc1 demo.fc2 Other TE modules - - input + + input rowwise (fprop) - - MXFP8 - - BF16 - - MXFP8 + + MXFP8 + + BF16 + + MXFP8 columnwise (wgrad) - - BF16 (original) - - BF16 - - MXFP8 + + BF16 (original) + + BF16 + + MXFP8 - - weight + + weight rowwise (fprop) - - MXFP8 - - BF16 - - MXFP8 + + MXFP8 + + BF16 + + MXFP8 columnwise (dgrad) - - NVFP4 - - BF16 - - MXFP8 + + NVFP4 + + BF16 + + MXFP8 - - grad_output + + grad_output rowwise (dgrad) - - NVFP4 - - BF16 - - MXFP8 + + NVFP4 + + BF16 + + MXFP8 columnwise (wgrad) - - BF16 (original) - - BF16 - - MXFP8 + + BF16 (original) + + BF16 + + MXFP8 - + MXFP8 - + NVFP4 - + BF16 (high precision) diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg index 982dc3913e..0e17ffdfa4 100644 --- a/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg +++ b/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg @@ -4,57 +4,56 @@ 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 + format: MXFP8 + + input.rowwise + MXFP8 × - - weight.rowwise - MXFP8 + + weight.rowwise + MXFP8 - + dgrad - format: NVFP4 - - grad_output.rowwise - NVFP4 + format: NVFP4 + + grad_output.rowwise + NVFP4 × - - weight.columnwise - NVFP4 + + weight.columnwise + NVFP4 - + wgrad - format: BF16 - - input.columnwise - BF16 + format: BF16 + + input.columnwise + BF16 × - - grad_output.columnwise - BF16 + + grad_output.columnwise + BF16 From b0e79a8c0fe937dfb2eebcc3297b4fa52a395ea2 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:06:00 +0200 Subject: [PATCH 3/5] Rename mixed-format quantization docs to fine-grained quantization recipes Signed-off-by: Pawel Gadzinski --- docs/_static/css/diagram-colors.css | 2 +- docs/api/pytorch.rst | 4 ++-- ...orch_fine_grained_quantization_example.py} | 10 ++++---- .../fine_grained_quantization.rst} | 23 ++++++++++--------- .../img/fine_grained_assignments.svg} | 0 .../img/fine_grained_linear_mapping.svg} | 0 .../img/hybrid_columnwise_source.svg | 0 .../img/hybrid_quantizer.svg | 0 .../features/low_precision_training/index.rst | 2 +- 9 files changed, 21 insertions(+), 20 deletions(-) rename docs/examples/{mixed_format_quantization/pytorch_mixed_format_quantization_example.py => fine_grained_quantization/pytorch_fine_grained_quantization_example.py} (94%) rename docs/features/low_precision_training/{mixed_format_quantization/mixed_format_quantization.rst => fine_grained_quantization/fine_grained_quantization.rst} (95%) rename docs/features/low_precision_training/{mixed_format_quantization/img/mixed_format_assignments.svg => fine_grained_quantization/img/fine_grained_assignments.svg} (100%) rename docs/features/low_precision_training/{mixed_format_quantization/img/mixed_format_linear_mapping.svg => fine_grained_quantization/img/fine_grained_linear_mapping.svg} (100%) rename docs/features/low_precision_training/{mixed_format_quantization => fine_grained_quantization}/img/hybrid_columnwise_source.svg (100%) rename docs/features/low_precision_training/{mixed_format_quantization => fine_grained_quantization}/img/hybrid_quantizer.svg (100%) diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index b96487302f..9ee5827bd1 100644 --- a/docs/_static/css/diagram-colors.css +++ b/docs/_static/css/diagram-colors.css @@ -280,7 +280,7 @@ html[data-theme="dark"] .memory-label { fill: #e0e0e0; } html[data-theme="dark"] .connector { stroke: #bdbdbd; } -/* mixed_format_quantization diagrams */ +/* 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; } diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5ac54c007a..87a81e8c65 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -112,8 +112,8 @@ Communication-computation overlap :members: FP8, NONE -Mixed-format quantization recipes ----------------------------------- +Fine-grained quantization recipes +--------------------------------- .. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="") diff --git a/docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py similarity index 94% rename from docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py rename to docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py index 0471a248b8..b511fc41dc 100644 --- a/docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py +++ b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py @@ -3,7 +3,7 @@ # # See LICENSE for license information. -"""Runnable mixed-format quantization recipe example. +"""Runnable fine-grained quantization recipe example. The factory assigns one precision to each ``demo.fc1`` Linear GEMM: @@ -16,8 +16,8 @@ Run from the Transformer Engine repository root:: - python docs/examples/mixed_format_quantization/\ - pytorch_mixed_format_quantization_example.py + python docs/examples/fine_grained_quantization/\ + pytorch_fine_grained_quantization_example.py """ from __future__ import annotations @@ -46,7 +46,7 @@ def require_supported_hardware() -> None: require_supported_hardware() -# START_MIXED_FORMAT_QUANTIZATION_EXAMPLE +# START_FINE_GRAINED_QUANTIZATION_EXAMPLE from typing import Optional @@ -119,7 +119,7 @@ def quantizer_factory(role: Optional[te.QuantizerRole]): loss = outputs.float().square().mean() loss.backward() -# END_MIXED_FORMAT_QUANTIZATION_EXAMPLE +# 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) diff --git a/docs/features/low_precision_training/mixed_format_quantization/mixed_format_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst similarity index 95% rename from docs/features/low_precision_training/mixed_format_quantization/mixed_format_quantization.rst rename to docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 8c636ca388..b16a4a66fa 100644 --- a/docs/features/low_precision_training/mixed_format_quantization/mixed_format_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -5,27 +5,28 @@ .. _fine-grained-quantization-recipes: .. _heterogeneous-quantization-recipes: -.. _mixed-format-quantization-recipes: -Mixed-format 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. Mixed-format +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:: - Mixed-format recipes are currently available only in the PyTorch API of + Fine-grained recipes are currently available only in the PyTorch API of TE. .. warning:: - Mixed-format recipes and their construction APIs are experimental: API, + 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. @@ -33,11 +34,11 @@ quantizer for each slot TE asks about, and pass it via Example: mixing MXFP8, NVFP4, and BF16 -------------------------------------- -The `runnable example `__ +The `runnable example `__ makes the following assignments: .. raw:: html - :file: img/mixed_format_assignments.svg + :file: img/fine_grained_assignments.svg *Figure 1. Precision assignments per module and GEMM used throughout this guide.* @@ -93,13 +94,13 @@ TE autocast path: output = model(inputs) The complete, runnable version is available -`on GitHub `__ +`on GitHub `__ (requires Blackwell or later); run it from the repository root after installing TE: .. code-block:: bash - python docs/examples/mixed_format_quantization/pytorch_mixed_format_quantization_example.py + python docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py CustomRecipe and quantizer factory ---------------------------------- @@ -309,7 +310,7 @@ For the example assignments (fprop in MXFP8, dgrad in NVFP4, wgrad in BF16): grad_output = HybridQuantizer(rowwise=NVFP4, columnwise=BF16) # dgrad | wgrad .. raw:: html - :file: img/mixed_format_linear_mapping.svg + :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.* diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg similarity index 100% rename from docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_assignments.svg rename to docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg similarity index 100% rename from docs/features/low_precision_training/mixed_format_quantization/img/mixed_format_linear_mapping.svg rename to docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg similarity index 100% rename from docs/features/low_precision_training/mixed_format_quantization/img/hybrid_columnwise_source.svg rename to docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg diff --git a/docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg similarity index 100% rename from docs/features/low_precision_training/mixed_format_quantization/img/hybrid_quantizer.svg rename to docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst index 8e45d1558d..b9649c00a4 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 - mixed_format_quantization/mixed_format_quantization.rst + fine_grained_quantization/fine_grained_quantization.rst speedups.rst From 193df9eb297fc3839719d6d3d895e98da48a4869 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:11:33 +0200 Subject: [PATCH 4/5] Simplify quantizer factory paragraph Signed-off-by: Pawel Gadzinski --- .../fine_grained_quantization.rst | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) 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 index b16a4a66fa..8d9b20777d 100644 --- 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 @@ -181,13 +181,12 @@ It does not call the factory on every unchanged forward. function for the factory itself, so that launchers and checkpointing setups can import or pickle it. - TE ships ready-made base factories for its native quantizers in + TE provides factories for its native quantizers in ``transformer_engine.pytorch.custom_recipes.quantizer_factories`` - (``mxfp8_factory``, ``nvfp4_factory``, ...) — the examples on this page - use them as fallbacks and as ``HybridQuantizer`` children. More - specialized factories live in - ``transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo``; - they are illustrative examples rather than validated defaults. + (``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 From 3d9431eaa32e34af68603fdf460bd19ed96c38fb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 7 Sep 2026 16:15:48 +0200 Subject: [PATCH 5/5] Generalize the fallback-path description Signed-off-by: Pawel Gadzinski --- .../fine_grained_quantization.rst | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 index 8d9b20777d..301b655b3f 100644 --- 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 @@ -336,11 +336,11 @@ 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 slow** — the recipe executes, but parts of it take fallback - paths. For example, ``HybridQuantizer`` launches a separate kernel per - direction instead of producing both representations in one fused kernel, - and under tensor parallelism it currently gathers the high-precision tensor - instead of the quantized one. +* **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. @@ -352,7 +352,7 @@ Before adopting a recipe for a real workload, check that: * accuracy and convergence hold on the target model and distributed setup; * throughput and memory actually improve on the target workload. -The slow paths are still useful: accuracy and convergence experiments can run +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.