Skip to content

Added TorchAO FP8 weight-only and static quantization + user guides - #4589

Open
micwill755 wants to merge 1 commit into
mainfrom
torchao-workflows
Open

Added TorchAO FP8 weight-only and static quantization + user guides#4589
micwill755 wants to merge 1 commit into
mainfrom
torchao-workflows

Conversation

@micwill755

@micwill755 micwill755 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Description

Adds TorchAO FP8 support to the Dynamo path: weight-only (dequantize_affine → TensorRT IDequantizeLayer) and static FP8 (non-decomposed quantize_affine_float8 / dequantize_affine_float8). Constant folding is updated so those DQ ops stay in the graph and are not folded into dense BF16/FP16 weights.

This matches the existing ModelOpt INT8/FP8/FP4 QDQ flow, but for TorchAO tensor-subclass checkpoints. Examples cover a toy nn.Linear (WOQ and static) and FLUX.1-dev WOQ. User-guide, troubleshooting, and gallery docs are updated. HTML under docs/ is left to the post-merge docgen bot.

Dependencies: torchao (examples list it in examples/dynamo/requirements.txt). FP8 compile/runtime needs Hopper or newer (compute capability ≥ 9).

Replaces #4577 (same commits; reopened from pytorch/TensorRT branches so GitHub native Stacks work). INT4 follow-up is #4590 (2/2).

Fixes # (issue)

Type of change

  • New feature (non-breaking change which adds functionality)
  • This change requires a documentation update

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@meta-cla meta-cla Bot added the cla signed label Aug 26, 2026
@github-actions github-actions Bot added documentation Improvements or additions to documentation component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 26, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 26, 2026 08:29

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/conversion/test_cumsum_aten.py	2026-08-26 08:29:50.568560+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/conversion/test_cumsum_aten.py	2026-08-26 08:30:14.958586+00:00
@@ -31,13 +31,11 @@
                    immutable_weights=False,
                    use_dynamo_tracer=True,
                )
            return

-        self.run_test(
-            Cumsum(), inputs, immutable_weights=False, use_dynamo_tracer=True
-        )
+        self.run_test(Cumsum(), inputs, immutable_weights=False, use_dynamo_tracer=True)

    @parameterized.expand(
        [
            ((3, 1), 0),
            ((3, 1), 1),
@@ -106,14 +104,11 @@
        has_static_trip_count = (
            min_shape[positive_dim]
            == opt_shape[positive_dim]
            == max_shape[positive_dim]
        )
-        if (
-            has_static_trip_count
-            and not is_tensorrt_rtx_version_supported("1.7")
-        ):
+        if has_static_trip_count and not is_tensorrt_rtx_version_supported("1.7"):
            with self.assertRaises(UnsupportedOperatorException):
                self.run_test_with_dynamic_shape(
                    Cumsum(),
                    inputs,
                    immutable_weights=False,

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant