feat: Expose quantization API in torch_tensorrt.dynamo

Author: keehyuna

py/torch_tensorrt/dynamo/__init__.py

@@ -15,6 +15,7 @@
         save_cross_compiled_exported_program,
     )
     from ._exporter import export
+    from ._quantization import quantize
     from ._refit import refit_module_weights
     from ._settings import CompilationSettings
     from ._SourceIR import SourceIR

py/torch_tensorrt/dynamo/_quantization.py

@@ -0,0 +1,35 @@
+import logging
+from typing import Any, Callable
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def quantize(
+    model: torch.nn.Module,
+    quant_format: str,
+    calibrate_loop: Callable[[], Any],
+    debug: bool = False,
+) -> torch.nn.Module:
+    try:
+        import modelopt.torch.quantization as mtq
+
+        assert torch.ops.tensorrt.quantize_op.default
+    except Exception:
+        logger.warning(
+            "Unable to import quantization op. Please install modelopt library"
+        )
+
+    if quant_format == "fp8":
+        quant_cfg = mtq.FP8_DEFAULT_CFG
+    elif quant_format == "nvfp4":
+        quant_cfg = mtq.NVFP4_DEFAULT_CFG
+    else:
+        raise RuntimeError("Unsupported quantization format")
+
+    quantized_model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
+    if debug:
+        mtq.print_quant_summary(quantized_model)
+
+    return quantized_model

tools/llm/quantize_utils.py

@@ -4,6 +4,7 @@
 
 import huggingface_hub
 import torch
+import torch_tensorrt
 from huggingface_hub import snapshot_download
 
 logger = logging.getLogger(__name__)
@@ -51,17 +52,11 @@ def quantize_model(model, args, tokenizer):
         num_samples=512,
         device="cuda:0",
     )
-    if args.quant_format == "fp8":
-        quant_cfg = mtq.FP8_DEFAULT_CFG
-    elif args.quant_format == "nvfp4":
-        quant_cfg = mtq.NVFP4_DEFAULT_CFG
-    else:
-        raise RuntimeError("Unsupported quantization format")
-    calibrate_loop = create_forward_loop(dataloader=calib_dataloader)
 
-    model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
-    if args.debug:
-        mtq.print_quant_summary(model)
+    calibrate_loop = create_forward_loop(dataloader=calib_dataloader)
+    model = torch_tensorrt.dynamo.quantize(
+        model, args.quant_format, calibrate_loop, debug=args.debug
+    )
 
     return model