[实践记录]Jetson orin Nx部署量化Qwen3实录

一、最终目录关系

整个流程会产生三类模型目录：

Qwen3-0.6B/
├── config.json
├── model.safetensors
├── tokenizer.json
└── ...

这是从 Hugging Face 下载的原始模型，也是权重转换和量化的起点。

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

这是 TensorRT-LLM 的中间 checkpoint，不能直接用于推理。

Qwen3-0.6B-engine/
├── config.json
└── rank0.engine

这是最终运行时加载的 TensorRT engine。

完整关系：

Hugging Face 原始模型
Qwen3-0.6B
        ↓
convert_checkpoint.py
        ↓
TensorRT-LLM checkpoint
        ↓
trtllm-build
        ↓
TensorRT engine
        ↓
实际推理

二、下载 Qwen3-0.6B

创建模型目录：

mkdir -p ~/tesnorrt_model
cd ~/tesnorrt_model

使用 Git LFS 下载：

git lfs install
git clone https://huggingface.co/Qwen/Qwen3-0.6B

下载完成后检查：

ls ~/tesnorrt_model/Qwen3-0.6B

应当至少包含：

config.json
model.safetensors
tokenizer.json
tokenizer_config.json

三、准备 Qwen3 转换环境

当前使用的是：

TensorRT-LLM 0.12.0-jetson
Jetson Orin NX
CUDA 12.6
TensorRT 10.3
Python 3.10

激活 TensorRT-LLM 环境：

source ~/venvs/trtllm/bin/activate

Qwen3 需要较新的 Transformers 支持，因此安装：

pip install "transformers==4.51.3"

验证：

python3 - <<'PY'
import transformers
from transformers.models.qwen3 import Qwen3ForCausalLM

print("transformers:", transformers.__version__)
print("Qwen3 support: OK")
PY

为了确保 Python 使用修改后的 TensorRT-LLM 源码，而不是虚拟环境中旧的 wheel，设置：

export PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH

验证路径：

python3 - <<'PY'
import tensorrt_llm
import tensorrt_llm.models.qwen3.model as qwen3_model

print(tensorrt_llm.__file__)
print(qwen3_model.__file__)
PY

路径应指向：

/home/nx-aiq11/TensorRT-LLM/tensorrt_llm/

而不是：

/home/nx-aiq11/venvs/trtllm/lib/python3.10/site-packages/

四、为 TensorRT-LLM 0.12 适配 Qwen3

TensorRT-LLM 0.12 原生支持 Qwen2，但不支持 Qwen3，因此基于原来的 Qwen 实现创建：

tensorrt_llm/models/qwen3/
├── __init__.py
├── config.py
├── model.py
└── convert.py

同时创建：

examples/qwen3/convert_checkpoint.py

主要适配内容如下。

1. 独立 head_dim

Qwen3-0.6B 的配置是：

hidden_size = 1024
num_attention_heads = 16
num_key_value_heads = 8
head_dim = 128

不能使用：

hidden_size // num_attention_heads

因为：

1024 / 16 = 64

但模型实际的 head_dim 是 128。

因此配置必须优先读取：

head_size = hf_config.head_dim

2. GQA

Qwen3-0.6B 使用：

Q heads  = 16
KV heads = 8
head_dim = 128

所以 QKV 投影宽度为：

Q = 16 × 128 = 2048
K =  8 × 128 = 1024
V =  8 × 128 = 1024

QKV = 4096

正确的 QKV 权重 shape：

(4096, 1024)

而不是普通 MHA 假设下的：

16 × 3 × 128 = 6144

3. Q/K RMSNorm

Qwen3 在 Q、K 投影后增加：

q_norm
k_norm

在 TensorRT-LLM 中启用：

qk_layernorm=True
layernorm_type=LayerNormType.RmsNorm
eps=config.norm_epsilon

HF 权重：

model.layers.N.self_attn.q_norm.weight
model.layers.N.self_attn.k_norm.weight

映射到：

transformer.layers.N.attention.q_layernorm.weight
transformer.layers.N.attention.k_layernorm.weight

它们的 shape 均为：

(128,)

4. 模型注册

创建独立模型类：

QWen3ForCausalLM

checkpoint 中写入：

"architecture": "QWen3ForCausalLM",
"qwen_type": "qwen3"

并将该 architecture 注册到 TensorRT-LLM 模型映射中。

5. Attention 支持 GQA + Q/K Norm

旧版 Attention 的 Q/K Norm 分支默认 Q、K、V head 数相同，会尝试：

[num_tokens, 4096]
→
[num_tokens, 16, 3, 128]

导致体积不一致。

修改后按照真实尺寸拆分：

Q: [num_tokens, 2048]
K: [num_tokens, 1024]
V: [num_tokens, 1024]

然后分别 reshape：

Q: [num_tokens, 16, 128]
K: [num_tokens, 8, 128]

执行 Q/K RMSNorm 后，再拼回：

[num_tokens, 4096]

五、转换 FP16 TensorRT-LLM checkpoint

删除旧输出：

rm -rf ~/tesnorrt_model/Qwen3-0.6B-trtllm-checkpoint

运行转换：

cd ~/TensorRT-LLM/examples/qwen3

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
python3 convert_checkpoint.py \
  --model_dir ~/tesnorrt_model/Qwen3-0.6B \
  --output_dir ~/tesnorrt_model/Qwen3-0.6B-trtllm-checkpoint \
  --dtype float16

成功日志：

Weights loaded.
Total time of converting checkpoints: ...

六、验证 FP16 checkpoint

检查配置：

python3 - <<'PY'
import json

path = "/home/nx-aiq11/tesnorrt_model/Qwen3-0.6B-trtllm-checkpoint/config.json"

with open(path, "r", encoding="utf-8") as f:
    cfg = json.load(f)

for key in [
    "architecture",
    "qwen_type",
    "hidden_size",
    "num_attention_heads",
    "num_key_value_heads",
    "head_size",
]:
    print(f"{key}: {cfg.get(key)}")
PY

正确结果：

architecture: QWen3ForCausalLM
qwen_type: qwen3
hidden_size: 1024
num_attention_heads: 16
num_key_value_heads: 8
head_size: 128

检查权重：

python3 - <<'PY'
from safetensors.torch import load_file

path = "/home/nx-aiq11/tesnorrt_model/Qwen3-0.6B-trtllm-checkpoint/rank0.safetensors"
weights = load_file(path)

keys = [
    "transformer.layers.0.attention.qkv.weight",
    "transformer.layers.0.attention.q_layernorm.weight",
    "transformer.layers.0.attention.k_layernorm.weight",
]

for key in keys:
    print(key, tuple(weights[key].shape))
PY

正确结果：

qkv.weight             (4096, 1024)
q_layernorm.weight     (128,)
k_layernorm.weight     (128,)

七、构建 FP16 TensorRT engine

删除旧 engine：

rm -rf ~/tesnorrt_model/Qwen3-0.6B-engine

构建：

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
trtllm-build \
  --checkpoint_dir ~/tesnorrt_model/Qwen3-0.6B-trtllm-checkpoint \
  --output_dir ~/tesnorrt_model/Qwen3-0.6B-engine \
  --gemm_plugin float16

成功后检查：

ls -lh ~/tesnorrt_model/Qwen3-0.6B-engine

应生成：

config.json
rank0.engine

八、测试 FP16 engine

Qwen3 是对话模型，必须应用 chat template，不能只使用裸文本测试。

可以先构造 Prompt：

PROMPT=$(python3 - <<'PY'
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "/home/nx-aiq11/tesnorrt_model/Qwen3-0.6B"
)

messages = [
    {
        "role": "user",
        "content": "你好，请用一句话介绍一下你自己。"
    }
]

print(tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False,
))
PY
)

运行：

cd ~/TensorRT-LLM/examples/qwen3

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
python3 ../run.py \
  --input_text "$PROMPT" \
  --max_output_len 100 \
  --tokenizer_dir ~/tesnorrt_model/Qwen3-0.6B \
  --engine_dir ~/tesnorrt_model/Qwen3-0.6B-engine

FP16 能正常输出后，说明：

模型结构适配成功
权重转换成功
TensorRT engine 构建成功
Qwen3 推理成功

九、在 RTX 4090 主机上进行 Weight-Only 量化

由于 Jetson NX 的 ARM64 环境中 ModelOpt、PyTorch distributed 和部分依赖安装困难，因此将量化任务放到 x86 RTX 4090 主机上完成。

重要原则：

4090 负责生成量化 checkpoint
NX 负责构建最终 TensorRT engine

不要在 4090 上构建 engine 后直接复制到 NX。

两台主机必须使用相同版本的 Qwen3 backport 代码，建议比较：

sha256sum \
  tensorrt_llm/models/qwen3/config.py \
  tensorrt_llm/models/qwen3/model.py \
  tensorrt_llm/models/qwen3/convert.py \
  tensorrt_llm/models/qwen3/__init__.py \
  examples/qwen3/convert_checkpoint.py

十、生成 W8A16 checkpoint

在 4090 主机上运行：

cd ~/TensorRT-LLM

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
python3 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

含义：

W8 = 权重 INT8
A16 = 激活 FP16

十一、生成 W4A16 checkpoint

在 4090 主机上运行：

cd ~/TensorRT-LLM

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
python3 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

含义：

