Integrate Automated QDQ placement tool - part 4.4

Signed-off-by: Will Guo <willg@nvidia.com>

Author: willg-nv

modelopt/onnx/quantization/autotune/__main__.py

@@ -33,6 +33,40 @@
 DEFAULT_TIMING_CACHE = str(Path(tempfile.gettempdir()) / "trtexec_timing.cache")
 DEFAULT_WARMUP_RUNS = 5
 DEFAULT_TIMING_RUNS = 20
+MODE_PRESETS = {
+    "quick": {"schemes_per_region": 30, "warmup_runs": 10, "timing_runs": 50},
+    "default": {"schemes_per_region": 50, "warmup_runs": 50, "timing_runs": 100},
+    "extensive": {"schemes_per_region": 200, "warmup_runs": 50, "timing_runs": 200},
+}
+
+
+class _StoreWithExplicitFlag(argparse.Action):
+    """Store the value and set an 'explicit' flag on the namespace so mode presets do not override."""
+
+    def __init__(self, explicit_attr: str, *args, **kwargs):
+        self._explicit_attr = explicit_attr
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, values)
+        setattr(namespace, self._explicit_attr, True)
+
+
+def apply_mode_presets(args) -> None:
+    """Apply --mode preset to schemes_per_region, warmup_runs, timing_runs.
+
+    Only applies preset for an option when that option was not explicitly set on the
+    command line (explicit flags override the preset even when the value equals the default).
+    """
+    if args.mode not in MODE_PRESETS:
+        return
+    preset = MODE_PRESETS[args.mode]
+    if not getattr(args, "_explicit_num_schemes", False):
+        args.num_schemes = preset["schemes_per_region"]
+    if not getattr(args, "_explicit_warmup_runs", False):
+        args.warmup_runs = preset["warmup_runs"]
+    if not getattr(args, "_explicit_timing_runs", False):
+        args.timing_runs = preset["timing_runs"]
 
 
 def validate_file_path(path: str | None, description: str) -> Path | None:
