lightly-ai / lightly-ai/lightly-train

[FEAT] TensorRT INT8 quantization support + TRT 11.x compatibility

Open
#935 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
1.7k
Forks
116
Avg merge
2d 21h
Merged PRs (30d)
6

Description

## Summary

This issue proposes adding TensorRT INT8 post-training quantization (PTQ) support to LightlyTrain's export pipeline, and also addresses a pre-existing TRT 11.x compatibility regression in `tensorrt_helpers.py`. The existing export pipeline supports FP32 and FP16 precision for both ONNX and TensorRT engines, but has no INT8 path.

## Part 1: TRT 11.x Compatibility Fix (bug)

### Problem

`tensorrt_helpers.py` uses several TensorRT APIs that were removed in TRT 11.x:

| API | Status in TRT 11.x |
|-----|---------------------|
| `NetworkDefinitionCreationFlag.EXPLICIT_BATCH` | **Removed** (explicit batch is now the only mode) |
| `BuilderFlag.FP16` | **Removed** |
| `BuilderFlag.INT8` | **Removed** |
| `BuilderFlag.OBEY_PRECISION_CONSTRAINTS` | **Removed** |

In TRT 11.x, precision is controlled entirely at the ONNX graph level (tensor types and QDQ nodes) rather than through builder flags. The `STRONGLY_TYPED` network flag still exists.

### Reproduction

```python
import tensorrt as trt # TRT 11.2.1.2
trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH # AttributeError
trt.BuilderFlag.FP16 # AttributeError
```

### Proposed Fix

Guard removed APIs with `hasattr` checks so the code works on both TRT 10.x and 11.x:

```python
# EXPLICIT_BATCH (line 139)
network_flags = 0
if hasattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH"):
network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

# FP16 builder flags (lines 208-216)
if hasattr(trt.BuilderFlag, "FP16") and builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
if hasattr(trt.BuilderFlag, "OBEY_PRECISION_CONSTRAINTS"):
config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
```

Also: `builder.build_serialized_network()` returns `IHostMemory` in TRT 11 which doesn't support `len()` — use `bytes(engine)` instead.

## Part 2: INT8 Post-Training Quantization via NVIDIA ModelOpt (feature)

I evaluated two ModelOpt quantization paths

**Path A — PyTorch-level (`mtq.quantize()`):** Inserts custom `tensorrt.quantize_op` operators into the PyTorch model. These custom ops are **not supported** by `torch.onnx.export()` (dynamo exporter), which raises:

```
DispatchError: No ONNX function found for
```

This path is designed for LLM deployment via TensorRT-LLM/vLLM, not for vision model ONNX pipelines. **Not viable for LightlyTrain.**

**Path B — ONNX-level (`modelopt.onnx.quantization.quantize()`):** Quantizes an already-exported FP32 ONNX model by inserting standard `QuantizeLinear`/`DequantizeLinear` nodes directly into the ONNX graph. TRT's ONNX parser auto-detects these QDQ nodes and fuses them into native INT8 kernels. **This is a viable approach because:

- Plugs into the existing ONNX-first export pipeline seamlessly
- No changes needed to PyTorch model code or ONNX export flow
- Just adds a post-processing step between ONNX export and TRT engine build
- Uses standard ONNX QDQ nodes (not custom ops)
- Calibration uses numpy arrays, easy to generate from existing data loaders

Tested the full pipeline on an H200 MIG (CUDA 13.0, PyTorch 2.13.0, TRT 11.2.1.2, ModelOpt 0.45.0):

```
FP32 ONNX (via ONNXExportMixin) → modelopt.onnx.quantization.quantize() → QDQ ONNX → TRT engine build → INT8 inference
```

Results:

- QDQ nodes correctly inserted (QuantizeLinear + DequantizeLinear)
- TRT engine builds successfully from QDQ ONNX
- INT8 vs FP32 max accuracy difference: **0.000296** (well within expected INT8 tolerance)

### Proposed Integration

1. **New module `_export/quantization.py`**: `QuantizationArgs` config class + wrapper around `modelopt.onnx.quantization.quantize()` with LightlyTrain's data loading for calibration
2. **`tensorrt_helpers.py`**: Add `"int8"` to precision literal; when INT8, run ONNX-level quantization on the exported FP32 ONNX before passing to TRT builder
3. **`benchmark_types.py`**: Add `"int8"` to `TensorRTBackendArgs` precision literal
4. **Task models**: Add `calibration_data` parameter to `export_tensorrt()` methods

### Dependencies

Only lightweight additions needed (no torch/CUDA version changes):

- `nvidia-modelopt` (core, **not** `[all]`) — requires `torch>=2.8`
- `onnx-graphsurgeon>=0.6.1`
- `lief`

## Test Environment

- GPU: NVIDIA H200 MIG 2g.35gb (32.5 GB VRAM)
- CUDA 13.0, Driver 580.126.20
- PyTorch 2.13.0+cu130
- TensorRT 11.2.1.2 (tensorrt_cu13 bindings)
- nvidia-modelopt 0.45.0
- onnx 1.22.0, onnxruntime-gpu 1.23.2
- LightlyTrain 0.17.0 — existing test suite passes (60 passed, 3 skipped, 0 failures)

cc @liopeer

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the existing export flow in tensorrt_helpers.py and its precision handling. Then review benchmark_types.py, the proposed _export/quantization.py, and task-model export_tensorrt() methods, and run the existing export tests. Done means the INT8 path and calibration integration work alongside the stated TRT 10.x and 11.x compatibility behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.