diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5fac0a89a6..60556d6cac 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -109,6 +109,12 @@ Communication-computation overlap :members: FP8, NONE +Heterogeneous quantization recipes +---------------------------------- + +.. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="") + + Quantized tensors ----------------- @@ -126,6 +132,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensorStorage(*, rowwise_storage, columnwise_storage, quantizer, fake_dtype=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensorStorage(*, hp_data, fake_dtype=None, quantizer=None) + .. autoapiclass:: transformer_engine.pytorch.Float8Tensor(shape, dtype, data, fp8_scale_inv, fp8_dtype, requires_grad=False, data_transpose=None, quantizer=None) .. autoapiclass:: transformer_engine.pytorch.MXFP8Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer) @@ -134,6 +144,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensor(shape, dtype, *, rowwise_storage, columnwise_storage, quantizer, requires_grad=False, device=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensor(shape, dtype, *, hp_data, quantizer=None, requires_grad=False, device=None) + Quantizers ---------- @@ -150,6 +164,10 @@ Quantizers .. autoapiclass:: transformer_engine.pytorch.NVFP4Quantizer(fp4_dtype, *, rowwise=True, columnwise=True, **kwargs) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizer(*, rowwise_quantizer, columnwise_quantizer, columnwise_source="original") + +.. autoapiclass:: transformer_engine.pytorch.IdentityQuantizer(*, dtype=None, rowwise=True, columnwise=True) + Tensor saving and restoring functions ------------------------------------- diff --git a/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst b/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst new file mode 100644 index 0000000000..e4ec7b2be1 --- /dev/null +++ b/docs/examples/heterogeneous_quantization/heterogeneous_quantization.rst @@ -0,0 +1,121 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fine-grained-quantization-tutorial: +.. _heterogeneous-quantization-tutorial: + +Building a heterogeneous quantization recipe +============================================== + +This tutorial demonstrates factory composition, module and name targeting, +the general three-format GEMM mapping, and high-precision directions. See +:doc:`Heterogeneous quantization recipes +<../../features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization>` +for the API concepts and direction mapping. + +Constructing a factory +---------------------- + +A robust factory follows four construction rules: + +* Return a quantizer for every role. Return ``IdentityQuantizer`` for an + intentional high-precision slot; do not return ``None``. +* Constructing a fresh quantizer for every call is recommended. + ``HybridQuantizer`` owns and configures its rowwise and columnwise children. +* A module-level function is the most portable factory definition, especially + when a launcher or checkpointing setup needs to import or pickle it. +* Treat role strings as selectors, not a fixed enumeration. Preserve a base + factory fallback for roles the factory does not recognize. + +For example, compose a TE-native factory by keeping one named ``linear`` +module in high precision, using NVFP4 for every ``grouped_linear`` role, and +retaining MXFP8 as the global fallback: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + from typing import Optional + + import transformer_engine.pytorch as te + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, + ) + + def my_factory(role: Optional[te.QuantizerRole]): + if role is not None: + if role.module_type == "linear" and role.name == "decoder.39.fc2": + return te.IdentityQuantizer() + if role.module_type == "grouped_linear": + return nvfp4_factory(role) + return mxfp8_factory(role) + +Runnable example +---------------- + +The following synthetic example demonstrates base-factory composition, the +general three-format mapping, high-precision directions, and module/name +targeting. It uses only TE-native quantizers. MXFP8 and NVFP4 execution requires +supported hardware and software. + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: pytorch_heterogeneous_quantization_example.py + :language: python + :start-after: # START_HETEROGENEOUS_QUANTIZATION_EXAMPLE + :end-before: # END_HETEROGENEOUS_QUANTIZATION_EXAMPLE + +Run it from the repository root after installing TE: + +.. code-block:: bash + + python docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py + +Recipe starting points +---------------------- + +The runnable example above is deliberately synthetic: it demonstrates the +expressiveness of the API, not a recommended training recipe. + +The TE-native base factories in +``transformer_engine/pytorch/custom_recipes/quantizer_factories.py`` construct +standard TE quantizers. They can be used directly as factory fallbacks or +composed as children of a ``HybridQuantizer``. + +More specialized starting points are available in +``transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py``. Some +zoo factories encode externally described recipe structures or have specific +motivating evidence. They remain illustrative examples rather than official, +broadly validated defaults; read each factory's rationale before adapting it. + +Validating and optimizing a recipe +---------------------------------- + +A factory can describe assignments beyond current optimized kernel coverage. +Before adopting an assignment for a workload: + +* Confirm that the module, GEMM layout and shape, software version, and GPU can + execute it. +* Check whether the selected module path has an appropriate optimized kernel + and integration. +* Validate accuracy and convergence on the target model and distributed setup. +* Benchmark throughput and memory on the target workload. + +Fine-grained assignments provide a way to explore an accuracy/performance +slider by varying precision and quantization logic per GEMM. The recipes that +can be realized efficiently are constrained by available kernels and +integrations. Accuracy and convergence experiments can run on functionally +executable, non-optimized paths before dedicated kernels are available. diff --git a/docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py b/docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py new file mode 100644 index 0000000000..c2834d2624 --- /dev/null +++ b/docs/examples/heterogeneous_quantization/pytorch_heterogeneous_quantization_example.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Runnable heterogeneous quantization recipe example. + +The factory assigns one precision to each ``demo.fc1`` Linear GEMM: + +* fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` +* dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` +* wgrad: ``input.col(original BF16) x grad_output.col(original BF16)`` + +``demo.fc2`` runs every GEMM in high precision. ``demo.output`` is not +special-cased and therefore exercises the MXFP8 base-factory fallback. + +Run from the Transformer Engine repository root:: + + python docs/examples/heterogeneous_quantization/\ + pytorch_heterogeneous_quantization_example.py +""" + +from __future__ import annotations + +import torch +import transformer_engine.pytorch as te + + +def require_supported_hardware() -> None: + """Fail early with TE's reason when either required format is unavailable.""" + + if not torch.cuda.is_available(): + raise SystemExit("This example requires a CUDA-capable NVIDIA GPU.") + + failures = [] + for name, check in ( + ("MXFP8", te.is_mxfp8_available), + ("NVFP4", te.is_nvfp4_available), + ): + available, reason = check(return_reason=True) + if not available: + failures.append(f"{name}: {reason}") + if failures: + raise SystemExit("Required formats are unavailable: " + "; ".join(failures)) + + +require_supported_hardware() + +# START_HETEROGENEOUS_QUANTIZATION_EXAMPLE + +from typing import Optional + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, +) + + +THREE_FORMAT_MODULE = "demo.fc1" +HIGH_PRECISION_MODULE = "demo.fc2" +BASE_FACTORY = mxfp8_factory + + +def quantizer_factory(role: Optional[te.QuantizerRole]): + """Return a fresh quantizer for every role, including ``None``. + + ``BASE_FACTORY`` makes the factory total: unknown roles, future role values, + and untargeted modules all retain valid MXFP8 behavior. + """ + + if role is not None and role.name == THREE_FORMAT_MODULE: + # Constructing fresh child quantizers for every call is recommended. + if role.tensor_type == "input": + # Wgrad retains the original BF16 input. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # Dgrad uses NVFP4 quantized from the dequantized MXFP8 fprop weight. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # Dgrad uses NVFP4 while wgrad retains the original BF16 gradient. + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + + if role is not None and role.name == HIGH_PRECISION_MODULE: + return te.IdentityQuantizer() + + return BASE_FACTORY(role) + + +linear_options = {"bias": False, "params_dtype": torch.bfloat16, "device": "cuda"} +model = torch.nn.Sequential( + te.Linear(128, 256, name=THREE_FORMAT_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 256, name=HIGH_PRECISION_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 128, name="demo.output", **linear_options), +) +inputs = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) +recipe = CustomRecipe(qfactory=quantizer_factory) + +with te.autocast(enabled=True, recipe=recipe): + outputs = model(inputs) + +loss = outputs.float().square().mean() +loss.backward() + +# END_HETEROGENEOUS_QUANTIZATION_EXAMPLE + +gradients = [inputs.grad, *(parameter.grad for parameter in model.parameters())] +assert all(gradient is not None for gradient in gradients) +assert all(torch.isfinite(gradient).all() for gradient in gradients) + +print(f"GPU: {torch.cuda.get_device_name()}") +print(f"TE Linear names: {[model[index].name for index in (0, 2, 4)]}") +print(f"loss: {loss.item():.6f}; forward and backward completed") diff --git a/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst b/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst new file mode 100644 index 0000000000..e7087e422b --- /dev/null +++ b/docs/features/low_precision_training/heterogeneous_quantization/heterogeneous_quantization.rst @@ -0,0 +1,252 @@ +.. + 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 new file mode 100644 index 0000000000..0910629db0 --- /dev/null +++ b/docs/features/low_precision_training/heterogeneous_quantization/img/heterogeneous_linear_mapping.svg @@ -0,0 +1,106 @@ + + + Heterogeneous precision mapping for Linear GEMMs + Input, weight, and output-gradient tensors each provide rowwise and columnwise representations. Forward uses format F, input-gradient uses format D, and weight-gradient uses format W. + + + + + + + + Tensor representations and Linear GEMMs + Quantizer outputs by tensor role + + + + Input + + rowwise + F + + columnwise + W + + + + + Weight + + rowwise + F + + columnwise + D + + + + + Grad output + + rowwise + D + + columnwise + W + + + TE selects the required direction for each GEMM + + + FPROP + + weight.rowwise + format F + × + + input.rowwise + format F + + + Output + + + + DGRAD + + weight.columnwise + format D + × + + grad_output.rowwise + format D + + + Input grad. + + + + WGRAD + + input.columnwise + format W + × + + grad_output.columnwise + format W + + + Weight grad. + + diff --git a/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg new file mode 100644 index 0000000000..d7aa328fe7 --- /dev/null +++ b/docs/features/low_precision_training/heterogeneous_quantization/img/hybrid_columnwise_source.svg @@ -0,0 +1,79 @@ + + + Hybrid quantizer columnwise source choices + With original provenance, both quantizers consume the original high-precision tensor. With rowwise-dequantized provenance, the columnwise quantizer consumes the dequantized rowwise representation. + + + + + + + + Choosing the columnwise source + + + + columnwise_source="original" + + + High-precision tensor + + + + same original source + + + Rowwise quantizer + + Columnwise quantizer + + + + + Rowwise + representation + + Columnwise + representation + + + + + columnwise_source="rowwise_dequantized" + + + High-precision tensor + + + + Rowwise quantizer + + + + Rowwise + representation + + + + Dequantize + + + + Columnwise quantizer + + + Columnwise + representation + + diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst index 0a798f1364..8e963c9d1b 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 + heterogeneous_quantization/heterogeneous_quantization.rst speedups.rst diff --git a/docs/index.rst b/docs/index.rst index fcd15a7a11..414c4d87b3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -60,6 +60,7 @@ Transformer Engine documentation examples/onnx/onnx_export.ipynb examples/te_jax_integration.rst examples/op_fuser/op_fuser.rst + examples/heterogeneous_quantization/heterogeneous_quantization.rst examples/gemm_profiling/gemm_profiling.rst .. toctree:: diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index a89ddba917..bdbaf224bd 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -635,13 +635,18 @@ class CustomRecipe(Recipe): ---------- qfactory : Callable Factory callable that returns a quantizer instance *or* a - ``QuantizerRequest`` subclass for a given ``QuantizerRole``. + ``QuantizerRequest`` subclass for a given optional ``QuantizerRole``. The callable is invoked as:: qfactory( - role: QuantizerRole, + role: Optional[QuantizerRole], ) -> Union[Quantizer, QuantizerRequest] + Boundary slots may provide ``None`` or a role with empty fields. The + factory must return a valid object for every call. Return an + ``IdentityQuantizer`` for an intentional high-precision slot instead + of returning ``None``. + ``QuantizerRole`` is a frozen dataclass with the following fields: - ``module_type`` (str): module type (empty string when not set), e.g. @@ -659,7 +664,8 @@ class CustomRecipe(Recipe): See ``transformer_engine.pytorch.quantization.QuantizerRole`` and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` - for full documentation. + for API details. See :ref:`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 dc65c9894b..4b9564ee23 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -19,6 +19,10 @@ class HybridQuantizer(Quantizer): """Quantizer that composes rowwise and columnwise representations. + .. warning:: + **EXPERIMENTAL**: ``HybridQuantizer`` is under active development and + its API is subject to change without notice. + When both representations are requested, applies ``rowwise_quantizer`` to produce the rowwise representation and ``columnwise_quantizer`` to produce the columnwise representation. The results are wrapped in a diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 9fb980a755..e21aafc3d9 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -26,6 +26,10 @@ class IdentityQuantizer(Quantizer): """Quantizer that produces a high-precision passthrough representation. + .. warning:: + **EXPERIMENTAL**: ``IdentityQuantizer`` is under active development and + its API is subject to change without notice. + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) holding the tensor directly, without a low-precision encoding. ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it @@ -174,6 +178,21 @@ class IdentityTensor(IdentityTensorStorage, QuantizedTensor): Presents as a standard tensor of its nominal dtype; internally it just holds data directly in that dtype, without a low-precision encoding. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Logical tensor datatype. + hp_data : torch.Tensor + Held high-precision data. + quantizer : IdentityQuantizer, optional + Quantizer that produced the tensor. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + device : torch.device, optional + Device containing the tensor. """ def __repr__(self, *, tensor_contents=None):