diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css index f5dc7da4dd..9ee5827bd1 100644 --- a/docs/_static/css/diagram-colors.css +++ b/docs/_static/css/diagram-colors.css @@ -279,3 +279,18 @@ html[data-theme="dark"] .subtitle, html[data-theme="dark"] .memory-label { fill: #e0e0e0; } html[data-theme="dark"] .connector { stroke: #bdbdbd; } + +/* fine_grained_quantization diagrams */ +html[data-theme="dark"] .fmt-mxfp8 { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .fmt-nvfp4 { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .fmt-bf16 { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .fmt-mxfp8-text { fill: #90caf9; } +html[data-theme="dark"] .fmt-nvfp4-text { fill: #ffcc80; } +html[data-theme="dark"] .fmt-bf16-text { fill: #a5d6a7; } +html[data-theme="dark"] .source { fill: #3a2f5c; stroke: #b39ddb; } +html[data-theme="dark"] .quantizer { fill: #1e4620; stroke: #81c784; } +html[data-theme="dark"] .representation { fill: #10375c; stroke: #64b5f6; } +html[data-theme="dark"] .dequantize { fill: #5c3a10; stroke: #ffb74d; } +html[data-theme="dark"] .rowlabel, +html[data-theme="dark"] .legend, +html[data-theme="dark"] .op { fill: #e0e0e0; } diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 54981c9086..a6afa2d0cc 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -6,7 +6,7 @@ PyTorch ======= -.. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs) +.. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, **kwargs) :members: forward, set_tensor_parallel_group .. autoapiclass:: transformer_engine.pytorch.GroupedLinear(in_features, out_features, bias=True, **kwargs) @@ -112,6 +112,12 @@ Communication-computation overlap :members: FP8, NONE +Fine-grained quantization recipes +--------------------------------- + +.. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="") + + Quantized tensors ----------------- @@ -129,6 +135,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensorStorage(*, rowwise_storage, columnwise_storage, quantizer, fake_dtype=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensorStorage(*, hp_data, fake_dtype=None, quantizer=None) + .. autoapiclass:: transformer_engine.pytorch.Float8Tensor(shape, dtype, data, fp8_scale_inv, fp8_dtype, requires_grad=False, data_transpose=None, quantizer=None) .. autoapiclass:: transformer_engine.pytorch.MXFP8Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer) @@ -137,6 +147,10 @@ Quantized tensors .. autoapiclass:: transformer_engine.pytorch.NVFP4Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensor(shape, dtype, *, rowwise_storage, columnwise_storage, quantizer, requires_grad=False, device=None) + +.. autoapiclass:: transformer_engine.pytorch.IdentityTensor(shape, dtype, *, hp_data, quantizer=None, requires_grad=False, device=None) + Quantizers ---------- @@ -153,6 +167,10 @@ Quantizers .. autoapiclass:: transformer_engine.pytorch.NVFP4Quantizer(fp4_dtype, *, rowwise=True, columnwise=True, **kwargs) +.. autoapiclass:: transformer_engine.pytorch.HybridQuantizer(*, rowwise_quantizer, columnwise_quantizer, columnwise_source="original") + +.. autoapiclass:: transformer_engine.pytorch.IdentityQuantizer(*, dtype=None, rowwise=True, columnwise=True) + Tensor saving and restoring functions ------------------------------------- diff --git a/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py new file mode 100644 index 0000000000..b511fc41dc --- /dev/null +++ b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Runnable fine-grained quantization recipe example. + +The factory assigns one precision to each ``demo.fc1`` Linear GEMM: + +* fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` +* dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` +* wgrad: ``input.col(original BF16) x grad_output.col(original BF16)`` + +``demo.fc2`` runs every GEMM in high precision. ``demo.output`` is not +special-cased and therefore exercises the MXFP8 base-factory fallback. + +Run from the Transformer Engine repository root:: + + python docs/examples/fine_grained_quantization/\ + pytorch_fine_grained_quantization_example.py +""" + +from __future__ import annotations + +import torch +import transformer_engine.pytorch as te + + +def require_supported_hardware() -> None: + """Fail early with TE's reason when either required format is unavailable.""" + + if not torch.cuda.is_available(): + raise SystemExit("This example requires a CUDA-capable NVIDIA GPU.") + + failures = [] + for name, check in ( + ("MXFP8", te.is_mxfp8_available), + ("NVFP4", te.is_nvfp4_available), + ): + available, reason = check(return_reason=True) + if not available: + failures.append(f"{name}: {reason}") + if failures: + raise SystemExit("Required formats are unavailable: " + "; ".join(failures)) + + +require_supported_hardware() + +# START_FINE_GRAINED_QUANTIZATION_EXAMPLE + +from typing import Optional + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, +) + + +THREE_FORMAT_MODULE = "demo.fc1" +HIGH_PRECISION_MODULE = "demo.fc2" +BASE_FACTORY = mxfp8_factory + + +def quantizer_factory(role: Optional[te.QuantizerRole]): + """Return a fresh quantizer for every role, including ``None``. + + ``BASE_FACTORY`` makes the factory total: unknown roles, future role values, + and untargeted modules all retain valid MXFP8 behavior. + """ + + if role is not None and role.name == THREE_FORMAT_MODULE: + # Constructing fresh child quantizers for every call is recommended. + if role.tensor_type == "input": + # Wgrad retains the original BF16 input. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # Dgrad uses NVFP4 quantized from the dequantized MXFP8 fprop weight. + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # Dgrad uses NVFP4 while wgrad retains the original BF16 gradient. + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + + if role is not None and role.name == HIGH_PRECISION_MODULE: + return te.IdentityQuantizer() + + return BASE_FACTORY(role) + + +linear_options = {"bias": False, "params_dtype": torch.bfloat16, "device": "cuda"} +model = torch.nn.Sequential( + te.Linear(128, 256, name=THREE_FORMAT_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 256, name=HIGH_PRECISION_MODULE, **linear_options), + torch.nn.GELU(), + te.Linear(256, 128, name="demo.output", **linear_options), +) +inputs = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) +recipe = CustomRecipe(qfactory=quantizer_factory) + +with te.autocast(enabled=True, recipe=recipe): + outputs = model(inputs) + +loss = outputs.float().square().mean() +loss.backward() + +# END_FINE_GRAINED_QUANTIZATION_EXAMPLE + +gradients = [inputs.grad, *(parameter.grad for parameter in model.parameters())] +assert all(gradient is not None for gradient in gradients) +assert all(torch.isfinite(gradient).all() for gradient in gradients) + +print(f"GPU: {torch.cuda.get_device_name()}") +print(f"TE Linear names: {[model[index].name for index in (0, 2, 4)]}") +print(f"loss: {loss.item():.6f}; forward and backward completed") diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst new file mode 100644 index 0000000000..301b655b3f --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -0,0 +1,364 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _fine-grained-quantization-recipes: +.. _heterogeneous-quantization-recipes: + +Fine-grained quantization recipes +================================= + +Standard TE recipes quantize the whole model the same way. That is often too +coarse: one sensitive layer may need BF16 while the rest runs in MXFP8, or a +gradient GEMM may tolerate a cheaper format than the forward pass. Fine-grained +recipes lift this restriction: you write a small factory function that picks a +quantizer for each slot TE asks about, and pass it via +:class:`~transformer_engine.common.recipe.CustomRecipe` to the usual +:class:`~transformer_engine.pytorch.autocast`. +"Fine-grained" refers to the granularity of that choice (per module, tensor +role, and GEMM direction), not to the block size of the scaling factors. + +.. warning:: + + Fine-grained recipes are currently available only in the PyTorch API of + TE. + +.. warning:: + + Fine-grained recipes and their construction APIs are experimental: API, + validation, and kernel coverage may change without notice. This guide does + not define a supported recipe or an expected accuracy/performance ordering. + + +Example: mixing MXFP8, NVFP4, and BF16 +-------------------------------------- + +The `runnable example `__ +makes the following assignments: + +.. raw:: html + :file: img/fine_grained_assignments.svg + +*Figure 1. Precision assignments per module and GEMM used throughout this +guide.* + +A minimal factory implementing these assignments, plugged into the standard +TE autocast path: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import transformer_engine.pytorch as te + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + mxfp8_factory, + nvfp4_factory, + ) + + + def quantizer_factory(role): + if role is not None and role.name == "demo.fc1": + if role.tensor_type == "input": + # wgrad keeps the original BF16 input + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role.tensor_type == "weight": + # fprop in MXFP8, dgrad in NVFP4 + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="rowwise_dequantized", + ) + if role.tensor_type == "grad_output": + # dgrad in NVFP4, wgrad keeps the original BF16 gradient + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="original", + ) + if role is not None and role.name == "demo.fc2": + return te.IdentityQuantizer() # whole module stays in BF16 + return mxfp8_factory(role) # every other TE module in MXFP8 + + + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + +The complete, runnable version is available +`on GitHub `__ +(requires Blackwell or later); run it from the repository root after +installing TE: + +.. code-block:: bash + + python docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py + +CustomRecipe and quantizer factory +---------------------------------- + +:class:`~transformer_engine.common.recipe.CustomRecipe` is used like any +other TE recipe (``DelayedScaling``, ``MXFP8BlockScaling``, ...), but carries +no quantization logic of its own: TE asks your ``qfactory`` for a quantizer +whenever a module needs one. + +Each TE module defines an ordered role list for the forward and backward +quantizer slots it needs. When module recipe state is initialized or rebuilt, +a ``CustomRecipe`` calls ``qfactory(role)`` once for every slot in that list. +It does not call the factory on every unchanged forward. + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + # QuantizerRole describes the slot being configured (fields below): + # + # @dataclasses.dataclass(frozen=True) + # class QuantizerRole: + # module_type: str = "" + # tensor_type: str = "" + # name: str = "" + + + def quantizer_factory(role: Optional[te.QuantizerRole]): + # construct a fresh quantizer on every call + ... + # Boundary slots may pass role=None or a role with empty fields, so + # always end with a default that covers every remaining role. + return mxfp8_factory(role) + + + # The factory plugs into the standard TE autocast path: + recipe = CustomRecipe(qfactory=quantizer_factory) + + with te.autocast(enabled=True, recipe=recipe): + output = model(inputs) + + **Module type** + + The kind of TE module that owns the slot, filled in by TE itself: + + * ``"linear"`` — ``Linear``, ``LayerNormLinear``, ``fc1``/``fc2`` in + ``LayerNormMLP``, ``qkv``/``proj`` in ``MultiheadAttention``; + * ``"grouped_linear"`` — ``GroupedLinear``; + * ``"dpa"`` — ``DotProductAttention``. + + **Tensor type** + + Which tensor of that module the quantizer will process, also filled in + by TE. For ``"linear"`` and ``"grouped_linear"``: + + * ``"input"`` — the activation (fprop, wgrad); + * ``"weight"`` — (fprop, dgrad); + * ``"grad_output"`` — the incoming gradient (dgrad, wgrad). + + For ``"dpa"``: + + * ``"qkv"`` — the query/key/value tensor; + * ``"s"`` — the softmax output; + * ``"do"`` — the output gradient; + * ``"dp"`` — the gradient of ``"s"``. + + **Name** + + The identity of one concrete module instance, supplied by the caller: + ``te.Linear(..., name="decoder.39.fc2")``. Composite TE modules may + append suffixes such as ``.fc1``, ``.fc2``, and ``.proj``. + + The role vocabulary is experimental and may grow between releases — + one more reason to end the factory with a total default. Treat the role + strings as selectors, not a fixed enumeration. Prefer a module-level + function for the factory itself, so that launchers and checkpointing + setups can import or pickle it. + + TE provides factories for its native quantizers in + ``transformer_engine.pytorch.custom_recipes.quantizer_factories`` + (``mxfp8_factory``, ``nvfp4_factory``, ...). They can be used as + defaults or to construct ``HybridQuantizer`` children. Additional + specialized recipes are available in + ``transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo``. + + The factory is not limited to TE-native quantizers: it may return your + own :class:`~transformer_engine.pytorch.Quantizer` subclass, and custom + quantizers can also serve as ``HybridQuantizer`` children. The GEMMs + still need to receive representations in formats they support. + +HybridQuantizer +--------------- + +During training, each tensor of a ``Linear`` or ``GroupedLinear`` layer feeds +two different GEMMs: its rowwise representation feeds one, its columnwise +representation the other (the exact operand layout is described in the +:doc:`Introduction <../introduction/introduction>`). Since those two GEMMs may +want different formats, the tensor needs a quantizer per direction: +:class:`~transformer_engine.pytorch.HybridQuantizer` composes a rowwise and a +columnwise quantizer, and its output, +:class:`~transformer_engine.pytorch.HybridQuantizedTensor`, composes the +corresponding representations. + +.. tabs:: + + .. tab:: PyTorch + + The following is pseudocode illustrating the composition: + + .. code-block:: text + + quantizer = te.HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3), + columnwise_quantizer=NVFP4Quantizer(), + columnwise_source="original", # or "rowwise_dequantized" + ) + + # Quantization yields a HybridQuantizedTensor whose rowwise + # representation is MXFP8 and columnwise representation is NVFP4; + # each GEMM consumes the representation it needs. + qtensor = quantizer(tensor) + +.. raw:: html + :file: img/hybrid_quantizer.svg + +*Figure 2. HybridQuantizer composes a rowwise and a columnwise quantizer; each +representation of the result feeds a different GEMM.* + +Choosing the columnwise source +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``columnwise_source`` is a separate numerical recipe choice that controls the +source for the columnwise representation: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Value + - Columnwise source + * - ``"original"`` + - The original high-precision tensor. + * - ``"rowwise_dequantized"`` + - Dequantized rowwise representation. + +.. raw:: html + :file: img/hybrid_columnwise_source.svg + +*Figure 3. The columnwise representation can be derived from the original +high-precision tensor or from the dequantized rowwise representation.* + +For forward inputs and weights, ``"rowwise_dequantized"`` derives the backward +representation from the value consumed in the forward direction. This +can improve forward/backward numerical consistency and may affect convergence. +It does not recover information discarded by rowwise quantization. +``"original"`` instead derives both representations from the original tensor. +Choose the provenance as part of the numerical recipe. + +IdentityQuantizer +----------------- + +:class:`~transformer_engine.pytorch.IdentityQuantizer` stores its input in the +held compute dtype, typically BF16, FP16, or FP32. It can keep a complete slot +in high precision or act as one child of a ``HybridQuantizer``: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + # whole slot in high precision (e.g. a module kept in BF16) + quantizer = te.IdentityQuantizer() + + # one direction in high precision, the other quantized + quantizer = te.HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + +Note that in the second example the columnwise direction is high precision but +holds the value reconstructed from MXFP8, not the original input — see +`Choosing the columnwise source`_ above. + +Example: one format per GEMM +---------------------------- + +A natural way to design a recipe is to pick one format for each GEMM. To +translate that into quantizers, look at what each GEMM consumes — both of its +operands must be in that GEMM's format: + +* **fprop** consumes ``input.rowwise`` and ``weight.rowwise``; +* **dgrad** consumes ``grad_output.rowwise`` and ``weight.columnwise``; +* **wgrad** consumes ``input.columnwise`` and ``grad_output.columnwise``. + +Reading the same table per tensor gives the ``HybridQuantizer`` for each role. +For the example assignments (fprop in MXFP8, dgrad in NVFP4, wgrad in BF16): + +.. code-block:: text + + input = HybridQuantizer(rowwise=MXFP8, columnwise=BF16) # fprop | wgrad + weight = HybridQuantizer(rowwise=MXFP8, columnwise=NVFP4) # fprop | dgrad + grad_output = HybridQuantizer(rowwise=NVFP4, columnwise=BF16) # dgrad | wgrad + +.. raw:: html + :file: img/fine_grained_linear_mapping.svg + +*Figure 4. Each GEMM consumes one representation of each of its two operand +tensors; giving both operands the same format sets that GEMM's precision.* + +If two directions use the same quantizer configuration, a plain quantizer may +replace the corresponding hybrid; one factory may return both plain and hybrid +quantizers. +The two operands of each GEMM still need a combination supported by that GEMM +backend. TE may reject incompatible quantizer pairs or unsupported layouts. + +.. note:: + + On supported hardware these recipes run TE's regular quantized kernels: the + tensors are quantized on the GPU and the GEMMs execute in the selected + low-precision formats. TE does not fall back to fake quantization + (quantize-dequantize followed by a high-precision GEMM). + + +Validating and optimizing a recipe +---------------------------------- + +The factory API can express more recipes than TE has kernels for, so any +assignment lands in one of three buckets: + +* **Fast** — quantization hits TE's fused kernels and every GEMM runs a + native low-precision implementation. +* **Correct but potentially unoptimized** — the recipe executes, but some + selected paths may not have fused or optimized implementations in the + current TE release. For example, ``HybridQuantizer`` may produce its rowwise + and columnwise representations in separate kernel launches; future releases + may fuse this work. +* **Rejected** — the two operands of some GEMM end up in a combination of + formats or layouts that no GEMM backend supports, and TE raises an error. + This can happen with plain and hybrid quantizers alike. + +Before adopting a recipe for a real workload, check that: + +* it executes at all on the target GPU, software version, and modules; +* it runs on optimized kernels rather than fallback paths; +* accuracy and convergence hold on the target model and distributed setup; +* throughput and memory actually improve on the target workload. + +The unoptimized paths are still useful: accuracy and convergence experiments can run +on them before dedicated kernels exist, so the precision of each GEMM can be +treated as an accuracy/performance trade-off to explore. + +API reference +------------- + +See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``, +``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types. +See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``. diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg new file mode 100644 index 0000000000..b9aa48333d --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg @@ -0,0 +1,95 @@ + + + Precision assignment by tensor role and module + Each tensor role provides a rowwise and a columnwise representation, consumed by fprop, dgrad, and wgrad GEMMs. demo.fc1: input is MXFP8 rowwise and original BF16 columnwise; weight is MXFP8 rowwise and NVFP4 columnwise; grad_output is NVFP4 rowwise and original BF16 columnwise. demo.fc2 keeps every tensor in BF16. Other TE modules use MXFP8 everywhere. + + + + + Precision assignment by tensor role and module + + demo.fc1 + demo.fc2 + Other TE modules + + + + input + rowwise (fprop) + + MXFP8 + + BF16 + + MXFP8 + + columnwise (wgrad) + + BF16 (original) + + BF16 + + MXFP8 + + + + + weight + rowwise (fprop) + + MXFP8 + + BF16 + + MXFP8 + + columnwise (dgrad) + + NVFP4 + + BF16 + + MXFP8 + + + + + grad_output + rowwise (dgrad) + + NVFP4 + + BF16 + + MXFP8 + + columnwise (wgrad) + + BF16 (original) + + BF16 + + MXFP8 + + + + + MXFP8 + + NVFP4 + + BF16 (high precision) + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg new file mode 100644 index 0000000000..0e17ffdfa4 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg @@ -0,0 +1,59 @@ + + + Per-GEMM formats and the operands each GEMM consumes + Three GEMM cards. Fprop consumes input.rowwise and weight.rowwise, both MXFP8. Dgrad consumes weight.columnwise and grad_output.rowwise, both NVFP4. Wgrad consumes input.columnwise and grad_output.columnwise, both BF16. + + + + + + + fprop + format: MXFP8 + + input.rowwise + MXFP8 + × + + weight.rowwise + MXFP8 + + + + + dgrad + format: NVFP4 + + grad_output.rowwise + NVFP4 + × + + weight.columnwise + NVFP4 + + + + + wgrad + format: BF16 + + input.columnwise + BF16 + × + + grad_output.columnwise + BF16 + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg new file mode 100644 index 0000000000..ebc077eb04 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg @@ -0,0 +1,74 @@ + + + Hybrid quantizer columnwise source choices + With original provenance, both quantizers consume the original high-precision tensor. With rowwise-dequantized provenance, the columnwise quantizer consumes the dequantized rowwise representation. + + + + + + + + Choosing the columnwise source + + + + columnwise_source="original" + + + High-precision tensor + + + + same original source + + + Rowwise quantizer + + Columnwise quantizer + + + + + Rowwise + representation + + Columnwise + representation + + + + + columnwise_source="rowwise_dequantized" + + + High-precision tensor + + + + Rowwise quantizer + + + + Rowwise + representation + + + + Dequantize + + + + Columnwise quantizer + + + Columnwise + representation + + diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg new file mode 100644 index 0000000000..6e542306c7 --- /dev/null +++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg @@ -0,0 +1,58 @@ + + + HybridQuantizer data flow + A high-precision tensor enters a HybridQuantizer whose rowwise child is an MXFP8 quantizer and columnwise child is an NVFP4 quantizer. The result is a HybridQuantizedTensor with an MXFP8 rowwise representation and an NVFP4 columnwise representation, each consumed by a different GEMM. + + + + + + + + + + tensor + high precision (BF16) + + + + + + + HybridQuantizer + + rowwise_quantizer + MXFP8Quantizer + + columnwise_quantizer + NVFP4Quantizer + + + + + + + HybridQuantizedTensor + + rowwise + MXFP8 + + columnwise + NVFP4 + + + + + GEMM 1 + GEMM 2 + diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst index 48d17db8d5..557a4b09e8 100644 --- a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst @@ -180,6 +180,37 @@ Blackwell and later (SM >= 10.0) – the recipe is emulated with MXFP8. Note tha ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Blockwise scaling uses + :class:`~transformer_engine.pytorch.Float8BlockQuantizer`. Each block of + the tensor gets its own power-of-two scale: ``block_scaling_dim=1`` + scales 1x128 blocks, ``block_scaling_dim=2`` (the default) scales + 128x128 blocks. This recipe is not available in TE/JAX. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=1, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + Developer Notes --------------- diff --git a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst index cac3792194..2436a07566 100644 --- a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst +++ b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst @@ -164,6 +164,57 @@ Here's how to use FP8 Current Scaling recipe in PyTorch and JAX: ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Current scaling uses + :class:`~transformer_engine.pytorch.Float8CurrentScalingQuantizer`. It + needs no external state: at each call it computes the amax of the input + tensor, derives the scale from it, and then quantizes. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8CurrentScalingQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + device="cuda", + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + Current scaling uses ``CurrentScaleQuantizer``. At each call it computes + the amax of the input tensor, derives the scale from it, and then + quantizes. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- @@ -177,4 +228,4 @@ On Blackwell and later, rowwise and columnwise tensors share the same memory lay so all-gather of columnwise tensors is directly supported. For Hopper and Ada, all-gather of transposed FP8 tensors is not supported. -The rowwise tensor is gathered first, then transposed to columnwise format. \ No newline at end of file +The rowwise tensor is gathered first, then transposed to columnwise format. diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst index d39787f6f5..99a379eed1 100644 --- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst +++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst @@ -160,4 +160,69 @@ However, amax reduction works slightly differently in different frameworks. Supported devices ----------------- -Ada and later (SM 8.9+) \ No newline at end of file +Ada and later (SM 8.9+) + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + Delayed scaling uses + :class:`~transformer_engine.pytorch.Float8Quantizer`. It does not + compute the scaling factor from the current tensor: one-element + ``scale`` and ``amax`` buffers are supplied at construction. + Quantization applies the given scale and records the tensor's amax into + the ``amax`` buffer. + + During training both buffers are views into the recipe state: ``scale`` + into its per-quantizer scale vector, ``amax`` into the current row of + its ``(amax_history_len, num_quantizers)`` amax history. At the end of + each step the recipe state computes a new scale from the history (its + max or most recent entry, per ``amax_compute_algo``), rolls the history + by one slot, and zeroes the current row — all in place, so the views + held by the quantizer stay valid for the whole training run. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.Float8Quantizer( + scale=torch.ones(1, device="cuda"), + amax=torch.zeros(1, device="cuda"), + fp8_dtype=te.DType.kFloat8E4M3, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + Delayed scaling uses ``DelayedScaleQuantizer``. The ``scale`` and the + ``amax_history`` (1024 entries by default) are fields of the quantizer + itself, carried through JAX transformations as its pytree state. Each + ``quantize()`` call applies the current ``scale``, then updates the + state: the tensor's amax is written into the history, a new scale is + computed from the history (max or most-recent entry, per + ``amax_compute_algo``), and the history is rolled by one slot. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.DELAYED_TENSOR_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst index 0a798f1364..b9649c00a4 100644 --- a/docs/features/low_precision_training/index.rst +++ b/docs/features/low_precision_training/index.rst @@ -15,4 +15,5 @@ Low precision training fp8_blockwise_scaling/fp8_blockwise_scaling.rst mxfp8/mxfp8.rst nvfp4/nvfp4.rst + fine_grained_quantization/fine_grained_quantization.rst speedups.rst diff --git a/docs/features/low_precision_training/introduction/introduction.rst b/docs/features/low_precision_training/introduction/introduction.rst index fba7796ece..2255308b04 100644 --- a/docs/features/low_precision_training/introduction/introduction.rst +++ b/docs/features/low_precision_training/introduction/introduction.rst @@ -283,3 +283,71 @@ so GEMM with tensors ``A`` and ``B`` returns ``B * A^T``. :file: img/fp8_linear_flow.svg *Figure 4: Forward pass of a Linear layer with low precision data flow.* + +Quantizers +---------- + +Every recipe implements its quantization logic in a **quantizer** — an object +that converts a high-precision tensor into a quantized one. TE modules create +and use quantizers internally according to the active recipe, but a quantizer +can also be used directly: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + qtensor = quantizer(tensor) # quantize + roundtrip = qtensor.dequantize() # back to high precision + + The main parts of the interface are: + + * ``quantize(tensor)`` — quantizes a high-precision tensor and returns a + ``QuantizedTensor``; calling the quantizer (``quantizer(tensor)``) is + a shorthand; + * ``update_quantized(src, dst)`` — quantizes ``src`` in place into an + already-allocated quantized tensor ``dst``; + * ``make_empty(shape)`` — allocates an uninitialized quantized tensor to + be filled later; + * ``rowwise_usage`` / ``columnwise_usage`` — flags selecting which of + the two GEMM-oriented representations the produced tensor holds; + * the returned ``QuantizedTensor`` supports ``dequantize()`` back to + high precision. + + .. tab:: JAX + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + + The main parts of the interface are: + + * ``quantize(x, is_rowwise=..., is_colwise=...)`` — quantizes a tensor + and returns a ``ScaledTensor`` holding the requested representations + (the default comes from the quantizer's ``q_layout``); + * the returned ``ScaledTensor`` supports ``dequantize()`` back to high + precision; + * quantizers are registered pytrees, so they can be passed through JAX + transformations. + +Each recipe section ends with a short description of that recipe's quantizer. diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst index 1fbcc43af9..1827d42cb6 100644 --- a/docs/features/low_precision_training/mxfp8/mxfp8.rst +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -152,6 +152,57 @@ SM 10.0, SM 10.3 ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + MXFP8 uses :class:`~transformer_engine.pytorch.MXFP8Quantizer`. Every + 32-element block shares one power-of-two (E8M0) scale, computed from the + block's amax at quantization time; no external state is needed. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + MXFP8 uses ``BlockScaleQuantizer`` — the JAX quantizer for block-based + scaling, selected by ``ScalingMode.MXFP8_1D_SCALING``. Instead of one + scale per tensor, the tensor is split along the quantization axis into + 32-element blocks and each block gets its own power-of-two (E8M0) scale, + computed from that block's amax at quantization time. Because the scale + is derived from the current data, no external state (scale buffers or + amax history) is needed. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- @@ -210,4 +261,4 @@ All-gather of columnwise tensors All-gather of columnwise tensors is supported and necessary because: - columnwise quantized tensors cannot be computed from rowwise quantized ones, -- gathering high-precision tensors is avoided in most cases for performance reasons. \ No newline at end of file +- gathering high-precision tensors is avoided in most cases for performance reasons. diff --git a/docs/features/low_precision_training/nvfp4/nvfp4.rst b/docs/features/low_precision_training/nvfp4/nvfp4.rst index 900789b0d3..26d798651d 100644 --- a/docs/features/low_precision_training/nvfp4/nvfp4.rst +++ b/docs/features/low_precision_training/nvfp4/nvfp4.rst @@ -250,6 +250,60 @@ Supported devices ---- + +Quantizer +--------- + +.. tabs:: + + .. tab:: PyTorch + + NVFP4 uses :class:`~transformer_engine.pytorch.NVFP4Quantizer`. It + implements the two-level scaling described above: an FP8 (E4M3) scale + per 16-element block plus one FP32 scale per tensor. Further keyword + options select the recipe variations from this page (random Hadamard + transforms, stochastic rounding, 2D weight scaling); they are internal + knobs and may change without notice. + + .. code-block:: python + + import torch + import transformer_engine.pytorch as te + + tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + + quantizer = te.NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + ) + + qtensor = quantizer(tensor) + roundtrip = qtensor.dequantize() + + .. tab:: JAX + + NVFP4 uses its own ``NVFP4Quantizer``, with the same two-level scaling. + ``ScalingMode.NVFP4_1D_SCALING`` selects per-block scaling only, + ``ScalingMode.NVFP4_2D_SCALING`` adds 2D weight scaling. + + .. code-block:: python + + import jax.numpy as jnp + from transformer_engine.jax.quantize import ( + QuantizerFactory, ScalingMode, QuantizeLayout, + ) + + x = jnp.ones((256, 256), dtype=jnp.bfloat16) + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.NVFP4_1D_SCALING, + q_dtype=jnp.float4_e2m1fn, + q_layout=QuantizeLayout.ROWWISE, + ) + qtensor = quantizer.quantize(x) + roundtrip = qtensor.dequantize() + Developer Notes --------------- diff --git a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst index 2c21799dd6..afa5d16c86 100644 --- a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst +++ b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst @@ -143,6 +143,61 @@ Transformer Engine chooses the best possible fusion internally taking the recipe *Figure 3: Three scenarios of producing quantized tensors in rowwise and columnwise usages.* +**Usages in the quantizer API** + +The usages are visible directly in the quantizer API: + +.. tabs:: + + .. tab:: PyTorch + + At quantization time, the quantizer's ``rowwise_usage`` and + ``columnwise_usage`` flags select which representations ``quantize()`` + produces; when both are set, the representations are computed together + in one fused kernel (scenario 1 above). + + After quantization, ``update_usage()`` on the quantized tensor removes a + representation or, when supported by the format, generates a missing one. + Passing ``rowwise_usage=False`` after the forward pass frees the rowwise + data while keeping the columnwise data for backward. Some formats also + support ``columnwise_usage=True`` to create the columnwise representation + from the data already present (e.g. by a transpose on Hopper — scenario 3 + above); unsupported requests raise an error. Arguments left as ``None`` + preserve the current state. + + .. code-block:: python + + quantizer = te.MXFP8Quantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + + qtensor = quantizer(tensor) # both representations, one fused kernel + + qtensor.update_usage(rowwise_usage=False) # drop rowwise, keep columnwise + + .. tab:: JAX + + The usages are selected when the tensor is quantized: the quantizer's + ``q_layout`` (``QuantizeLayout.ROWWISE``, ``COLWISE``, or + ``ROWWISE_COLWISE``) sets the default, and ``quantize()`` accepts + ``is_rowwise``/``is_colwise`` overrides. Requesting both usages returns + a ``ScaledTensor2x`` holding the two representations. There is no + in-place ``update_usage()``: JAX arrays are immutable, so a + representation is not added or dropped later — unneeded ones are simply + not requested and get dropped by XLA's dead-code elimination. + + .. code-block:: python + + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=QuantizeLayout.ROWWISE_COLWISE, + ) + + qtensor = quantizer.quantize(x) # ScaledTensor2x, both representations + rowwise_only = quantizer.quantize(x, is_rowwise=True, is_colwise=False) Memory usage @@ -470,4 +525,3 @@ Actual behavior depends on the recipe and module configuration. *Figure 5: All-gather of quantized tensors for input and gradient tensors. This is one possible scenario — actual behavior varies depending on the recipe and module configuration.* - diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index b6fc8f7794..1ca4121a4f 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -612,11 +612,6 @@ def _check_configs(self): pytest.skip( "B1SS, BHSS and 11SS bias shapes are only supported for non-padding mask" ) - elif self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: - pytest.skip( - "B1SS, BHSS and 11SS bias shapes are only supported for " - "the F16_arbitrary_seqlen backend." - ) def _setup_inputs(self): self._check_configs() diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 08e762045a..fc482ce5b2 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1733,6 +1733,10 @@ def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): ), f"missing hybrid current-scaling directions: {checked}" +@pytest.mark.skipif( + not te.is_fp8_available(), + reason=te.is_fp8_available(return_reason=True)[1], +) def test_fused_adam_hybrid_identity_fp8_master_weights(): """FSDP2 + FusedAdam with Hybrid(FP8 current rowwise, Identity columnwise). diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 39d825e701..c2a9df765e 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -510,6 +510,10 @@ def _hybrid_param_count(): _check_fp8_fsdp2_allgather(model, tols=dict(atol=5e-4, rtol=5e-3)) +@pytest.mark.skipif( + not te.is_fp8_available(), + reason=te.is_fp8_available(return_reason=True)[1], +) def test_distributed_hybrid_identity_all(): """FSDP2 training/all-gather with an all-Identity CustomRecipe. diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index d733286093..c314bf5bb6 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -12,6 +12,7 @@ import pathlib import subprocess import sys +import tempfile from typing import Optional import pytest @@ -58,7 +59,8 @@ def world_group() -> torch.distributed.ProcessGroup: torch.cuda.set_device(rank) group = torch.distributed.init_process_group( "nccl", - init_method="file:///tmp/rdzv", + # Each parallel job must use a fresh FileStore shared by only its ranks. + init_method=f"file://{os.environ['NVTE_TEST_RDZV_PATH']}", world_size=world_size, rank=rank, ) @@ -1053,10 +1055,10 @@ def test_distributed_fuser_ops(world_size: int) -> None: current_file, "--parallel", ] - result = subprocess.run( - command, - check=True, - ) + with tempfile.TemporaryDirectory(prefix="te-test-fusible-ops-") as temp_dir: + env = dict(os.environ) + env["NVTE_TEST_RDZV_PATH"] = str(pathlib.Path(temp_dir) / "rdzv") + subprocess.run(command, check=True, env=env) def main() -> None: diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index 07dffebf5f..38f49a96cb 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -12,6 +12,7 @@ import pathlib import subprocess import sys +import tempfile import pytest import torch @@ -107,7 +108,8 @@ def world_group() -> torch.distributed.ProcessGroup: torch.cuda.set_device(local_rank) group = torch.distributed.init_process_group( "nccl", - init_method="file:///tmp/rdzv", + # Each parallel job must use a fresh FileStore shared by only its ranks. + init_method=f"file://{os.environ['NVTE_TEST_RDZV_PATH']}", world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{local_rank}"), @@ -471,7 +473,9 @@ def test_fuser_ops_with_userbuffers( env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" # Launch parallel job - run_distributed(command, env=env) + with tempfile.TemporaryDirectory(prefix="te-test-fusible-ops-userbuffers-") as temp_dir: + env["NVTE_TEST_RDZV_PATH"] = str(pathlib.Path(temp_dir) / "rdzv") + run_distributed(command, env=env) def main() -> None: diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 5a848dc0e8..8f85f57f32 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -5,6 +5,8 @@ from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest import copy +import gc +import weakref import torch from transformer_engine.pytorch import ( @@ -994,6 +996,206 @@ def hook(module: torch.nn.Module) -> None: ] +def test_ordered_warmup_releases_consumed_outputs() -> None: + """Ordered warmup should only retain outputs until their corresponding backward.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.previous_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + is_warmup = not torch.cuda.is_current_stream_capturing() + if is_warmup and self.previous_output is not None: + assert self.previous_output() is None + output = input_ * 2 + if is_warmup: + self.previous_output = weakref.ref(output) + return output + + module = OutputLifetimeModule() + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2)) + graphed_callables = make_graphed_callables( + (module,), + sample_args, + num_warmup_iters=2, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + ) + assert module.previous_output is not None + assert module.previous_output() is None + reset_graphs(graphed_callables) + + +def test_unordered_warmup_releases_consumed_outputs() -> None: + """Unordered warmup should release each output after its corresponding backward.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self, output_refs: list, module_idx: int) -> None: + super().__init__() + self.output_refs = output_refs + self.module_idx = module_idx + self.capture_started = False + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.capture_started = True + else: + self.output_refs[self.module_idx] = weakref.ref(output) + return output + + output_refs = [None, None] + modules = tuple(OutputLifetimeModule(output_refs, module_idx) for module_idx in range(2)) + + def first_module_backward_pre_hook(_module: torch.nn.Module) -> None: + if not modules[0].capture_started: + assert output_refs[1] is not None + assert output_refs[1]() is None + + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in modules), + num_warmup_iters=2, + capture_time_hooks=[ + {"backward_pre_hooks": {0: first_module_backward_pre_hook}}, + None, + ], + ) + assert all(output_ref is not None and output_ref() is None for output_ref in output_refs) + reset_graphs(graphed_callables) + + +def test_inference_warmup_does_not_retain_outputs() -> None: + """Inference warmup should release outputs as soon as each forward returns.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self, previous_output: list) -> None: + super().__init__() + self.previous_output = previous_output + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + is_warmup = not torch.cuda.is_current_stream_capturing() + if is_warmup and self.previous_output[0] is not None: + assert self.previous_output[0]() is None + output = input_ * 2 + if is_warmup: + self.previous_output[0] = weakref.ref(output) + return output + + previous_output = [None] + modules = tuple(OutputLifetimeModule(previous_output).eval() for _ in range(2)) + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, 8, device="cuda"),) for _ in modules), + num_warmup_iters=2, + ) + assert previous_output[0] is not None + assert previous_output[0]() is None + reset_graphs(graphed_callables) + + +def test_reused_capture_buffers_release_outputs_after_backward() -> None: + """Capture locals must not keep weak-refed output buffers alive.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.previous_capture_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + if ( + torch.cuda.is_current_stream_capturing() + and self.previous_capture_output is not None + ): + assert self.previous_capture_output() is None + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.previous_capture_output = weakref.ref(output) + return output + + module = OutputLifetimeModule() + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2)) + graphed_callables = make_graphed_callables( + (module,), + sample_args, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + ) + assert module.previous_capture_output is not None + assert module.previous_capture_output() is None + reset_graphs(graphed_callables) + + +def test_reset_releases_only_the_selected_callable() -> None: + """Reset releases one callable's graph state without retaining its peers.""" + + class CaptureOutputModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.capture_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.capture_output = weakref.ref(output) + return output + + modules = tuple(CaptureOutputModule().cuda() for _ in range(2)) + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, device="cuda", requires_grad=True),) for _ in modules), + ) + capture_outputs = tuple(module.capture_output for module in modules) + assert all(output is not None and output() is not None for output in capture_outputs) + + graphed_callables[0].reset() + graphed_callables[0].reset() + gc.collect() + assert capture_outputs[0]() is None + assert capture_outputs[1]() is not None + + output = graphed_callables[1](torch.randn(4, device="cuda", requires_grad=True)) + output.sum().backward() + del output + graphed_callables[1].reset() + gc.collect() + assert capture_outputs[1]() is None + + +@pytest.mark.parametrize("with_order", (False, True)) +def test_reset_rejects_all_replay_entry_points(with_order: bool) -> None: + """Reset is idempotent and terminal for forward and backward replay.""" + + class TestModule(torch.nn.Module): + def forward(self, input_: torch.Tensor) -> torch.Tensor: + return input_ * 2 + + module = TestModule().cuda() + sample_input = torch.ones(4, device="cuda", requires_grad=True) + graph_options = {} + if with_order: + graph_options = {"_order": [1, -1], "_num_layers_per_chunk": [1]} + graphed_callable = make_graphed_callables(module, (sample_input,), **graph_options) + output = graphed_callable(torch.randn_like(sample_input, requires_grad=True)) + torch.cuda.synchronize() + + graphed_callable.reset() + graphed_callable.reset() + if not with_order: + # The eager fallback for a different training state is invalid after reset too. + graphed_callable.eval() + + error = "has been reset and can no longer be used" + with pytest.raises(RuntimeError, match=error): + graphed_callable(torch.randn_like(sample_input, requires_grad=True)) + with pytest.raises(RuntimeError, match=error): + graphed_callable.backward_dw() + with pytest.raises(RuntimeError, match=error): + output.sum().backward() + + @pytest.mark.parametrize("with_order", (False, True)) def test_make_graphed_callables_with_capture_time_hooks(with_order: bool) -> None: """Test capture-time hooks around warmup and graph capture.""" diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 6cd1fc3065..2adce717c9 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -146,6 +146,173 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +def test_operation_fuser_caches_plans_by_grad_requirement(monkeypatch) -> None: + """Cache and restore fusion plans for checkpoint forward and recompute.""" + + # Count fusion-plan construction without depending on any particular real + # fusion implementation. Each distinct fusion configuration invokes this + # hook once, while a cache hit must bypass it entirely. + fusion_calls = 0 + + def track_fusion(ops, *, recipe): # pylint: disable=unused-argument + nonlocal fusion_calls + fusion_calls += 1 + # Preserve the operation list so this hook observes plan construction + # without changing the topology under test. + return ops + + # The fusion registries are class attributes shared by every OperationFuser. + # pytest's monkeypatch fixture restores all three after the test, preventing + # this synthetic fusion function from leaking into other tests. Keep only a + # joint forward-backward fusion hook so each plan build has one countable + # callback and no registered TE fusion can affect the result. + monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion]) + monkeypatch.setattr(OperationFuser, "forward_fusion_functions", []) + monkeypatch.setattr(OperationFuser, "backward_fusion_functions", []) + + # One Identity op is enough to exercise the cache. With one basic op, + # first_op_requiring_backward has an intentionally simple interpretation: + # 0: backward starts at the Identity op; + # 1: the boundary is past the only op, so no backward work is required. + fuser = OperationFuser([te_ops.Identity()]) + x = torch.ones(1, requires_grad=True) + # maybe_fuse_ops expects one extra-input collection per basic op. Identity + # has no extra inputs, so its collection is an empty tuple. + extra_inputs = [()] + + # Phase 1: the original checkpointed forward runs with grad disabled. This + # is the first invocation, so the fuser must construct and cache the no-grad + # configuration. The runtime backward boundary is past the only op. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 1 + assert fuser.first_op_requiring_backward == 1 + no_grad_forward_ops = fuser._forward_ops + no_grad_backward_ops = fuser._backward_ops + + # Phase 2: backward replays the checkpointed region with grad enabled. The + # backward boundary is part of the fusion key, allowing future fusion rules + # to choose a training-specific topology. The first grad-enabled invocation + # therefore constructs and caches a second configuration. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 0 + grad_forward_ops = fuser._forward_ops + grad_backward_ops = fuser._backward_ops + assert grad_forward_ops is not no_grad_forward_ops + assert grad_backward_ops is not no_grad_backward_ops + + # Phase 3: the next checkpointed forward must select the exact no-grad lists + # cached in phase 1. Before the cache was added, every boundary transition + # rebuilt the fused operations and called track_fusion again. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 1 + assert fuser._forward_ops is no_grad_forward_ops + assert fuser._backward_ops is no_grad_backward_ops + + # Phase 4: another recomputation must likewise restore the grad-enabled + # lists from phase 2. The full alternating sequence has built only the two + # configurations represented by its two fusion keys. + fuser.maybe_fuse_ops(True, None, x, extra_inputs) + assert fusion_calls == 2 + assert fuser.first_op_requiring_backward == 0 + assert fuser._forward_ops is grad_forward_ops + assert fuser._backward_ops is grad_backward_ops + + +def test_operation_fuser_resets_recipe_state_independently_from_plan_cache(monkeypatch) -> None: + """Track recipe-state resets independently from fusion-plan construction.""" + + fusion_calls = 0 + + def track_fusion(ops, *, recipe): # pylint: disable=unused-argument + nonlocal fusion_calls + fusion_calls += 1 + return ops + + # Replace the process-wide fusion registries so one callback corresponds to + # one plan construction. monkeypatch restores the registries after the test. + monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion]) + monkeypatch.setattr(OperationFuser, "forward_fusion_functions", []) + monkeypatch.setattr(OperationFuser, "backward_fusion_functions", []) + + op = te_ops.Identity() + reset_recipes = [] + first_forward_calls = 0 + + def track_recipe_reset(*, recipe): + reset_recipes.append(recipe) + + def track_first_forward(): + nonlocal first_forward_calls + first_forward_calls += 1 + + # Identity has no quantizers, so replace its state hooks with counters. This + # keeps the test CPU-only and isolates OperationFuser's reset decisions. + monkeypatch.setattr(op, "reset_recipe_state", track_recipe_reset) + monkeypatch.setattr(op, "pre_first_fuser_forward", track_first_forward) + + fuser = OperationFuser([op]) + x = torch.ones(1) + extra_inputs = [()] + + current_scaling = transformer_engine.common.recipe.Float8CurrentScaling(backward_override=None) + fuser.maybe_fuse_ops(False, current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 1 + + # A fresh but equivalent recipe does not invalidate state or the plan. + equivalent_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling( + backward_override=None + ) + fuser.maybe_fuse_ops(False, equivalent_current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 1 + + # Backward override affects both recipe state and fusion topology, so it + # triggers one reset and constructs a distinct cached plan. + overridden_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling( + backward_override="high_precision" + ) + fuser.maybe_fuse_ops(False, overridden_current_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling, overridden_current_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 2 + + delayed_scaling = transformer_engine.common.recipe.DelayedScaling( + amax_history_len=8, + backward_override=None, + ) + fuser.maybe_fuse_ops(False, delayed_scaling, x, extra_inputs) + assert reset_recipes == [current_scaling, overridden_current_scaling, delayed_scaling] + assert first_forward_calls == 1 + assert fusion_calls == 3 + + # Amax history length only affects delayed-scaling recipe state. Reset that + # state, but restore the existing DelayedScaling fusion plan from the cache. + resized_delayed_scaling = transformer_engine.common.recipe.DelayedScaling( + amax_history_len=16, + backward_override=None, + ) + fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs) + assert reset_recipes == [ + current_scaling, + overridden_current_scaling, + delayed_scaling, + resized_delayed_scaling, + ] + assert first_forward_calls == 1 + assert fusion_calls == 3 + + # Repeating the exact recipe parameters performs neither operation again. + fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs) + assert len(reset_recipes) == 4 + assert first_forward_calls == 1 + assert fusion_calls == 3 + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d48f7afae6..15c1ff6c51 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable +import contextlib import functools import os import math @@ -1067,6 +1068,46 @@ def train_step( class TestGroupedMLPFusedOp: """Tests for grouped MLP fused op""" + def test_fusion_requires_supported_grad_output_format(self, monkeypatch) -> None: + """Fuse E4M3 MXFP8 and NVFP4, but decline MXFP8 with an E5M2 backward.""" + from transformer_engine.common.recipe import Format, MXFP8BlockScaling, NVFP4BlockScaling + + fused_op_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + monkeypatch.setattr(fused_op_cls, "is_supported", classmethod(lambda cls: True)) + + fc1 = te.ops.GroupedLinear(1, 64, 128, bias=False, device="cuda") + activation = te.ops.ScaledSwiGLU(glu_interleave_size=32) + fc2 = te.ops.GroupedLinear(1, 64, 64, bias=False, device="cuda") + ops = [fc1, activation, fc2] + + def fuse(recipe): + return grouped_mlp_module.fuse_grouped_mlp_ops( + ops, + recipe=recipe, + fused_op_cls=fused_op_cls, + ) + + def assert_fused(recipe): + fused_ops = fuse(recipe) + assert len(fused_ops) == 1 + fused_op = fused_ops[0] + assert isinstance(fused_op, fused_op_cls) + assert list(fused_op.basic_ops) == ops + + hybrid = MXFP8BlockScaling(fp8_format=Format.HYBRID) + assert fuse(hybrid) is ops + + e4m3 = MXFP8BlockScaling(fp8_format=Format.E4M3) + assert_fused(e4m3) + + # NVFP4 quantizes gradients to FP4, so the FP8 format must not gate it. Forcing the + # lookup to E5M2 is what an NVFP4 recipe would hit if the check were not MXFP8-only. + monkeypatch.setattr( + grouped_mlp_module, "get_fp8_torch_dtype", lambda *_, **__: torch.float8_e5m2 + ) + nvfp4 = NVFP4BlockScaling(disable_rht=False) + assert_fused(nvfp4) + @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @@ -2805,6 +2846,194 @@ def train_step( assert_close(graph_grad, param.grad, **tols) +class TestGroupedMLPDeterminism: + """Determinism coverage for the CuTe DSL fused grouped MLP. + + Only the dSReLU wrapper can make ``dprob`` bit-exact, and only from cuDNN FE 1.28.0 on. + Anything else must refuse a determinism request rather than run non-deterministically. + """ + + @pytest.fixture + def _restore_torch_determinism(self): + """``use_deterministic_algorithms`` is process-global, so put it back.""" + previous = torch.are_deterministic_algorithms_enabled() + yield + torch.use_deterministic_algorithms(previous) + + @pytest.mark.parametrize( + "allow_nondeterministic,torch_flag,expected", + ( + (None, False, False), # default: non-deterministic algorithms are allowed + ("1", False, False), + ("0", False, True), # the TE variable alone + (None, True, True), # the torch flag alone, which TE must not ignore + ("1", True, True), # ... including when the TE variable says otherwise + ("0", True, True), + ), + ) + def test_either_knob_requests_determinism( + self, + monkeypatch, + _restore_torch_determinism, + *, + allow_nondeterministic: Optional[str], + torch_flag: bool, + expected: bool, + ) -> None: + """``=1`` is the absence of a request, not a request for non-determinism.""" + if allow_nondeterministic is None: + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + else: + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", allow_nondeterministic) + torch.use_deterministic_algorithms(torch_flag) + assert grouped_mlp_module._deterministic_algorithms_required() is expected + + def test_only_the_srelu_path_can_be_deterministic(self) -> None: + """The capability belongs to the wrapper, not the environment. Needs no GPU.""" + glu = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + unary = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + assert glu.grouped_gemm_dactivation_is_deterministic() is False + assert isinstance(unary.grouped_gemm_dactivation_is_deterministic(), bool) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + @pytest.mark.parametrize( + "activation,fused_cls", + ( + ("scaled_srelu", grouped_mlp_module.GroupedMLP_CuTeGEMMUnary), + ("scaled_swiglu", grouped_mlp_module.GroupedMLP_CuTeGEMMGLU), + ), + ) + def test_determinism_either_runs_or_refuses( + self, monkeypatch, *, activation, fused_cls + ) -> None: + """A request TE cannot honor must fail loudly; one it can must still be correct.""" + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + expectation = ( + contextlib.nullcontext() + if fused_cls.grouped_gemm_dactivation_is_deterministic() + else pytest.raises(RuntimeError, match="dprob") + ) + with expectation: + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=False, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation=activation, + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_scale_bias_refuses_under_the_torch_flag( + self, monkeypatch, _restore_torch_determinism + ) -> None: + """``scale_bias`` finishes ``dprob`` in a Triton kernel that reads only the env var. + + So the torch flag alone is the combination that used to pass this op's own check and + then reduce nondeterministically anyway, on a front-end new enough to say yes. + """ + fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + # warn_only so torch's own enforcement cannot raise first and mask what TE does. + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.raises(RuntimeError, match="dprob"): + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=True, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation="scaled_srelu", + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: + """Repeated identical runs must give a bit-identical ``dprob``. + + An ulp of reordering passes every tolerance in this file, so only an exact + comparison across runs can see it. + """ + fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + if not fused_cls.grouped_gemm_dactivation_is_deterministic(): + pytest.skip("dSReLU determinism needs cuDNN frontend 1.28.0 or later") + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + + device = torch.device("cuda") + dtype = torch.bfloat16 + # Measured on GB300, determinism off, 8 launches per shape (job 538058): this shape + # gives 7/7 runs differing from run 0, so the assertion below can actually fail. + # Shapes matter more than they look -- l=8 with the same n and tokens/group varies + # only 2/7, which an 8-run sample reports as stable often enough to be useless, and + # cudnn-frontend#521 measured its own l=4 / [256]*4 / n=512 as never varying. + group_size = 16 + hidden_size = 2048 + tokens_per_group = 1024 + split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) + num_tokens = tokens_per_group * group_size + + recipe = make_recipe("mxfp8") + + # Plain random tensors, not make_reference_and_test_tensors: this test compares two + # runs against each other, never against a reference, so the fp64 companion and the + # MXFP8 representability round-trip would both be allocated and thrown away. + def _rand(*shape, requires_grad=True) -> torch.Tensor: + out = torch.empty(shape, dtype=dtype, device=device).uniform_(-0.25, 0.25) + return out.requires_grad_() if requires_grad else out + + x = _rand(num_tokens, hidden_size) + dy = _rand(num_tokens, hidden_size, requires_grad=False) + probs = _rand(num_tokens) + + # No bias, or probs.grad comes from the Triton dbias kernel instead of cuDNN. + with te.quantized_model_init(enabled=True, recipe=recipe): + module = te.ops.Sequential( + te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ), + te.ops.ScaledSReLU(), + te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ), + ) + + def _run() -> torch.Tensor: + x.grad = None + probs.grad = None + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + y.backward(dy) + return probs.grad.detach().clone() + + runs = [_run()] + # Without the fusion there is no cuDNN dprob and the comparison proves nothing. + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], fused_cls) + # More than two, as cudnn-frontend#521 does: the cross-CTA order that determinism + # removes is set by the scheduler, so two runs can agree by luck. + runs += [_run() for _ in range(int(os.getenv("NVTE_TEST_DETERMINISM_REPEATS", "4")) - 1)] + torch.cuda.synchronize() + + assert torch.isfinite(runs[0]).all(), "dprob is not finite; the comparison would be moot" + # Bytes, not values: torch.equal calls +0.0 and -0.0 equal, and a change in reduction + # order can produce exactly that. Weight grads are excluded from the comparison -- + # the CuTe DSL wgrad kernel has its own K-split atomics, which this change leaves. + for index, later in enumerate(runs[1:], start=1): + assert torch.equal( + runs[0].contiguous().view(torch.uint8), later.contiguous().view(torch.uint8) + ), ( + f"dprob differs between run 0 and run {index} under determinism; max |delta| =" + f" {(runs[0].float() - later.float()).abs().max().item()}" + ) + + def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: if not mxfp8_available: pytest.skip(reason_for_no_mxfp8) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..9252051387 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -4,10 +4,12 @@ import abc import contextlib +import dataclasses import os import re import sys import warnings +from typing import Union import pytest import torch @@ -22,7 +24,7 @@ from torch._opaque_base import OpaqueBaseMeta from torch._library.opaque_object import ( get_opaque_type_name, - register_opaque_type, + register_custom_class, MemberType, ) @@ -37,11 +39,17 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear +from transformer_engine.pytorch.ops.fuser import OperationFuser +from transformer_engine.pytorch.ops.op import BasicOperation from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer -from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer +from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensor, + QuantizedTensorStorage, + Quantizer, +) from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec from transformer_engine.pytorch import ( is_fp8_available, @@ -215,9 +223,9 @@ def __fx_repr__(self): {"ToyQuantizer": ToyQuantizer}, ) - register_opaque_type( + register_custom_class( ToyQuantizer, - typ="value", + typ="constant", members={ "__setattr__": MemberType.USE_REAL, "set_usage": MemberType.USE_REAL, @@ -2339,3 +2347,320 @@ def fn(inp): "Unexpected recompilation(s) across different batch sizes: " f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled" ) + + +# --------------------------------------------------------------------------- # +# transformer_engine.pytorch.ops under torch.compile +# --------------------------------------------------------------------------- # + + +@dataclasses.dataclass(slots=True) +class _ScaleFwdArgs: + """Flat, ``self``-free inputs to the test operation's forward.""" + + input_: torch.Tensor + scale: torch.Tensor + + +@dataclasses.dataclass(slots=True) +class _ScaleBwdArgs: + """Flat inputs to the test operation's backward.""" + + grad_output: torch.Tensor = None + saved_input: torch.Tensor = None + scale: torch.Tensor = None + + +class _ScaleOp(BasicOperation): + """Test-only operation: multiply by a learnable scalar. + + Exists so the fuser's compiled path can be exercised without depending on + which real operations happen to have a custom op. It is the + smallest operation that still has a parameter gradient and a saved tensor. + """ + + fwd_args_type = _ScaleFwdArgs + bwd_args_type = _ScaleBwdArgs + + def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype)) + + @classmethod + def forward_impl(cls, args): + return args.input_ * args.scale, () + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), () + + @classmethod + def backward_impl(cls, args): + dy = args.grad_output + return dy * args.scale, (dy * args.saved_input).sum() + + @classmethod + def backward_fake(cls, args): + dy = args.grad_output + return ( + TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device), + TensorSpec(shape=(), dtype=dy.dtype, device=dy.device), + ) + + def setup_context(self, ctx, args, aux): + del aux + ctx.save_for_backward(args.input_, args.scale) + + def resolve_fwd_args( + self, + input_, + *, + requires_grad, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + ): + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + return _ScaleFwdArgs(input_=input_, scale=self.scale) + + def resolve_bwd_args(self, ctx, grad_output): + x, scale = ctx.saved_tensors + return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=scale) + + +class _BackwardScalePair(te.ops.FusedOperation): + """Backward-only fusion for the compile gate test.""" + + def fuser_backward(self, basic_op_ctxs, grad_output, **unused): + dx, grad_params_1 = self.basic_ops[1].op_backward(basic_op_ctxs[1], grad_output) + dx, grad_params_0 = self.basic_ops[0].op_backward(basic_op_ctxs[0], dx) + return dx, [grad_params_0, grad_params_1], [(), ()] + + +def _fuse_backward_scale_pair(ops, **unused): + if len(ops) == 2 and all(isinstance(op, _ScaleOp) for op in ops): + return [_BackwardScalePair(ops)] + return ops + + +@dataclasses.dataclass(slots=True) +class _ScaleKwargsFwdArgs: + """Flat inputs to the kwarg-taking test operation's forward.""" + + input_: torch.Tensor + scale: torch.Tensor + extra_scale: float + offset: Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclasses.dataclass(slots=True) +class _ScaleKwargsBwdArgs: + """Flat inputs to the kwarg-taking test operation's backward.""" + + grad_output: torch.Tensor = None + saved_input: torch.Tensor = None + scale: torch.Tensor = None + extra_scale: float = 1.0 + + +class _ScaleWithKwargsOp(BasicOperation): + """Test-only operation taking forward kwargs: a value and a tensor. + + ``offset`` is declared as tensor-or-quantized, so a quantized kwarg crosses + the op boundary as its inner buffers. Neither kwarg carries a gradient -- + that is what "read-only" means here. + """ + + fwd_args_type = _ScaleKwargsFwdArgs + bwd_args_type = _ScaleKwargsBwdArgs + + def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__() + self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype)) + + @classmethod + def forward_impl(cls, args): + offset = args.offset + if isinstance(offset, QuantizedTensor): + offset = offset.dequantize() + out = args.input_ * args.scale * args.extra_scale + offset + return out, () + + @classmethod + def forward_fake(cls, args): + x = args.input_ + return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), () + + @classmethod + def backward_impl(cls, args): + dy = args.grad_output + return ( + dy * args.scale * args.extra_scale, + (dy * args.saved_input).sum() * args.extra_scale, + ) + + @classmethod + def backward_fake(cls, args): + dy = args.grad_output + return ( + TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device), + TensorSpec(shape=(), dtype=dy.dtype, device=dy.device), + ) + + def setup_context(self, ctx, args, aux): + del aux + ctx.save_for_backward(args.input_, args.scale) + ctx.extra_scale = args.extra_scale + + def resolve_fwd_args( + self, + input_, + *, + requires_grad, + prev_op_grad_output_quantizer=None, + next_op_input_quantizer=None, + extra_scale=1.0, + offset=None, + ): + del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer + if offset is None: + offset = torch.zeros((), device=input_.device, dtype=input_.dtype) + return _ScaleKwargsFwdArgs( + input_=input_, + scale=self.scale, + extra_scale=extra_scale, + offset=offset, + ) + + def resolve_bwd_args(self, ctx, grad_output): + x, scale = ctx.saved_tensors + return _ScaleKwargsBwdArgs( + grad_output=grad_output, + saved_input=x, + scale=scale, + extra_scale=ctx.extra_scale, + ) + + +def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,)): + """Run a Sequential eagerly and compiled on identical inputs; compare both + the output and every parameter gradient. + + Each pass gets its own freshly built model, so the compiled one is traced on + a first run: nothing has built the module groups, resolved the fusions or run + ``pre_first_fuser_forward`` on it beforehand. ``make_model`` must therefore + build deterministically identical models. + + Several ``op_kwargs`` are run in order on the same pair of models, which is + what exercises Dynamo's guards on a kwarg value. + """ + eager_model = make_model() + compiled_model = make_model() + compiled = torch.compile(compiled_model, fullgraph=True) + + for op_kwargs in op_kwargs_seq: + call_kwargs = {} if op_kwargs is None else {"op_kwargs": op_kwargs} + + inp_eager = base.detach().clone().requires_grad_(True) + eager_model.zero_grad(set_to_none=True) + out_eager = eager_model(inp_eager, **call_kwargs) + out_eager.sum().backward() + ref_out = out_eager.detach().clone() + ref_igrad = inp_eager.grad.detach().clone() + ref_pgrads = [p.grad.detach().clone() for p in eager_model.parameters()] + + inp_compiled = base.detach().clone().requires_grad_(True) + compiled_model.zero_grad(set_to_none=True) + out_compiled = compiled(inp_compiled, **call_kwargs).clone() + out_compiled.sum().backward() + + torch.testing.assert_close(out_compiled, ref_out) + torch.testing.assert_close(inp_compiled.grad, ref_igrad) + for got, expected in zip(compiled_model.parameters(), ref_pgrads): + torch.testing.assert_close(got.grad, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_single_op_group_compiles(): + """``fullgraph=True`` over an ``OperationFuser`` group holding one operation. + + The pipeline-level ``autograd.Function`` is traced as a higher-order op and + calls the operation's custom ops inside, so forward and backward both end up + in the graph. + """ + torch._dynamo.reset() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_multi_op_group_uses_eager_implementations(): + """A group of several operations is gated onto the eager implementations.""" + torch._dynamo.reset() + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + with pytest.warns(UserWarning, match="several operations"): + _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp(), _ScaleOp()), base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_backward_fusion_uses_eager_implementations(): + """A backward fusion prevents the group from using basic-op custom ops.""" + te.ops.register_backward_fusion(_fuse_backward_scale_pair, prepend=True) + try: + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + with pytest.warns(UserWarning, match="backward fusion"): + _assert_sequential_matches_eager( + lambda: te.ops.Sequential(_ScaleOp(), _ScaleOp()), base + ) + finally: + OperationFuser.backward_fusion_functions.remove(_fuse_backward_scale_pair) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_te_ops_unsupported_group_still_compiles_eagerly(): + """An operation without a custom op runs its eager implementation. + + Note that this is not a fallback: under ``fullgraph=True`` there is no + leaving the graph, so the pipeline is traced either way and only the choice + of implementation changes. That is why the tracing constraints -- no + mutation of anything from an enclosing scope -- have to hold on both paths. + """ + torch._dynamo.reset() + assert te.ops.Identity().compile_unsupported_reason() is not None + + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_te_ops_forward_kwargs_compile(): + """A tensor forward kwarg reaches the operation through its custom op. + + The tensor is quantized, so it crosses the op boundary as its inner buffers, + and it changes between calls, which a graph input absorbs without a + recompilation. The last call adds a value kwarg: that one is gated onto the + eager implementation, since Dynamo turns a changed scalar into a symbol that + cannot be carried as opaque config. + """ + torch._dynamo.reset() + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cuda"), + ) + + def offset(value): + return quantizer(torch.full((64,), value, dtype=torch.bfloat16, device="cuda")) + + base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + _assert_sequential_matches_eager( + lambda: te.ops.Sequential(_ScaleWithKwargsOp()), + base, + op_kwargs_seq=( + {0: {"offset": offset(0.5)}}, + {0: {"offset": offset(1.5)}}, + # No quantized offset here: this call runs the eager implementation, + # which is traced directly, and dequantize() is not traceable. + {0: {"extra_scale": 3.0}}, + ), + ) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index e47a6f11b2..155ceb9baf 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -16,6 +16,7 @@ import sys import sysconfig from typing import Optional, Tuple +import warnings @functools.lru_cache(maxsize=None) @@ -193,6 +194,31 @@ def load_framework_extension(framework: str) -> None: sys.modules[module_name] = solib spec.loader.exec_module(solib) + # Plugin system: set NVTE_PLUGIN= to let plugin stub take over + # transformer_engine_torch and register original pybind as _nv for CUDA backend. + # Only applies to the PyTorch extension — JAX has no plugin stub. + _nvte_plugin = os.environ.get("NVTE_PLUGIN") + if _nvte_plugin and framework == "torch": + _original_module = sys.modules.get(module_name) + try: + # Register _nv alias BEFORE importing the plugin, because the + # plugin module may import transformer_engine_torch_nv at top level. + sys.modules[module_name + "_nv"] = solib + _plugin = importlib.import_module(_nvte_plugin) + _plugin.load_plugins() + except Exception as e: # pylint: disable=broad-exception-caught + # Rollback to pre-plugin state if plugin failed to fully initialize + sys.modules.pop(module_name + "_nv", None) + if _original_module is not None: + sys.modules[module_name] = _original_module + else: + sys.modules.pop(module_name, None) + warnings.warn( + f"NVTE_PLUGIN={_nvte_plugin} but plugin loading failed: {e}", + RuntimeWarning, + stacklevel=2, + ) + def sanity_checks_for_pypi_installation() -> None: """Ensure that package is installed correctly if using PyPI.""" diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh index 61218d654a..a919b0a5ee 100644 --- a/transformer_engine/common/cast/core/grouped_tma.cuh +++ b/transformer_engine/common/cast/core/grouped_tma.cuh @@ -53,51 +53,59 @@ inline bool dimensions_supported_by_TMA(const Tensor *const t) { return cols % alignment_requirement == 0; } -__device__ __forceinline__ unsigned char *align_smem_ptr_per_TMA_requirements(unsigned char *p) { - size_t addr = reinterpret_cast(p); - addr = (addr + TMA_SHMEM_ALIGNMENT - 1) & ~(TMA_SHMEM_ALIGNMENT - 1); - return reinterpret_cast(addr); -} - -// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index +// Copy the base tensor map to shared memory, modify it, and publish it to global memory. This +// function must be called convergently by every thread in a one-warp CTA. __device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, CUtensorMap *global_tensor_map, const uintptr_t global_data_ptr, const size_t global_dim_Y, const size_t global_dim_X, const size_t data_type_size_bytes) { - __shared__ CUtensorMap shared_tensor_map; - shared_tensor_map = base_tensor_map; // Copy the base tensor map into shmem + __shared__ alignas(128) CUtensorMap shared_tensor_map; constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; if constexpr (is_blackwell) { - const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; - if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); - } - if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); + const uint32_t shared_tensor_map_ptr = __cvta_generic_to_shared(&shared_tensor_map); + if (threadIdx.x == 0) { + shared_tensor_map = base_tensor_map; + + const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; + if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); + } + if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); + } + + asm volatile( + "tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;\n\t" + "tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %2;\n\t" + "tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %3;\n\t" + "tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %4;\n" + : + : "r"(shared_tensor_map_ptr), "l"(global_data_ptr), + "r"(static_cast(global_dim_Y)), "r"(static_cast(global_dim_X)), + "l"(static_cast(global_stride_bytes)) + : "memory"); } + // tensormap.cp_fenceproxy is a warp-collective instruction. Besides copying the complete + // 128-byte descriptor, its GPU-scope release makes the update visible to tensor-map proxy + // accesses that perform a matching acquire in a consumer CTA. + __syncwarp(); + const uintptr_t global_tensor_map_ptr = reinterpret_cast(global_tensor_map); asm volatile( - "{\n\t" - ".reg.b64 tensor_map_ptr; \n\t" - "mov.b64 tensor_map_ptr, %0; \n\t" - "tensormap.replace.tile.global_address.b1024.b64 [tensor_map_ptr], %1; \n\t" - "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 1, %2; \n\t" // DIM Y - "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 0, %3; \n\t" // DIM X - "tensormap.replace.tile.global_stride.b1024.b64 [tensor_map_ptr], 0, %4; \n" - "}\n" ::"l"(reinterpret_cast(&shared_tensor_map)), - "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), - "r"(static_cast(global_dim_X)), "l"(static_cast(global_stride_bytes)) + "tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned " + "[%0], [%1], 128;" + : + : "l"(global_tensor_map_ptr), "r"(shared_tensor_map_ptr) : "memory"); - *global_tensor_map = shared_tensor_map; } else { NVTE_DEVICE_ERROR("tensormap.replace is architecture-specific. "); } } template -__global__ void __launch_bounds__(1) +__global__ void __launch_bounds__(THREADS_PER_WARP) update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input, const __grid_constant__ CUtensorMap base_tensor_map_act_input, const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise, @@ -118,9 +126,11 @@ __global__ void __launch_bounds__(1) const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); const size_t offset_elts = offsets_ptr[tensor_id]; - g_tensor_maps.rows[tensor_id] = rows; - g_tensor_maps.cols[tensor_id] = cols; - g_tensor_maps.offsets[tensor_id] = offset_elts; + if (threadIdx.x == 0) { + g_tensor_maps.rows[tensor_id] = rows; + g_tensor_maps.cols[tensor_id] = cols; + g_tensor_maps.offsets[tensor_id] = offset_elts; + } // Zero-sized groups: skip TMA descriptor update. The main kernel already returns // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension @@ -162,7 +172,8 @@ __global__ void __launch_bounds__(1) __device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map)); + // The descriptor updater and consumer execute in different CTAs, so CTA scope is insufficient. + asm volatile("fence.proxy.tensormap::generic.acquire.gpu [%0], 128;" ::"l"(tensor_map)); #else NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+."); #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh index 203f569471..31feaf833d 100644 --- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh +++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh @@ -314,9 +314,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma // Dynamic smem holds the IType input tile (TMA dest, must be 128 B aligned). // warp_amaxes and tma_mbar are static smem. - extern __shared__ unsigned char smem_raw_2d_tma[]; - IType(*smem_in)[kTileDim] = reinterpret_cast( - common::align_smem_ptr_per_TMA_requirements(smem_raw_2d_tma)); + extern __shared__ char smem_raw_2d_tma[]; + IType(*smem_in)[kTileDim] = + reinterpret_cast(align_up(smem_raw_2d_tma, TMA_SHMEM_ALIGNMENT)); __shared__ CType warp_amaxes[kNumWarps]; __shared__ size_t warp_offset_partials[kNumWarps]; @@ -603,8 +603,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke // Dynamic smem: IType[kTileDim][kTileDim], 128 B aligned for TMA. Static smem // (smem_T when CW, tma_mbar) lives outside the dynamic region. - extern __shared__ unsigned char smem_raw_1d_tma[]; - unsigned char* smem_base = common::align_smem_ptr_per_TMA_requirements(smem_raw_1d_tma); + extern __shared__ char smem_raw_1d_tma[]; + char* smem_base = align_up(smem_raw_1d_tma, TMA_SHMEM_ALIGNMENT); IType(*smem)[kTileDim] = reinterpret_cast(smem_base); __shared__ uint64_t tma_mbar; diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index dad8d18d6f..141efabbaf 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -89,7 +89,6 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_ const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, const int64_t *const __restrict__ last_dims_ptr) { - const bool leading_thread = (threadIdx.x == 0); const size_t tensor_id = blockIdx.x; const size_t rows = @@ -105,7 +104,7 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_ return; } - if (leading_thread && (tensor_id < num_tensors)) { + if (tensor_id < num_tensors) { { const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id], @@ -469,7 +468,7 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, const IType *const input_dptr = reinterpret_cast(input_data.dptr); OType *const output_dptr = reinterpret_cast(output->data.dptr); - update_tma_descriptors<<>>( + update_tma_descriptors<<>>( tensor_map_input, tensor_map_output, input_dptr, output_dptr, shape_rep, num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr); diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 980a77db0a..bd8c3052f6 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -133,6 +133,13 @@ struct LaunchConfig { dim3 grid; }; +struct alignas(8) DirectVaryingFirstMapperStorage { + size_t active_elements; + size_t tensor_id; + size_t rows; + size_t tensor_start_offset; +}; + template LaunchConfig get_launch_config(const size_t first_logical_dim, const size_t last_logical_dim, const size_t elts_total, const size_t num_tensors) { @@ -659,6 +666,22 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ constexpr bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); const bool leading_thread = (threadIdx.x == 0); + // Keep mapper metadata separate from dynamic shared memory. The latter is the destination of + // asynchronous TMA loads and may be overwritten before every warp has consumed the metadata. + __shared__ DirectVaryingFirstMapperStorage direct_mapper_storage; + + if constexpr (use_direct_varying_first_mapper) { + static_assert(CastTraits::THREADS_PER_CHUNK >= MAX_SUPPORTED_TENSOR_DESCRIPTORS, + "The first CTA must have enough threads to validate every tensor."); + // The direct mapper relies on every tensor boundary being TILE_DIM_Y-aligned. The legacy + // mapper validated this while decoding each job, but the direct path no longer calls it. + // Validate all per-tensor row counts once, using the first CTA. num_tensors is bounded by + // MAX_SUPPORTED_TENSOR_DESCRIPTORS (64), which is smaller than the CTA size. + if (blockIdx.x == 0 && threadIdx.x < num_tensors) { + get_tensor_rows_num(threadIdx.x, first_logical_dim, + first_dims_ptr, num_tensors); + } + } // Decode the linear direct-mapper grid once per CTA. Valid CUDA grid extents fit in uint, which // also keeps this one-time coordinate calculation in 32-bit arithmetic. @@ -704,8 +727,8 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem); + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned IType *sIn_ptr = reinterpret_cast(dshmem); @@ -719,26 +742,26 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ if constexpr (use_direct_varying_first_mapper) { // logical_shape may describe graph-safe capacity beyond the active tensors. Resolve the - // active tail once per CTA and reject it before initializing TMA barriers. Reuse the same - // temporary storage for the exceptional colwise-swizzled tensor metadata. - size_t *const mapper_storage = reinterpret_cast(dshmem); + // active tail once per CTA and reject it before initializing TMA barriers. Cache the + // exceptional colwise-swizzled tensor metadata in dedicated static shared memory. const size_t block_offset_Y = direct_block_id_Y * CHUNK_DIM_Y; const size_t tensor_offset = block_offset_Y * last_logical_dim; if (leading_thread) { const size_t active_elements = static_cast(offsets_ptr[num_tensors]); - mapper_storage[0] = active_elements; + direct_mapper_storage.active_elements = active_elements; if constexpr (WITH_GEMM_SWIZZLED_SCALES && COLWISE_SCALING) { if (tensor_offset < active_elements) { const size_t mapped_tensor_id = find_tensor_from_offsets(offsets_ptr, num_tensors, tensor_offset); - mapper_storage[1] = mapped_tensor_id; - mapper_storage[2] = static_cast(first_dims_ptr[mapped_tensor_id]); - mapper_storage[3] = static_cast(offsets_ptr[mapped_tensor_id]); + direct_mapper_storage.tensor_id = mapped_tensor_id; + direct_mapper_storage.rows = static_cast(first_dims_ptr[mapped_tensor_id]); + direct_mapper_storage.tensor_start_offset = + static_cast(offsets_ptr[mapped_tensor_id]); } } } __syncthreads(); - if (tensor_offset >= mapper_storage[0]) { + if (tensor_offset >= direct_mapper_storage.active_elements) { return; } } @@ -839,10 +862,10 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_ if constexpr (WITH_GEMM_SWIZZLED_SCALES && COLWISE_SCALING) { // Colwise GEMM-swizzled scale indices restart at each tensor and depend on M_i. // The leading thread decoded this exceptional metadata before barrier initialization. - size_t *const mapper_storage = reinterpret_cast(dshmem); - tensor_id = mapper_storage[1]; - rows = mapper_storage[2]; - tensor_start_offset = mapper_storage[3]; + tensor_id = direct_mapper_storage.tensor_id; + rows = direct_mapper_storage.rows; + tensor_start_offset = direct_mapper_storage.tensor_start_offset; + tensor_offset_Y = block_offset_Y - tensor_start_offset / cols; } } else { block_id_Y = current_block_id / fixed_blocks_X; @@ -1289,12 +1312,13 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations use_colwise_scaling ? reinterpret_cast(output->columnwise_data.dptr) : nullptr; - update_tma_descriptors<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, input_dptr, act_input_dptr, - output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors, - first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, - last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT); + update_tma_descriptors + <<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, input_dptr, act_input_dptr, + output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors, + first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, + last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT); } TRANSFORMER_ENGINE_SWITCH_CONDITION( diff --git a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh index 24f84fa359..878fc93107 100644 --- a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh @@ -209,8 +209,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_scaled_swiglu_mxfp8_k DIVUP_TO_MULTIPLE(CHUNK_DIM_Y * sizeof(float), TMA_SHMEM_ALIGNMENT); // shmem layout: [act input][gate input][colwise output][prob] - extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem); + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); IType *sInAct_ptr = reinterpret_cast(dshmem); IType *sInGate_ptr = reinterpret_cast(dshmem + buff_size_aligned_in); diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index cdd0d4916a..58af1f7938 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -410,8 +410,8 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = common::align_smem_ptr_per_TMA_requirements(dynamic_shmem); + extern __shared__ char dynamic_shmem[]; + char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT); IType *sIn_ptr = reinterpret_cast(dshmem); fp4e2m1x2 *sOut_ptr = reinterpret_cast(dshmem + in_mem); diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 5d5ce1f6cf..128e8280bb 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -639,13 +639,18 @@ class CustomRecipe(Recipe): ---------- qfactory : Callable Factory callable that returns a quantizer instance *or* a - ``QuantizerRequest`` subclass for a given ``QuantizerRole``. + ``QuantizerRequest`` subclass for a given optional ``QuantizerRole``. The callable is invoked as:: qfactory( - role: QuantizerRole, + role: Optional[QuantizerRole], ) -> Union[Quantizer, QuantizerRequest] + Boundary slots may provide ``None`` or a role with empty fields. The + factory must return a valid object for every call. Return an + ``IdentityQuantizer`` for an intentional high-precision slot instead + of returning ``None``. + ``QuantizerRole`` is a frozen dataclass with the following fields: - ``module_type`` (str): module type (empty string when not set), e.g. @@ -663,7 +668,8 @@ class CustomRecipe(Recipe): See ``transformer_engine.pytorch.quantization.QuantizerRole`` and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` - for full documentation. + for API details. See :ref:`heterogeneous-quantization-recipes` for + construction rules and direction mapping. backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index de082c8fae..cd087fd642 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -408,6 +408,18 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: _dpa_fp8ds_reduce_amax = os.getenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") == "1" +# Plugin system: override FlashAttention and get_attention_backend if enabled +if os.environ.get("NVTE_PLUGIN"): + import transformer_engine_torch as tex + + _FlashAttentionNative = FlashAttention + FlashAttention = getattr(tex, "flash_attention", _FlashAttentionNative) + _plugin_get_attention_backend = getattr(tex, "get_attention_backend", None) + if _plugin_get_attention_backend is not None: + dpa_utils._original_get_attention_backend = dpa_utils.get_attention_backend + dpa_utils.get_attention_backend = _plugin_get_attention_backend + + __all__ = ["DotProductAttention"] diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index c74d4ddb6d..97bc70bddd 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -59,6 +59,21 @@ std::atomic g_zero_copy_enabled{false}; // is not symm-mem-backed; the backend treats it as "no window, use staged copy". constexpr NVTECommWindow kNoWindow = {nullptr, 0}; +#ifdef NCCL_HAS_SYMMEM_SUPPORT +// Offset of a symm-mem allocation relative to the start of its NCCL window. +// Newer torch places the signal pad at the front of the allocation and exposes +// get_window_offset() for it; on older torch the window starts at the buffer +// base, where get_offset() is already window-relative. +template +auto symm_mem_window_offset(T* sm, int) -> decltype(sm->get_window_offset()) { + return sm->get_window_offset(); +} +template +size_t symm_mem_window_offset(T* sm, ...) { + return sm->get_offset(); +} +#endif + // Resolve ``t`` to an NCCL symm-mem window for the zero-copy one-sided path. // Returns ``kNoWindow`` when symm-mem support isn't compiled in, zero-copy is // disabled, no group is set, or ``t`` isn't symm-mem-backed; callers pass the @@ -78,11 +93,11 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) { NVTE_CHECK(nccl_sm != nullptr, "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to " "\"NCCL\" before allocating EP payload buffers."); - // NCCL EP consumes window-relative offsets (the NCCL window starts at the signal pad, - // not at the buffer base). get_window_offset() = buffer_offset + get_offset(); add - // ``t``'s own storage offset for slice/view positioning. + // NCCL EP consumes window-relative offsets. Add ``t``'s own storage offset so a + // slice/view of a symm-mem allocation (e.g. the scale region carved from a shared + // recv buffer) resolves to its true position in the window. const uint64_t offset = - static_cast(nccl_sm->get_window_offset()) + + static_cast(symm_mem_window_offset(nccl_sm, 0)) + static_cast(t.storage_offset()) * static_cast(t.element_size()); return NVTECommWindow{static_cast(nccl_sm->get_window()), offset}; #else diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py index e42eb8f9f6..042c209c1f 100644 --- a/transformer_engine/pytorch/dynamo/__init__.py +++ b/transformer_engine/pytorch/dynamo/__init__.py @@ -6,7 +6,11 @@ from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer from .tensor_spec import TensorSpec, to_tensor_spec -from .custom_op import register_custom_op, TensorOrQuantized +from .custom_op import ( + register_custom_op, + register_custom_op_with_autograd, + TensorOrQuantized, +) __all__ = [ "register_value_opaque_quantizer", @@ -14,5 +18,6 @@ "TensorSpec", "to_tensor_spec", "register_custom_op", + "register_custom_op_with_autograd", "TensorOrQuantized", ] diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 00846d615a..8d80eebf61 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/transformer_engine/pytorch/dynamo/custom_op.py @@ -6,11 +6,16 @@ Registers TE modules' eager forward/backward as ``torch.library`` custom ops so ``torch.compile(fullgraph=True)`` traces them as single graph nodes. -``register_custom_op`` is the entry point; ``module/linear.py`` is the first user. +``register_custom_op_with_autograd`` is the entry point for a module that +wires autograd on the op itself (``module/linear.py`` is the first user); +``register_custom_op`` hands back the forward and backward ops separately, for a +caller that drives autograd at a higher level (``ops/fuser.py``). A TE forward/backward implementation takes one dataclass argument (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix tensors, quantized tensors, quantizers, process groups and plain Python values. +The autograd-free forward returns an ``(output, aux)`` tuple; the autograd-wired +API keeps its saved-tensor and context-metadata contract. A ``torch.library`` custom op is narrower: it only accepts flat schema slots (tensors plus opaque objects) and returns a flat ``Tensor[]``. @@ -41,9 +46,9 @@ only when its value is trivial (``None`` / all-``None``) at call time. What runs where. Each op registers a data-free fake (``register_fake``) so it -traces under ``torch.compile`` without allocating. ``register_custom_op`` returns -``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward -call through it: +traces under ``torch.compile`` without allocating. +``register_custom_op_with_autograd`` returns ``forward_fn`` -- the drop-in for +the eager ``autograd.Function.apply``. A forward call through it: * runs the fake ``fwd_fake_impl`` on ``TensorSpec`` descriptors (data-free; see ``tensor_spec.py``) and parses its result into an ``_OutputPlan`` -- the @@ -111,6 +116,7 @@ _TE_OP_NAMESPACE = "transformer_engine_compile" + # Annotation for an op arg field that may hold a plain tensor, a quantized # tensor subclass or a *bare* ``QuantizedTensorStorage`` (the internal-quantizer # optimization). Matched exactly by ``_TensorOrQuantizedAdapter``. @@ -180,7 +186,7 @@ def is_simple_value(cls, value: Any) -> bool: return True if isinstance(value, type): return True - if _is_opaque_value_type is not None and _is_opaque_value_type(type(value)): + if _is_opaque_constant_type is not None and _is_opaque_constant_type(type(value)): return True if isinstance(value, dict): return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items()) @@ -223,7 +229,7 @@ def _fmt_simple(cls, value: Any) -> str: if isinstance(value, tuple): body = ", ".join(cls._fmt_simple(v) for v in value) return f"({body},)" if len(value) == 1 else f"({body})" - if _is_opaque_value_type(type(value)): + if _is_opaque_constant_type(type(value)): return value.__fx_repr__()[0] # repr(float('inf')) is 'inf', which is not an evaluable literal. if isinstance(value, float) and not math.isfinite(value): @@ -285,7 +291,7 @@ def _collect(value: Any) -> None: return if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES): return - if _is_opaque_value_type(type(value)): + if _is_opaque_constant_type(type(value)): _, extra = value.__fx_repr__() globals_.update(extra) @@ -297,18 +303,18 @@ def _collect(value: Any) -> None: try: from torch._library.opaque_object import ( get_opaque_type_name, - is_opaque_value_type as _is_opaque_value_type, - register_opaque_type, + is_opaque_constant_type as _is_opaque_constant_type, + register_custom_class, ) - register_opaque_type(OpaqueValueBundle, typ="value") + register_custom_class(OpaqueValueBundle, typ="constant") _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle) # Older torch without opaque_object support. except Exception as e: # pylint: disable=broad-exception-caught # pragma: no cover record_compile_disabled( f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build" ) - _is_opaque_value_type = None + _is_opaque_constant_type = None _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None try: @@ -461,8 +467,8 @@ def _is_simple_annot(annot: Any) -> bool: return True if ( isinstance(annot, type) - and _is_opaque_value_type is not None - and _is_opaque_value_type(annot) + and _is_opaque_constant_type is not None + and _is_opaque_constant_type(annot) ): return True if get_origin(annot) in (tuple, list): @@ -824,14 +830,16 @@ def _pack_fwd_result(result: Any) -> List[torch.Tensor]: return flat -def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]: +def _pack_bwd_result( + grads: Any, num_grad_inputs: Optional[int], op_qualname: str +) -> List[torch.Tensor]: """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload. - Each grad occupies exactly one slot (validated against ``num_grad_inputs``); - a :class:`TensorSpec` grad is materialized into a single tensor. + Each grad occupies exactly one slot (validated against ``num_grad_inputs`` + when given); a :class:`TensorSpec` grad is materialized into a single tensor. """ grads = list(grads) - if len(grads) != num_grad_inputs: + if num_grad_inputs is not None and len(grads) != num_grad_inputs: raise RuntimeError( f"{op_qualname} expected bwd_impl to return {num_grad_inputs} grads " f"(one per input_tensors_for_grad entry), got {len(grads)}" @@ -930,7 +938,7 @@ def _slice_user_grads( # --------------------------------------------------------------------------- # -# Op registration +# Op registration: base and wrapper ops, autograd wiring, one op # --------------------------------------------------------------------------- # @@ -1166,7 +1174,238 @@ def _all_quantized_tensor_subclasses() -> List[type]: return found +@dataclasses.dataclass(frozen=True) +class _RegisteredOp: + """One registered custom op: its arg plan and the base / wrapper definitions.""" + + plan: _ArgPlan + base_def: Any + base_op: Any + wrapper_def: Any + wrapper_op: Any + + def __call__(self, args: Any) -> List[torch.Tensor]: + """Pack the args dataclass into slots and call the wrapper op.""" + kwargs = self.plan.pack(args) + return self.wrapper_op(*[kwargs[name] for name in self.plan.slot_names]) + + +def _register_op( + *, + name: str, + arg_type: type, + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + pack_result: Callable[[Any], List[torch.Tensor]], + flatten_in_body: bool, +) -> _RegisteredOp: + """Define one two-tier custom op: the base kernel, the wrapper op that lets + ``QuantizedTensor`` subclasses be inputs, and the passthrough registrations. + + ``flatten_in_body`` also flattens subclass inputs inside the wrapper body, + not only through the ``register_torch_dispatch`` rules. + """ + plan = _parse_arg_type(arg_type) + schema = f"{plan.schema_str} -> Tensor[]" + subclasses = _all_quantized_tensor_subclasses() + slot_offsets = plan.tensor_or_quantized_offsets() + namespace = getattr(torch.ops, _TE_OP_NAMESPACE) + + base_def = _register_base_op( + op_name=f"{name}_base", + schema_str=schema, + plan=plan, + impl=impl, + fake_impl=fake_impl, + pack_result=pack_result, + ) + base_op = getattr(namespace, f"{name}_base") + wrapper_def = _register_wrapper_op( + wrapper_op_name=name, + schema_str=schema, + base_op=base_op, + slot_offsets=slot_offsets if flatten_in_body else (), + subclasses=subclasses if flatten_in_body else (), + ) + wrapper_op = getattr(namespace, name) + + rule = _make_dispatch_rule(_make_slot_forwarder(base_op, slot_offsets, subclasses)) + for sub in subclasses: + wrapper_def.register_torch_dispatch(sub, rule) + _quantized_tensor_passthrough_ops.update((base_op.default, wrapper_op.default)) + + return _RegisteredOp( + plan=plan, + base_def=base_def, + base_op=base_op, + wrapper_def=wrapper_def, + wrapper_op=wrapper_op, + ) + + +def _register_forward_op( + *, name: str, arg_type: type, impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any] +) -> _RegisteredOp: + return _register_op( + name=name, + arg_type=arg_type, + impl=impl, + fake_impl=fake_impl, + pack_result=_pack_fwd_result, + flatten_in_body=True, + ) + + +def _register_backward_op( + *, + name: str, + arg_type: type, + impl: Callable[[Any], Any], + fake_impl: Callable[[Any], Any], + num_grad_inputs: Optional[int], +) -> _RegisteredOp: + # Pass-through body: a subclass input reaches the base op through the + # dispatch rule, never through the wrapper body. + qualname = f"{_TE_OP_NAMESPACE}::{name}_base" + return _register_op( + name=name, + arg_type=arg_type, + impl=impl, + fake_impl=fake_impl, + pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, qualname), + flatten_in_body=False, + ) + + +def _run_forward( + fwd_op: _RegisteredOp, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], fwd_args: Any +) -> Tuple[_OutputPlan, List[torch.Tensor]]: + """Run the forward op on ``fwd_args``: its output plan and flat payload.""" + spec_obj = _spec_view(fwd_args, fwd_op.plan.tensor_field_names()) + out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) + return out_plan, fwd_op(fwd_args) + + +# --------------------------------------------------------------------------- # +# Op registration: the autograd-free pair, and the autograd-wired variant +# --------------------------------------------------------------------------- # + + def register_custom_op( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: Optional[int] = None, +) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]: + """Register an op's forward and backward as two independent custom ops. + + Autograd is the caller's: it decides how the two are wired, which is what + lets a pipeline-level ``torch.autograd.Function`` -- traced by Dynamo as a + higher-order op -- group the forward and backward passes differently, as + ``ops.OperationFuser`` does. :func:`register_custom_op_with_autograd` builds + on this and wires them the usual way instead. + + Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through + without dequantization. + + Callable contracts: + + * ``fwd_impl(fwd_args) -> (output, aux)`` -- ``aux`` is a tuple of fresh tensors + * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec` + * ``bwd_impl(bwd_args) -> tuple`` of gradients (``num_grad_inputs`` of them, + if given) + * ``bwd_fake_impl`` -- its data-free twin + + Returns ``(forward_fn, backward_fn)``: + + * ``forward_fn(fwd_args) -> (output, aux)`` -- ``aux`` contains only fresh + tensors produced by the custom op. The caller decides which tensors and + metadata to persist for backward. + * ``backward_fn(bwd_args) -> tuple`` of gradients. + + Returns ``None`` if registration fails (recorded once), so callers can fall + back to eager rather than breaking import. + """ + try: + return _register_custom_op_impl( + op_name=op_name, + fwd_arg_type=fwd_arg_type, + fwd_impl=fwd_impl, + fwd_fake_impl=fwd_fake_impl, + bwd_arg_type=bwd_arg_type, + bwd_impl=bwd_impl, + bwd_fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + except (ImportError, AttributeError, RuntimeError, TypeError) as e: + record_compile_disabled( + f"could not register the autograd-free custom ops '{op_name}' ({type(e).__name__}: {e})" + ) + return None + + +def _register_custom_op_impl( + *, + op_name: str, + fwd_arg_type: type, + fwd_impl: Callable[[Any], Any], + fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + bwd_arg_type: type, + bwd_impl: Callable[[Any], Any], + bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], + num_grad_inputs: Optional[int], +) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]: + """Body of :func:`register_custom_op`; see it for semantics.""" + + def adapt_forward(impl): + def wrapped(args): + result = impl(args) + if not isinstance(result, tuple) or len(result) != 2: + raise TypeError( + "autograd-free fwd impl must return an (output, aux) tuple, got" + f" {type(result).__name__}" + ) + output, aux = result + return output, tuple(aux), None + + return wrapped + + adapted_fwd_fake_impl = adapt_forward(fwd_fake_impl) + fwd_op = _register_forward_op( + name=op_name, + arg_type=fwd_arg_type, + impl=adapt_forward(fwd_impl), + fake_impl=adapted_fwd_fake_impl, + ) + bwd_op = _register_backward_op( + name=f"{op_name}_backward", + arg_type=bwd_arg_type, + impl=bwd_impl, + fake_impl=bwd_fake_impl, + num_grad_inputs=num_grad_inputs, + ) + + def forward_fn(fwd_args): + out_plan, payload = _run_forward(fwd_op, adapted_fwd_fake_impl, fwd_args) + outputs = out_plan.user_outputs(payload) + aux = out_plan.saved_tensors(payload) + return outputs[0], tuple(aux) + + def backward_fn(bwd_args): + # Unlike the forward payload, each grad occupies exactly one slot + # (``_pack_bwd_result`` materializes a TensorSpec grad), so there is + # nothing to reassemble. + return tuple(_decode_none(t) for t in bwd_op(bwd_args)) + + return forward_fn, backward_fn + + +def register_custom_op_with_autograd( *, op_name: str, input_tensors_for_grad: List[str], @@ -1231,7 +1470,7 @@ def register_custom_op( ``torch.compile`` (a graph break) rather than breaking import. """ try: - return _register_custom_op_impl( + return _register_custom_op_with_autograd_impl( op_name=op_name, input_tensors_for_grad=input_tensors_for_grad, fwd_arg_type=fwd_arg_type, @@ -1249,7 +1488,7 @@ def register_custom_op( return None -def _register_custom_op_impl( +def _register_custom_op_with_autograd_impl( *, op_name: str, input_tensors_for_grad: List[str], @@ -1261,7 +1500,7 @@ def _register_custom_op_impl( fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], bwd_fake_impl: Callable[[Any], Tuple[Any, ...]], ) -> Callable[..., Any]: - """Body of :func:`register_custom_op`; see it for semantics.""" + """Body of :func:`register_custom_op_with_autograd`; see it for semantics.""" # Existence check at the API boundary: every ``input_tensors_for_grad`` name # must be an actual field of ``fwd_arg_type`` (differentiability -- whether # that field can carry a gradient -- is checked later, in @@ -1271,96 +1510,32 @@ def _register_custom_op_impl( if missing: raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}") - wrapper_fwd_name = op_name - wrapper_bwd_name = f"{op_name}_backward" - base_fwd_name = f"{op_name}_base" - base_bwd_name = f"{wrapper_bwd_name}_base" - subclass_list = _all_quantized_tensor_subclasses() - - fwd_plan = _parse_arg_type(fwd_arg_type) - bwd_plan = _parse_arg_type(bwd_arg_type) - - num_grad_inputs = len(input_tensors_for_grad) - grad_targets = fwd_plan.resolve_grad_targets(input_tensors_for_grad) - - fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]" - bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]" - - base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}" - - base_fwd_def = _register_base_op( - op_name=base_fwd_name, - schema_str=fwd_schema, - plan=fwd_plan, - impl=fwd_impl, - fake_impl=fwd_fake_impl, - pack_result=_pack_fwd_result, + fwd_op = _register_forward_op( + name=op_name, arg_type=fwd_arg_type, impl=fwd_impl, fake_impl=fwd_fake_impl ) - _register_base_op( - op_name=base_bwd_name, - schema_str=bwd_schema, - plan=bwd_plan, + bwd_op = _register_backward_op( + name=f"{op_name}_backward", + arg_type=bwd_arg_type, impl=bwd_impl, fake_impl=bwd_fake_impl, - pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname), - ) - - base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name) - base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name) - - fwd_slot_offsets = fwd_plan.tensor_or_quantized_offsets() - bwd_slot_offsets = bwd_plan.tensor_or_quantized_offsets() - - wrapper_fwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_fwd_name, - schema_str=fwd_schema, - base_op=base_fwd_op, - slot_offsets=fwd_slot_offsets, - subclasses=subclass_list, - ) - # Pass-through: a subclass input reaches the base op through the dispatch - # rule below, never through the wrapper body. - wrapper_bwd_def = _register_wrapper_op( - wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op + num_grad_inputs=len(input_tensors_for_grad), ) autograd_common = { - "fwd_plan": fwd_plan, - "bwd_plan": bwd_plan, - "grad_targets": grad_targets, + "fwd_plan": fwd_op.plan, + "bwd_plan": bwd_op.plan, + "grad_targets": fwd_op.plan.resolve_grad_targets(input_tensors_for_grad), "setup_context_user": setup_context, "fwd_fake_impl": fwd_fake_impl, } - wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name) - wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name) - - _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common) - _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common) - - _fwd_rule = _make_dispatch_rule( - _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list) + _register_autograd_for_op(fwd_op=fwd_op.base_def, bwd_op=bwd_op.base_op, **autograd_common) + _register_autograd_for_op( + fwd_op=fwd_op.wrapper_def, bwd_op=bwd_op.wrapper_op, **autograd_common ) - _bwd_rule = _make_dispatch_rule( - _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list) - ) - - for sub in subclass_list: - wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule) - wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule) - - _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default) - _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default) - _quantized_tensor_passthrough_ops.add(base_fwd_op.default) - _quantized_tensor_passthrough_ops.add(base_bwd_op.default) def forward_fn(fwd_args): - spec_obj = _spec_view(fwd_args, fwd_plan.tensor_field_names()) - out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj)) - kwargs = fwd_plan.pack(fwd_args) - flat_in = [kwargs[name] for name in fwd_plan.slot_names] - result = wrapper_fwd_op(*flat_in) - - outputs = out_plan.user_outputs(result) + out_plan, payload = _run_forward(fwd_op, fwd_fake_impl, fwd_args) + outputs = out_plan.user_outputs(payload) if len(outputs) == 1: return outputs[0] return tuple(outputs) diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index dc689258d1..da595b1070 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py @@ -108,15 +108,15 @@ def register_value_opaque_quantizer(cls: type) -> None: "the field in ``_rebuild_derived_state`` instead." ) cls._value_field_names = tuple(fields) - # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the + # ``register_custom_class`` requires ``__fx_repr__`` to already exist on the # class, so attach it before registering. if "__fx_repr__" not in cls.__dict__: cls.__fx_repr__ = _quantizer_fx_repr try: from torch._library.opaque_object import ( # pylint: disable=import-outside-toplevel - register_opaque_type, - is_opaque_value_type, + register_custom_class, + is_opaque_constant_type, ) except (ImportError, AttributeError) as e: # Older PyTorch without the opaque-object API: eager value semantics @@ -127,8 +127,8 @@ def register_value_opaque_quantizer(cls: type) -> None: return try: - if not is_opaque_value_type(cls): - register_opaque_type(cls, typ="value") + if not is_opaque_constant_type(cls): + register_custom_class(cls, typ="constant") except (RuntimeError, TypeError) as e: # Keep TE importable: neither the opaque-type query nor the registration # must crash the import, e.g. on PyTorch versions with only partial / diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b298b3d8ff..04fa56721d 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -8,7 +8,7 @@ import gc import warnings from math import ceil -from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, TypeVar, Union import torch from torch.utils._pytree import tree_flatten as _tree_flatten @@ -44,6 +44,13 @@ ) +class _GraphedCallableHelpers(NamedTuple): + """Lifecycle helpers owned by one graphed callable invocation.""" + + ensure_not_reset: Callable[[], None] + release_static_state: Callable[[], None] + + def set_capture_start() -> None: """Record beginning of `make_graphed_callables`.""" global _IS_GRAPH_CAPTURING @@ -633,13 +640,17 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): warmup_outputs = [] for func_idx, func in zip(warmup_func_idx, warmup_func): outputs = _run_warmup_forward(func_idx, func, func_idx) - warmup_outputs.append((func_idx, func, outputs)) - if is_training: - for func_idx, func, outputs in reversed(warmup_outputs): - _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + if is_training: + warmup_outputs.append((func_idx, func, outputs)) + else: + del outputs + while warmup_outputs: + func_idx, func, outputs = warmup_outputs.pop() + _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + del outputs else: # Follow _order exactly, mirroring the capture phase. - per_fwd_outputs = {} # per_callable_fwd_idx -> flattened outputs + per_fwd_outputs = {} # per_callable_fwd_idx -> outstanding flattened outputs fwd_idx = [0] * num_model_chunks bwd_idx = [0] * num_model_chunks for c_id in _order: @@ -653,7 +664,10 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): ) + (fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) func = callables[callable_idx] outputs = _run_warmup_forward(per_callable_fwd_idx, func, callable_idx) - per_fwd_outputs[per_callable_fwd_idx] = outputs + if is_training: + per_fwd_outputs[per_callable_fwd_idx] = outputs + else: + del outputs fwd_idx[m_chunk] += 1 elif ceil(c_id) == c_id: # Backward pass for chunk -c_id. @@ -665,10 +679,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): _prefix_num_layers[m_chunk] * num_microbatches ) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) func = callables[callable_idx] - outputs = per_fwd_outputs[per_callable_bwd_idx] + outputs = per_fwd_outputs.pop(per_callable_bwd_idx) _run_warmup_backward( per_callable_bwd_idx, func, outputs, warmup_iter, callable_idx ) + del outputs bwd_idx[m_chunk] += 1 if post_warmup_hook is not None: @@ -729,6 +744,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs) per_callable_output_unflatten_spec[per_callable_fwd_idx] = spec graph_callables[per_callable_fwd_idx] = func + del outputs, flatten_outputs fwd_idx[m_chunk] += 1 else: # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1] @@ -917,6 +933,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): per_callable_static_grad_inputs[idx] ) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx + + # The per-callable containers now own all tensors that must survive + # capture. Drop local strong references so weak-refed graph buffers can + # be returned to the shared CUDA graph pool before the next capture. + del static_outputs, static_grad_inputs, grad_inputs if ceil(c_id) == c_id: bwd_idx[m_chunk] += 1 else: @@ -1028,12 +1049,22 @@ def make_graphed_autograd_function( static_grad_inputs, returned_param_grad_clone_slots, ): + is_reset = False + + def ensure_not_reset(): + """Reject replay after this callable's graph state has been released.""" + if is_reset: + raise RuntimeError( + "This graphed callable has been reset and can no longer be used." + ) + class Graphed(torch.autograd.Function): """Autograd function for graph replay.""" @staticmethod def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *inputs): # pylint: disable=missing-function-docstring + ensure_not_reset() # Set flag for whether to update FP8 weight updates ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module() @@ -1071,6 +1102,7 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i @torch.autograd.function.once_differentiable def backward(ctx, *grads): # pylint: disable=missing-function-docstring + ensure_not_reset() # Replay backward graph if len(grads) != len(static_grad_outputs): @@ -1119,6 +1151,7 @@ def backward(ctx, *grads): return (None, None, None) + tuple(grad_inputs) def functionalized(*user_args, **user_kwargs): + ensure_not_reset() # Decide whether to update FP8 weights skip_fp8_weight_update = None @@ -1170,16 +1203,43 @@ def functionalized(*user_args, **user_kwargs): ) return _tree_unflatten(out, output_unflatten_spec) - return functionalized - - def make_graphed_attribute_functions(graph_idx): - # Get te modules for current graph + def release_static_state(): + """Release per-callable state captured by replay closures.""" + nonlocal fwd_graph, bwd_graph, is_reset + nonlocal module_params + nonlocal static_input_surface, static_outputs + nonlocal static_grad_outputs, static_grad_inputs + + is_reset = True + + # Drop the per-callable references that can own graph-pool storage. + fwd_graph = None + bwd_graph = None + module_params = () + static_input_surface = () + static_outputs = () + static_grad_outputs = () + static_grad_inputs = () + + helpers = _GraphedCallableHelpers( + ensure_not_reset=ensure_not_reset, + release_static_state=release_static_state, + ) + return functionalized, helpers + + def make_graphed_attribute_functions(graph_idx, helpers): + # Snapshot per-callable state so returned closures do not retain the outer lists. + fwd_graph = fwd_graphs[graph_idx] + bwd_graph = bwd_graphs[graph_idx] + bwd_dw_graph = bwd_dw_graphs[graph_idx] + need_bwd_dw = need_bwd_dw_graph.get(graph_idx, False) te_modules = visited_te_modules.get(graph_idx, set()) # Attach backward_dw as an attribute to the graphed callable. def backward_dw(): - if need_bwd_dw_graph.get(graph_idx, False): - bwd_dw_graphs[graph_idx].replay() + helpers.ensure_not_reset() + if need_bwd_dw: + bwd_dw_graph.replay() # Trigger the grad accumulation hook for wgrad graphs. for module in te_modules: @@ -1191,16 +1251,24 @@ def backward_dw(): # Attach reset as an attribute to the graphed callable. def reset(): - fwd_graphs[graph_idx].reset() - bwd_graphs[graph_idx].reset() - bwd_dw_graphs[graph_idx].reset() + nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules + + for graph in (fwd_graph, bwd_graph, bwd_dw_graph): + if graph is not None: + graph.reset() + + fwd_graph = None + bwd_graph = None + bwd_dw_graph = None + te_modules = () + helpers.release_static_state() return backward_dw, reset # Put together the final graphed callables ret = [] for i in range(len(sample_args)): - graphed = make_graphed_autograd_function( + graphed, helpers = make_graphed_autograd_function( fwd_graphs[i], bwd_graphs[i], per_callable_module_params[i], @@ -1218,8 +1286,17 @@ def reset(): te_modules = visited_te_modules.get(i, set()) if isinstance(func, torch.nn.Module): - def make_graphed_forward(func, graph_training_state, graphed, orig_fwd, te_modules): + def make_graphed_forward( + func, + graph_training_state, + graphed, + orig_fwd, + te_modules, + helpers, + ): def new_fwd(*user_args, **user_kwargs): + helpers.ensure_not_reset() + # If the module's training-or-eval state matches what we graphed, # run the graph, otherwise run the original forward method if func.training == graph_training_state: @@ -1264,7 +1341,14 @@ def new_fwd(*user_args, **user_kwargs): return new_fwd - forward = make_graphed_forward(func, func.training, graphed, func.forward, te_modules) + forward = make_graphed_forward( + func, + func.training, + graphed, + func.forward, + te_modules, + helpers, + ) if _order is None: func.forward = forward ret.append(func) @@ -1273,7 +1357,10 @@ def new_fwd(*user_args, **user_kwargs): else: ret.append(graphed) - backward_dw_func, reset_func = make_graphed_attribute_functions(i) + backward_dw_func, reset_func = make_graphed_attribute_functions( + i, + helpers, + ) setattr(ret[-1], "backward_dw", backward_dw_func) setattr(ret[-1], "reset", reset_func) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 55fc69ef7f..b6a25cc37f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -88,7 +88,7 @@ from ..dynamo import ( TensorSpec, TensorOrQuantized, - register_custom_op, + register_custom_op_with_autograd, is_value_opaque_quantizer, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer @@ -1788,7 +1788,7 @@ def _linear_backward_fake( # Custom op used under ``torch.compile``. -_linear_op = register_custom_op( +_linear_op = register_custom_op_with_autograd( op_name="linear", input_tensors_for_grad=["weight", "inp", "bias"], fwd_arg_type=LinearFwdArgs, @@ -1901,7 +1901,13 @@ def _linear_eager( class Linear(TransformerEngineBaseModule): """Applies a linear transformation to the incoming data :math:`y = xA^T + b` - On NVIDIA GPUs it is a drop-in replacement for ``torch.nn.Linear``. + On NVIDIA GPUs, this module implements the same linear transformation as + ``torch.nn.Linear``. + + .. note:: + + Its constructor signature differs from ``torch.nn.Linear``. Pass optional + arguments, including ``bias``, by keyword. Parameters ---------- diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 5c33c08b44..26d5261b13 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -16,7 +16,7 @@ from ...constants import DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer -from ...utils import clear_tensor_data +from ...utils import _compile_safe_warn, clear_tensor_data from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize @@ -404,14 +404,14 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + extra_input = basic_op_extra_inputs[0][0] + if self.activation_recompute_in_mlp: - raise RuntimeError( - f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " - "fused grouped MLP path." + _compile_safe_warn( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported " + "in the fused grouped MLP path." ) - extra_input = basic_op_extra_inputs[0][0] - if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") elif isinstance(input_, torch.Tensor): @@ -447,18 +447,18 @@ def fuser_backward( ]: del basic_op_grad_extra_outputs - if self.activation_recompute_in_mlp: - raise RuntimeError( - f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " - "fused grouped MLP path." - ) - ctx = basic_op_ctxs[0] x, scales = ctx.saved_tensors x = maybe_dequantize(x.contiguous(), ctx.dtype) scales = maybe_dequantize(scales, ctx.dtype) grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype) + if self.activation_recompute_in_mlp: + _compile_safe_warn( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported " + "in the fused grouped MLP path." + ) + grad_input, grad_extra_input = self._scaled_unary_backward( grad_output, x, @@ -483,7 +483,9 @@ class ScaledSReLU(_ScaledUnary): ---------- activation_recompute_in_mlp : bool, default = ``False`` Enable fused grouped MLP kernels to recompute activation outputs - during backward when supported instead of saving them. + during backward when supported instead of saving them. Outside the + fused grouped MLP path this option has no effect and a warning is + emitted. """ def _scaled_unary_forward( diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 61f80b9d9f..bd8b82d3fa 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -17,7 +17,7 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex -from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed_weight import ( @@ -27,7 +27,7 @@ finalize_weight_grads, ) from ...module.base import _2X_ACC_WGRAD -from ...quantization import Recipe +from ...quantization import Recipe, get_fp8_torch_dtype from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor @@ -170,6 +170,17 @@ def _cudnn_frontend_supports_single_group_runtime_offsets( ) and _cudnn_frontend_version_at_least("1.27.0") +def _deterministic_algorithms_required() -> bool: + """Whether bit-exact reproducibility was asked for. Same union as ``DotProductAttention``. + + Uncached: both knobs can change during the process. + """ + return ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() + ) + + def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, @@ -228,10 +239,11 @@ def _group_quantize_for_grouped_mlp( split_sizes: Optional[torch.Tensor], *, tensor_offsets: Optional[torch.Tensor] = None, + use_dense_single_group: bool = False, ) -> GroupedTensor: """Quantize into grouped storage.""" - if num_groups != 1 or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)): + if not use_dense_single_group: return tex.group_quantize( tensor, quantizer, @@ -259,6 +271,7 @@ def _group_quantize_with_amax_for_grouped_mlp( columnwise_amax: torch.Tensor, *, tensor_offsets: Optional[torch.Tensor] = None, + use_dense_single_group: bool = False, ) -> GroupedTensor: """Quantize with precomputed NVFP4 amaxes into grouped storage.""" if not isinstance(quantizer, NVFP4Quantizer): @@ -268,9 +281,10 @@ def _group_quantize_with_amax_for_grouped_mlp( num_groups, split_sizes, tensor_offsets=tensor_offsets, + use_dense_single_group=use_dense_single_group, ) - if num_groups != 1: + if not use_dense_single_group: return tex.nvfp4_group_quantize_with_amax( tensor, quantizer, @@ -679,6 +693,7 @@ def _compute_grad_params( scale_view_dtype, sf_vec_size, offsets, + use_dense_single_group, ): """Compute weight gradients and build grad_params for a GroupedLinear layer. Returns the grad_params list in parameter registration order. @@ -746,7 +761,7 @@ def _compute_grad_params( "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False." ) if ( - num_groups == 1 + use_dense_single_group and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage)) and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage)) and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer)) @@ -877,6 +892,12 @@ def fuse_grouped_mlp_ops( # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT. if recipe.nvfp4() and recipe.disable_rht: return ops + # The fused MXFP8 backward reinterprets the grad output's storage as E4M3, so an E5M2 + # backward format would have its gradients misread rather than converted. This declines + # MXFP8 with Format.HYBRID. fp8_format does not describe NVFP4 gradients, so NVFP4 is + # excluded from the check rather than relying on its value. + if recipe.mxfp8() and get_fp8_torch_dtype(recipe, fprop_tensor=False) != torch.float8_e4m3fn: + return ops if activation_op_types is None: activation_op_types = [ScaledSwiGLU, ScaledClampedQGeGLU] if _cudnn_frontend_supports_grouped_gemm_situglu(): @@ -950,6 +971,11 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: """Fused kernel for grouped GEMM, activation backward, and scale grad.""" raise NotImplementedError + @classmethod + def grouped_gemm_dactivation_is_deterministic(cls) -> bool: + """Whether this op's dactivation kernel can produce a bit-exact ``dprob``.""" + return False + @classmethod @functools.lru_cache(maxsize=None) def grouped_gemm_quant_kernel(cls) -> Callable: @@ -1103,6 +1129,9 @@ def fuser_forward( raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).") num_groups = fc1_op.num_groups + use_dense_single_group = num_groups == 1 and ( + _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) + ) fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 @@ -1173,7 +1202,7 @@ def fuser_forward( # Older cuDNN frontends do not expose this specialization, so use the # live generic offset calculation rather than caching CUDA metadata. use_offsetless_metadata = ( - num_groups == 1 + use_dense_single_group and unit_activation_scale and isinstance(fc1_input_quantizer, MXFP8Quantizer) and supports_single_group_runtime_offsets @@ -1254,6 +1283,7 @@ def fuser_forward( fc1_weight_quantizer, num_groups, None, + use_dense_single_group=use_dense_single_group, ) else: fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] @@ -1291,6 +1321,7 @@ def fuser_forward( fc2_weight_quantizer, num_groups, None, + use_dense_single_group=use_dense_single_group, ) else: fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] @@ -1338,6 +1369,7 @@ def fuser_forward( num_groups, split_sizes, tensor_offsets=fc1_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) use_nvfp4 = isinstance(fc1_input_quantizer, NVFP4Quantizer) or isinstance( @@ -1398,7 +1430,7 @@ def fuser_forward( alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(device.index).cuda_stream fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op) fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) @@ -1473,7 +1505,7 @@ def fuser_forward( fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4 if supports_single_group_runtime_offsets: - fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + fc1_activation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group if self._pass_geglu_runtime_params: fc1_activation_kwargs.update( linear_offset=self._cudnn_linear_offset, @@ -1492,7 +1524,7 @@ def fuser_forward( if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. fc1_weight_for_gemm = grouped_fc1_weight.copy() - use_single_group_weight_swizzle = num_groups == 1 + use_single_group_weight_swizzle = use_dense_single_group if use_single_group_weight_swizzle: fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm) fc1_weight_single._columnwise_data = None @@ -1529,7 +1561,7 @@ def fuser_forward( fc1_activation_kwargs["b_tensor"] = fc1_w_data fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales else: - use_single_discrete_weight = num_groups == 1 + use_single_discrete_weight = use_dense_single_group if use_single_discrete_weight: fc1_weight_single = grouped_fc1_weight[0] original_rowwise_scale = fc1_weight_single._rowwise_scale_inv @@ -1631,6 +1663,7 @@ def fuser_forward( fc1_kernel_out["amax_tensor"].view(-1), fc1_kernel_out["post_rht_amax_tensor"].view(-1), tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) else: grouped_fc2_x = _group_quantize_for_grouped_mlp( @@ -1639,11 +1672,12 @@ def fuser_forward( num_groups, split_sizes, tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) if ( - num_groups == 1 + use_dense_single_group and grouped_fc2_x.columnwise_data is not None and grouped_fc2_x.columnwise_scale_inv is not None ): @@ -1701,7 +1735,7 @@ def fuser_forward( with_gemm_swizzled_scales=True, ) - use_single_group_dense_fc2 = num_groups == 1 + use_single_group_dense_fc2 = use_dense_single_group fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) if use_single_group_dense_fc2: fc2_out = _single_group_fc2_gemm( @@ -1742,7 +1776,7 @@ def fuser_forward( } fc2_quant_kernel = self.grouped_gemm_quant_kernel() if supports_single_group_runtime_offsets: - fc2_quant_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + fc2_quant_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -1913,6 +1947,9 @@ def fuser_backward( grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) out_shape = list(grad_output.size()) num_groups = fc1_op.num_groups + use_dense_single_group = num_groups == 1 and ( + _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) + ) fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 device = fc1_weight_param.device @@ -2019,6 +2056,7 @@ def fuser_backward( num_groups, split_sizes, tensor_offsets=fc2_out_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) use_nvfp4 = ( @@ -2026,6 +2064,13 @@ def fuser_backward( or isinstance(fc1_weight_param, NVFP4Tensor) or isinstance(fc2_weight_param, NVFP4Tensor) ) + if not use_nvfp4 and fc2_grad_output_quantizer.dtype != DType.kFloat8E4M3: + # The pack below reinterprets the grad output's storage as E4M3 rather than + # converting it, so anything else would be read as the wrong format. + raise RuntimeError( + "Fused grouped MLP backward requires an E4M3 grad output, but the recipe " + f"produced {fc2_grad_output_quantizer.dtype}." + ) data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn scale_view_dtype = torch.float8_e4m3fn if use_nvfp4 else torch.float8_e8m0fnu sf_vec_size = NVFP4_BLOCK_SCALING_SIZE if use_nvfp4 else MXFP8_BLOCK_SCALING_SIZE @@ -2082,9 +2127,31 @@ def fuser_backward( # Kernel scaling factors alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) - current_stream = torch.cuda.current_stream().cuda_stream + current_stream = torch.cuda.current_stream(device.index).cuda_stream unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) + # A unit activation scale produces no dprob, so there is nothing to make deterministic. + deterministic_dactivation = ( + not unit_activation_scale and _deterministic_algorithms_required() + ) + if deterministic_dactivation: + # Two kernels write dprob and both have to be exact. The cuDNN dactivation + # epilogue produces it below; then, when scale_bias is set, it is passed to + # compute_grouped_dbias_dscales as the ``dscales`` accumulator and atomically + # added into (see triton/grouped_dbias_dscales.py). That Triton kernel is never + # deterministic, so scale_bias rules out a bit-exact dprob on its own. + dprob_is_deterministic = ( + self.grouped_gemm_dactivation_is_deterministic() and not scale_bias + ) + if not dprob_is_deterministic: + raise RuntimeError( + "Deterministic execution was requested" + " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" + " torch.use_deterministic_algorithms), but the scale gradient (dprob) is" + " accumulated with nondeterministic atomics on this configuration." + " A bit-exact dprob requires the scaled-SReLU activation," + " nvidia-cudnn-frontend 1.28.0 or later, and an FC2 without scale_bias." + ) scales_f32 = None scales_tensor = None dscales_tensor = None @@ -2139,8 +2206,11 @@ def fuser_backward( "use_dynamic_sched": True, } dactivation_kernel = self.grouped_gemm_dactivation_kernel() + if deterministic_dactivation: + # Never passed to a wrapper that would reject it -- the check above raises first. + fc2_dactivation_kwargs["deterministic"] = True if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): - fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func @@ -2194,7 +2264,7 @@ def fuser_backward( fc2_dactivation_kwargs["b_tensor"] = fc2_w_data fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: - use_single_discrete_weight = num_groups == 1 + use_single_discrete_weight = use_dense_single_group if use_single_discrete_weight: fc2_weight_single = grouped_fc2_weight[0] original_rowwise_data = fc2_weight_single._rowwise_data @@ -2299,6 +2369,7 @@ def fuser_backward( num_groups, split_sizes, tensor_offsets=fc2_x_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) else: sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") @@ -2376,6 +2447,7 @@ def fuser_backward( num_groups, split_sizes, tensor_offsets=fc1_dy_tensor_offsets, + use_dense_single_group=use_dense_single_group, ) else: grouped_fc1_dy = GroupedTensor( @@ -2412,6 +2484,7 @@ def fuser_backward( scale_view_dtype=scale_view_dtype, sf_vec_size=sf_vec_size, offsets=split_points, + use_dense_single_group=use_dense_single_group, ) # Clear FC2 input tensor if possible @@ -2437,7 +2510,7 @@ def fuser_backward( if is_distributed_weight(fc1_leader): grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) - use_single_group_dense_dgrad = num_groups == 1 + use_single_group_dense_dgrad = use_dense_single_group if use_single_group_dense_dgrad: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) _single_group_dgrad_gemm( @@ -2491,7 +2564,7 @@ def fuser_backward( } fc1_dgrad_kernel = self.grouped_gemm_quant_kernel() if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): - fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -2568,6 +2641,7 @@ def fuser_backward( scale_view_dtype=scale_view_dtype, sf_vec_size=sf_vec_size, offsets=split_points, + use_dense_single_group=use_dense_single_group, ) # Clear FC1 input tensor if possible @@ -2669,6 +2743,19 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: return grouped_gemm_dsrelu_wrapper_sm100 + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_dactivation_is_deterministic(cls) -> bool: + """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+).""" + try: + kernel = cls.grouped_gemm_dactivation_kernel() + except ImportError: + return False + try: + return "deterministic" in inspect.signature(kernel).parameters + except (TypeError, ValueError): + return False + def fuse_ops( ops: list[FusibleOperation], diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fd66529ba8..86174ca27a 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -6,13 +6,15 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Sequence +import copy import itertools from typing import Any, Optional, TypeAlias import torch -from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling +from ..quantization import FP8GlobalStateManager, Recipe, _has_delayed_scaling_state from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx +from ..utils import warn_compile_eager_fallback from .op import ( BasicOperation, FusibleOperation, @@ -48,6 +50,8 @@ def _is_graph_capturing() -> bool: OperationFusionFunction: TypeAlias = ( "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]" ) +_FusedOpList: TypeAlias = list[tuple[FusibleOperation, list[int]]] +_FusionParams: TypeAlias = tuple[type, int, Optional[str]] class _OperationFuserAutogradFunction(torch.autograd.Function): @@ -66,6 +70,7 @@ def forward( fuser: OperationFuser, basic_op_kwargs: list[dict[str, Any]], set_output_requires_grad: bool, + use_custom_ops: bool, *params_and_extra_inputs: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Forward pass @@ -82,6 +87,9 @@ def forward( Keyword arguments to BasicOperation set_output_requires_grad: bool Whether to set ``requires_grad`` flags on returned tensors + use_custom_ops: bool + Whether to call the operations' custom ops instead of tracing their + eager implementations. Decided once per group by ``OperationFuser``. *params_and_extra_inputs: torch.Tensor Other tensor inputs to include in autograd graph. Consists of parameter tensors, followed by extra operation inputs. @@ -98,9 +106,14 @@ def forward( # Operation autograd contexts basic_op_ctxs = [OperationContext() for _ in range(fuser._num_basic_ops)] - # Mark input tensors as not deletable in backward - for tensor in (input_,) + params_and_extra_inputs: - tensor._do_not_clear = True + # Mark input tensors as not deletable in backward. Skipped whenever this + # is being traced -- not merely when the custom ops are used: these + # tensors are created outside this function, and a higher-order op may + # not mutate anything from an enclosing scope. Under fullgraph there is + # no falling back out of the graph, so the constraint holds either way. + if not torch.compiler.is_compiling(): + for tensor in (input_,) + params_and_extra_inputs: + tensor._do_not_clear = True # Place user provided extra inputs into their basic-op slots. Slots bound to # internal channels are filled lazily as their producers execute. @@ -153,14 +166,24 @@ def forward( if next_op is not None: next_op_input_quantizer = next_op.get_input_quantizer() - x, fused_op_extra_outputs = op.fuser_forward( - [basic_op_ctxs[idx] for idx in basic_op_idxs], - x, - basic_op_extra_inputs=extra_inputs, - prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, - next_op_input_quantizer=next_op_input_quantizer, - basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], - ) + if use_custom_ops: + x = op.compiled_op_forward( + basic_op_ctxs[basic_op_idxs[0]], + x, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **basic_op_kwargs[basic_op_idxs[0]], + ) + fused_op_extra_outputs = [()] + else: + x, fused_op_extra_outputs = op.fuser_forward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + x, + basic_op_extra_inputs=extra_inputs, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs], + ) if len(fused_op_extra_outputs) != len(basic_op_idxs): raise RuntimeError( f"Expected {type(op).__name__} to generate extra outputs for " @@ -227,9 +250,13 @@ def forward( func_ctx.save_for_backward(*tensors_to_save) func_ctx.tensor_objects = tensor_objects - # Whether to perform recipe update in backward pass + # Whether to perform recipe update in backward pass. Skipped under + # compile: this reads and flips global FP8 state, and delayed + # scaling -- the only recipe it serves -- is gated out anyway. is_first_module = False - if fuser.first_op_requiring_backward < fuser._num_basic_ops: + if not torch.compiler.is_compiling() and ( + fuser.first_op_requiring_backward < fuser._num_basic_ops + ): is_first_module = FP8GlobalStateManager.is_first_fp8_module() # Other context @@ -244,15 +271,20 @@ def forward( func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources func_ctx.is_first_module = is_first_module - - # Mark output tensors as not deletable in backward - for tensor in itertools.chain( - (x,), - (y for ys in extra_outputs for y in ys if y is not None), - ): - tensor._do_not_clear = True - - if set_output_requires_grad: + func_ctx.use_custom_ops = use_custom_ops + + # Mark output tensors as not deletable in backward (eager only; see above) + if not torch.compiler.is_compiling(): + for tensor in itertools.chain( + (x,), + (y for ys in extra_outputs for y in ys if y is not None), + ): + tensor._do_not_clear = True + + # Autograd marks the outputs of an ``apply`` itself, so this is only + # needed on the eager path -- and AOTAutograd's functionalization drops + # a requires_grad_() applied to a graph output anyway. + if set_output_requires_grad and not torch.compiler.is_compiling(): x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops) if extra_outputs_flat: @@ -277,7 +309,12 @@ def backward( # Restore saved tensors saved_tensors = restore_from_func_ctx(func_ctx) - # Unflatten list of saved tensors + # Unflatten list of saved tensors. Under compile the contexts were + # created in the forward, which is a different subgraph, so writing to + # them here would be a side effect on an enclosing scope; copy them into + # this one instead. The copy carries the attributes the forward set. + if torch.compiler.is_compiling(): + basic_op_ctxs = [copy.copy(ctx) for ctx in basic_op_ctxs] for ctx in basic_op_ctxs: ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)] ctx._saved_tensors_range = None @@ -327,11 +364,16 @@ def backward( channel_grad if output_grad is None else output_grad + channel_grad ) grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs] - dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( - [basic_op_ctxs[idx] for idx in basic_op_idxs], - dx, - basic_op_grad_extra_outputs=grad_extra_outputs, - ) + if func_ctx.use_custom_ops: + dx, grad_params_one = op.compiled_op_backward(basic_op_ctxs[basic_op_idxs[0]], dx) + fused_op_grad_params = [grad_params_one] + fused_op_grad_extra_inputs = [()] + else: + dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward( + [basic_op_ctxs[idx] for idx in basic_op_idxs], + dx, + basic_op_grad_extra_outputs=grad_extra_outputs, + ) for idx, dparams in zip(basic_op_idxs, fused_op_grad_params): grad_params[idx] = dparams basic_op_ctxs[idx].saved_tensors = None @@ -392,6 +434,7 @@ def backward( None, # fuser None, # basic_op_kwargs None, # set_output_requires_grad + None, # use_custom_ops *grad_params_flat, *grad_extra_inputs_flat, ) @@ -535,15 +578,22 @@ def __init__( op._lock_extra_tensor_channels() # Ops for forward and backward pass, will be populated in maybe_fuse_ops - self._forward_ops: list[tuple[FusibleOperation, list[int]]] - self._backward_ops: list[tuple[FusibleOperation, list[int]]] + self._forward_ops: _FusedOpList + self._backward_ops: _FusedOpList + + # Fused operation configurations are reusable wrappers around the basic + # ops, so cache each configuration by the state that selected it. + self._fused_ops_cache: dict[_FusionParams, tuple[_FusedOpList, _FusedOpList]] = {} # Cache and detect change of state relevant for fusing operations self.recipe_type = None - self.first_op_requiring_backward = 0 self.backward_override = None self._last_amax_history_len = 0 + # Runtime backward boundary. Full activation recompute alternates this + # between the checkpointed forward and the grad-enabled recomputation. + self.first_op_requiring_backward = 0 + # Flatten list of parameters self._basic_op_params = [list(op.parameters()) for op in self._basic_ops] self._basic_op_num_params = list(map(len, self._basic_op_params)) @@ -626,36 +676,48 @@ def maybe_fuse_ops( first_op_requiring_backward = op_idx break - # Early exit if fusion parameters haven't changed - need_reset = False + # Update the runtime backward boundary on every invocation, including + # paths that reuse a cached fused operation configuration. + self.first_op_requiring_backward = first_op_requiring_backward + + # Check if recipe parameters don't match cached values. In this case, + # the recipe state in the basic ops might be invalid, so reset it. recipe_type = type(recipe) + need_to_reset_recipe_state = self.recipe_type != recipe_type + backward_override = recipe.backward_override if recipe is not None else None - fusion_params = (recipe_type, first_op_requiring_backward, backward_override) - if fusion_params != ( - self.recipe_type, - self.first_op_requiring_backward, - self.backward_override, - ): - # Recipe type, backward override, or grad requirements have changed - need_reset = True - elif ( + if backward_override != self.backward_override: + self.backward_override = backward_override + need_to_reset_recipe_state = True + + if ( recipe is not None and recipe.delayed() and self._last_amax_history_len != recipe.amax_history_len ): - # FP8 delayed scaling has changed amax history length - need_reset = True - if not need_reset: - return - - # Reset recipe state - for op in self._basic_ops: - op.reset_recipe_state(recipe=recipe) + self._last_amax_history_len = recipe.amax_history_len + need_to_reset_recipe_state = True - # Check if this is the first iteration - if self.recipe_type is None: + if need_to_reset_recipe_state: for op in self._basic_ops: - op.pre_first_fuser_forward() + op.reset_recipe_state(recipe=recipe) + + # Check if this is the first iteration + if self.recipe_type is None: + for op in self._basic_ops: + op.pre_first_fuser_forward() + + self.recipe_type = recipe_type + + # Training and inference may support different fusions. Keep the + # backward boundary in the key, but pay construction cost only once for + # each configuration. Full recompute therefore builds at most one + # no-grad plan and one grad-enabled plan for a stable recipe. + fusion_params = (recipe_type, first_op_requiring_backward, backward_override) + cached_ops = self._fused_ops_cache.get(fusion_params) + if cached_ops is not None: + self._forward_ops, self._backward_ops = cached_ops + return # Apply joint forward-backward fusions first joint_ops = OperationFuser._apply_fusions( @@ -682,14 +744,59 @@ def maybe_fuse_ops( self._basic_ops, ) - # Save current fusion params - self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params - - # Save amax history length - if isinstance(recipe, DelayedScaling): - self._last_amax_history_len = recipe.amax_history_len - else: - self._last_amax_history_len = 0 + # The FusedOperation contract excludes parameters and per-invocation + # state, so the mapped lists can be selected directly on cache hits. + self._fused_ops_cache[fusion_params] = (self._forward_ops, self._backward_ops) + + def _custom_ops_unsupported_reason( + self, basic_op_kwargs: list[dict[str, Any]] + ) -> Optional[str]: + """Why this group may not run through its operations' custom ops.""" + for mode, ops in (("forward", self._forward_ops), ("backward", self._backward_ops)): + if len(ops) != self._num_basic_ops or any( + op is not self._basic_ops[idx] or basic_op_idxs != [idx] + for idx, (op, basic_op_idxs) in enumerate(ops) + ): + return f"a {mode} fusion" + if self._num_basic_ops != 1: + return "a group of several operations" + for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True): + # A kwarg an operation declares is resolved into its args container + # like any other config. Anything else -- notably the preallocated + # buffers of the grouped operations -- is written to by the op, and a + # custom op may not mutate a tensor from an enclosing scope. + undeclared = sorted(name for name in kwargs if name not in op.fwd_kwarg_names) + if undeclared: + return f"{type(op).__name__} with undeclared keyword arguments {undeclared}" + # Only tensors. The other fields of an args container are values read + # off the module, constant across calls and baked into the graph; a + # kwarg changes per call, and on the second value Dynamo hands over a + # symbolic scalar, which cannot go into an opaque value bundle. Pass a + # 0-d tensor instead -- it is a graph input, so it does not recompile. + values = sorted(name for name, v in kwargs.items() if not isinstance(v, torch.Tensor)) + if values: + return f"{type(op).__name__} with non-tensor keyword arguments {values}" + for op in self._basic_ops: + if op.num_extra_inputs or op.num_extra_outputs: + return f"{type(op).__name__} with extra tensor inputs or outputs" + reason = op.compile_unsupported_reason() + if reason is not None: + return reason + return None + + def _use_custom_ops(self, basic_op_kwargs: list[dict[str, Any]]) -> bool: + """Whether this group runs through its operations' custom ops. + + Decided once for the whole group: a pipeline compiles as a whole, so one + unsupported operation sends all of them to eager. + """ + if not torch.compiler.is_compiling(): + return False + reason = self._custom_ops_unsupported_reason(basic_op_kwargs) + if reason is None: + return True + warn_compile_eager_fallback(reason) + return False def __call__( self, @@ -727,17 +834,29 @@ def __call__( # Initialization before forward for idx, op in enumerate(self._basic_ops): + if torch.compiler.is_compiling() and op._fp8_metas is not None: + if any( + meta is not None and _has_delayed_scaling_state(meta) + for meta in op._fp8_metas.values() + ): + raise RuntimeError( + "Delayed scaling is not supported under torch.compile in OperationFuser, " + "including CustomRecipe with DelayedScalingRequest." + ) op.pre_fuser_forward(requires_grad=idx >= self.first_op_requiring_backward) # Fuser forward pass # Note: We call forward directly when is_grad_enabled=False, # which can expose non-leaf tensors to the inner ops. Avoid # problems in this case by passing set_output_requires_grad=False. + use_custom_ops = self._use_custom_ops(basic_op_kwargs) + args = ( input, self, basic_op_kwargs, is_grad_enabled, # set_output_requires_grad + use_custom_ops, *self._flat_basic_op_params, *extra_inputs, ) diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index d057d46816..d50cd86637 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -8,8 +8,9 @@ import abc from collections.abc import Iterable, Sequence import dataclasses +import inspect import pickle -from typing import Any, Optional +from typing import Any, Callable, Optional import torch @@ -22,6 +23,7 @@ autocast, ) from ..tensor import Quantizer +from ..dynamo import is_value_opaque_quantizer, register_custom_op @dataclasses.dataclass @@ -186,6 +188,45 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): # Number of extra tensor outputs num_extra_outputs: int = 0 + # torch.compile support: an operation declares its argument containers and + # implements forward_impl / backward_impl and their fakes (see op_forward); + # its custom op is registered once per class. + fwd_args_type: Optional[type] = None + bwd_args_type: Optional[type] = None + # Forward kwargs accepted by resolve_fwd_args. Read-only, no gradient. + fwd_kwarg_names: tuple[str, ...] = () + # (forward_fn, backward_fn), or None if the operation cannot be compiled. + compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if cls.fwd_args_type is not None and cls.bwd_args_type is not None: + cls._register_compile_ops() + + @classmethod + def _register_compile_ops(cls) -> None: + for name in ("fwd_args_type", "bwd_args_type"): + if not dataclasses.is_dataclass(getattr(cls, name)): + raise TypeError(f"{cls.__name__}.{name} must be a dataclass") + params = inspect.signature(cls.resolve_fwd_args).parameters + if any(p.kind is p.VAR_KEYWORD for p in params.values()): + raise TypeError(f"{cls.__name__}.resolve_fwd_args must name its keyword arguments") + base_params = inspect.signature(BasicOperation.resolve_fwd_args).parameters + cls.fwd_kwarg_names = tuple( + name + for name, p in params.items() + if p.kind is p.KEYWORD_ONLY and name not in base_params + ) + cls.compile_ops = register_custom_op( + op_name=cls.__name__.lower(), + fwd_arg_type=cls.fwd_args_type, + fwd_impl=cls.forward_impl, + fwd_fake_impl=cls.forward_fake, + bwd_arg_type=cls.bwd_args_type, + bwd_impl=cls.backward_impl, + bwd_fake_impl=cls.backward_fake, + ) + def __init__(self) -> None: super().__init__() @@ -508,7 +549,6 @@ def _load_fp8_metas(self, fp8_metas: Optional[dict[str, Any]]) -> None: self._fp8_metas[mode][fp8_meta_key].scale.copy_(scale) self._fp8_metas[mode][fp8_meta_key].amax_history.copy_(amax_history) - @abc.abstractmethod def op_forward( self, ctx: OperationContext, @@ -520,6 +560,8 @@ def op_forward( ) -> torch.Tensor: """Forward pass + Operations with a custom op inherit this; the rest override it. + Parameters ---------- ctx: OperationContext @@ -537,8 +579,15 @@ def op_forward( Output tensor """ + return self._forward( + self.forward_impl, + ctx, + input_, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **kwargs, + ) - @abc.abstractmethod def op_backward( self, ctx: OperationContext, @@ -546,6 +595,8 @@ def op_backward( ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: """Backward pass + Operations with a custom op inherit this; the rest override it. + Parameters ---------- ctx: OperationContext @@ -561,6 +612,118 @@ def op_backward( Loss gradients w.r.t. parameters """ + return self._backward(self.backward_impl, ctx, grad_output) + + def compiled_op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + **kwargs: Any, + ) -> torch.Tensor: + """:meth:`op_forward` through this operation's custom op.""" + return self._forward( + self.compile_ops[0], + ctx, + input_, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + **kwargs, + ) + + def compiled_op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + """:meth:`op_backward` through this operation's custom op.""" + return self._backward(self.compile_ops[1], ctx, grad_output) + + def _forward( + self, + compute: Callable[[Any], tuple[torch.Tensor, tuple]], + ctx: OperationContext, + input_: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + args = self.resolve_fwd_args(input_, requires_grad=ctx.requires_grad, **kwargs) + output, aux = compute(args) + if ctx.requires_grad: + self.setup_context(ctx, args, aux) + return output + + def _backward( + self, + compute: Callable[[Any], tuple], + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + grads = compute(self.resolve_bwd_args(ctx, grad_output)) + grad_input = grad_output if grads[0] is None else grads[0] + return grad_input, tuple(grads[1:]) + + # Custom op implementation: *_impl runs eagerly and as the op body, *_fake + # on TensorSpec. Neither may read self, global state, or mutate its + # arguments. A None grad_input means "grad_output, unchanged". + + @classmethod + def forward_impl(cls, args: Any) -> tuple[torch.Tensor, tuple]: + """Forward over ``args``; returns ``(output, aux)``, aux being fresh tensors.""" + raise NotImplementedError + + @classmethod + def forward_fake(cls, args: Any) -> tuple[Any, tuple]: + """Shape-only twin of :meth:`forward_impl`.""" + raise NotImplementedError + + @classmethod + def backward_impl(cls, args: Any) -> tuple: + """Backward over ``args``; returns grad_input, then parameter grads.""" + raise NotImplementedError + + @classmethod + def backward_fake(cls, args: Any) -> tuple: + """Shape-only twin of :meth:`backward_impl`.""" + raise NotImplementedError + + def resolve_fwd_args( + self, + input_: torch.Tensor, + *, + requires_grad: bool, + prev_op_grad_output_quantizer: Optional[Quantizer] = None, + next_op_input_quantizer: Optional[Quantizer] = None, + ) -> Any: + """Build ``fwd_args_type`` from the input, module state and forward kwargs. + + Forward kwargs are declared as additional keyword-only parameters + with defaults. + """ + raise NotImplementedError + + def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any: + """Build ``bwd_args_type`` from the saved context.""" + raise NotImplementedError + + def setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None: + """Save backward state from the forward args and the fresh aux tensors.""" + del args + ctx.save_for_backward(*aux) + + def compile_unsupported_reason(self) -> Optional[str]: + """Why this operation cannot run through its custom op, or ``None``.""" + if self.compile_ops is None: + return f"{self.__class__.__name__} without a custom op" + for mode in ("forward", "backward"): + for index in range(self.num_quantizers(mode)): + quantizer = self.get_quantizer(mode, index) + if quantizer is not None and not is_value_opaque_quantizer(quantizer): + return ( + f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)" + ) + return None def fuser_forward( self, diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 8df2ec8b4b..26d0798b92 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -19,6 +19,10 @@ class HybridQuantizer(Quantizer): """Quantizer that composes rowwise and columnwise representations. + .. warning:: + **EXPERIMENTAL**: ``HybridQuantizer`` is under active development and + its API is subject to change without notice. + When both representations are requested, applies ``rowwise_quantizer`` to produce the rowwise representation and ``columnwise_quantizer`` to produce the columnwise representation. The results are wrapped in a diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 8310afc653..ec171564fe 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -26,6 +26,10 @@ class IdentityQuantizer(Quantizer): """Quantizer that produces a high-precision passthrough representation. + .. warning:: + **EXPERIMENTAL**: ``IdentityQuantizer`` is under active development and + its API is subject to change without notice. + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) holding the tensor directly, without a low-precision encoding. ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it @@ -174,6 +178,21 @@ class IdentityTensor(IdentityTensorStorage, QuantizedTensor): Presents as a standard tensor of its nominal dtype; internally it just holds data directly in that dtype, without a low-precision encoding. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Logical tensor datatype. + hp_data : torch.Tensor + Held high-precision data. + quantizer : IdentityQuantizer, optional + Quantizer that produced the tensor. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + device : torch.device, optional + Device containing the tensor. """ def __repr__(self, *, tensor_contents=None): diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index c4c6024f47..00bc627031 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -131,6 +131,8 @@ def clear_tensor_data(*tensors: Tuple[Optional[torch.Tensor], ...]) -> None: Must be used carefully. """ + if torch.compiler.is_compiling(): + return for t in tensors: if t is not None: