[实践记录]4090上使用tensorRT量化qwen3

1. 任务目标

在 RTX 4090 主机上，将 Hugging Face 格式的 Qwen3-0.6B 模型转换为 TensorRT-LLM 的 Weight-Only 量化 checkpoint，然后将量化结果传输到 Jetson Orin NX 上构建 TensorRT Engine。

整体流程：

Qwen3-0.6B Hugging Face 模型
              ↓
RTX 4090 执行权重量化
              ↓
TensorRT-LLM W4A16/W8A16 checkpoint
              ↓
传输到 Jetson Orin NX
              ↓
使用 trtllm-build 构建 Engine

---

2. 环境信息

4090 主机环境：

操作系统：Ubuntu 22.04
GPU：NVIDIA RTX 4090
Python：3.10
CUDA：12.6
TensorRT：10.3
TensorRT-LLM：0.12.0
Transformers：4.51.3

原始模型路径：

/home/scjs/trt_models/Qwen3-0.6B

TensorRT-LLM 源码路径：

/home/scjs/TensorRT-LLM

---

3. 激活量化环境

进入 Conda 环境：

conda activate trtllm

设置 CUDA 和 TensorRT 环境变量：

export CUDA_HOME=/usr/local/cuda-12.6
export PATH="$CUDA_HOME/bin:$PATH"

export TRT_ROOT="$HOME/Downloads/TensorRT-10.3.0.26"

export LD_LIBRARY_PATH="$CUDA_HOME/lib64:$TRT_ROOT/lib:$TRT_ROOT/targets/x86_64-linux-gnu/lib:$LD_LIBRARY_PATH"

进入 TensorRT-LLM 源码目录：

cd ~/TensorRT-LLM

---

4. 检查 TensorRT-LLM 环境

验证 TensorRT-LLM 是否能够正常导入：

PYTHONPATH=$PWD python -c \
"import tensorrt_llm; print(tensorrt_llm.__version__)"

正常输出：

[TensorRT-LLM] TensorRT-LLM version: 0.12.0
0.12.0

验证 Qwen3 模型代码：

PYTHONPATH=$PWD python -c \
"from tensorrt_llm.models.qwen3.model import QWen3ForCausalLM; print('Qwen3 OK')"

正常输出：

Qwen3 OK

验证 Weight-Only 量化算子：

PYTHONPATH=$PWD python - <<'PY'
import torch
import tensorrt_llm

print(
    hasattr(
        torch.ops.trtllm,
        "symmetric_quantize_last_axis_of_batched_matrix"
    )
)
PY

正常输出：

True

---

5. 执行 W4A16 量化

W4A16 表示：

权重：INT4
激活：FP16

执行命令：

cd ~/TensorRT-LLM

PYTHONPATH=$PWD python examples/qwen3/convert_checkpoint.py \
    --model_dir /home/scjs/trt_models/Qwen3-0.6B \
    --output_dir /home/scjs/trt_models/Qwen3-0.6B-trtllm-w4a16 \
    --dtype float16 \
    --use_weight_only \
    --weight_only_precision int4

量化成功时会输出：

Weights loaded. Total time: 00:00:05
Total time of converting checkpoints: 00:00:07

生成目录：

/home/scjs/trt_models/Qwen3-0.6B-trtllm-w4a16

主要文件：

Qwen3-0.6B-trtllm-w4a16/
├── config.json
└── rank0.safetensors

其中：

• config.json：TensorRT-LLM 模型配置和量化配置。

• rank0.safetensors：INT4 Weight-Only 模型权重。

---

6. 可选：执行 W8A16 量化

W8A16 表示：

权重：INT8
激活：FP16

执行命令：

cd ~/TensorRT-LLM

PYTHONPATH=$PWD python examples/qwen3/convert_checkpoint.py \
    --model_dir /home/scjs/trt_models/Qwen3-0.6B \
    --output_dir /home/scjs/trt_models/Qwen3-0.6B-trtllm-w8a16 \
    --dtype float16 \
    --use_weight_only \
    --weight_only_precision int8

输出目录：

/home/scjs/trt_models/Qwen3-0.6B-trtllm-w8a16

---

7. 检查量化结果

检查 W4A16 checkpoint：

ls -lh /home/scjs/trt_models/Qwen3-0.6B-trtllm-w4a16

检查 W8A16 checkpoint：

ls -lh /home/scjs/trt_models/Qwen3-0.6B-trtllm-w8a16

正常情况下应当包含：

config.json
rank0.safetensors

也可以检查配置中的量化信息：

cat /home/scjs/trt_models/Qwen3-0.6B-trtllm-w4a16/config.json

W4A16 配置中应当包含类似：

{
  "quantization": {
    "quant_algo": "W4A16"
  }
}

W8A16 配置中应当包含类似：

{
  "quantization": {
    "quant_algo": "W8A16"
  }
}

---

8. 将量化模型传输到 Orin NX

可以使用 scp 将整个量化目录传输到 NX：

scp -r \
    /home/scjs/trt_models/Qwen3-0.6B-trtllm-w4a16 \
    nx-aiq11@NX_IP:/home/nx-aiq11/tesnorrt_model/

传输完成后，在 NX 上检查：

ls -lh \
/home/nx-aiq11/tesnorrt_model/Qwen3-0.6B-trtllm-w4a16

确认包含：

config.json
rank0.safetensors

---

9. 后续操作

RTX 4090 主机负责完成：

Hugging Face FP16 模型
        ↓
W4A16/W8A16 TensorRT-LLM checkpoint

Jetson Orin NX 负责完成：

TensorRT-LLM checkpoint
        ↓
trtllm-build
        ↓
TensorRT Engine
        ↓
模型推理服务

需要注意，4090 上生成的是 TensorRT-LLM checkpoint，并不是能够直接运行的 TensorRT Engine。

最终 Engine 应当在 Jetson Orin NX 上，根据 NX 的 CUDA、TensorRT、TensorRT-LLM 和 GPU 架构进行构建。

---

10. 总结

本次在 RTX 4090 上完成了 Qwen3-0.6B 的 TensorRT-LLM Weight-Only 量化。

W4A16 量化命令的核心参数为：

--dtype float16 \
--use_weight_only \
--weight_only_precision int4

最终输出：

config.json
rank0.safetensors

完整主干流程为：

准备 Qwen3 Hugging Face 模型
              ↓
配置 CUDA、TensorRT 和 TensorRT-LLM 环境
              ↓
验证 Qwen3 与 INT4 量化算子
              ↓
运行 convert_checkpoint.py
              ↓
生成 W4A16/W8A16 checkpoint
              ↓
将 checkpoint 传输到 Orin NX
              ↓
在 NX 上构建 TensorRT Engine