@@ -94,6 +128,7 @@ def run_autotune() -> int:
         - 130: Interrupted by user (Ctrl+C)
     """
     args = _get_autotune_parser().parse_args()
+    apply_mode_presets(args)
     model_path = validate_file_path(args.onnx_path, "Model file")
     validate_file_path(args.qdq_baseline, "QDQ baseline model")
     output_dir = Path(args.output_dir)
@@ -167,6 +202,12 @@ def _get_autotune_parser() -> argparse.ArgumentParser:
   # Basic usage
   python -m modelopt.onnx.quantization.autotune --onnx_path model.onnx
 
+  # Quick mode (fewer schemes and benchmark runs for fast iteration)
+  python -m modelopt.onnx.quantization.autotune --onnx_path model.onnx --mode quick
+
+  # Extensive mode (more schemes and runs for thorough tuning)
+  python -m modelopt.onnx.quantization.autotune --onnx_path model.onnx --mode extensive
+
   # Import patterns from QDQ baseline model
   python -m modelopt.onnx.quantization.autotune \\
       --onnx_path model.onnx --qdq_baseline baseline.onnx
@@ -198,13 +239,26 @@ def _get_autotune_parser() -> argparse.ArgumentParser:
 
     # Autotuning Strategy
     strategy_group = parser.add_argument_group("Autotuning Strategy")
+    strategy_group.add_argument(
+        "--mode",
+        type=str,
+        default="default",
+        choices=["quick", "default", "extensive"],
+        help="Preset for schemes_per_region, warmup_runs, and timing_runs. "
+        "'quick': fewer schemes/runs for fast iteration; "
+        "'default': balanced; "
+        "'extensive': more schemes/runs for thorough tuning. "
+        "Explicit --schemes_per_region, --warmup_runs, --timing_runs override the preset.",
+    )
     strategy_group.add_argument(
         "--schemes_per_region",
         "-s",
         type=int,
         default=DEFAULT_NUM_SCHEMES,
         dest="num_schemes",
-        help=f"Number of schemes to test per region (default: {DEFAULT_NUM_SCHEMES})",
+        action=_StoreWithExplicitFlag,
+        explicit_attr="_explicit_num_schemes",
+        help=f"Number of schemes to test per region (default: {DEFAULT_NUM_SCHEMES}; overridden by --mode)",
     )
     strategy_group.add_argument(
         "--pattern_cache",
@@ -268,13 +322,17 @@ def _get_autotune_parser() -> argparse.ArgumentParser:
         "--warmup_runs",
         type=int,
         default=DEFAULT_WARMUP_RUNS,
-        help=f"Number of warmup runs (default: {DEFAULT_WARMUP_RUNS})",
+        action=_StoreWithExplicitFlag,
+        explicit_attr="_explicit_warmup_runs",
+        help=f"Number of warmup runs (default: {DEFAULT_WARMUP_RUNS}; overridden by --mode)",
     )
     trt_group.add_argument(
         "--timing_runs",
         type=int,
         default=DEFAULT_TIMING_RUNS,
-        help=f"Number of timing runs (default: {DEFAULT_TIMING_RUNS})",
+        action=_StoreWithExplicitFlag,
+        explicit_attr="_explicit_timing_runs",
+        help=f"Number of timing runs (default: {DEFAULT_TIMING_RUNS}; overridden by --mode)",
     )
     trt_group.add_argument(
         "--plugin_libraries",

pyproject.toml

@@ -48,6 +48,7 @@ dependencies = [
 [project.optional-dependencies]
 onnx = [
     "cppimport",
+    "cuda-python",
     "cupy-cuda12x; platform_machine != 'aarch64' and platform_system != 'Darwin'",
     "lief",
     "ml_dtypes",

tests/_test_utils/onnx/quantization/autotune/models.py

@@ -25,9 +25,11 @@
 
 def _create_simple_conv_onnx_model():
     """Build ONNX model: Input -> Conv -> Relu -> Output (minimal for autotuner tests)."""
-    input_tensor = helper.make_tensor_value_info("input", onnx.TensorProto.FLOAT, [32, 3, 224, 224])
+    input_tensor = helper.make_tensor_value_info(
+        "input", onnx.TensorProto.FLOAT, [64, 32, 224, 224]
+    )
     output_tensor = helper.make_tensor_value_info(
-        "output", onnx.TensorProto.FLOAT, [32, 64, 224, 224]
+        "output", onnx.TensorProto.FLOAT, [64, 64, 224, 224]
     )
     conv_node = helper.make_node(
         "Conv", inputs=["input", "conv_weight"], outputs=["conv_out"], name="conv"
@@ -40,7 +42,7 @@ def _create_simple_conv_onnx_model():
         [output_tensor],
         initializer=[
             helper.make_tensor(
-                "conv_weight", onnx.TensorProto.FLOAT, [64, 3, 3, 3], [0.1] * (64 * 3 * 3 * 3)
+                "conv_weight", onnx.TensorProto.FLOAT, [64, 32, 3, 3], [0.1] * (64 * 32 * 3 * 3)
             )
         ],
     )

tests/unit/onnx/quantization/autotune/test_autotune_config.py

@@ -14,11 +14,17 @@
 # limitations under the License.
 
 """
-Tests for the Config class in the autotuner.
+Tests for the Config class and CLI mode presets in the autotuner.
 
