From 79a5d1d0142fe73c0918a465ed41c24fe66c33ae Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:11:11 +0000 Subject: [PATCH 1/2] Drop DataLoader workers from the torch_onnx calibration loaders load_calibration_data builds its DataLoaders over lists of tensors that are already on the GPU, so the four workers load nothing. They only add cost: the script forced the "spawn" start method (CUDA cannot be used in forked workers), so every loader pass spawned four processes that re-imported torch, timm and modelopt and received the dataset over CUDA IPC. That is about 14 s per pass and about 300 s per torch_onnx CI job. Use num_workers=0 in both loaders and drop the module-level mp.set_start_method("spawn", force=True). The latter also stops importing the script from changing the start method of the importing process (two tests import it inside the pytest process). Nothing else in the script uses multiprocessing. The --evaluate path keeps the CPU-only DataLoader workers of examples/onnx_ptq/evaluation.py; they now run with the default start method, exactly as in the onnx_ptq example that calls the same function. Testing: py_compile and ruff 0.15.20 (check and format) on the script, and read the file for other users of mp and of the start method (none). Not run: tests/examples/torch_onnx needs a GPU, TensorRT and model downloads, so GPU CI validation is still needed. Signed-off-by: Shengliang Xu --- examples/torch_onnx/torch_quant_to_onnx.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index bd935146e0e..a08338d053e 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -27,7 +27,6 @@ import timm import torch -import torch.multiprocessing as mp import torch.nn.functional as F from datasets import load_dataset from download_example_onnx import export_to_onnx @@ -56,9 +55,6 @@ """ -mp.set_start_method("spawn", force=True) # Needed for data loader with multiple workers - - _FP8_CONV_OVERRIDE: list = [ { "parent_class": "nn.Conv2d", @@ -311,16 +307,17 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals calib_tensor = [transforms(img) for img in images] calib_tensor = [t.to(device) for t in calib_tensor] + # The samples already sit on ``device``; worker processes would only add spawn and import cost. if with_labels: labels = dataset["train"][:data_size]["label"] labels = torch.tensor(labels, device=device) calib_dataset = [{"image": img, "label": lbl} for img, lbl in zip(calib_tensor, labels)] return torch.utils.data.DataLoader( - calib_dataset, batch_size=batch_size, shuffle=True, num_workers=4 + calib_dataset, batch_size=batch_size, shuffle=True, num_workers=0 ) else: return torch.utils.data.DataLoader( - calib_tensor, batch_size=batch_size, shuffle=True, num_workers=4 + calib_tensor, batch_size=batch_size, shuffle=True, num_workers=0 ) From cae95d32277059753e579931cdc7f4c14b594716 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:11:23 +0000 Subject: [PATCH 2/2] Skip swinv2 int8/mxfp8/auto on PRs and run trtexec for one iteration Two cuts to the torch_onnx example lane. swinv2_tiny runs all five formats on top of swin_tiny, which already covers every format through the same exporter and TensorRT path; swinv2 only adds its cosine attention and post-norm blocks. On PRs (GITHUB_REF under refs/heads/pull-request/) skip swinv2 int8, mxfp8 and auto and keep fp8 and nvfp4, which still exercise those blocks. Scheduled and local runs keep every combination, and test ids are unchanged. swinv2_tiny[auto] alone averages 121 s over 304 CI runs (the 22nd slowest example test); the three together cost about 110 s of every PR run. build_trt_engine let trtexec run its default benchmark after the build (200 ms warm-up, then at least 3 s and 10 iterations), and the output is discarded on success. Run a single iteration instead (--warmUp=0 --duration=0 --iterations=1): the engine still executes once, but about 3 s of timing per engine goes away, roughly 40 s per lane run. --trt_build users see no difference besides it returning sooner. Testing: ruff 0.15.20 (check and format), py_compile and pytest --collect-only on CPU (24 tests, ids unchanged). With the example command stubbed, GITHUB_REF=refs/heads/pull-request/123 skips exactly swinv2 int8, mxfp8 and auto, and any other ref skips none. The two in-process mapping tests that import the script pass on CPU, and build_trt_engine passes the three flags to a fake trtexec. Not run: the lane itself (needs a GPU, TensorRT and model downloads) and the flags against a real trtexec, so GPU CI validation is still needed. Signed-off-by: Shengliang Xu --- examples/torch_onnx/torch_quant_to_onnx.py | 4 ++++ tests/examples/torch_onnx/test_torch_quant_to_onnx.py | 7 +++++++ 2 files changed, 11 insertions(+) diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index a08338d053e..a9a441e1a6d 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -758,6 +758,10 @@ def build_trt_engine(onnx_path, optimization_level=4): f"--onnx={onnx_path}", "--stronglyTyped", f"--builderOptimizationLevel={optimization_level}", + # One inference iteration shows the engine runs; the default 3 s timed benchmark is not shown. + "--warmUp=0", + "--duration=0", + "--iterations=1", ] print(f"\nBuilding TensorRT engine: {' '.join(cmd)}") try: diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index 28de24e9429..9fa4aa958d1 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -15,6 +15,7 @@ import importlib +import os from collections import defaultdict from pathlib import Path @@ -38,6 +39,10 @@ "resnet50": ("resnet50", None), } +# PRs run swinv2 for fp8 and nvfp4 only (swin_tiny covers every format); scheduled runs keep all. +_PR_CI = os.environ.get("GITHUB_REF", "").startswith("refs/heads/pull-request/") +_SKIPPED_ON_PR = {("swinv2_tiny", "int8"), ("swinv2_tiny", "mxfp8"), ("swinv2_tiny", "auto")} + def _assert_residual_inputs_are_quantized(onnx_save_path): model = onnx.load(onnx_save_path) @@ -74,6 +79,8 @@ def _assert_residual_inputs_are_quantized(onnx_save_path): def test_torch_onnx(tmp_path, model_key, qformat): if model_key == "resnet50" and qformat not in _RESNET_RECIPE_QFORMATS: pytest.skip("Only FP8 and INT8 quantization are supported for ResNet") + if _PR_CI and (model_key, qformat) in _SKIPPED_ON_PR: + pytest.skip("Covered by swin_tiny on PRs; runs in the scheduled CI") timm_model_name, model_kwargs = _MODELS[model_key] onnx_save_path = tmp_path / f"{model_key}.{qformat}.onnx"