[OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models

tests/_test_utils/onnx/lib_test_models.py

@@ -260,12 +260,12 @@ def find_init(onnx_model: onnx.ModelProto, init_name: str) -> np.ndarray:
 #           conv_4
 #             |
 #
-def build_conv_concat_model():
+def build_conv_concat_model(channels=128, height=240, width=64):
     # Define your model inputs and outputs
     input_names = ["input_0"]
     output_names = ["conv4_conv/Conv2D:0"]
-    input_shapes = [(1, 128, 240, 64)]
-    output_shapes = [(1, 256, 240, 64)]
+    input_shapes = [(1, channels, height, width)]
+    output_shapes = [(1, 2 * channels, height, width)]
 
     inputs = [
         helper.make_tensor_value_info(input_name, onnx.TensorProto.FLOAT, input_shape)
@@ -340,26 +340,26 @@ def build_conv_concat_model():
         helper.make_tensor(
             name="ReadVariableOp:0",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(128, 128, 3, 3),
-            vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3),
+            dims=(channels, channels, 3, 3),
+            vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3),
         ),
         helper.make_tensor(
             name="ReadVariableOp:1",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(128, 128, 3, 3),
-            vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3),
+            dims=(channels, channels, 3, 3),
+            vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3),
         ),
         helper.make_tensor(
             name="ReadVariableOp:2",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(128, 128, 3, 3),
-            vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3),
+            dims=(channels, channels, 3, 3),
+            vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3),
         ),
         helper.make_tensor(
             name="ReadVariableOp:3",
             data_type=onnx.TensorProto.FLOAT,
-            dims=(256, 384, 1, 1),
-            vals=np.random.uniform(low=0.5, high=1.0, size=256 * 384 * 1 * 1),
+            dims=(2 * channels, 3 * channels, 1, 1),
+            vals=np.random.uniform(low=0.5, high=1.0, size=2 * channels * 3 * channels),
         ),
     ]
 

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

@@ -25,13 +25,13 @@
 from onnx import helper
 
 
-def _create_simple_conv_onnx_model():
+def _create_simple_conv_onnx_model(batch_size=64, spatial_size=224):
     """Build ONNX model: Input -> Conv -> Relu -> Output (minimal for autotuner tests)."""
     input_tensor = helper.make_tensor_value_info(
-        "input", onnx.TensorProto.FLOAT, [64, 32, 224, 224]
+        "input", onnx.TensorProto.FLOAT, [batch_size, 32, spatial_size, spatial_size]
     )
     output_tensor = helper.make_tensor_value_info(
-        "output", onnx.TensorProto.FLOAT, [64, 64, 224, 224]
+        "output", onnx.TensorProto.FLOAT, [batch_size, 64, spatial_size, spatial_size]
     )
     conv_node = helper.make_node(
         "Conv",

tests/gpu/onnx/quantization/autotune/test_benchmark.py

@@ -46,8 +46,8 @@
 
 @pytest.fixture
 def simple_conv_model_bytes():
-    """ONNX model bytes: Input -> Conv -> Relu -> Output (from _test_utils)."""
-    model = _create_simple_conv_onnx_model()
+    """ONNX model bytes: Input -> Conv -> Relu -> Output (from _test_utils), small input."""
+    model = _create_simple_conv_onnx_model(batch_size=1, spatial_size=32)
     return model.SerializeToString()
 
 
@@ -431,7 +431,12 @@ def test_run_with_path(self, simple_conv_model_path):
         with tempfile.NamedTemporaryFile(suffix=".cache", delete=False) as f:
             cache_path = f.name
         try:
-            benchmark = TrtExecBenchmark(timing_cache_file=cache_path, warmup_runs=1, timing_runs=2)
+            benchmark = TrtExecBenchmark(
+                timing_cache_file=cache_path,
+                warmup_runs=1,
+                timing_runs=2,
+                trtexec_args=["--duration=0"],
+            )
             latency_ms = benchmark.run(simple_conv_model_path)
             assert isinstance(latency_ms, float)
             assert latency_ms > 0
@@ -445,7 +450,12 @@ def test_run_with_bytes(self, simple_conv_model_bytes):
         with tempfile.NamedTemporaryFile(suffix=".cache", delete=False) as f:
             cache_path = f.name
         try:
-            benchmark = TrtExecBenchmark(timing_cache_file=cache_path, warmup_runs=1, timing_runs=2)
+            benchmark = TrtExecBenchmark(
+                timing_cache_file=cache_path,
+                warmup_runs=1,
+                timing_runs=2,
+                trtexec_args=["--duration=0"],
+            )
             latency_ms = benchmark.run(simple_conv_model_bytes)
             assert isinstance(latency_ms, float)
             assert latency_ms > 0

tests/gpu/onnx/quantization/autotune/test_workflow.py

@@ -59,7 +59,10 @@ def test_export_quantized_model(use_trtexec, simple_conv_model):
     output_path = output_dir + ".quant.onnx"
 
     try:
-        init_benchmark_instance(use_trtexec=use_trtexec, timing_runs=100)
+        # trtexec measures for at least 3 s by default; --duration=0 stops after timing_runs iterations
+        init_benchmark_instance(
+            use_trtexec=use_trtexec, timing_runs=100, trtexec_args=["--duration=0"]
+        )
         autotuner = region_pattern_autotuning_workflow(baseline_model_path, Path(output_dir))
 
         # Export model with Q/DQ insertion

tests/gpu/onnx/quantization/test_concat_elim.py

@@ -23,10 +23,19 @@
 
 from modelopt.onnx.quantization.quantize import quantize
 
+# Keep every Conv above the 16-channel limit for INT8/FP8 quantization (graph_selection.py) so the
+# same nodes are quantized as with the default 128-channel model, but calibrate a small input.
+SMALL_MODEL_SHAPE = {"channels": 32, "height": 16, "width": 16}
+
 
 def _check_concat_qdq_status(onnx_path, quantize_mode):
-    # Quantize the input model
-    quantize(onnx_path, quantize_mode=quantize_mode, passes="concat_elimination")
+    # Quantize the input model. Only the graph rewrite is checked, so calibrate on CPU.
+    quantize(
+        onnx_path,
+        quantize_mode=quantize_mode,
+        passes="concat_elimination",
+        calibration_eps=["cpu"],
+    )
 
     # Output model should be produced in the same tmpdir
     output_onnx_path = onnx_path.replace(".onnx", ".quant.onnx")
@@ -61,7 +70,7 @@ def _check_concat_qdq_status(onnx_path, quantize_mode):
 
 def test_concat_elim_int8(tmp_path):
     # Copy the test model to the tmp_path
-    model = build_conv_concat_model()
+    model = build_conv_concat_model(**SMALL_MODEL_SHAPE)
     this_function_name = sys._getframe().f_code.co_name
     onnx_path = os.path.join(tmp_path, f"{this_function_name}.onnx")
     onnx.save(model, onnx_path)
@@ -70,7 +79,7 @@ def test_concat_elim_int8(tmp_path):
 
 def test_concat_elim_fp8(tmp_path):
     # Copy the test model to the tmp_path
-    model = build_conv_concat_model()
+    model = build_conv_concat_model(**SMALL_MODEL_SHAPE)
     this_function_name = sys._getframe().f_code.co_name
     onnx_path = os.path.join(tmp_path, f"{this_function_name}.onnx")
     onnx.save(model, onnx_path)