W4 = 权重 INT4
A16 = 激活 FP16

量化配置中应出现：

"quantization": {
    "quant_algo": "W4A16",
    "kv_cache_quant_algo": null
}

十二、将量化 checkpoint 拷贝回 NX

例如使用 SCP：

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

W8A16 同理。

最终 NX 上目录：

/home/nx-aiq11/tesnorrt_model/Qwen3-0.6B-trtllm-w8a16

/home/nx-aiq11/tesnorrt_model/Qwen3-0.6B-trtllm-w4a16

十三、在 NX 上构建 W8A16 engine

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
trtllm-build \
  --checkpoint_dir ~/tesnorrt_model/Qwen3-0.6B-trtllm-w8a16 \
  --output_dir ~/tesnorrt_model/Qwen3-0.6B-engine-w8a16 \
  --gemm_plugin float16

虽然权重是 INT8，但这里仍使用：

--gemm_plugin float16

因为 W8A16 中激活和计算仍然是 FP16。

十四、在 NX 上构建 W4A16 engine

PYTHONPATH=$HOME/TensorRT-LLM:$PYTHONPATH \
trtllm-build \
  --checkpoint_dir ~/tesnorrt_model/Qwen3-0.6B-trtllm-w4a16 \
  --output_dir ~/tesnorrt_model/Qwen3-0.6B-engine-w4a16 \
  --gemm_plugin float16

同样：

权重：INT4
激活：FP16
GEMM 计算：FP16

十五、量化结果验证

最初量化模型曾出现重复输出：

我是你吗？我是你吗？我是你吗？……

后来通过固定以下条件重新测试：

相同 tokenizer
相同 Qwen3 chat template
相同 input_ids
相同 decoding 参数

固定参数：

max_output_len = 64
temperature = 1.0
top_k = 1
top_p = 0.0
random_seed = 0

三种模型使用的 input_ids 完全一致。

最终测试结果：

FP16：
我是AI助手，专注于帮助用户解决问题和提供支持。

W8A16：
我是AI助手，专注于帮助用户解决问题和提供支持。

W4A16：
能够输出正常、连贯文本，但相比 FP16/W8A16 有一定质量退化。

同时验证量化确实生效：

W8A16：
torch.int8 weight
+ per_channel_scale

W4A16：
packed INT4 weight，以 torch.int8 保存
+ per_channel_scale

量化模型实际实例化为：

attention.qkv   → WeightOnlyQuantLinear
attention.dense → WeightOnlyQuantRowLinear
mlp.gate        → WeightOnlyQuantLinear
mlp.fc          → WeightOnlyQuantLinear
mlp.proj        → WeightOnlyQuantRowLinear

因此最终结论：

FP16    正常
W8A16   正常
W4A16   可用，但存在一定精度退化

十六、最终模型目录

原始模型：

~/tesnorrt_model/Qwen3-0.6B

FP16 checkpoint：

~/tesnorrt_model/Qwen3-0.6B-trtllm-checkpoint

FP16 engine：

~/tesnorrt_model/Qwen3-0.6B-engine

W8A16 checkpoint：

~/tesnorrt_model/Qwen3-0.6B-trtllm-w8a16

W8A16 engine：

~/tesnorrt_model/Qwen3-0.6B-engine-w8a16

W4A16 checkpoint：

~/tesnorrt_model/Qwen3-0.6B-trtllm-w4a16

W4A16 engine：

~/tesnorrt_model/Qwen3-0.6B-engine-w4a16

十七、完整流程总结

1. 从 Hugging Face 下载 Qwen3-0.6B
        ↓
2. 在 TensorRT-LLM 0.12 中 backport Qwen3
        ↓
3. 适配 head_dim、GQA、q_norm/k_norm
        ↓
4. 转换 FP16 TRT-LLM checkpoint
        ↓
5. 验证 config 和权重 shape
        ↓
6. 在 NX 上构建 FP16 engine
        ↓
7. 使用 Qwen3 chat template 验证 FP16 推理
        ↓
8. 在 RTX 4090 上生成 W8A16/W4A16 checkpoint
        ↓
9. 将量化 checkpoint 复制回 NX
        ↓
10. 在 NX 上分别构建 W8A16/W4A16 engine
        ↓
11. 使用相同 input_ids 和解码参数进行 A/B 测试
        ↓
12. 确认 W8A16 正常、W4A16 可用

核心原则：

原始 safetensors 是转换和量化起点
TRT-LLM checkpoint 是中间产物
TensorRT engine 才是真正运行的模型

量化在 4090 上完成
最终 engine 在 NX 上重新构建

posted @ 2026-07-24 14:26 阿基米德的澡盆阅读(29) 评论(0) 收藏举报

刷新页面返回顶部

阿基米德的澡盆