diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 43b869800d7..61c169ad0ed 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -126,6 +126,7 @@ repos: examples/llm_eval/lm_eval_hf.py| examples/llm_eval/mmlu.py| examples/llm_eval/modeling.py| + examples/onnx_ptq/far3d/evaluate.py| examples/llm_qat/train.py| examples/llm_sparsity/weight_sparsity/finetune.py| examples/specdec_bench/specdec_bench/models/specbench_medusa.py| diff --git a/LICENSE b/LICENSE index 0450e9ebffd..a894d488493 100644 --- a/LICENSE +++ b/LICENSE @@ -223,6 +223,7 @@ the following copyright holders, licensed under the Apache License, Version 2.0 Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li Copyright (c) 2024 Heming Xia Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team + Copyright (c) OpenMMLab. All rights reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use these files except in compliance with the License. You may obtain a copy diff --git a/examples/onnx_ptq/README.md b/examples/onnx_ptq/README.md index 24d3d97b410..3cee5535a84 100644 --- a/examples/onnx_ptq/README.md +++ b/examples/onnx_ptq/README.md @@ -129,6 +129,10 @@ The top5 accuracy of the model is Inference latency of the model is ms ``` +### FAR3D 3D object detection + +The [FAR3D example](./far3d/) demonstrates an end-to-end workflow that exports and quantizes the FAR3D ONNX image encoder, builds TensorRT engines, and evaluates 3D object detection mAP on the Argoverse 2 validation set. + ## Advanced Features ### Per node calibration of ONNX models diff --git a/examples/onnx_ptq/far3d/Dockerfile b/examples/onnx_ptq/far3d/Dockerfile new file mode 100644 index 00000000000..65beaf464bd --- /dev/null +++ b/examples/onnx_ptq/far3d/Dockerfile @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +FROM nvcr.io/nvidia/pytorch:25.06-py3 + +ENV LD_LIBRARY_PATH=/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + libgl1 \ + libglib2.0-0 && \ + rm -rf /var/lib/apt/lists/* + +ENV UV_PYTHON_INSTALL_DIR=/opt/python + +RUN python -m pip install --no-cache-dir uv && \ + uv python install 3.8 && \ + uv venv --seed --python 3.8 /opt/far3d + +RUN env -u PIP_CONSTRAINT /opt/far3d/bin/python -m pip install --no-cache-dir \ + torch==1.13.1+cu117 \ + torchvision==0.14.1+cu117 \ + --extra-index-url https://download.pytorch.org/whl/cu117 + +RUN env -u PIP_CONSTRAINT /opt/far3d/bin/python -m pip install --no-cache-dir \ + "numpy<1.24" \ + opencv-python==4.5.5.64 \ + yapf==0.32.0 \ + mmcv-full==1.7.0 \ + -f https://download.openmmlab.com/mmcv/dist/cu117/torch1.13.0/index.html + +RUN env -u PIP_CONSTRAINT /opt/far3d/bin/python -m pip install --no-cache-dir \ + av2==0.2.1 \ + einops \ + "ipython<9" \ + kornia==0.6.12 \ + mmdet==2.28.2 \ + mmsegmentation==0.30.0 \ + onnx \ + onnx-graphsurgeon==0.6.1 \ + onnxruntime \ + onnxsim \ + refile \ + tensorrt-cu12-bindings==10.11.0.33 \ + --extra-index-url https://pypi.nvidia.com && \ + mkdir -p /opt/far3d/lib/python3.8/site-packages/tensorrt && \ + cp /opt/far3d/lib/python3.8/site-packages/tensorrt_bindings/__init__.py \ + /opt/far3d/lib/python3.8/site-packages/tensorrt/__init__.py && \ + cp /opt/far3d/lib/python3.8/site-packages/tensorrt_bindings/tensorrt.so \ + /opt/far3d/lib/python3.8/site-packages/tensorrt/tensorrt.so + +RUN env -u PIP_CONSTRAINT /opt/far3d/bin/python -m pip install --no-cache-dir \ + "setuptools<81" && \ + env -u PIP_CONSTRAINT /opt/far3d/bin/python -m pip install --no-cache-dir \ + --no-build-isolation \ + mmdet3d==1.0.0rc6 + +COPY . /opt/modelopt +RUN env -u PIP_CONSTRAINT python -m pip install --no-cache-dir \ + lief \ + ml_dtypes \ + omegaconf \ + "onnx-graphsurgeon>=0.6.1" \ + "onnx~=1.21.0" \ + "onnxruntime-gpu~=1.24.2" \ + onnxscript \ + "onnxslim>=0.1.76" \ + "setuptools>=80" && \ + env -u PIP_CONSTRAINT python -m pip install --no-cache-dir --no-deps -e /opt/modelopt diff --git a/examples/onnx_ptq/far3d/README.md b/examples/onnx_ptq/far3d/README.md new file mode 100644 index 00000000000..8c4e0e3f0a5 --- /dev/null +++ b/examples/onnx_ptq/far3d/README.md @@ -0,0 +1,148 @@ +# FAR3D ONNX PTQ and Argoverse 2 evaluation + +This example quantizes the FAR3D image encoder and decoder to INT8 or FP8 with Model Optimizer and evaluates the complete pipeline on the Argoverse 2 validation set. It follows the [NVIDIA DL4AGX FAR3D workflow](https://github.com/NVIDIA/DL4AGX/tree/master/AV-Solutions/far3d-trt). + +FAR3D uses a legacy PyTorch/MMCV environment that is incompatible with the current Model Optimizer Python dependencies. The provided image uses `nvcr.io/nvidia/pytorch:25.06-py3` for the complete workflow and isolates the legacy FAR3D packages in a Python 3.8 virtual environment. + +## 1. Prepare FAR3D and Argoverse 2 + +Clone DL4AGX, initialize its submodules, and apply its FAR3D patch: + +```bash +git clone https://github.com/NVIDIA/DL4AGX.git +cd DL4AGX +git submodule update --init --recursive +cd AV-Solutions/far3d-trt/dependencies/Far3D +git apply ../../patch/far3d.patch +git apply /path/to/Model-Optimizer/examples/onnx_ptq/far3d/far3d_pytorch_25_06.patch +cd ../.. +``` + +The second patch makes the unused CUDA 11-only FlashAttention implementation optional; the reference configuration uses MMCV `MultiheadAttention`. + +Download the [Argoverse 2 sensor validation set](https://www.argoverse.org/av2.html), the [reference FAR3D checkpoint](https://github.com/NVIDIA/DL4AGX/tree/master/AV-Solutions/far3d-trt#pytorch-model-to-onnx), and its configuration. The remaining commands assume: + +```text +far3d-trt/ +├── data/av2/val/ +├── dependencies/Far3D/projects/configs/far3d.py +└── weights/iter_82548.pth +``` + +Build the example image from the Model Optimizer checkout: + +```bash +docker build \ + -f /path/to/Model-Optimizer/examples/onnx_ptq/far3d/Dockerfile \ + -t far3d-modelopt \ + /path/to/Model-Optimizer +``` + +Start the image and mount the FAR3D checkout: + +```bash +docker run --rm -it --network=host --gpus=all --shm-size=80G --privileged \ + -v /data/av2:/data/av2 \ + -v /path/to/far3d-trt:/workspace/far3d-trt \ + far3d-modelopt +``` + +Use `/opt/far3d/bin/python` for data preparation, export, and evaluation. It selects the isolated legacy FAR3D environment: + +```bash +export PYTHONPATH=/workspace/far3d-trt/dependencies/Far3D +cd /workspace/far3d-trt +/opt/far3d/bin/python /opt/modelopt/examples/onnx_ptq/far3d/prepare_metadata.py data/av2 +``` + +## 2. Export the ONNX models + +```bash +/opt/far3d/bin/python tools/export_onnx.py \ + dependencies/Far3D/projects/configs/far3d.py \ + weights/iter_82548.pth +``` + +This produces `far3d.encoder.onnx` and `far3d.decoder.onnx`. + +## 3. Prepare calibration batches + +Build temporary engines from the exported models. They run the reference pipeline while collecting representative encoder and decoder inputs: + +```bash +trtexec \ + --onnx=far3d.encoder.onnx \ + --saveEngine=far3d.encoder.fp16.engine \ + --fp16 \ + --skipInference +trtexec \ + --onnx=far3d.decoder.onnx \ + --saveEngine=far3d.decoder.fp16.engine \ + --stronglyTyped \ + --skipInference +``` + +Extract 512 batches sampled every 20 frames from the Argoverse 2 validation loader: + +```bash +/opt/far3d/bin/python /opt/modelopt/examples/onnx_ptq/far3d/prepare_calibration.py \ + dependencies/Far3D/projects/configs/far3d.py \ + data/far3d_calibration \ + --encoder-engine far3d.encoder.fp16.engine \ + --decoder-engine far3d.decoder.fp16.engine \ + --num-samples 512 \ + --sample-skip-interval 20 +``` + +The calibration directory contains separate `encoder/` and `decoder/` batches. Decoder batches include the image features, camera geometry, and temporal state seen by the reference decoder. + +## 4. Quantize the models + +Use the base Python environment for Model Optimizer: + +```bash +python /opt/modelopt/examples/onnx_ptq/far3d/quantize.py \ + far3d.encoder.onnx \ + far3d.decoder.onnx \ + data/far3d_calibration +``` + +Both models use max calibration. INT8 is the default; use `--quantization-mode fp8` to produce `far3d.encoder.fp8.onnx` and `far3d.decoder.fp8.onnx` instead. FP8 deployment requires an FP8-capable GPU. + +The quantizer preserves the accuracy-sensitive exclusions used by the DL4AGX reference: the `OSA4_5` block and nodes downstream of `lateral_convs` remain in high precision. + +To keep the decoder in its original mixed FP16/FP32 precision, add `--fp16-decoder`; decoder calibration batches are not required in that mode. This flag can be combined with either quantization mode. + +Build both engines in the same container. Serialized TensorRT engines are not portable across TensorRT versions or GPU architectures. + +Set the precision to the quantization mode used above: + +```bash +precision=int8 # Use fp8 for FP8 models. +trtexec \ + --onnx=far3d.encoder.${precision}.onnx \ + --saveEngine=far3d.encoder.${precision}.engine \ + --stronglyTyped \ + --skipInference +trtexec \ + --onnx=far3d.decoder.${precision}.onnx \ + --saveEngine=far3d.decoder.${precision}.engine \ + --stronglyTyped \ + --skipInference +``` + +When using `--fp16-decoder`, build `far3d.decoder.onnx` as `far3d.decoder.fp16.engine` instead. + +## 5. Evaluate accuracy + +```bash +precision=int8 # Use fp8 for FP8 models. +/opt/far3d/bin/python /opt/modelopt/examples/onnx_ptq/far3d/evaluate.py \ + dependencies/Far3D/projects/configs/far3d.py \ + far3d.encoder.${precision}.engine \ + far3d.decoder.${precision}.engine +``` + +The evaluator reports the dataset metrics, including mAP. The DL4AGX reference reports 0.230 mAP for its INT8 encoder and mixed FP16/FP32 decoder, compared with 0.232 mAP for FP16 encoder and decoder. It reports neither an INT8 decoder nor FP8 results; validate the selected precision on the target TensorRT version and GPU. + +Use `--max-samples N` for an inference smoke test. Dataset metrics are skipped when only part of the validation set is processed. diff --git a/examples/onnx_ptq/far3d/evaluate.py b/examples/onnx_ptq/far3d/evaluate.py new file mode 100644 index 00000000000..fae6093b1fd --- /dev/null +++ b/examples/onnx_ptq/far3d/evaluate.py @@ -0,0 +1,308 @@ +# Adapted from https://github.com/NVIDIA/DL4AGX/blob/9f7b29104c253d5bc68334e7b83b3eecb72d4572/AV-Solutions/far3d-trt/tools/test_tensorrt.py +# which was modified from https://github.com/megvii-research/Far3D/blob/5efb9d73a246c39fac79b3cf8c20a8e059611c3f/tools/test.py. +# Copyright (c) OpenMMLab. All rights reserved. +# Modified by Zhiqi Li. +# +# SPDX-FileCopyrightText: Copyright (c) 2023-2024, 2026 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. + +import argparse +import importlib +import os +import warnings + +import tensorrt as trt +import torch +from mmcv import Config, DictAction +from mmdet.apis import set_random_seed +from mmdet3d.core.bbox.structures.lidar_box3d import LiDARInstance3DBoxes +from mmdet3d.datasets import build_dataset +from projects.mmdet3d_plugin.datasets.builder import build_dataloader +from tqdm import tqdm + +TRT_TO_TORCH = { + trt.DataType.FLOAT: torch.float32, + trt.DataType.HALF: torch.float16, + trt.DataType.INT8: torch.int8, + trt.DataType.INT32: torch.int32, + trt.DataType.BOOL: torch.bool, + trt.DataType.UINT8: torch.uint8, +} +if int(trt.__version__.split(".")[0]) >= 10: + TRT_TO_TORCH[trt.DataType.INT64] = torch.int64 + +TRT_LOGGER = trt.Logger(trt.Logger.WARNING) +trt.init_libnvinfer_plugins(TRT_LOGGER, "") + + +def aligned_tensor(shape, dtype, device, alignment=256): + element_size = torch.empty((), dtype=dtype).element_size() + element_count = int(torch.tensor(shape).prod().item()) + storage = torch.empty(element_count + alignment // element_size, dtype=dtype, device=device) + offset_bytes = (-storage.data_ptr()) % alignment + offset = offset_bytes // element_size + return storage[offset : offset + element_count].view(shape) + + +class TensorRTRunner: + def __init__(self, engine_path, state_names=()): + with open(engine_path, "rb") as engine_file: + engine_bytes = engine_file.read() + self.engine = trt.Runtime(TRT_LOGGER).deserialize_cuda_engine(engine_bytes) + if self.engine is None: + raise RuntimeError(f"Failed to deserialize {engine_path}") + self.context = self.engine.create_execution_context() + if self.context is None: + raise RuntimeError(f"Failed to create an execution context for {engine_path}") + self.tensor_names = [ + self.engine.get_tensor_name(index) for index in range(self.engine.num_io_tensors) + ] + self.input_shapes = {} + self.output_shapes = {} + self.tensor_dtypes = {} + for name in self.tensor_names: + shape = tuple(self.engine.get_tensor_shape(name)) + dtype = TRT_TO_TORCH[self.engine.get_tensor_dtype(name)] + self.tensor_dtypes[name] = dtype + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + self.input_shapes[name] = shape + else: + self.output_shapes[name] = shape + + self.state = {} + for base_name in state_names: + name = self.resolve_name(base_name) + if name in self.input_shapes: + tensor = aligned_tensor(self.input_shapes[name], self.tensor_dtypes[name], "cuda") + tensor.zero_() + self.state[name] = tensor + self.context.set_tensor_address(name, tensor.data_ptr()) + if self.state: + torch.cuda.synchronize() + + def resolve_name(self, base_name): + if base_name in self.tensor_names: + return base_name + suffixed_name = f"{base_name}.1" + return suffixed_name if suffixed_name in self.tensor_names else base_name + + def reset_state(self): + for tensor in self.state.values(): + tensor.zero_() + + def prepare_input(self, name, inputs): + shape = self.input_shapes[name] + base_name = name.rsplit(".1", maxsplit=1)[0] if name.endswith(".1") else name + if base_name not in inputs: + raise KeyError(f"Missing TensorRT input {base_name}") + value = inputs[base_name].to(device="cuda", dtype=self.tensor_dtypes[name]) + if tuple(value.shape) != shape: + if tuple(value.shape[1:]) == shape: + value = value.squeeze(0) + elif tuple(shape[1:]) == tuple(value.shape): + value = value.unsqueeze(0) + else: + raise ValueError( + f"Input {base_name} has shape {tuple(value.shape)}, expected {shape}" + ) + return value + + def __call__(self, stream, **inputs): + input_buffers = {} + for name, shape in self.input_shapes.items(): + if name in self.state: + continue + value = self.prepare_input(name, inputs) + buffer = aligned_tensor(shape, value.dtype, value.device) + buffer.copy_(value) + input_buffers[name] = buffer + self.context.set_tensor_address(name, buffer.data_ptr()) + + outputs = {} + for name, shape in self.output_shapes.items(): + output = aligned_tensor(shape, self.tensor_dtypes[name], "cuda") + outputs[name] = output + self.context.set_tensor_address(name, output.data_ptr()) + + if not self.context.execute_async_v3(stream.cuda_stream): + raise RuntimeError("TensorRT execution failed") + stream.synchronize() + return outputs + + +STATE_NAMES = ( + "memory_embedding", + "memory_reference_point", + "memory_egopose", + "memory_velo", + "memory_timestamp", +) + + +class Far3DDecoderRunner(TensorRTRunner): + def __init__(self, engine_path, input_callback=None): + super().__init__(engine_path, STATE_NAMES) + self.input_callback = input_callback + self.scene_token = None + self.timestamp_offset = None + + def __call__(self, stream, img_metas, timestamp, **inputs): + scene_token = img_metas[0].data[0][0]["scene_token"] + new_scene = self.scene_token != scene_token + if new_scene: + self.reset_state() + self.scene_token = scene_token + self.timestamp_offset = timestamp.clone() + prev_exists_name = self.resolve_name("prev_exists") + if prev_exists_name in self.input_shapes: + inputs["prev_exists"] = torch.full( + self.input_shapes[prev_exists_name], + not new_scene, + dtype=self.tensor_dtypes[prev_exists_name], + device="cuda", + ) + inputs["timestamp"] = (timestamp - self.timestamp_offset).float() + if self.input_callback: + calibration_inputs = {} + for name in self.input_shapes: + base_name = name.rsplit(".1", maxsplit=1)[0] if name.endswith(".1") else name + if name in self.state: + value = self.state[name] + else: + value = self.prepare_input(name, inputs) + calibration_inputs[base_name] = value + self.input_callback(calibration_inputs) + outputs = super().__call__(stream, **inputs) + for base_name in STATE_NAMES: + input_name = self.resolve_name(base_name) + output_name = f"{base_name}_out" + if input_name in self.state and output_name in outputs: + state_length = self.state[input_name].shape[1] + self.state[input_name].copy_(outputs[output_name][:, :state_length]) + return outputs + + +class Far3DPipeline: + def __init__(self, encoder_engine, decoder_engine, decoder_input_callback=None): + self.encoder = TensorRTRunner(encoder_engine) + self.decoder = Far3DDecoderRunner(decoder_engine, decoder_input_callback) + + @staticmethod + def unpack(data): + lidar2img = data["lidar2img"][0].data[0][0].unsqueeze(0).cuda() + return { + "img": data["img"][0].data[0].flip(2).permute(0, 1, 3, 4, 2).contiguous().cuda(), + "intrinsics": data["intrinsics"][0].data[0][0].unsqueeze(0).cuda(), + "extrinsics": data["extrinsics"][0].data[0][0].unsqueeze(0).cuda(), + "lidar2img": lidar2img, + "img2lidar": lidar2img.inverse(), + "ego_pose": data["ego_pose"][0].data[0][0].unsqueeze(0).cuda(), + "ego_pose_inv": data["ego_pose_inv"][0].data[0][0].unsqueeze(0).cuda(), + "pad_shape": torch.tensor(data["img_metas"][0].data[0][0]["pad_shape"][0]).cuda(), + "timestamp": torch.tensor(data["timestamp"][0].data[0][0]).cuda(), + } + + def __call__(self, stream, data): + with torch.cuda.stream(stream): + inputs = self.unpack(data) + image_features = self.encoder(stream, **inputs) + decoder_inputs = dict(data) + decoder_inputs.update(inputs) + decoder_inputs.update(image_features) + return self.decoder(stream, **decoder_inputs) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Evaluate FAR3D TensorRT engines on Argoverse 2") + parser.add_argument("config", help="Path to the FAR3D configuration file") + parser.add_argument("encoder_engine") + parser.add_argument("decoder_engine") + parser.add_argument("--cfg-options", nargs="+", action=DictAction) + parser.add_argument("--eval-options", nargs="+", action=DictAction) + parser.add_argument("--options", nargs="+", action=DictAction) + parser.add_argument("--max-samples", type=int) + args = parser.parse_args() + if args.max_samples is not None and args.max_samples < 1: + raise ValueError("--max-samples must be positive") + if args.options and args.eval_options: + raise ValueError("--options and --eval-options cannot both be specified") + if args.options: + warnings.warn("--options is deprecated; use --eval-options", stacklevel=2) + args.eval_options = args.options + return args + + +def import_plugin(cfg): + plugin_dir = os.path.dirname(cfg.plugin_dir).split("/") + importlib.import_module(".".join(plugin_dir)) + + +def main(): + args = parse_args() + cfg = Config.fromfile(args.config) + if args.cfg_options: + cfg.merge_from_dict(args.cfg_options) + if cfg.get("custom_imports"): + from mmcv.utils import import_modules_from_strings + + import_modules_from_strings(**cfg.custom_imports) + import_plugin(cfg) + + cfg.model.pretrained = None + cfg.data.test.test_mode = True + set_random_seed(0, deterministic=False) + dataset = build_dataset(cfg.data.test) + data_loader = build_dataloader( + dataset, + samples_per_gpu=1, + workers_per_gpu=cfg.data.workers_per_gpu, + dist=False, + shuffle=False, + nonshuffler_sampler=cfg.data.nonshuffler_sampler, + ) + + pipeline = Far3DPipeline(args.encoder_engine, args.decoder_engine) + stream = torch.cuda.Stream() + outputs = [] + for data in tqdm(data_loader): + result = pipeline(stream, data) + boxes = LiDARInstance3DBoxes(result["bboxes"].cpu()) + outputs.append( + { + "pts_bbox": { + "boxes_3d": boxes, + "scores_3d": result["scores"].cpu(), + "labels_3d": result["labels"].cpu(), + } + } + ) + if args.max_samples is not None and len(outputs) == args.max_samples: + break + + if len(outputs) < len(dataset): + print(f"Processed {len(outputs)} samples; skipping dataset metrics") + return + + eval_kwargs = cfg.get("evaluation", {}).copy() + for key in ("interval", "tmpdir", "start", "gpu_collect", "save_best", "rule"): + eval_kwargs.pop(key, None) + if args.eval_options: + eval_kwargs.update(args.eval_options) + print(dataset.evaluate(outputs, **eval_kwargs)) + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("fork") + main() diff --git a/examples/onnx_ptq/far3d/far3d_pytorch_25_06.patch b/examples/onnx_ptq/far3d/far3d_pytorch_25_06.patch new file mode 100644 index 00000000000..22d775a7e17 --- /dev/null +++ b/examples/onnx_ptq/far3d/far3d_pytorch_25_06.patch @@ -0,0 +1,17 @@ +--- a/projects/mmdet3d_plugin/models/utils/petr_transformer.py ++++ b/projects/mmdet3d_plugin/models/utils/petr_transformer.py +@@ -17,3 +17,6 @@ + from torch.nn import ModuleList +-from .attention import FlashMHA ++try: ++ from .attention import FlashMHA ++except ImportError: ++ FlashMHA = None + import torch.utils.checkpoint as cp +@@ -65,3 +68,5 @@ + self.batch_first = True +- ++ if FlashMHA is None: ++ raise ImportError("flash-attn is required for PETRMultiheadFlashAttention") ++ + self.attn = FlashMHA(embed_dims, num_heads, attn_drop, dtype=torch.float16, device='cuda', diff --git a/examples/onnx_ptq/far3d/prepare_calibration.py b/examples/onnx_ptq/far3d/prepare_calibration.py new file mode 100644 index 00000000000..86a17f88588 --- /dev/null +++ b/examples/onnx_ptq/far3d/prepare_calibration.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +import argparse +from pathlib import Path + +import numpy as np +import torch +from evaluate import Far3DPipeline +from mmcv import Config +from mmdet.datasets import replace_ImageToTensor +from mmdet3d.datasets import build_dataset +from projects.mmdet3d_plugin.datasets.builder import build_dataloader +from torch.utils.data import Subset + + +def parse_args(): + parser = argparse.ArgumentParser(description="Prepare FAR3D calibration batches") + parser.add_argument("config", help="Path to the FAR3D configuration file") + parser.add_argument("output_dir", type=Path) + parser.add_argument("--encoder-engine") + parser.add_argument("--decoder-engine") + parser.add_argument("--num-samples", type=int, default=512) + parser.add_argument("--sample-skip-interval", type=int, default=20) + return parser.parse_args() + + +def build_validation_loader(config_path, num_samples, sample_skip_interval): + cfg = Config.fromfile(config_path) + samples_per_gpu = 1 + if isinstance(cfg.data.test, dict): + cfg.data.test.test_mode = True + samples_per_gpu = cfg.data.test.pop("samples_per_gpu", 1) + if samples_per_gpu > 1: + cfg.data.test.pipeline = replace_ImageToTensor(cfg.data.test.pipeline) + else: + for dataset_cfg in cfg.data.test: + dataset_cfg.test_mode = True + samples_per_gpu = max( + dataset_cfg.pop("samples_per_gpu", 1) for dataset_cfg in cfg.data.test + ) + if samples_per_gpu > 1: + for dataset_cfg in cfg.data.test: + dataset_cfg.pipeline = replace_ImageToTensor(dataset_cfg.pipeline) + + dataset = build_dataset(cfg.data.test) + sample_indices = range( + sample_skip_interval - 1, + min(len(dataset), num_samples * sample_skip_interval), + sample_skip_interval, + ) + dataset = Subset(dataset, sample_indices) + return build_dataloader( + dataset, + samples_per_gpu=samples_per_gpu, + workers_per_gpu=cfg.data.workers_per_gpu, + dist=False, + shuffle=False, + nonshuffler_sampler=cfg.data.nonshuffler_sampler, + ) + + +class DecoderCalibrationWriter: + def __init__(self, output_dir): + self.output_dir = output_dir + self.saved = 0 + + def __call__(self, inputs): + batch = {name: value.detach().cpu().numpy() for name, value in inputs.items()} + np.savez(self.output_dir / f"batch_{self.saved:04d}.npz", **batch) + self.saved += 1 + + +def main(): + args = parse_args() + if args.num_samples < 1: + raise ValueError("--num-samples must be positive") + if args.sample_skip_interval < 1: + raise ValueError("--sample-skip-interval must be positive") + if bool(args.encoder_engine) != bool(args.decoder_engine): + raise ValueError("--encoder-engine and --decoder-engine must be specified together") + + encoder_dir = args.output_dir / "encoder" + encoder_dir.mkdir(parents=True, exist_ok=True) + if any(encoder_dir.glob("*.npy")): + raise FileExistsError( + f"{encoder_dir} already contains calibration batches; use an empty directory" + ) + + decoder_writer = pipeline = None + if args.encoder_engine: + decoder_dir = args.output_dir / "decoder" + decoder_dir.mkdir(parents=True, exist_ok=True) + if any(decoder_dir.glob("*.npz")): + raise FileExistsError( + f"{decoder_dir} already contains calibration batches; use an empty directory" + ) + decoder_writer = DecoderCalibrationWriter(decoder_dir) + pipeline = Far3DPipeline( + args.encoder_engine, + args.decoder_engine, + decoder_input_callback=decoder_writer, + ) + stream = torch.cuda.Stream() + + saved = 0 + data_loader = build_validation_loader(args.config, args.num_samples, args.sample_skip_interval) + for data in data_loader: + images = data["img"][0].data[0].cpu().permute(0, 1, 3, 4, 2).numpy() + np.save(encoder_dir / f"batch_{saved:04d}.npy", images) + if pipeline: + pipeline(stream, data) + saved += 1 + if saved == args.num_samples: + break + + if saved < args.num_samples: + raise RuntimeError( + f"Only prepared {saved} of {args.num_samples} requested calibration batches" + ) + if decoder_writer and decoder_writer.saved != saved: + raise RuntimeError(f"Prepared {saved} encoder and {decoder_writer.saved} decoder batches") + print(f"Saved {saved} encoder calibration batches to {encoder_dir}") + if decoder_writer: + print( + f"Saved {decoder_writer.saved} decoder calibration batches to {decoder_writer.output_dir}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/onnx_ptq/far3d/prepare_metadata.py b/examples/onnx_ptq/far3d/prepare_metadata.py new file mode 100644 index 00000000000..4daaee7b516 --- /dev/null +++ b/examples/onnx_ptq/far3d/prepare_metadata.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +import argparse +from pathlib import Path + +import pandas as pd +from av2.utils.io import read_feather +from tools.create_infos_av2.create_av2_infos import create_av2_infos + + +def parse_args(): + parser = argparse.ArgumentParser(description="Prepare FAR3D Argoverse 2 validation metadata") + parser.add_argument("dataset_dir", type=Path, help="Argoverse 2 root containing val/") + return parser.parse_args() + + +def main(): + args = parse_args() + info_path = args.dataset_dir / "av2_val_infos.pkl" + annotation_path = args.dataset_dir / "val_anno.feather" + for output_path in (info_path, annotation_path): + if output_path.exists(): + raise FileExistsError(f"Refusing to overwrite {output_path}") + + create_av2_infos(dataset_dir=args.dataset_dir, split="val", out_dir=args.dataset_dir) + generated_info_path = args.dataset_dir / "av2_val_infos_mini.pkl" + generated_info_path.replace(info_path) + + annotations = [] + for path in sorted((args.dataset_dir / "val").glob("*/annotations.feather")): + frame = read_feather(path) + frame["log_id"] = path.parent.name + annotations.append(frame) + if not annotations: + raise RuntimeError(f"No validation annotations found under {args.dataset_dir / 'val'}") + pd.concat(annotations).reset_index().to_feather(annotation_path) + print(f"Saved {info_path} and {annotation_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/onnx_ptq/far3d/quantize.py b/examples/onnx_ptq/far3d/quantize.py new file mode 100644 index 00000000000..0a98e2ebf53 --- /dev/null +++ b/examples/onnx_ptq/far3d/quantize.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +import argparse +from pathlib import Path + +import numpy as np +import onnx +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.quantization import quantize +from modelopt.onnx.utils import topologically_sort_graph_nodes + + +class FileCalibrationReader(CalibrationDataReader): + def __init__(self, calibration_dir, pattern): + self.batch_paths = sorted(Path(calibration_dir).glob(pattern)) + if not self.batch_paths: + raise ValueError(f"No {pattern} calibration batches found in {calibration_dir}") + self.rewind() + + def get_next(self): + batch_path = next(self._iterator, None) + return None if batch_path is None else self.load(batch_path) + + def get_first(self): + return self.load(self.batch_paths[0]) + + def rewind(self): + self._iterator = iter(self.batch_paths) + + def load(self, batch_path): + raise NotImplementedError + + +class EncoderCalibrationReader(FileCalibrationReader): + def __init__(self, calibration_dir): + super().__init__(calibration_dir, "*.npy") + + def load(self, batch_path): + return {"img": np.load(batch_path)} + + +class DecoderCalibrationReader(FileCalibrationReader): + def __init__(self, calibration_dir, onnx_path): + graph = onnx.load(onnx_path, load_external_data=False).graph + self.input_dtypes = { + value.name: onnx.helper.tensor_dtype_to_np_dtype(value.type.tensor_type.elem_type) + for value in graph.input + } + super().__init__(calibration_dir, "*.npz") + + def load(self, batch_path): + with np.load(batch_path) as batch: + missing = self.input_dtypes.keys() - batch.files + if missing: + raise ValueError(f"{batch_path} is missing decoder inputs: {sorted(missing)}") + return { + name: batch[name].astype(dtype, copy=False) + for name, dtype in self.input_dtypes.items() + } + + +def find_encoder_nodes_to_exclude(onnx_path): + graph = onnx.load(onnx_path, load_external_data=False).graph + topologically_sort_graph_nodes(graph) + + excluded = set() + downstream_tensors = set() + for node in graph.node: + is_osa = "OSA4_5" in node.name + is_downstream = any(name in downstream_tensors for name in node.input) + if is_osa or is_downstream: + excluded.add(node.name) + if "lateral_convs" in node.name or (is_downstream and not is_osa): + downstream_tensors.update(node.output) + return sorted(excluded) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Quantize the FAR3D ONNX models") + parser.add_argument("encoder_onnx", help="Path to far3d.encoder.onnx") + parser.add_argument("decoder_onnx", help="Path to far3d.decoder.onnx") + parser.add_argument("calibration_dir", help="Directory created by prepare_calibration.py") + parser.add_argument("--quantization-mode", choices=("int8", "fp8"), default="int8") + parser.add_argument("--encoder-output") + parser.add_argument("--decoder-output") + parser.add_argument( + "--fp16-decoder", + action="store_true", + help="Skip decoder quantization and use the original mixed-precision decoder", + ) + return parser.parse_args() + + +def quantize_encoder(args): + encoder_dir = Path(args.calibration_dir) + if (encoder_dir / "encoder").is_dir(): + encoder_dir /= "encoder" + excluded_nodes = find_encoder_nodes_to_exclude(args.encoder_onnx) + print(f"Excluding {len(excluded_nodes)} accuracy-sensitive nodes from quantization") + quantize( + onnx_path=args.encoder_onnx, + quantize_mode=args.quantization_mode, + calibration_data_reader=EncoderCalibrationReader(encoder_dir), + calibration_method="max", + calibration_eps=["cuda:0", "cpu"], + nodes_to_exclude=excluded_nodes, + high_precision_dtype="fp16", + output_path=args.encoder_output, + ) + + +def quantize_decoder(args): + decoder_dir = Path(args.calibration_dir) / "decoder" + quantize( + onnx_path=args.decoder_onnx, + quantize_mode=args.quantization_mode, + calibration_data_reader=DecoderCalibrationReader(decoder_dir, args.decoder_onnx), + calibration_method="max", + calibration_eps=["cuda:0", "cpu"], + high_precision_dtype="fp16" if args.quantization_mode == "fp8" else "fp32", + output_path=args.decoder_output, + ) + + +def main(): + args = parse_args() + if args.encoder_output is None: + args.encoder_output = f"far3d.encoder.{args.quantization_mode}.onnx" + if args.decoder_output is None: + args.decoder_output = f"far3d.decoder.{args.quantization_mode}.onnx" + quantize_encoder(args) + if args.fp16_decoder: + print("Skipping decoder quantization; use the original mixed-precision decoder ONNX") + else: + quantize_decoder(args) + + +if __name__ == "__main__": + main()