-Tests configuration parameter validation and defaults.
+Tests configuration parameter validation, defaults, and CLI --mode preset
+selection and explicit-flag precedence.
 """
 
+from modelopt.onnx.quantization.autotune.__main__ import (
+    MODE_PRESETS,
+    _get_autotune_parser,
+    apply_mode_presets,
+)
 from modelopt.onnx.quantization.autotune.common import Config
 
 
@@ -95,3 +101,96 @@ def test_pattern_cache_params(self):
 
         assert config.pattern_cache_minimum_distance == 3
         assert config.pattern_cache_max_entries_per_pattern == 10
+
+
+class TestModePresets:
+    """Test --mode preset selection and explicit-flag precedence."""
+
+    @staticmethod
+    def _parse_cli(argv):
+        """Parse argv with the autotune CLI parser and apply mode presets."""
+        parser = _get_autotune_parser()
+        args = parser.parse_args(argv)
+        apply_mode_presets(args)
+        return args
+
+    def test_mode_quick_applies_preset_when_no_explicit_flags(self):
+        """With --mode quick and no explicit schemes/warmup/timing, preset values are used."""
+        args = self._parse_cli(["--onnx_path", "model.onnx", "--mode", "quick"])
+        preset = MODE_PRESETS["quick"]
+        assert args.num_schemes == preset["schemes_per_region"]
+        assert args.warmup_runs == preset["warmup_runs"]
+        assert args.timing_runs == preset["timing_runs"]
+
+    def test_mode_default_applies_preset_when_no_explicit_flags(self):
+        """With --mode default and no explicit flags, preset values are used."""
+        args = self._parse_cli(["--onnx_path", "model.onnx", "--mode", "default"])
+        preset = MODE_PRESETS["default"]
+        assert args.num_schemes == preset["schemes_per_region"]
+        assert args.warmup_runs == preset["warmup_runs"]
+        assert args.timing_runs == preset["timing_runs"]
+
+    def test_mode_extensive_applies_preset_when_no_explicit_flags(self):
+        """With --mode extensive and no explicit flags, preset values are used."""
+        args = self._parse_cli(["--onnx_path", "model.onnx", "--mode", "extensive"])
+        preset = MODE_PRESETS["extensive"]
+        assert args.num_schemes == preset["schemes_per_region"]
+        assert args.warmup_runs == preset["warmup_runs"]
+        assert args.timing_runs == preset["timing_runs"]
+
+    def test_explicit_schemes_per_region_overrides_mode_preset(self):
+        """Explicit --schemes_per_region is kept even when it differs from preset."""
+        args = self._parse_cli(
+            ["--onnx_path", "model.onnx", "--mode", "default", "--schemes_per_region", "99"]
+        )
+        assert args.num_schemes == 99
+        assert args.warmup_runs == MODE_PRESETS["default"]["warmup_runs"]
+        assert args.timing_runs == MODE_PRESETS["default"]["timing_runs"]
+
+    def test_explicit_default_value_not_overridden_by_mode(self):
+        """Explicit --schemes_per_region 30 (parser default) is not overridden by --mode default."""
+        args = self._parse_cli(
+            ["--onnx_path", "model.onnx", "--mode", "default", "--schemes_per_region", "30"]
+        )
+        assert args.num_schemes == 30
+
+    def test_explicit_warmup_runs_overrides_mode_preset(self):
+        """Explicit --warmup_runs is kept and not overridden by preset."""
+        args = self._parse_cli(
+            ["--onnx_path", "model.onnx", "--mode", "extensive", "--warmup_runs", "3"]
+        )
+        assert args.warmup_runs == 3
+        assert args.num_schemes == MODE_PRESETS["extensive"]["schemes_per_region"]
+        assert args.timing_runs == MODE_PRESETS["extensive"]["timing_runs"]
+
+    def test_explicit_timing_runs_overrides_mode_preset(self):
+        """Explicit --timing_runs is kept and not overridden by preset."""
+        args = self._parse_cli(
+            ["--onnx_path", "model.onnx", "--mode", "quick", "--timing_runs", "7"]
+        )
+        assert args.timing_runs == 7
+        assert args.num_schemes == MODE_PRESETS["quick"]["schemes_per_region"]
+        assert args.warmup_runs == MODE_PRESETS["quick"]["warmup_runs"]
+
+    def test_multiple_explicit_overrides_mode_preset(self):
+        """Multiple explicit flags override only their respective preset values."""
+        args = self._parse_cli(
+            [
+                "--onnx_path",
+                "model.onnx",
+                "--mode",
+                "extensive",
+                "--schemes_per_region",
+                "10",
+                "--timing_runs",
+                "5",
+            ]
+        )
+        assert args.num_schemes == 10
+        assert args.timing_runs == 5
+        assert args.warmup_runs == MODE_PRESETS["extensive"]["warmup_runs"]
+
+    def test_short_flag_schemes_per_region_overrides_mode(self):
+        """Short form -s for schemes_per_region is treated as explicit and overrides preset."""
+        args = self._parse_cli(["--onnx_path", "model.onnx", "--mode", "default", "-s", "25"])
+        assert args.num_schemes == 25