From 9099fe5e9f64e863de9a427a1e30a8708b655cee Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 25 Aug 2026 17:44:27 -0700 Subject: [PATCH] Added TorchAO FP8 weight-only and static quantization + user guides --- docsrc/conf.py | 2 +- docsrc/debugging/troubleshooting.rst | 8 + docsrc/user_guide/shapes_precision/index.rst | 5 +- .../shapes_precision/quantization.rst | 91 ++++++++- examples/dynamo/README.rst | 3 + examples/dynamo/requirements.txt | 1 + examples/dynamo/torchao/README.rst | 22 +++ .../torchao/quantize_linear_fp8_static.py | 119 ++++++++++++ .../dynamo/torchao/quantize_linear_fp8_woq.py | 84 +++++++++ examples/dynamo/torchao/static_fp8_utils.py | 174 ++++++++++++++++++ .../torchao/torch_export_flux_fp8_woq.py | 161 ++++++++++++++++ examples/dynamo/torchao/utils.py | 59 ++++++ .../dynamo/conversion/aten_ops_converters.py | 104 ++++++++++- .../dynamo/conversion/impl/quantize.py | 125 ++++++++++++- .../lowering/passes/constant_folding.py | 27 ++- .../py/dynamo/models/test_torchao_fp8_woq.py | 172 +++++++++++++++++ 16 files changed, 1148 insertions(+), 9 deletions(-) create mode 100644 examples/dynamo/torchao/README.rst create mode 100644 examples/dynamo/torchao/quantize_linear_fp8_static.py create mode 100644 examples/dynamo/torchao/quantize_linear_fp8_woq.py create mode 100644 examples/dynamo/torchao/static_fp8_utils.py create mode 100644 examples/dynamo/torchao/torch_export_flux_fp8_woq.py create mode 100644 examples/dynamo/torchao/utils.py create mode 100644 tests/py/dynamo/models/test_torchao_fp8_woq.py diff --git a/docsrc/conf.py b/docsrc/conf.py index bd12a18854..cca631a7f3 100644 --- a/docsrc/conf.py +++ b/docsrc/conf.py @@ -104,7 +104,7 @@ "tutorials/_rendered_examples/distributed_inference", ], # Exclude pure utility modules that aren't standalone runnable examples. - "ignore_pattern": r"(utils\.py|rotary_embedding\.py|tensor_parallel_initialize_dist\.py)", + "ignore_pattern": r"(utils\.py|static_fp8_utils\.py|rotary_embedding\.py|tensor_parallel_initialize_dist\.py)", } # Setup the breathe extension diff --git a/docsrc/debugging/troubleshooting.rst b/docsrc/debugging/troubleshooting.rst index 68ba01fa2f..e9abcb6aa8 100644 --- a/docsrc/debugging/troubleshooting.rst +++ b/docsrc/debugging/troubleshooting.rst @@ -215,6 +215,14 @@ Accuracy / Performance Issues See `NVIDIA ModelOpt documentation `_ for the full list of built-in quantization configs and customization options. + For **FP8 weight-only** quantization with TorchAO (no calibration loop), see + :ref:`quantize_linear_fp8_woq` and :ref:`torch_export_flux_fp8_woq`. Those + examples keep ``dequantize_affine`` in the graph so TensorRT can emit an + ``IDequantizeLayer`` on an FP8 weight constant. + + For **static FP8** (activations and weights, with a calibration loop), see + :ref:`quantize_linear_fp8_static`. + ---- Distributed / Tensor-Parallel Issues diff --git a/docsrc/user_guide/shapes_precision/index.rst b/docsrc/user_guide/shapes_precision/index.rst index de0883e6e8..195e40bb1d 100644 --- a/docsrc/user_guide/shapes_precision/index.rst +++ b/docsrc/user_guide/shapes_precision/index.rst @@ -2,7 +2,7 @@ Precision & Quantization ========================= Control numerical precision with FP16, BF16, and mixed-precision autocast, -and reduce model size with INT8/FP8/FP4 quantization via ModelOpt. +and reduce model size with INT8/FP8/FP4 quantization via ModelOpt or TorchAO. .. toctree:: :maxdepth: 1 @@ -12,3 +12,6 @@ and reduce model size with INT8/FP8/FP4 quantization via ModelOpt. quantization ../../tutorials/_rendered_examples/dynamo/vgg16_ptq Example: ViT FP8 Quantization <../../tutorials/_rendered_examples/dynamo/quantize_vit_fp8> + Example: TorchAO FP8 WOQ Linear <../../tutorials/_rendered_examples/dynamo/torchao/quantize_linear_fp8_woq> + Example: TorchAO Static FP8 Linear <../../tutorials/_rendered_examples/dynamo/torchao/quantize_linear_fp8_static> + Example: FLUX.1-dev FP8 WOQ <../../tutorials/_rendered_examples/dynamo/torchao/torch_export_flux_fp8_woq> diff --git a/docsrc/user_guide/shapes_precision/quantization.rst b/docsrc/user_guide/shapes_precision/quantization.rst index e703e6f372..ade862fb7c 100644 --- a/docsrc/user_guide/shapes_precision/quantization.rst +++ b/docsrc/user_guide/shapes_precision/quantization.rst @@ -5,8 +5,10 @@ Quantization (INT8 / FP8 / FP4) Torch-TensorRT supports post-training quantization (PTQ) with **INT8**, **FP8**, and **FP4** precisions via NVIDIA's -`ModelOpt `_ library. ModelOpt -inserts quantize/dequantize (QDQ) nodes into the model graph; Torch-TensorRT then +`ModelOpt `_ library, and +**FP8 weight-only** and **static FP8** quantization via +`TorchAO `_. Quantizers insert +quantize/dequantize (QDQ) nodes into the model graph; Torch-TensorRT then converts those nodes into TRT quantization layers and sets the appropriate builder flags. ---- @@ -14,11 +16,12 @@ converts those nodes into TRT quantization layers and sets the appropriate build Prerequisites ------------- -Install ModelOpt (requires ``nvidia-modelopt``): +Install ModelOpt (requires ``nvidia-modelopt``) and/or TorchAO: .. code-block:: bash pip install nvidia-modelopt + pip install torchao Hardware requirements: @@ -152,6 +155,73 @@ Quantization also works with ``torch.compile``: ---- +TorchAO FP8 Weight-Only Quantization +------------------------------------ + +TorchAO ``Float8WeightOnlyConfig`` quantizes Linear weights to FP8 (e4m3) while +leaving activations in BF16/FP16. Unlike ModelOpt PTQ, no calibration dataset is +required. + +Default TorchAO ``Float8Tensor.dequantize`` decomposes into primitive ops. The +examples promote weights to a ``Float8TensorNonDecomposed`` subclass so export +emits ``torch.ops.torchao.dequantize_affine``, which Torch-TensorRT converts to +``IDequantizeLayer``. + +.. code-block:: python + + from torchao.quantization import Float8WeightOnlyConfig, quantize_ + + quantize_(model, Float8WeightOnlyConfig()) + model = pre_process_model_for_export(model) # emit dequantize_affine + + with exclude_dq_from_constant_folding(): + exp_program = torch.export.export(model, (example_input,), strict=True) + + trt_model = torch_tensorrt.dynamo.compile( + exp_program, + inputs=[example_input], + min_block_size=1, + use_explicit_typing=True, + require_full_compilation=True, + ) + +The intended engine keeps an FP8 weight constant plus a DQ prologue into GEMM. +On **Blackwell**, Myelin can fuse that prologue into the matmul. On other GPUs +DQ + GEMM may run as two kernels — that is still correct as long as the FP8 +weight is not constant-folded into a dense high-precision weight. + +See :ref:`quantize_linear_fp8_woq` for a toy Linear model and +:ref:`torch_export_flux_fp8_woq` for FLUX.1-dev. + +---- + +TorchAO Static FP8 Quantization +-------------------------------- + +Static FP8 quantizes **activations and weights**. Activation scales are chosen +offline with min/max observers (per-tensor activations, per-channel weights), +then Linear layers are rewritten so export emits +``quantize_affine_float8_non_decomposed`` and +``dequantize_affine_float8_non_decomposed``. Torch-TensorRT maps those ops to +``IQuantizeLayer`` / ``IDequantizeLayer``. After fusion, GEMMs can run in FP8. + +.. code-block:: python + + quantize_static_fp8(model, (example_input,), calibration_steps=10) + exp_program = torch.export.export(model, (example_input,), strict=True) + trt_model = torch_tensorrt.dynamo.compile( + exp_program, + inputs=[example_input], + enabled_precisions={torch.float8_e4m3fn}, + min_block_size=1, + require_full_compilation=True, + ) + +This needs a calibration loop (unlike weight-only). See +:ref:`quantize_linear_fp8_static`. + +---- + How QDQ Nodes Are Converted ----------------------------- @@ -161,6 +231,15 @@ graph (inserted by ModelOpt), the ``IDequantizeLayer`` pairs. The TRT builder then fuses these with adjacent compute layers (e.g. Conv, Linear) to produce INT8 or FP8 kernel variants. +For TorchAO weight-only graphs, ``torch.ops.torchao.dequantize_affine.default`` is +mapped to ``IDequantizeLayer`` (weight constant stays FP8/INT, activations stay +high precision). + +For TorchAO static FP8 graphs, +``quantize_affine_float8_non_decomposed`` / ``dequantize_affine_float8_non_decomposed`` +map to ``IQuantizeLayer`` / ``IDequantizeLayer`` so both activations and weights +participate in FP8 GEMM fusion. + For FP4, ``torch.ops.tensorrt.dynamic_block_quantize_op.default`` nodes are converted via the dynamic block quantize converter, which uses TRT's ``add_dynamic_quantize`` API (TRT ≥ 10.8). @@ -228,3 +307,9 @@ Troubleshooting **FP4 "requires TRT ≥ 10.8" error** Upgrade TensorRT. FP4 uses ``add_dynamic_quantize`` which is only available in TRT 10.8 and newer. + +**TorchAO FP8 weights folded to BF16/FP16** + Export or Torch-TensorRT constant folding removed ``dequantize_affine``. + Promote weights with ``pre_process_model_for_export`` and wrap export in + ``exclude_dq_from_constant_folding`` as in :ref:`quantize_linear_fp8_woq`. + Install ``torchao`` so the converter and constant-folding exclusion register. diff --git a/examples/dynamo/README.rst b/examples/dynamo/README.rst index 4b27ad92e6..847781cb39 100644 --- a/examples/dynamo/README.rst +++ b/examples/dynamo/README.rst @@ -25,6 +25,9 @@ Model Zoo * :ref:`_torch_export_llama2`: Compiling a Llama2 model using AOT workflow (`ir=dynamo`) * :ref:`_torch_export_sam2`: Compiling SAM2 model using AOT workflow (`ir=dynamo`) * :ref:`_torch_export_flux_dev`: Compiling FLUX.1-dev model using AOT workflow (`ir=dynamo`) +* :ref:`quantize_linear_fp8_woq`: TorchAO FP8 weight-only quantization of a Linear layer (``examples/dynamo/torchao``) +* :ref:`quantize_linear_fp8_static`: TorchAO static FP8 (act + weight) quantization of a Linear layer (``examples/dynamo/torchao``) +* :ref:`torch_export_flux_fp8_woq`: Compiling FLUX.1-dev with TorchAO FP8 weight-only quantization (``examples/dynamo/torchao``) * :ref:`debugger_example`: Debugging Torch-TensorRT Compilation * :ref:`torch_export_3d_rope`: Compiling a 3D RoPE video-transformer block with complex numerics support * :ref:`engine_converter_binding_names`: Naming input / output bindings when emitting a raw serialized TRT engine \ No newline at end of file diff --git a/examples/dynamo/requirements.txt b/examples/dynamo/requirements.txt index 09f1438780..cd95e6234d 100644 --- a/examples/dynamo/requirements.txt +++ b/examples/dynamo/requirements.txt @@ -6,3 +6,4 @@ matplotlib pandas huggingface_hub opencv-python +torchao diff --git a/examples/dynamo/torchao/README.rst b/examples/dynamo/torchao/README.rst new file mode 100644 index 0000000000..cb76b403d4 --- /dev/null +++ b/examples/dynamo/torchao/README.rst @@ -0,0 +1,22 @@ +""" +TorchAO quantization +==================== + +Compile models quantized with `TorchAO `_ using +the Torch-TensorRT Dynamo backend. + +Weight-only FP8 keeps activations in BF16/FP16. Export emits +``dequantize_affine``, which Torch-TensorRT maps to TensorRT +``IDequantizeLayer`` so the engine can keep an FP8 weight constant. + +Static FP8 also quantizes activations. Calibrate observers, then export +``quantize_affine_float8_non_decomposed`` / ``dequantize_affine_float8_non_decomposed`` +so TensorRT can fuse Q/DQ into FP8 GEMMs. + +.. code-block:: bash + + pip install -r ../requirements.txt + python quantize_linear_fp8_woq.py + python quantize_linear_fp8_static.py + python torch_export_flux_fp8_woq.py +""" diff --git a/examples/dynamo/torchao/quantize_linear_fp8_static.py b/examples/dynamo/torchao/quantize_linear_fp8_static.py new file mode 100644 index 0000000000..6f1cb0d04c --- /dev/null +++ b/examples/dynamo/torchao/quantize_linear_fp8_static.py @@ -0,0 +1,119 @@ +""" +.. _quantize_linear_fp8_static: + +TorchAO Static FP8 Quantization (Linear) +======================================== + +This example calibrates a two-layer Linear model with TorchAO observers, then +rewrites those layers so **activations and weights** are quantized to FP8 +(e4m3). Export keeps explicit ``quantize_affine_float8_non_decomposed`` / +``dequantize_affine_float8_non_decomposed`` nodes, which Torch-TensorRT maps +to ``IQuantizeLayer`` / ``IDequantizeLayer``. + +Contrast with :ref:`quantize_linear_fp8_woq`, which quantizes **weights only** +and needs no calibration. Static FP8 can run the GEMM itself in FP8 (Tensor +Cores) after Q/DQ fusion. + +Graph after export (one Linear):: + + BF16 act ──► Q ──► FP8 act ──► DQ ──┐ + ▼ + aten.linear + ▲ + FP8 weight ──► DQ ──────────────────┘ + +Requirements: + +* NVIDIA GPU with FP8 support (Hopper or newer) +* ``torchao`` +* ``torch-tensorrt`` with the TorchAO float8_non_decomposed converters + +""" + +# %% +# Imports +# ^^^^^^^ +# This example lives in ``examples/dynamo/torchao/``. Move that directory off +# the front of ``sys.path`` so ``import torchao`` resolves the PyPI package +# instead of this folder. + +import sys +from pathlib import Path + +_EXAMPLE_DIR = str(Path(__file__).resolve().parent) +if sys.path and Path(sys.path[0]).resolve() == Path(_EXAMPLE_DIR): + sys.path.pop(0) + +import torch +import torch_tensorrt as torchtrt + +sys.path.insert(0, _EXAMPLE_DIR) +from static_fp8_utils import quantize_static_fp8 + + +def sqnr(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Signal-to-quantization-noise ratio in dB (higher is closer).""" + a = a.float().flatten() + b = b.float().flatten() + signal = torch.norm(a) + noise = torch.norm(a - b) + return 20 * torch.log10(signal / noise.clamp_min(1e-12)) + + +# %% +# Define a small two-layer Linear model +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + +class LinearModel(torch.nn.Module): + def __init__(self, in_features=256, hidden=512, out_features=128): + super().__init__() + self.linear1 = torch.nn.Linear(in_features, hidden, bias=False) + self.linear2 = torch.nn.Linear(hidden, out_features, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear2(self.linear1(x)) + + +model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda") +example_input = torch.randn(32, 256, dtype=torch.bfloat16, device="cuda") + +# %% +# Calibrate, then insert static FP8 Q/DQ +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# Observers record per-tensor activation min/max and per-channel weight min/max +# over a few forward passes. Those scales are baked into ``QuantizedLinearQDQ``. + +with torch.no_grad(): + fp_out = model(example_input) + +quantize_static_fp8(model, (example_input,), calibration_steps=10) + +with torch.no_grad(): + quant_out = model(example_input) +print(f"eager SQNR after static FP8: {sqnr(fp_out, quant_out):.2f} dB") + +# %% +# Export and compile +# ^^^^^^^^^^^^^^^^^^ +# The exported graph should contain +# ``quantize_affine_float8_non_decomposed`` (activations) and +# ``dequantize_affine_float8_non_decomposed`` (activations and weights). + +exp_program = torch.export.export(model, (example_input,), strict=True) +exp_program.graph_module.print_readable() + +trt_model = torchtrt.dynamo.compile( + exp_program, + inputs=[example_input], + enabled_precisions={torch.float8_e4m3fn}, + min_block_size=1, + require_full_compilation=True, +) + +with torch.no_grad(): + trt_out = trt_model(example_input) + if isinstance(trt_out, (list, tuple)): + trt_out = trt_out[0] +print(f"TRT SQNR vs quantized eager: {sqnr(quant_out, trt_out):.2f} dB") +print(trt_out) diff --git a/examples/dynamo/torchao/quantize_linear_fp8_woq.py b/examples/dynamo/torchao/quantize_linear_fp8_woq.py new file mode 100644 index 0000000000..f7f74982a6 --- /dev/null +++ b/examples/dynamo/torchao/quantize_linear_fp8_woq.py @@ -0,0 +1,84 @@ +""" +.. _quantize_linear_fp8_woq: + +TorchAO FP8 Weight-Only Quantization (Linear) +============================================= + +This example quantizes a toy ``nn.Linear`` with TorchAO +``Float8WeightOnlyConfig``, keeps an explicit ``dequantize_affine`` in the +exported graph, and compiles with the Torch-TensorRT Dynamo backend. + +The intended engine keeps an FP8 weight constant plus a DQ prologue into GEMM. +On Blackwell, Myelin can fuse that prologue into the matmul. On other GPUs the +DQ + GEMM may run as two kernels — that is still correct as long as the FP8 +weight is **not** constant-folded into a dense FP16/BF16 weight. + +Requirements: + +* NVIDIA GPU with FP8 support (Hopper or newer). Blackwell for DQ fusion. +* ``torchao`` +* ``torch-tensorrt`` with the ``torchao.dequantize_affine`` converter + +""" + +# %% +# Imports +# ^^^^^^^ +# This example lives in ``examples/dynamo/torchao/``. Move that directory off +# the front of ``sys.path`` so ``import torchao`` resolves the PyPI package +# instead of this folder. + +import sys +from pathlib import Path + +_EXAMPLE_DIR = str(Path(__file__).resolve().parent) +if sys.path and Path(sys.path[0]).resolve() == Path(_EXAMPLE_DIR): + sys.path.pop(0) + +import torch +import torch_tensorrt as torchtrt +from torchao.quantization import Float8WeightOnlyConfig, quantize_ + +sys.path.insert(0, _EXAMPLE_DIR) +from utils import exclude_dq_from_constant_folding, pre_process_model_for_export + +# %% +# Define a linear model and quantize weights to FP8 +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + +class LinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(3072, 4096) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + +model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda") +example_input = torch.randn(32, 3072, dtype=torch.bfloat16, device="cuda") + +quantize_(model, Float8WeightOnlyConfig()) +processed_model = pre_process_model_for_export(model) + +# %% +# Export and compile +# ^^^^^^^^^^^^^^^^^^ +# Wrap export in ``exclude_dq_from_constant_folding`` so inductor does not fold +# ``dequantize_affine`` before Torch-TensorRT lowering. Torch-TensorRT also +# marks that op impure in its own constant-folding pass. + +with exclude_dq_from_constant_folding(): + exp_program = torch.export.export(processed_model, (example_input,), strict=True) + +trt_model = torchtrt.dynamo.compile( + exp_program, + inputs=[example_input], + min_block_size=1, + use_explicit_typing=True, + require_full_compilation=True, +) + +output = trt_model(example_input) +print(output) diff --git a/examples/dynamo/torchao/static_fp8_utils.py b/examples/dynamo/torchao/static_fp8_utils.py new file mode 100644 index 0000000000..1be42cd407 --- /dev/null +++ b/examples/dynamo/torchao/static_fp8_utils.py @@ -0,0 +1,174 @@ +"""Helpers for TorchAO static FP8 (activation + weight) quantization examples. + +Static FP8 calibrates activation and weight ranges with observers, then rewrites +Linear layers to explicit ``quantize_affine_float8_non_decomposed`` / +``dequantize_affine_float8_non_decomposed`` so Torch-TensorRT can emit +``IQuantizeLayer`` / ``IDequantizeLayer`` pairs. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import Tensor +from torchao.core.config import AOBaseConfig +from torchao.quantization import quantize_ +from torchao.quantization.granularity import PerAxis, PerTensor +from torchao.quantization.observer import AffineQuantizedMinMaxObserver +from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter +from torchao.quantization.quant_primitives import ( + MappingType, + _dequantize_affine_float8_non_decomposed, + _quantize_affine_float8_non_decomposed, +) +from torchao.quantization.transform_module import register_quantize_module_handler + + +class ObservedLinear(torch.nn.Linear): + """Linear that records activation and weight ranges during calibration.""" + + def __init__( + self, + in_features, + out_features, + act_obs, + weight_obs, + bias=True, + device=None, + dtype=None, + ): + super().__init__(in_features, out_features, bias, device, dtype) + self.act_obs = act_obs + self.weight_obs = weight_obs + + def forward(self, input: Tensor): + observed_input = self.act_obs(input) + observed_weight = self.weight_obs(self.weight) + return F.linear(observed_input, observed_weight, self.bias) + + @classmethod + def from_float(cls, float_linear, act_obs, weight_obs): + observed = cls( + float_linear.in_features, + float_linear.out_features, + act_obs, + weight_obs, + bias=float_linear.bias is not None, + device=float_linear.weight.device, + dtype=float_linear.weight.dtype, + ) + observed.weight = float_linear.weight + observed.bias = float_linear.bias + return observed + + +class QuantizedLinearQDQ(torch.nn.Module): + """Linear with explicit FP8 Q/DQ on activations and a pre-quantized FP8 weight.""" + + def __init__( + self, + act_obs, + weight_obs, + weight: torch.Tensor, + bias: torch.Tensor | None, + target_dtype: torch.dtype, + ): + super().__init__() + assert target_dtype == torch.float8_e4m3fn + self.act_scale, _ = act_obs.calculate_qparams() + weight_scale, _ = weight_obs.calculate_qparams() + self.target_dtype = target_dtype + self.bias = bias + self.output_dtype = weight.dtype + weight_scale_2d = ( + weight_scale.view(-1, 1) if weight_scale.dim() == 1 else weight_scale + ) + self.register_buffer( + "weight_fp8", + _quantize_affine_float8_non_decomposed( + weight, weight_scale_2d, target_dtype + ), + ) + self.register_buffer("weight_scale", weight_scale_2d) + + def forward(self, input: Tensor): + input_fp8 = _quantize_affine_float8_non_decomposed( + input, self.act_scale, self.target_dtype + ) + input_hp = _dequantize_affine_float8_non_decomposed( + input_fp8, self.act_scale, self.output_dtype + ) + weight_hp = _dequantize_affine_float8_non_decomposed( + self.weight_fp8, self.weight_scale, self.output_dtype + ) + return F.linear(input_hp, weight_hp, self.bias) + + @classmethod + def from_observed(cls, observed_linear, target_dtype): + return cls( + observed_linear.act_obs, + observed_linear.weight_obs, + observed_linear.weight, + observed_linear.bias, + target_dtype, + ) + + +@dataclass +class StaticQuantConfigQDQ(AOBaseConfig): + target_dtype: torch.dtype + + +@register_quantize_module_handler(StaticQuantConfigQDQ) +def _apply_static_quant_qdq_transform(module, config): + return QuantizedLinearQDQ.from_observed(module, config.target_dtype) + + +def insert_observers_(model, act_obs, weight_obs): + def replacement_fn(m): + return ObservedLinear.from_float( + m, copy.deepcopy(act_obs), copy.deepcopy(weight_obs) + ) + + _replace_with_custom_fn_if_matches_filter( + model, + replacement_fn, + lambda m, fqn: isinstance(m, torch.nn.Linear), + ) + + +def create_fp8_observers(): + common_kwargs = dict( + mapping_type=MappingType.SYMMETRIC, + target_dtype=torch.float8_e4m3fn, + eps=torch.finfo(torch.float32).eps, + scale_dtype=torch.float32, + zero_point_dtype=torch.float32, + ) + act_obs = AffineQuantizedMinMaxObserver(granularity=PerTensor(), **common_kwargs) + weight_obs = AffineQuantizedMinMaxObserver( + granularity=PerAxis(axis=0), **common_kwargs + ) + return act_obs, weight_obs + + +def quantize_static_fp8( + model: torch.nn.Module, + example_inputs, + calibration_steps: int = 10, +) -> torch.nn.Module: + """Calibrate Linear layers and replace them with explicit FP8 Q/DQ modules.""" + act_obs, weight_obs = create_fp8_observers() + insert_observers_(model, act_obs, weight_obs) + with torch.no_grad(): + for _ in range(calibration_steps): + model(*example_inputs) + quantize_( + model, + StaticQuantConfigQDQ(torch.float8_e4m3fn), + lambda m, fqn: isinstance(m, ObservedLinear), + ) + return model diff --git a/examples/dynamo/torchao/torch_export_flux_fp8_woq.py b/examples/dynamo/torchao/torch_export_flux_fp8_woq.py new file mode 100644 index 0000000000..092a1d56d3 --- /dev/null +++ b/examples/dynamo/torchao/torch_export_flux_fp8_woq.py @@ -0,0 +1,161 @@ +""" +.. _torch_export_flux_fp8_woq: + +Compiling FLUX.1-dev with TorchAO FP8 weight-only quantization +============================================================== + +This example quantizes the ``transformer`` of +`FLUX.1-dev `_ with TorchAO +``Float8WeightOnlyConfig``, then compiles it with the Torch-TensorRT Dynamo +backend. + +Weight-only FP8 keeps activations in BF16 and stores Linear weights as FP8 plus +per-channel scales. After export, Torch-TensorRT maps ``dequantize_affine`` to +TensorRT ``IDequantizeLayer`` so the engine can keep an FP8 weight constant +instead of folding DQ into a dense BF16 GEMM. + +You need access to FLUX.1-dev on Hugging Face and a GPU with enough memory to +load and compile the 12B transformer (the unquantized Flux compile path wants +>80GB; FP8 WOQ reduces the weight footprint). + +.. code-block:: bash + + pip install torchao diffusers transformers accelerate sentencepiece protobuf + +On Blackwell, set Myelin prologue-fusion flags to encourage DQ+GEMM fusion:: + + export __LUNOWUD='-log:level=1 -log:dump=on -trace:use_id=on -mlir:prologue_fusion=1 -mlir:fusion_profit_threshold=0.01' + +""" + +# %% +# Imports +# ------- +# This example lives in ``examples/dynamo/torchao/``. Move that directory off +# the front of ``sys.path`` so ``import torchao`` resolves the PyPI package +# instead of this folder. + +import gc +import os +import sys +from pathlib import Path + +_EXAMPLE_DIR = str(Path(__file__).resolve().parent) +if sys.path and Path(sys.path[0]).resolve() == Path(_EXAMPLE_DIR): + sys.path.pop(0) + +import torch +import torch_tensorrt +from diffusers import FluxPipeline +from torchao.quantization import Float8WeightOnlyConfig, quantize_ + +sys.path.insert(0, _EXAMPLE_DIR) +from utils import exclude_dq_from_constant_folding, pre_process_model_for_export + +DEVICE = "cuda:0" +MODEL_ID = os.environ.get("FLUX_MODEL_ID", "black-forest-labs/FLUX.1-dev") + +# %% +# Load FLUX.1-dev and quantize the transformer +# -------------------------------------------- +# Only the transformer is quantized and compiled. Text encoders and the VAE stay +# BF16. ``Float8WeightOnlyConfig`` is weight-only, so no calibration dataset is +# required. + +pipe = FluxPipeline.from_pretrained( + MODEL_ID, + torch_dtype=torch.bfloat16, +) +config = pipe.transformer.config +backbone = pipe.transformer.to(DEVICE) + +quantize_(backbone, Float8WeightOnlyConfig()) +backbone = pre_process_model_for_export(backbone) + +# %% +# Export the quantized transformer +# -------------------------------- +# Dummy inputs match the Flux transformer signature used by the BF16 Flux +# example. ``exclude_dq_from_constant_folding`` keeps ``dequantize_affine`` in +# the graph during ``torch.export``. + +dummy_inputs = { + "hidden_states": torch.randn(1, 4096, 64, dtype=torch.bfloat16, device=DEVICE), + "encoder_hidden_states": torch.randn( + 1, 512, 4096, dtype=torch.bfloat16, device=DEVICE + ), + "pooled_projections": torch.randn(1, 768, dtype=torch.bfloat16, device=DEVICE), + "timestep": torch.randn(1, dtype=torch.bfloat16, device=DEVICE), + "guidance": torch.randn(1, dtype=torch.float32, device=DEVICE), + "img_ids": torch.randn(4096, 3, dtype=torch.bfloat16, device=DEVICE), + "txt_ids": torch.randn(512, 3, dtype=torch.bfloat16, device=DEVICE), + "joint_attention_kwargs": {}, + "return_dict": False, +} + +with exclude_dq_from_constant_folding(): + exp_program = torch.export.export( + backbone, + args=(), + kwargs=dummy_inputs, + strict=True, + ) + +# %% +# Compile with Torch-TensorRT +# --------------------------- +# .. note:: +# Compilation of the 12B transformer takes on the order of 20–30 minutes on +# an H100. ``offload_module_to_cpu`` frees PyTorch weights after they are +# ingested by TensorRT. + +trt_gm = torch_tensorrt.dynamo.compile( + exp_program, + inputs=dummy_inputs, + truncate_double=True, + min_block_size=1, + use_explicit_typing=True, + require_full_compilation=True, + immutable_weights=False, + offload_module_to_cpu=True, +) + +# %% +# Swap the compiled transformer into the pipeline +# ----------------------------------------------- + +pipe.transformer = None +pipe.to(DEVICE) +pipe.transformer = trt_gm +pipe.transformer.config = config +trt_gm.device = torch.device("cuda") +del exp_program, backbone +gc.collect() +torch.cuda.empty_cache() + +# %% +# Generate an image +# ----------------- + + +def generate_image(pipe, prompt, image_name): + seed = 42 + with torch.no_grad(): + image = pipe( + prompt, + output_type="pil", + num_inference_steps=20, + generator=torch.Generator("cuda").manual_seed(seed), + ).images[0] + image.save(f"{image_name}.png") + print(f"Image generated using {image_name} model saved as {image_name}.png") + + +generate_image( + pipe, + [ + "Baroque style, a lavish palace interior with ornate gilded ceilings, " + "intricate tapestries, and dramatic lighting over a grand staircase." + ], + "flux_fp8_woq", +) diff --git a/examples/dynamo/torchao/utils.py b/examples/dynamo/torchao/utils.py new file mode 100644 index 0000000000..6d99e1b25e --- /dev/null +++ b/examples/dynamo/torchao/utils.py @@ -0,0 +1,59 @@ +"""Helpers for TorchAO FP8 weight-only quantization examples. + +TorchAO's default ``Float8Tensor.dequantize`` decomposes into primitive ops. +``Float8TensorNonDecomposed`` keeps an explicit ``dequantize_affine`` in the +exported graph so Torch-TensorRT can map it to ``IDequantizeLayer``. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Iterator + +import torch +from torchao.quantization import dequantize_affine +from torchao.quantization.quantize_.workflows import Float8Tensor + + +class Float8TensorNonDecomposed(Float8Tensor): + """``Float8Tensor`` that dequantizes via explicit ``dequantize_affine``.""" + + def dequantize(self, output_dtype=None): + if output_dtype is None: + output_dtype = torch.bfloat16 + return dequantize_affine( + self.qdata, + self.block_size, + self.scale, + None, + self.qdata.dtype, + output_dtype=output_dtype, + ) + + +def pre_process_model_for_export(model: torch.nn.Module) -> torch.nn.Module: + """Promote ``Float8Tensor`` parameters so export emits ``dequantize_affine``.""" + for param in model.parameters(): + if isinstance(param, Float8Tensor) and not isinstance( + param, Float8TensorNonDecomposed + ): + param.__class__ = Float8TensorNonDecomposed + param.requires_grad_(False) + return model + + +@contextmanager +def exclude_dq_from_constant_folding() -> Iterator[None]: + """Keep ``dequantize_affine`` out of inductor constant folding during export.""" + from torch._inductor.constant_folding import ( + _dont_constant_fold, + add_dont_constant_fold, + ) + + op = torch.ops.torchao.dequantize_affine.default + add_dont_constant_fold(op) + try: + yield + finally: + if op in _dont_constant_fold: + _dont_constant_fold.remove(op) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index b23205398a..598780d2d2 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1030,6 +1030,108 @@ def aten_ops_dynamic_block_quantize_op( ) +try: + import torchao # noqa: F401 + + assert torch.ops.torchao.dequantize_affine.default +except Exception: + _LOGGER.debug( + "torchao not available; skipping torchao.dequantize_affine converter. " + "Install torchao to compile TorchAO weight-only quantized models." + ) +else: + + @dynamo_tensorrt_converter( + torch.ops.torchao.dequantize_affine.default, + supports_dynamic_shapes=False, + ) + def aten_ops_torchao_dequantize_affine( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, + ) -> Union[TRTTensor, Sequence[TRTTensor]]: + # dequantize_affine(input, block_size, scale, zero_point, input_dtype, + # quant_min=None, quant_max=None, output_dtype=...) + output_dtype = kwargs.get("output_dtype") + if output_dtype is None: + output_dtype = args[7] if len(args) > 7 else torch.float16 + input_dtype = args[4] if len(args) > 4 else None + return impl.quantize.dequantize_affine( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args[1], + args[2], + output_dtype, + input_dtype=input_dtype, + ) + + +try: + from torchao.quantization.quant_primitives import ( # noqa: F401 + _dequantize_affine_float8_non_decomposed, + _quantize_affine_float8_non_decomposed, + ) + + assert torch.ops.torchao.quantize_affine_float8_non_decomposed.default + assert torch.ops.torchao.dequantize_affine_float8_non_decomposed.default +except Exception: + _LOGGER.debug( + "torchao float8_non_decomposed ops not available; skipping static FP8 " + "Q/DQ converters. Import torchao quantization primitives to compile " + "TorchAO static FP8 graphs." + ) +else: + + @dynamo_tensorrt_converter( + torch.ops.torchao.quantize_affine_float8_non_decomposed.default, + supports_dynamic_shapes=True, + ) + def aten_ops_torchao_quantize_affine_float8( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, + ) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.quantize.quantize_affine_float8( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args[1], + ) + + @dynamo_tensorrt_converter( + torch.ops.torchao.dequantize_affine_float8_non_decomposed.default, + supports_dynamic_shapes=True, + ) + def aten_ops_torchao_dequantize_affine_float8( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, + ) -> Union[TRTTensor, Sequence[TRTTensor]]: + output_dtype = ( + args[2] if len(args) > 2 else kwargs.get("output_dtype", torch.bfloat16) + ) + return impl.quantize.dequantize_affine_float8( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args[1], + output_dtype, + ) + + @dynamo_tensorrt_converter(torch.ops.aten.squeeze.dim, supports_dynamic_shapes=True) @dynamo_tensorrt_converter(torch.ops.aten.squeeze.dims, supports_dynamic_shapes=True) def aten_ops_squeeze( @@ -1227,7 +1329,7 @@ def _index_copy_kv_eligible( if len(node.args) < 4: return False if input_node is None: - input_node = node.args[0] # type: ignore[assignment] + input_node = node.args[0] dim, _index_node, src_node = node.args[1:4] if not isinstance(input_node, Node) or input_node.op != "placeholder": diff --git a/py/torch_tensorrt/dynamo/conversion/impl/quantize.py b/py/torch_tensorrt/dynamo/conversion/impl/quantize.py index 8ec0895832..0258f0a099 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/quantize.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/quantize.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Optional, Sequence, Union import numpy as np import tensorrt as trt @@ -6,6 +6,7 @@ from tensorrt import ITensor as TRTTensor from torch.fx.experimental.proxy_tensor import unset_fake_temporarily from torch.fx.node import Target +from torch_tensorrt import _enums from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext @@ -138,3 +139,125 @@ def quantize( dq_output = dequantize_layer.get_output(0) return dq_output + + +def _block_size_as_ints(block_size: Sequence[object]) -> list[int]: + dims: list[int] = [] + for bs in block_size: + if isinstance(bs, torch.Tensor): + dims.append(int(bs.item())) + elif isinstance(bs, (int, float)): + dims.append(int(bs)) + else: + raise TypeError(f"Unsupported block_size dim type {type(bs)}: {bs!r}") + return dims + + +def dequantize_affine( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + qdata: Union[torch.Tensor, TRTTensor], + block_size: Sequence[object], + scale: Union[np.ndarray, torch.Tensor, TRTTensor], + output_dtype: torch.dtype, + input_dtype: Optional[torch.dtype] = None, +) -> TRTTensor: + """Map ``torchao.dequantize_affine`` to TensorRT ``IDequantizeLayer``. + + Used by TorchAO weight-only quantization (FP8/INT8/INT4). The quantized + weight stays a low-precision constant and is dequantized at the GEMM + boundary so Myelin can fuse DQ into the matmul prologue instead of + folding it into a dense high-precision weight. + """ + qdata_trt = get_trt_tensor(ctx, qdata, f"{name}_qdata", dtype=input_dtype) + + axis = None + if isinstance(scale, torch.Tensor): + scale_for_trt = scale.squeeze() + if scale_for_trt.numel() != 1: + # Per-channel axis is the dimension whose block size is 1 + # (quantized independently per slice). Example: weight (3072, 64) + # with block_size [3072, 1] → axis 1. + bs = _block_size_as_ints(block_size) + try: + axis = next(i for i, dim in enumerate(bs) if dim == 1) + except StopIteration as exc: + raise ValueError( + f"Unable to derive IDequantizeLayer axis from block_size={bs} " + f"and scale shape {tuple(scale.shape)}" + ) from exc + else: + scale_for_trt = scale + + scale_trt = get_trt_tensor(ctx, scale_for_trt, f"{name}_scale", dtype=torch.float32) + trt_output_dtype = _enums.dtype._from(output_dtype).to(trt.DataType) + + dequantize_layer = ctx.net.add_dequantize( + qdata_trt, + scale_trt, + output_type=trt_output_dtype, + ) + if axis is not None: + dequantize_layer.axis = axis + set_layer_name(dequantize_layer, target, f"{name}_dequantize", source_ir) + return dequantize_layer.get_output(0) + + +def _fp8_scale_and_axis( + ctx: ConversionContext, + scale: Union[np.ndarray, torch.Tensor, TRTTensor], + name: str, +) -> tuple[TRTTensor, Optional[int]]: + axis = None + if isinstance(scale, torch.Tensor): + scale_for_trt = scale.squeeze() + if scale_for_trt.numel() != 1: + axis = 0 + else: + scale_for_trt = scale + scale_trt = get_trt_tensor(ctx, scale_for_trt, f"{name}_scale", dtype=torch.float32) + return scale_trt, axis + + +def quantize_affine_float8( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_tensor: Union[torch.Tensor, TRTTensor], + scale: Union[np.ndarray, torch.Tensor, TRTTensor], +) -> TRTTensor: + """Map TorchAO ``quantize_affine_float8_non_decomposed`` to ``IQuantizeLayer``.""" + input_trt = get_trt_tensor(ctx, input_tensor, f"{name}_input") + scale_trt, axis = _fp8_scale_and_axis(ctx, scale, name) + quantize_layer = ctx.net.add_quantize(input_trt, scale_trt, trt.DataType.FP8) + if axis is not None: + quantize_layer.axis = axis + set_layer_name(quantize_layer, target, f"{name}_quantize", source_ir) + return quantize_layer.get_output(0) + + +def dequantize_affine_float8( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_tensor: Union[torch.Tensor, TRTTensor], + scale: Union[np.ndarray, torch.Tensor, TRTTensor], + output_dtype: torch.dtype = torch.bfloat16, +) -> TRTTensor: + """Map TorchAO ``dequantize_affine_float8_non_decomposed`` to ``IDequantizeLayer``.""" + input_trt = get_trt_tensor(ctx, input_tensor, f"{name}_input") + scale_trt, axis = _fp8_scale_and_axis(ctx, scale, name) + trt_output_dtype = _enums.dtype._from(output_dtype).to(trt.DataType) + dequantize_layer = ctx.net.add_dequantize( + input_trt, + scale_trt, + output_type=trt_output_dtype, + ) + if axis is not None: + dequantize_layer.axis = axis + set_layer_name(dequantize_layer, target, f"{name}_dequantize", source_ir) + return dequantize_layer.get_output(0) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index a88aed4a9c..f7d999144c 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -107,8 +107,8 @@ def replace_node_with_constant( class _TorchTensorRTConstantFolder(ConstantFolder): # type: ignore[misc] def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - # Set of known quantization ops to be excluded from constant folding. - # Currently, we exclude all quantization ops coming from modelopt library. + # Quantization ops excluded from constant folding so TRT sees QDQ + # (ModelOpt tensorrt.quantize_op and TorchAO dequantize_affine). self.quantization_ops: Set[torch._ops.OpOverload] = set() try: # modelopt import ensures torch.ops.tensorrt.quantize_op.default is registered @@ -123,6 +123,29 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: except Exception as e: pass + try: + import torchao # noqa: F401 + + assert torch.ops.torchao.dequantize_affine.default + self.quantization_ops.add(torch.ops.torchao.dequantize_affine.default) + except Exception: + pass + + try: + from torchao.quantization.quant_primitives import ( # noqa: F401 + _dequantize_affine_float8_non_decomposed, + _quantize_affine_float8_non_decomposed, + ) + + self.quantization_ops.add( + torch.ops.torchao.quantize_affine_float8_non_decomposed.default + ) + self.quantization_ops.add( + torch.ops.torchao.dequantize_affine_float8_non_decomposed.default + ) + except Exception: + pass + # TODO: Update this function when quantization is added def is_impure(self, node: torch.fx.node.Node) -> bool: diff --git a/tests/py/dynamo/models/test_torchao_fp8_woq.py b/tests/py/dynamo/models/test_torchao_fp8_woq.py new file mode 100644 index 0000000000..989bfc98cc --- /dev/null +++ b/tests/py/dynamo/models/test_torchao_fp8_woq.py @@ -0,0 +1,172 @@ +# type: ignore +import importlib +import unittest + +import pytest +import torch +import torch_tensorrt as torchtrt +from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity + +assertions = unittest.TestCase() + + +def _has_fp8_gpu() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major >= 9 + + +def _promote_float8_woq(model: torch.nn.Module) -> torch.nn.Module: + from torchao.quantization import dequantize_affine + from torchao.quantization.quantize_.workflows import Float8Tensor + + class Float8TensorNonDecomposed(Float8Tensor): + def dequantize(self, output_dtype=None): + if output_dtype is None: + output_dtype = torch.bfloat16 + return dequantize_affine( + self.qdata, + self.block_size, + self.scale, + None, + self.qdata.dtype, + output_dtype=output_dtype, + ) + + for param in model.parameters(): + if isinstance(param, Float8Tensor): + param.__class__ = Float8TensorNonDecomposed + param.requires_grad_(False) + return model + + +@pytest.mark.unit +@unittest.skipIf(importlib.util.find_spec("torchao") is None, "torchao not installed") +@unittest.skipIf(not _has_fp8_gpu(), "FP8 GPU (compute capability >= 9.0) is required") +def test_linear_fp8_woq(): + from torch._inductor.constant_folding import ( + _dont_constant_fold, + add_dont_constant_fold, + ) + from torchao.quantization import Float8WeightOnlyConfig, quantize_ + + class LinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(64, 128) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda") + example_input = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda") + quantize_(model, Float8WeightOnlyConfig()) + model = _promote_float8_woq(model) + + op = torch.ops.torchao.dequantize_affine.default + add_dont_constant_fold(op) + try: + exp_program = torch.export.export(model, (example_input,), strict=True) + finally: + if op in _dont_constant_fold: + _dont_constant_fold.remove(op) + + dq_nodes = [ + n + for n in exp_program.graph.nodes + if n.target == torch.ops.torchao.dequantize_affine.default + ] + assertions.assertTrue( + len(dq_nodes) >= 1, + msg="Exported graph is missing torchao.dequantize_affine", + ) + + trt_mod = torchtrt.dynamo.compile( + exp_program, + inputs=[example_input], + min_block_size=1, + use_explicit_typing=True, + require_full_compilation=True, + cache_built_engines=False, + reuse_cached_engines=False, + ) + + with torch.no_grad(): + eager_out = model(example_input) + trt_out = trt_mod(example_input) + if isinstance(trt_out, (list, tuple)): + trt_out = trt_out[0] + cos_sim = cosine_similarity(eager_out, trt_out) + assertions.assertTrue( + cos_sim > COSINE_THRESHOLD, + msg=( + "TorchAO FP8 WOQ TRT outputs don't match eager. " + f"Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}" + ), + ) + torch._dynamo.reset() + + +@pytest.mark.unit +@unittest.skipIf(importlib.util.find_spec("torchao") is None, "torchao not installed") +@unittest.skipIf(not _has_fp8_gpu(), "FP8 GPU (compute capability >= 9.0) is required") +def test_linear_fp8_static(): + import sys + from pathlib import Path + + example_dir = ( + Path(__file__).resolve().parents[4] / "examples" / "dynamo" / "torchao" + ) + sys.path.insert(0, str(example_dir)) + from static_fp8_utils import quantize_static_fp8 + + class LinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = torch.nn.Linear(32, 64, bias=False) + self.linear2 = torch.nn.Linear(64, 16, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear2(self.linear1(x)) + + model = LinearModel().eval().to(dtype=torch.bfloat16, device="cuda") + example_input = torch.randn(4, 32, dtype=torch.bfloat16, device="cuda") + quantize_static_fp8(model, (example_input,), calibration_steps=3) + + exp_program = torch.export.export(model, (example_input,), strict=True) + graph_str = str(exp_program.graph) + assertions.assertIn( + "quantize_affine_float8_non_decomposed", + graph_str, + "Exported graph is missing static FP8 quantize", + ) + assertions.assertIn( + "dequantize_affine_float8_non_decomposed", + graph_str, + "Exported graph is missing static FP8 dequantize", + ) + + trt_mod = torchtrt.dynamo.compile( + exp_program, + inputs=[example_input], + enabled_precisions={torch.float8_e4m3fn}, + min_block_size=1, + require_full_compilation=True, + cache_built_engines=False, + reuse_cached_engines=False, + ) + with torch.no_grad(): + eager_out = model(example_input) + trt_out = trt_mod(example_input) + if isinstance(trt_out, (list, tuple)): + trt_out = trt_out[0] + cos_sim = cosine_similarity(eager_out, trt_out) + assertions.assertTrue( + cos_sim > COSINE_THRESHOLD, + msg=( + "TorchAO static FP8 TRT outputs don't match eager. " + f"Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}" + ), + ) + torch._dynamo.reset()