From 9eb2e6f1816d66b93553d70dda57bb3baf973da6 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 17 Apr 2026 20:13:23 +0000 Subject: [PATCH 1/4] Add FP8 MHA quantization support for HuggingFace ViT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enables TensorRT attention-v2 fusion for HuggingFace ViT (and similar transformer vision models) when exported to ONNX with FP8 Q/DQ. - fp8_exporter: rewrite attention-scaling Mul and K Transpose to the Q-side so DQ feeds MatMul directly, pre-transpose weight constants, insert FP8 Q/DQ on Softmax outputs for MHA-v2 fusion. Scale dtype now matches the graph's float dtype to keep strongly-typed builds consistent. - onnx/utils: fold Cast(FP16<->FP32) nodes that convert_float_to_float16 inserts around Q/DQ by rewriting scale initializers to FP16, so TRT fuses DQ into the downstream GEMM/MatMul kernel. - torch/quantization/export_onnx: keep FP8 Q/DQ scale in the native input dtype so no Cast is injected between graph and Q/DQ. - torch/quantization/nn: register nn.LayerNorm in QuantModuleRegistry so LayerNorm output quantizers are honored. - torch/quantization/plugins/huggingface: skip attention wrappers whose children are also "*Attention" to avoid double-patching eager_attention_forward (e.g. ViTAttention vs ViTSelfAttention). Example: examples/torch_onnx/vit_mha_quantization.py shows a ViT-FP8 config (extends FP8_DEFAULT_CFG with LayerNorm output quantizer, disabled input quantizers on LayerNorm-followed layers, and *_bmm_quantizer entries) plus accuracy + TRT-latency comparison against an FP16 baseline. Measured on ViT-base-patch16-224 (RTX 6000 Ada, batch=1): - Top-1 / top-5 on 5k ImageNet-val: 81.16% / 95.50% (FP16) vs 80.96% / 95.44% (torch FP8) — -0.20% / -0.06% - TRT latency: 0.721 ms (FP16) vs 0.646 ms (torch FP8) — 1.12x speedup Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 4 + examples/torch_onnx/torch_quant_to_onnx.py | 24 +- modelopt/onnx/export/fp8_exporter.py | 289 +++++++++++++++++- modelopt/onnx/utils.py | 78 ++++- modelopt/torch/_deploy/utils/torch_onnx.py | 6 + modelopt/torch/quantization/export_onnx.py | 45 +-- modelopt/torch/quantization/nn/__init__.py | 1 + .../nn/modules/quant_layernorm.py | 28 ++ .../torch/quantization/plugins/huggingface.py | 19 ++ .../quantization/test_fp8_mha_exporter.py | 118 +++++++ tests/unit/onnx/test_fold_casts.py | 99 ++++++ .../plugins/test_nested_attention_skip.py | 44 +++ .../quantization/test_quant_layernorm.py | 31 ++ 13 files changed, 744 insertions(+), 42 deletions(-) create mode 100644 modelopt/torch/quantization/nn/modules/quant_layernorm.py create mode 100644 tests/unit/onnx/quantization/test_fp8_mha_exporter.py create mode 100644 tests/unit/onnx/test_fold_casts.py create mode 100644 tests/unit/torch/quantization/plugins/test_nested_attention_skip.py create mode 100644 tests/unit/torch/quantization/test_quant_layernorm.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6d70548aee8..3441790f9a1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,10 @@ Changelog - ``from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import ...`` → ``from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import ...`` - ``from modelopt.torch.sparsity.attention_sparsity.kernels import ...`` → ``from modelopt.torch.kernels.sparsity.attention import ...`` +**New Features** + +- Add FP8 MHA quantization support for vision transformers. Adds an attention-aware ONNX post-processing pass (scale Mul / K-transpose move before Q, Q→DQ insertion on softmax output) in :class:`FP8QuantExporter `, per-instance nested-attention-wrapper skipping in the HF plugin, and ``nn.LayerNorm`` registration in ``QuantModuleRegistry`` so BMM input quantizers and LayerNorm output quantizers defined in FP8_DEFAULT_CFG are honored end-to-end. See `examples/torch_onnx/torch_quant_to_onnx.py `_ for the general timm-model quantize→ONNX workflow. + 0.44 (2026-05-xx) ^^^^^^^^^^^^^^^^^ diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 98daa2c13d4..37a97392485 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -88,6 +88,22 @@ }, ] +# FP8 MHA-aware config entries: quantize LayerNorm output so TRT can fuse the shared +# Q/DQ across all downstream Q/K/V/FC consumers. Softmax-output Q/DQ is handled by the +# FP8 ONNX exporter's post-processing pass (fixed 1/448 scale, data-independent). +_FP8_MHA_OVERRIDE: list = [ + { + "parent_class": "nn.LayerNorm", + "quantizer_name": "*output_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + }, + { + "parent_class": "nn.LayerNorm", + "quantizer_name": "*input_quantizer", + "enable": False, + }, +] + # Auto-quantize format configs that use block quantization and need Conv2d overrides for TRT. # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. _NEEDS_FP8_CONV_OVERRIDE: set[str] = { @@ -102,11 +118,16 @@ def get_quant_config(quantize_mode): """Get quantization config, overriding Conv2d for TRT compatibility. TensorRT only supports FP8 and INT8 for Conv layers. + - For FP8: add MHA-aware LayerNorm output quantizer so TRT fuses shared Q/DQ into + downstream attention matmuls. Softmax-output Q/DQ is inserted by the FP8 ONNX + exporter's post-processing (fixed 1/448 scale, no calibration needed). - For MXFP8, NVFP4: override Conv2d to FP8 - For INT4_AWQ: override Conv2d to INT8 """ config: dict = copy.deepcopy(QUANT_CONFIG_DICT[quantize_mode]) - if quantize_mode in ("mxfp8", "nvfp4"): + if quantize_mode == "fp8": + config["quant_cfg"].extend(_FP8_MHA_OVERRIDE) + elif quantize_mode in ("mxfp8", "nvfp4"): warnings.warn( f"TensorRT only supports FP8/INT8 for Conv layers. " f"Overriding Conv2d quantization to FP8 for '{quantize_mode}' mode." @@ -458,6 +479,7 @@ def main(): # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8 # quantizers require calibration data. config = get_quant_config(args.quantize_mode) + data_loader = load_calibration_data( model, args.calibration_data_size, diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index dcae618dd0a..427a7791f3b 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -17,6 +17,7 @@ import time +import numpy as np import onnx import onnx_graphsurgeon as gs import torch @@ -26,6 +27,11 @@ from .base_exporter import ONNXQuantExporter +# FP8 E4M3 max representable magnitude; softmax output in [0, 1] saturates exactly at 1.0 +# when using 1/448 as the Q scale (single fixed value — softmax range is data-independent). +_FP8_E4M3_MAX = 448.0 +_FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -62,6 +68,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: # Fold constants is required since the scale is not constant yet. graph.cleanup().toposort().fold_constants().cleanup() + n_t_folded = 0 + for node in graph.nodes: if node.op == "TRT_FP8QuantizeLinear": # Should not remove input QDQ (only process weight quantization) @@ -74,9 +82,44 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: torch_scale = torch.from_numpy(scale.values) quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] - assert dq_op.op == "TRT_FP8DequantizeLinear", ( - f"QDQ does not occur in pairs. You reached {dq_op.op}" - ) + if dq_op.op != "TRT_FP8DequantizeLinear": + raise RuntimeError(f"QDQ does not occur in pairs. You reached {dq_op.op}") + + # Pre-transpose constant weights if DQ feeds ``Transpose → MatMul`` (or + # ``Cast → Transpose → MatMul`` after fp16 conversion) so TRT sees DQ→MatMul. + # Control flow: scan candidates; a Cast-wrapped candidate is accepted only if it + # leads to a Transpose; a bare Transpose whose all consumers are MatMul wins and + # breaks the loop. Any other shape defaults `cast_to_remove` back to None and + # continues scanning. + transpose_to_remove = None + cast_to_remove = None + for candidate in list(dq_op.outputs[0].outputs): + if candidate.op == "Cast": + cast_to_remove = candidate + candidate = next( + (c for c in candidate.outputs[0].outputs if c.op == "Transpose"), + None, + ) + if candidate is None: + cast_to_remove = None + continue + if candidate.op != "Transpose": + cast_to_remove = None + continue + t_consumers = list(candidate.outputs[0].outputs) + # Only fold the transpose when every downstream consumer is MatMul; otherwise + # non-MatMul consumers would observe the un-transposed weights. + if t_consumers and all(c.op == "MatMul" for c in t_consumers): + perm = candidate.attrs.get("perm", None) + torch_weights = ( + torch_weights.permute(*perm).contiguous() + if perm is not None + else torch_weights.T.contiguous() + ) + transpose_to_remove = candidate + else: + cast_to_remove = None + break # Replace it with Dequantize with FP8 weights. This is a WAR because numpy does not support fp8. numpy_weights = ( @@ -94,9 +137,23 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: dq_op.inputs[0] = onnx_weights_fp8 dq_op.op = "DequantizeLinear" dq_op.outputs[0].dtype = dq_op.inputs[1].dtype + dq_op.outputs[0].shape = list(numpy_weights.shape) + + if transpose_to_remove is not None: + t_out = transpose_to_remove.outputs[0] + for consumer in list(t_out.outputs): + for i, inp in enumerate(consumer.inputs): + if inp is t_out: + consumer.inputs[i] = dq_op.outputs[0] + transpose_to_remove.outputs.clear() + if cast_to_remove is not None: + cast_to_remove.outputs.clear() + n_t_folded += 1 graph.cleanup().toposort() end_time = time.time() + if n_t_folded > 0: + logger.info(f"Folded {n_t_folded} weight Transpose nodes during weight compression") print(f"fp8 qdq replaced with only dq completed in {end_time - start_time}s.") return gs.export_onnx(graph) @@ -121,7 +178,6 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: Returns: Number of Conv weight DQ nodes inserted. """ - fp8_max = 448.0 count = 0 for node in list(graph.nodes): @@ -142,7 +198,7 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: amax = torch_weights.abs().max().float() if amax == 0: continue - scale_val = (amax / fp8_max).item() + scale_val = (amax / _FP8_E4M3_MAX).item() # Quantize weights to FP8 (WAR: numpy doesn't support fp8) fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() @@ -155,8 +211,6 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: ) # Scale in FP16 — DQ output type matches scale dtype, must match activation type - import numpy as np - scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", np.array(scale_val, dtype=np.float16), @@ -175,13 +229,220 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: return count + @staticmethod + def _move_mul_before_qdq(graph: gs.Graph) -> int: + """Move attention-scaling Mul(const) from after DQ to before Q for TRT MatMul fusion. + + Handles both ``DQ → Mul → MatMul`` and ``DQ → Transpose → Mul → MatMul`` (K path). + """ + count = 0 + for mul_node in list(graph.nodes): + if mul_node.op != "Mul": + continue + + const_input = next( + (i for i in mul_node.inputs if isinstance(i, gs.Constant) and i.values.size == 1), + None, + ) + tensor_input = next( + (i for i in mul_node.inputs if not isinstance(i, gs.Constant)), None + ) + if const_input is None or tensor_input is None: + continue + if not (isinstance(tensor_input, gs.Variable) and len(tensor_input.inputs) == 1): + continue + + producer = tensor_input.inputs[0] + transpose_node = producer if producer.op == "Transpose" else None + dq_node = producer if producer.op == "DequantizeLinear" else None + if transpose_node is not None: + t_input = transpose_node.inputs[0] + if ( + isinstance(t_input, gs.Variable) + and len(t_input.inputs) == 1 + and t_input.inputs[0].op == "DequantizeLinear" + ): + dq_node = t_input.inputs[0] + if dq_node is None: + continue + + q_output = dq_node.inputs[0] + if ( + not isinstance(q_output, gs.Variable) + or len(q_output.inputs) != 1 + or q_output.inputs[0].op != "QuantizeLinear" + ): + continue + q_node = q_output.inputs[0] + q_input = q_node.inputs[0] + if not isinstance(q_input, gs.Variable): + continue + + mul_output = mul_node.outputs[0] + mul_consumers = list(mul_output.outputs) + # Require every consumer to be MatMul: rewiring all consumers to bypass the Mul + # would silently drop the scale for any non-MatMul branch. + if not mul_consumers or not all(c.op == "MatMul" for c in mul_consumers): + continue + + new_mul_output = gs.Variable( + q_input.name + "_scaled", dtype=q_input.dtype, shape=q_input.shape + ) + graph.nodes.append( + gs.Node( + op="Mul", + name=mul_node.name + "_moved", + inputs=[q_input, const_input], + outputs=[new_mul_output], + ) + ) + q_node.inputs[0] = new_mul_output + + replacement = ( + transpose_node.outputs[0] if transpose_node is not None else dq_node.outputs[0] + ) + for consumer in mul_consumers: + for i, inp in enumerate(consumer.inputs): + if inp is mul_output: + consumer.inputs[i] = replacement + mul_node.outputs.clear() + count += 1 + + graph.cleanup().toposort() + return count + + @staticmethod + def _move_transpose_before_qdq(graph: gs.Graph) -> int: + """Move Transpose from ``DQ → Transpose → MatMul`` to ``Transpose → Q → DQ → MatMul`` (K path).""" + count = 0 + for transpose_node in list(graph.nodes): + if transpose_node.op != "Transpose": + continue + + t_input = transpose_node.inputs[0] + if ( + not isinstance(t_input, gs.Variable) + or len(t_input.inputs) != 1 + or t_input.inputs[0].op != "DequantizeLinear" + ): + continue + dq_node = t_input.inputs[0] + + dq_input = dq_node.inputs[0] + if ( + not isinstance(dq_input, gs.Variable) + or len(dq_input.inputs) != 1 + or dq_input.inputs[0].op != "QuantizeLinear" + ): + continue + q_node = dq_input.inputs[0] + q_input = q_node.inputs[0] + if not isinstance(q_input, gs.Variable): + continue + + t_output = transpose_node.outputs[0] + t_consumers = list(t_output.outputs) + # Require every consumer to be MatMul: rewiring to dq_node.outputs[0] would drop + # the transpose for any non-MatMul branch, producing a wrong-shape tensor. + if not t_consumers or not all(c.op == "MatMul" for c in t_consumers): + continue + + new_t_output = gs.Variable(q_input.name + "_transposed", dtype=q_input.dtype) + graph.nodes.append( + gs.Node( + op="Transpose", + name=transpose_node.name + "_moved", + inputs=[q_input], + outputs=[new_t_output], + attrs=transpose_node.attrs, + ) + ) + q_node.inputs[0] = new_t_output + + for consumer in t_consumers: + for i, inp in enumerate(consumer.inputs): + if inp is t_output: + consumer.inputs[i] = dq_node.outputs[0] + transpose_node.outputs.clear() + count += 1 + + graph.cleanup().toposort() + return count + + @staticmethod + def _insert_qdq_after_softmax(graph: gs.Graph) -> int: + """Insert FP8 Q→DQ on Softmax outputs feeding MatMul (required by TRT MHA fusion). + + Softmax output is data-independently bounded to [0, 1], so we use a fixed scale + ``_FP8_E4M3_SOFTMAX_SCALE`` (1/448) that saturates exactly at 1.0 while covering + the full FP8 E4M3 representable range. No calibration is required. Only applied + when every Softmax consumer is a MatMul so we do not insert quantization error + on unrelated branches. + """ + count = 0 + for softmax_node in list(graph.nodes): + if softmax_node.op != "Softmax": + continue + softmax_output = softmax_node.outputs[0] + consumers = list(softmax_output.outputs) + if not consumers or not all(c.op == "MatMul" for c in consumers): + continue + if any(c.op == "QuantizeLinear" for c in consumers): + continue + + # Match scale dtype to the graph's current float dtype so TRT stronglyTyped + # sees consistent Q/DQ types with the surrounding compute. + scale_dtype = softmax_output.dtype if softmax_output.dtype is not None else np.float32 + scale_val = np.array(_FP8_E4M3_SOFTMAX_SCALE, dtype=scale_dtype) + scale_constant = gs.Constant(softmax_node.name + "/softmax_q_scale", scale_val) + dq_scale_constant = gs.Constant( + softmax_node.name + "/softmax_dq_scale", scale_val.copy() + ) + + zp_tensor = onnx.TensorProto() + zp_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN + zp_tensor.dims.extend([1]) + zp_tensor.raw_data = b"\x00" + zp_constant = gs.Constant( + softmax_node.name + "/softmax_q_zero_point", LazyValues(zp_tensor) + ) + + q_output = gs.Variable(softmax_node.name + "/q_output") + dq_output = gs.Variable(softmax_node.name + "/dq_output", dtype=softmax_output.dtype) + q_node = gs.Node( + op="QuantizeLinear", + name=softmax_node.name + "/QuantizeLinear", + inputs=[softmax_output, scale_constant, zp_constant], + outputs=[q_output], + attrs={"saturate": 1}, + ) + dq_node = gs.Node( + op="DequantizeLinear", + name=softmax_node.name + "/DequantizeLinear", + inputs=[q_output, dq_scale_constant], + outputs=[dq_output], + ) + graph.nodes.extend([q_node, dq_node]) + + for consumer in consumers: + if consumer is q_node: + continue + for i, inp in enumerate(consumer.inputs): + if inp is softmax_output: + consumer.inputs[i] = dq_output + count += 1 + + graph.cleanup().toposort() + return count + @staticmethod def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for FP8 quantization. - Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear and + Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear, adds FP8 weight DQ for Conv layers whose weight quantizers were disabled during - TorchScript export. + TorchScript export, and rewrites attention scaling / K-transpose / softmax-output + patterns so TRT can fuse DQ into the attention MatMul kernels. Args: onnx_model: The ONNX model containing TRT_FP8 quantization nodes. @@ -223,5 +484,15 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: if count > 0: logger.info(f"Inserted FP8 weight DequantizeLinear for {count} Conv nodes") + # Attention-aware rewrites so TRT can fuse DQ into the attention MatMuls. + n_mul = FP8QuantExporter._move_mul_before_qdq(graph) + n_t = FP8QuantExporter._move_transpose_before_qdq(graph) + n_sm = FP8QuantExporter._insert_qdq_after_softmax(graph) + if n_mul or n_t or n_sm: + logger.info( + f"Attention QDQ rewrites: moved {n_mul} Mul, {n_t} Transpose; " + f"inserted QDQ on {n_sm} Softmax outputs" + ) + graph.cleanup().toposort() return gs.export_onnx(graph) diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index ac93bc2a26c..7b1d7903c1a 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1415,6 +1415,70 @@ def _bypass_cast_node(model: onnx.ModelProto, node: onnx.NodeProto) -> None: consumer.input[i] = input_tensor +_DQ_OPS = {"DequantizeLinear", "TRT_FP8DequantizeLinear"} +_Q_OPS = {"QuantizeLinear", "TRT_FP8QuantizeLinear"} + + +def _scale_fp32_to_fp16(scale_init: onnx.TensorProto) -> None: + """Convert a scalar Q/DQ scale initializer in-place from FP32 to FP16. + + Warns if any non-zero scale saturates to 0/inf in FP16 (out of FP16 representable range). + """ + if scale_init.data_type != onnx.TensorProto.FLOAT: + return + scale_data = np.frombuffer(scale_init.raw_data, dtype=np.float32) + if not scale_data.size: + scale_data = np.array(scale_init.float_data, dtype=np.float32) + fp16_data = scale_data.astype(np.float16) + if np.any(np.isinf(fp16_data)) or (np.any(fp16_data == 0) and np.any(scale_data != 0)): + logger.warning(f"Q/DQ scale '{scale_init.name}' overflows or underflows when cast to FP16") + scale_init.data_type = onnx.TensorProto.FLOAT16 + scale_init.raw_data = fp16_data.tobytes() + del scale_init.float_data[:] + + +def fold_q_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + """Remove ``Cast(FP16→FP32) → Q`` patterns inserted by ``convert_float_to_float16``. + + The Q scale is rewritten to FP16 so Q consumes the FP16 graph directly. Skipped for + opsets below ``BASE_MIN_OPSET`` since FP16 Q scales require opset >= 19. + """ + if get_opset_version(onnx_model) < BASE_MIN_OPSET: + logger.debug( + f"Skipping fold_q_fp16_to_fp32_casts: opset < {BASE_MIN_OPSET} (FP16 Q scale unsupported)" + ) + return onnx_model + + consumer_map: dict[str, list[onnx.NodeProto]] = {} + for node in onnx_model.graph.node: + for inp in node.input: + consumer_map.setdefault(inp, []).append(node) + initializers = {init.name: init for init in onnx_model.graph.initializer} + + to_remove = [] + for node in onnx_model.graph.node: + if node.op_type != "Cast": + continue + cast_to = next((a.i for a in node.attribute if a.name == "to"), None) + if cast_to != onnx.TensorProto.FLOAT: + continue + consumers = consumer_map.get(node.output[0], []) + if not consumers or not all(c.op_type in _Q_OPS for c in consumers): + continue + + for q_node in consumers: + if len(q_node.input) >= 2 and q_node.input[1] in initializers: + _scale_fp32_to_fp16(initializers[q_node.input[1]]) + + _bypass_cast_node(onnx_model, node) + to_remove.append(node) + + logger.debug(f"Folded {len(to_remove)} Cast(FP16->FP32) -> Q patterns") + for node in to_remove: + onnx_model.graph.node.remove(node) + return onnx_model + + def _is_foldable_constant_cast_pattern(model: onnx.ModelProto, node: onnx.NodeProto) -> bool: """Check if a Constant -> Cast pattern can be folded.""" assert node.op_type == "Cast" @@ -1523,7 +1587,12 @@ def fold_dq_fp32_to_fp16_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: Returns: The ONNX model with Cast nodes removed and DQ outputs set to FP16. """ - import numpy as np + if get_opset_version(onnx_model) < BASE_MIN_OPSET: + logger.debug( + f"Skipping fold_dq_fp32_to_fp16_casts: opset < {BASE_MIN_OPSET} " + "(FP16 DQ scale unsupported)" + ) + return onnx_model dq_ops = {"DequantizeLinear", "TRT_FP8DequantizeLinear"} @@ -1623,6 +1692,13 @@ def fold_qdq_scale_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.Model Returns: The ONNX model with redundant scale-path casts removed. """ + if get_opset_version(onnx_model) < BASE_MIN_OPSET: + logger.debug( + f"Skipping fold_qdq_scale_fp16_to_fp32_casts: opset < {BASE_MIN_OPSET} " + "(FP16 Q/DQ scale unsupported)" + ) + return onnx_model + qdq_ops = { "QuantizeLinear", "DequantizeLinear", diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 9ec110b7887..01fb754bbae 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -48,6 +48,7 @@ change_casts_to_fp16, check_model_uses_external_data, fold_dq_fp32_to_fp16_casts, + fold_q_fp16_to_fp32_casts, fold_qdq_scale_fp16_to_fp32_casts, get_input_names, get_input_shapes, @@ -663,6 +664,11 @@ def get_onnx_bytes_and_metadata( onnx_opt_graph = remove_redundant_casts(onnx_opt_graph) + # Remove Cast nodes around Q/DQ for optimal TRT fusion + if is_fp8_quantized(model): + onnx_opt_graph = fold_q_fp16_to_fp32_casts(onnx_opt_graph) + onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) + # TensorRT expects all scales to be postive onnx_opt_graph = replace_zero_scale_with_smallest_nonzero(onnx_opt_graph) diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index 05efe48842f..7b42dce5782 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -216,56 +216,36 @@ def _fp8_quantize( g: "GraphContext", inputs: torch.Value, scale_inv: float, - trt_high_precision_dtype: str, ): """Helper Function for Quantization.""" + # Emit the scale in the native input dtype so no Cast is inserted between the + # graph and Q/DQ (Cast nodes block TRT from fusing DQ into the MatMul kernel). output_shape = sym_help._get_tensor_sizes(inputs) - - # TRT StronglyType only supports FP16 QDQs - # custom ops, so cast the input if needed. - input_type = inputs.type().scalarType() - assert trt_high_precision_dtype in (input_type, "Float"), ( - "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." - ) - if trt_high_precision_dtype != input_type: - inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) - scale = g.op( "Constant", - value_t=torch.tensor(scale_inv).to(torch_dtype_map[trt_high_precision_dtype]), + value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]), ) - q_op = g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( + return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) ) - return q_op def _fp8_dequantize( g: "GraphContext", inputs: torch.Value, scale_inv: float, - trt_high_precision_dtype: str, otype: str | None = None, ): """Helper Function for Dequantization.""" output_shape = sym_help._get_tensor_sizes(inputs) - assert trt_high_precision_dtype in (otype, "Float"), ( - "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." - ) scale = g.op( "Constant", value_t=torch.tensor(scale_inv, dtype=torch_dtype_map[otype]), # type: ignore[index] ) - out = g.op("trt::TRT_FP8DequantizeLinear", inputs, scale).setType( - inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) + return g.op("trt::TRT_FP8DequantizeLinear", inputs, scale).setType( + inputs.type().with_dtype(torch_dtype_map[otype]).with_sizes(output_shape) # type: ignore[index] ) - # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT - # custom ops, so cast the output if needed. - if trt_high_precision_dtype != otype: - out = g.op("Cast", out, to_i=onnx_dtype_map[otype]) # type: ignore[index] - return out - def export_fp8( g: "GraphContext", @@ -273,14 +253,17 @@ def export_fp8( amax: float, trt_high_precision_dtype: str | None, ): - """Export quantized model to FP8 ONNX.""" + """Export quantized model to FP8 ONNX. + + ``trt_high_precision_dtype`` is accepted for API compatibility but unused: Q/DQ now + emit scales in the native input dtype, so no intermediate Cast is required. + """ + del trt_high_precision_dtype scale = 1.0 if amax is None else 448.0 / float(amax) otype = inputs.type().scalarType() - if trt_high_precision_dtype is None: - trt_high_precision_dtype = otype - q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype) - return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype, otype) + q_tensor = _fp8_quantize(g, inputs, 1.0 / scale) + return _fp8_dequantize(g, q_tensor, 1.0 / scale, otype) def scaled_dot_product_attention( diff --git a/modelopt/torch/quantization/nn/__init__.py b/modelopt/torch/quantization/nn/__init__.py index ca7082eb1cb..af9490c8311 100644 --- a/modelopt/torch/quantization/nn/__init__.py +++ b/modelopt/torch/quantization/nn/__init__.py @@ -19,6 +19,7 @@ from .modules.quant_batchnorm import * from .modules.quant_conv import * from .modules.quant_instancenorm import * +from .modules.quant_layernorm import * from .modules.quant_linear import * from .modules.quant_module import * from .modules.quant_pooling import * diff --git a/modelopt/torch/quantization/nn/modules/quant_layernorm.py b/modelopt/torch/quantization/nn/modules/quant_layernorm.py new file mode 100644 index 00000000000..10e9ba47582 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_layernorm.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Registers ``torch.nn.LayerNorm`` with ``QuantInputBase``. + +Enables LayerNorm output quantizers to be honored during quantization. Required for FP8 +attention fusion where a single LayerNorm output QDQ is shared across all downstream +Q/K/V/FC consumers (instead of repeating it on each input), which enables TRT to fuse DQ +into the attention MatMul kernels. +""" + +import torch.nn as nn + +from .quant_module import QuantInputBase, QuantModuleRegistry + +QuantModuleRegistry.register({nn.LayerNorm: "nn.LayerNorm"})(QuantInputBase) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 92eaf12ece5..ba40128fa5e 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -274,6 +274,22 @@ def forward(self, *args, **kwargs): return super().forward(*args, **kwargs) +def _wraps_nested_attention(module): + """Return True when ``module`` contains another Attention child on this specific instance. + + Checked per-instance (not by class) so an attention class reused as both wrapper and + leaf is not dropped everywhere. In a 3-level hierarchy (Outer → Middle → Inner), both + Outer and Middle are treated as wrappers and only Inner is registered for KV-cache + quantization. Used to avoid double-patching ``eager_attention_forward`` when a wrapper + attention module (e.g. ``ViTAttention``) delegates to a nested self-attention child + (e.g. ``ViTSelfAttention``). + """ + return any( + child is not module and type(child).__name__.endswith("Attention") + for _, child in module.named_modules() + ) + + def register_hf_attentions_on_the_fly(model): """Find HF Attention modules in the model and register them for KV Cache quantization. @@ -286,9 +302,12 @@ def register_hf_attentions_on_the_fly(model): attention_cls = set() registered_attn_module = False + for name, module in model.named_modules(): # Only register attention classes that are from Huggingface transformers if type(module).__name__.endswith("Attention"): + if _wraps_nested_attention(module): + continue attention_type = _QuantAttention.get_attn_type(module) # Add modules to be registered only if they arent already registered if ( diff --git a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py new file mode 100644 index 00000000000..1f7251a9ad9 --- /dev/null +++ b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the attention-aware FP8 ONNX graph rewrites in ``FP8QuantExporter``.""" + +import numpy as np +import onnx_graphsurgeon as gs +import pytest + +from modelopt.onnx.export.fp8_exporter import FP8QuantExporter + + +def _var(name): + return gs.Variable(name, dtype=np.float32) + + +def _qdq(src): + """Build ``QuantizeLinear → DequantizeLinear`` and return [Q, DQ], dq_out.""" + scale = gs.Constant("scale", np.array(0.1, dtype=np.float32)) + q_out, dq_out = _var("q_out"), _var("dq_out") + return [ + gs.Node(op="QuantizeLinear", inputs=[src, scale], outputs=[q_out]), + gs.Node(op="DequantizeLinear", inputs=[q_out, scale], outputs=[dq_out]), + ], dq_out + + +def _graph(nodes, inputs, outputs): + return gs.Graph(nodes=nodes, inputs=inputs, outputs=outputs, opset=19) + + +def test_move_mul_before_qdq_rewrites_dq_mul_matmul_pattern(): + """``DQ → Mul(const) → MatMul`` collapses to ``Mul → Q → DQ → MatMul``.""" + x, k, y, mul_out = _var("x"), _var("k"), _var("y"), _var("mul_out") + qdq_nodes, dq_out = _qdq(x) + mul = gs.Node( + op="Mul", + inputs=[dq_out, gs.Constant("c", np.array(0.5, dtype=np.float32))], + outputs=[mul_out], + ) + mm = gs.Node(op="MatMul", inputs=[mul_out, k], outputs=[y]) + graph = _graph([*qdq_nodes, mul, mm], [x, k], [y]) + + assert FP8QuantExporter._move_mul_before_qdq(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert q.inputs[0].inputs[0].op == "Mul" + + +def test_move_transpose_before_qdq_rewrites_dq_transpose_matmul_pattern(): + """``DQ → Transpose → MatMul`` collapses to ``Transpose → Q → DQ → MatMul``.""" + k_in, q_in, scores, t_out = _var("k_in"), _var("q_in"), _var("scores"), _var("t_out") + qdq_nodes, dq_out = _qdq(k_in) + t = gs.Node(op="Transpose", inputs=[dq_out], outputs=[t_out], attrs={"perm": [0, 2, 1]}) + mm = gs.Node(op="MatMul", inputs=[q_in, t_out], outputs=[scores]) + graph = _graph([*qdq_nodes, t, mm], [k_in, q_in], [scores]) + + assert FP8QuantExporter._move_transpose_before_qdq(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert q.inputs[0].inputs[0].op == "Transpose" + + +def test_insert_qdq_after_softmax_adds_fixed_scale_q_dq(): + """Softmax → MatMul picks up ``Q → DQ`` with the fixed ``1/448`` scale.""" + scores, v, y, sm_out = _var("scores"), _var("v"), _var("y"), _var("sm_out") + sm = gs.Node(op="Softmax", inputs=[scores], outputs=[sm_out], attrs={"axis": -1}) + mm = gs.Node(op="MatMul", inputs=[sm_out, v], outputs=[y]) + graph = _graph([sm, mm], [scores, v], [y]) + + assert FP8QuantExporter._insert_qdq_after_softmax(graph) == 1 + q = next(n for n in graph.nodes if n.op == "QuantizeLinear") + assert np.isclose(float(q.inputs[1].values), 1.0 / 448.0) + + +@pytest.mark.parametrize( + "rewrite", ["_move_mul_before_qdq", "_move_transpose_before_qdq", "_insert_qdq_after_softmax"] +) +def test_rewrites_skip_when_non_matmul_consumer_exists(rewrite): + """Every MHA rewrite must skip when the candidate tensor fans out to a non-MatMul branch.""" + x, k, y_mm, y_side, shared = _var("x"), _var("k"), _var("y_mm"), _var("y_side"), _var("shared") + + if rewrite == "_move_mul_before_qdq": + qdq_nodes, dq_out = _qdq(x) + producer = gs.Node( + op="Mul", + inputs=[dq_out, gs.Constant("c", np.array(0.5, dtype=np.float32))], + outputs=[shared], + ) + prelude = [*qdq_nodes, producer] + elif rewrite == "_move_transpose_before_qdq": + qdq_nodes, dq_out = _qdq(x) + producer = gs.Node( + op="Transpose", inputs=[dq_out], outputs=[shared], attrs={"perm": [1, 0]} + ) + prelude = [*qdq_nodes, producer] + else: + prelude = [gs.Node(op="Softmax", inputs=[x], outputs=[shared], attrs={"axis": -1})] + + graph = _graph( + [ + *prelude, + gs.Node(op="MatMul", inputs=[shared, k], outputs=[y_mm]), + gs.Node(op="Relu", inputs=[shared], outputs=[y_side]), + ], + [x, k], + [y_mm, y_side], + ) + assert getattr(FP8QuantExporter, rewrite)(graph) == 0 diff --git a/tests/unit/onnx/test_fold_casts.py b/tests/unit/onnx/test_fold_casts.py new file mode 100644 index 00000000000..59a434d1206 --- /dev/null +++ b/tests/unit/onnx/test_fold_casts.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the FP16 Q/DQ scale cast-folding helpers in ``modelopt.onnx.utils``.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.utils import fold_dq_fp32_to_fp16_casts, fold_q_fp16_to_fp32_casts + + +def _dq_cast_model(opset): + """``DQ → Cast(FP32→FP16) → MatMul(x)`` with FP32 scale.""" + nodes = [ + helper.make_node("DequantizeLinear", ["w_q", "w_scale", "w_zp"], ["dq_out"], "dq"), + helper.make_node("Cast", ["dq_out"], ["cast_out"], "cast", to=TensorProto.FLOAT16), + helper.make_node("MatMul", ["x", "cast_out"], ["y"], "matmul"), + ] + inits = [ + numpy_helper.from_array(np.ones((4, 4), dtype=np.int8), "w_q"), + numpy_helper.from_array(np.array(0.1, dtype=np.float32), "w_scale"), + numpy_helper.from_array(np.array(0, dtype=np.int8), "w_zp"), + ] + return helper.make_model( + helper.make_graph( + nodes, + "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [None, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT16, [None, 4])], + initializer=inits, + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + + +def _cast_q_model(opset): + """``Cast(FP16→FP32) → Q → DQ → MatMul`` with FP32 scale.""" + nodes = [ + helper.make_node("Cast", ["x"], ["c_out"], "cast", to=TensorProto.FLOAT), + helper.make_node("QuantizeLinear", ["c_out", "scale", "zp"], ["q_out"], "q"), + helper.make_node("DequantizeLinear", ["q_out", "scale", "zp"], ["dq_out"], "dq"), + helper.make_node("MatMul", ["dq_out", "w"], ["y"], "matmul"), + ] + inits = [ + numpy_helper.from_array(np.ones((4, 4), dtype=np.float16), "w"), + numpy_helper.from_array(np.array(0.1, dtype=np.float32), "scale"), + numpy_helper.from_array(np.array(0, dtype=np.int8), "zp"), + ] + return helper.make_model( + helper.make_graph( + nodes, + "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [None, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [None, 4])], + initializer=inits, + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + + +@pytest.mark.parametrize( + ("fold_fn", "build_model", "scale_name"), + [ + (fold_dq_fp32_to_fp16_casts, _dq_cast_model, "w_scale"), + (fold_q_fp16_to_fp32_casts, _cast_q_model, "scale"), + ], +) +def test_fold_rewrites_cast_and_scale_at_opset_19(fold_fn, build_model, scale_name): + folded = fold_fn(build_model(opset=19)) + assert "Cast" not in {n.op_type for n in folded.graph.node} + scale = next(i for i in folded.graph.initializer if i.name == scale_name) + assert scale.data_type == TensorProto.FLOAT16 + + +@pytest.mark.parametrize( + ("fold_fn", "build_model", "scale_name"), + [ + (fold_dq_fp32_to_fp16_casts, _dq_cast_model, "w_scale"), + (fold_q_fp16_to_fp32_casts, _cast_q_model, "scale"), + ], +) +def test_fold_is_noop_below_min_opset(fold_fn, build_model, scale_name): + folded = fold_fn(build_model(opset=18)) + assert "Cast" in {n.op_type for n in folded.graph.node} + scale = next(i for i in folded.graph.initializer if i.name == scale_name) + assert scale.data_type == TensorProto.FLOAT diff --git a/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py b/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py new file mode 100644 index 00000000000..27dcc034f96 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the per-instance nested-attention skip in the HF quantization plugin.""" + +import torch.nn as nn + +from modelopt.torch.quantization.plugins.huggingface import _wraps_nested_attention + + +def _attn(name, child=None): + """Build a module whose class name ends with ``Attention`` and optionally wraps ``child``.""" + cls = type(name, (nn.Module,), {"__init__": lambda self: nn.Module.__init__(self)}) + m = cls() + if child is not None: + m.inner = child + return m + + +def test_wraps_nested_attention_flags_only_wrappers_per_instance(): + """Leaf attention is not a wrapper; wrappers (any level) are; same class reused is + checked per-instance.""" + leaf = _attn("SelfAttention") + wrapper = _attn("ViTAttention", child=_attn("ViTSelfAttention")) + outer = _attn("OuterAttention", child=wrapper) + reused_wrapper = _attn("ReusedAttention", child=_attn("ReusedAttention")) + + assert not _wraps_nested_attention(leaf) + assert _wraps_nested_attention(wrapper) + assert _wraps_nested_attention(outer) and _wraps_nested_attention(outer.inner) + assert _wraps_nested_attention(reused_wrapper) + assert not _wraps_nested_attention(reused_wrapper.inner) diff --git a/tests/unit/torch/quantization/test_quant_layernorm.py b/tests/unit/torch/quantization/test_quant_layernorm.py new file mode 100644 index 00000000000..1e20662524a --- /dev/null +++ b/tests/unit/torch/quantization/test_quant_layernorm.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ``torch.nn.LayerNorm`` being registered in ``QuantModuleRegistry``.""" + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.nn import QuantModuleRegistry + + +def test_layernorm_quant_wrapper_is_identity_when_quantizers_disabled(): + qln = QuantModuleRegistry.convert(nn.LayerNorm(8)) + qln.input_quantizer.disable() + qln.output_quantizer.disable() + + x = torch.randn(2, 8) + ref = nn.functional.layer_norm(x, (8,), qln.weight, qln.bias, eps=qln.eps) + assert torch.allclose(qln(x), ref, rtol=0, atol=0) From 473e9203cd72b70eddb62d3e5ec2897f2e53a4bd Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 22 Apr 2026 17:30:17 +0000 Subject: [PATCH 2/4] Fix torch_onnx PTQ for Swin/SwinV2 under MXFP8/NVFP4 and FP8 Two independent bugs surfaced by the parametrized matrix in tests/examples/torch_onnx/test_torch_quant_to_onnx.py: - MXFP8/NVFP4 lower input quantizers to TRT DynamicQuantize, which only supports 2D/3D input. Swin/SwinV2 keep the 4D (B, H, W, C) layout on per-block norm1, downsample.norm, and the top-level norm, causing trtexec (MXFP8) and the NVFP4 autocast TRT-shape-inference pre-pass to reject the graph. Added _disable_high_rank_input_quantizers which runs a forward-pass rank probe and disables quantizers on 4D+ inputs; gated on mxfp8 / nvfp4 / auto so FP8 and INT8 still quantize those layers (their Q/DQ has no rank constraint). Name-based alternatives would false-positive on ViT, whose same-named top-level norm is 3D. - swinv2_tiny-fp8 hit ZeroDivisionError in export_fp8 (448 / amax): timm's res-post-norm scheme zero-inits each SwinV2 block's norm1 / norm2 weight and bias, so under --no_pretrained those LayerNorm outputs are exactly zero, and the FP8 MHA override's output_quantizer calibrates to amax == 0. Added _disable_dead_quantizers to drop any quantizer whose calibrated amax is NaN or <= 0 before export. Full matrix (4 models x 5 modes) now passes: 20/20 in ~33 min. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_onnx/torch_quant_to_onnx.py | 76 ++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 37a97392485..b99eeebbdc0 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -147,6 +147,9 @@ def filter_func(name): ``downsample.reduction`` (Swin/SwinV2) is excluded because it operates on 4D tensors and TRT's DynamicQuantize layer (used for MXFP8/NVFP4) requires 2D/3D input. + Other 4D-input layers (e.g. Swin's ``norm1``, ``downsample.norm``, top-level ``norm``) + are handled dynamically by ``_disable_high_rank_input_quantizers`` via a forward-pass + rank probe — that avoids false positives on ViT, whose same-named ``norm`` sees 3D input. """ pattern = re.compile( r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|" @@ -156,6 +159,41 @@ def filter_func(name): return pattern.match(name) is not None +def _disable_high_rank_input_quantizers(model, input_shape, device): + """Disable quantizers on Linear/LayerNorm modules that receive 4D+ input. + + TRT's MXFP8/NVFP4 ``DynamicQuantize`` op only supports 2D/3D input, so Swin's + per-block ``norm1``, ``downsample.norm``, and top-level ``norm`` (all 4D in Swin + but 3D in ViT) must be skipped. A forward pass with hooks identifies them at + runtime, so this works across architectures without hardcoded paths. + """ + high_rank: set[str] = set() + handles = [] + for name, mod in model.named_modules(): + if isinstance(mod, (torch.nn.Linear, torch.nn.LayerNorm)): + + def hook(m, inp, out, _n=name): + if inp and hasattr(inp[0], "ndim") and inp[0].ndim > 3: + high_rank.add(_n) + + handles.append(mod.register_forward_hook(hook)) + + was_training = model.training + model.eval() + try: + with torch.no_grad(): + model(torch.randn(input_shape, device=device)) + finally: + for h in handles: + h.remove() + model.train(was_training) + + if not high_rank: + return + prefixes = tuple(n + "." for n in high_rank) + mtq.disable_quantizer(model, lambda n: n.startswith(prefixes)) + + def load_calibration_data(model, data_size, batch_size, device, with_labels=False): """Load and prepare calibration data. @@ -190,6 +228,28 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals ) +def _disable_dead_quantizers(model): + """Disable quantizers whose calibrated ``amax`` is non-positive or NaN. + + ``export_fp8`` computes ``scale = 448 / amax`` and blows up on ``amax == 0``. + This shows up on SwinV2 with ``--no_pretrained``: timm's ``res-post-norm`` scheme + zero-inits each block's ``norm1``/``norm2`` weight and bias, so those LayerNorm + outputs are exactly zero at init and the MHA override's output_quantizer + calibrates to ``amax == 0``. Disable such dead quantizers — they have nothing + meaningful to quantize and would otherwise break ONNX export. + """ + for _, mod in model.named_modules(): + for attr in ("input_quantizer", "output_quantizer", "weight_quantizer"): + q = getattr(mod, attr, None) + if q is None or not q.is_enabled: + continue + amax = q.amax + if amax is None or not torch.is_tensor(amax): + continue + if torch.any(torch.isnan(amax)) or torch.all(amax <= 0): + q.disable() + + def _calibrate_uncalibrated_quantizers(model, data_loader): """Calibrate FP8 quantizers that weren't calibrated by mtq.quantize(). @@ -240,6 +300,10 @@ def forward_loop(model): if data_loader is not None: _calibrate_uncalibrated_quantizers(quantized_model, data_loader) + # Drop quantizers whose calibration saw only zeros (e.g. SwinV2 zero-init norm1/norm2) + # so ``export_fp8`` doesn't divide by zero. + _disable_dead_quantizers(quantized_model) + return quantized_model @@ -324,6 +388,8 @@ def auto_quantize_model( # Disable quantization for specified layers mtq.disable_quantizer(quantized_model, filter_func) + _disable_dead_quantizers(quantized_model) + return quantized_model, search_state @@ -490,6 +556,16 @@ def main(): quantized_model = quantize_model(model, config, data_loader) + # MXFP8/NVFP4 lower their input quantizers to TRT DynamicQuantize (2D/3D only). + # Disable quantizers on 4D-input layers (Swin's norm1 / downsample.norm / top-level norm). + # Auto mode also needs this when an MXFP8/NVFP4 candidate format is in the search set. + uses_dynamic_quantize = args.quantize_mode in ("mxfp8", "nvfp4") or ( + args.quantize_mode == "auto" + and any(fmt in _NEEDS_FP8_CONV_OVERRIDE for fmt in args.auto_quantization_formats) + ) + if uses_dynamic_quantize: + _disable_high_rank_input_quantizers(quantized_model, input_shape, device) + # Print quantization summary print("\nQuantization Summary:") mtq.print_quant_summary(quantized_model) From fd40ee44b5ad54214590fb09be9a1a424dc78f4f Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 22 Apr 2026 18:11:56 +0000 Subject: [PATCH 3/4] Skip test_nested_attention_skip when transformers is missing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The test module imports from modelopt.torch.quantization.plugins.huggingface, which imports transformers at module scope. Under the partial-install (torch) CI job — which installs only torch, without transformers/onnx/diffusers — collection failed with ModuleNotFoundError, taking the whole unit-torch partial-install step down. Add pytest.importorskip("transformers") before the plugin import, matching the pattern used by the sibling test_fused_experts.py. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .../torch/quantization/plugins/test_nested_attention_skip.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py b/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py index 27dcc034f96..7a919799b86 100644 --- a/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py +++ b/tests/unit/torch/quantization/plugins/test_nested_attention_skip.py @@ -15,8 +15,11 @@ """Tests for the per-instance nested-attention skip in the HF quantization plugin.""" +import pytest import torch.nn as nn +pytest.importorskip("transformers") + from modelopt.torch.quantization.plugins.huggingface import _wraps_nested_attention From a9f87bf41288ac9cad6cad7d77f38fb31be823af Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 22 Apr 2026 19:20:06 +0000 Subject: [PATCH 4/4] Skip first-conv input_quantizer on 3-channel Conv2d for FP8 modes CI (Blackwell, compute capability 12.0) fails TRT engine build for resnet50 under fp8 / mxfp8 / nvfp4 / auto: Error Code 10: Could not find any implementation for node /conv1/input_quantizer/TRT_FP8QuantizeLinear ... [ElementWise] The node is ResNet50's top-level conv1 (7x7 stride-2, in_channels=3). TRT's Blackwell tactics for FP8 Q -> Conv fusion don't cover the raw-RGB (3-channel) first-layer pattern. Ada (compute capability 8.9, the local dev GPU) happens to have a tactic, which is why the matrix passed locally. Swin/ViT avoid this because their first conv (patch_embed.proj, also 3-channel) is already excluded by filter_func's patch_embed pattern. ResNet50's conv1 wasn't on any list. Add _disable_low_channel_conv_input_quantizers to disable the input_quantizer on any Conv2d with in_channels <= 3 for FP8-family modes. Weight quantization is preserved. This also aligns with standard quantization practice (leave first/last layers in higher precision). INT8 is unchanged - INT8 Q/DQ has broader TRT kernel coverage on Blackwell and built successfully in CI. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_onnx/torch_quant_to_onnx.py | 33 ++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index b99eeebbdc0..97bd0b60c18 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -194,6 +194,30 @@ def hook(m, inp, out, _n=name): mtq.disable_quantizer(model, lambda n: n.startswith(prefixes)) +def _disable_low_channel_conv_input_quantizers(model): + """Disable ``input_quantizer`` on Conv2d modules whose ``in_channels <= 3``. + + The first Conv2d of an image backbone (e.g. ResNet50's ``conv1``) consumes raw + RGB input, so ``in_channels == 3``. On Blackwell (compute capability 12.0) TRT + fails to find an FP8/MXFP8/NVFP4 tactic for this first-layer Q→Conv fusion: + + Error Code 10: Could not find any implementation for node + /conv1/input_quantizer/TRT_FP8QuantizeLinear ... [ElementWise] + + Ada (8.9) happens to have a tactic, which is why local runs pass. Disabling the + input quantizer on the raw-RGB conv is also standard quantization practice — + first/last layers are typically left in higher precision. Weight quantization + still applies. Swin/ViT's ``patch_embed.proj`` is already excluded via + ``filter_func``'s ``patch_embed`` pattern, so this helper is effectively the + ResNet-shaped analogue. + """ + for _, mod in model.named_modules(): + if isinstance(mod, torch.nn.Conv2d) and mod.in_channels <= 3: + q = getattr(mod, "input_quantizer", None) + if q is not None and q.is_enabled: + q.disable() + + def load_calibration_data(model, data_size, batch_size, device, with_labels=False): """Load and prepare calibration data. @@ -566,6 +590,15 @@ def main(): if uses_dynamic_quantize: _disable_high_rank_input_quantizers(quantized_model, input_shape, device) + # FP8-family modes emit TRT_FP8QuantizeLinear on the first-layer conv; Blackwell has + # no tactic for that 3-channel Q→Conv fusion. Skip for pure INT8 (unaffected). + uses_fp8_conv_input = args.quantize_mode in ("fp8", "mxfp8", "nvfp4") or ( + args.quantize_mode == "auto" + and any(fmt != "INT8_DEFAULT_CFG" for fmt in args.auto_quantization_formats) + ) + if uses_fp8_conv_input: + _disable_low_channel_conv_input_quantizers(quantized_model) + # Print quantization summary print("\nQuantization Summary:") mtq.print_quant_summary(quantized_model)