模型转换与优化：将 PyTorch 模型转换为 TensorRT-LLM 支持的格式

在完成 TensorRT-LLM 的环境部署后，我们面临的第一个核心任务便是将训练好的 PyTorch 模型转换为 TensorRT-LLM 能够识别并高效执行的格式。这一过程并非简单的文件格式变更，而是涉及权重重映射、网络结构解析、层融合以及量化策略嵌入的深度优化。理解这一转换链路，是后续实现流式推理与性能调优的基石。

4.1 核心原理：从 PyTorch 权重到 TensorRT-LLM 引擎的必经之路

TensorRT-LLM 并不直接读取 PyTorch 的 .pt 或 .pth 文件。其标准工作流分为两个阶段：权重转换与引擎构建。权重转换阶段将 PyTorch 的 state_dict 映射为 TensorRT-LLM 定义的、按层拆分的二进制权重文件（通常是 .safetensors 格式）。引擎构建阶段则利用这些权重文件，结合网络定义与编译选项，生成针对特定 GPU 架构优化的 .engine 文件。

该流程的核心价值在于：其一，解耦训练框架与推理框架，使得推理引擎不依赖 PyTorch 运行时；其二，引入图优化与内核自动调优，如算子融合、KV Cache 量化、Attention 变体选择等；其三，支持权重量化，将 FP16/BF16 甚至 INT8/INT4 的量化参数固化在权重文件中，避免运行时重复计算。

下图展示了从 PyTorch 模型到 TensorRT-LLM 引擎的完整转换链路与关键产物：

4.2 方法流程：以 LLaMA-7B 为例的完整转换实操

我们以 Hugging Face 格式的 LLaMA-7B 模型为例，演示标准转换流程。首先需要准备模型文件，并确保目录结构符合 TensorRT-LLM 的预期。

# 1. 设置环境变量
export MODEL_DIR=/models/llama-7b-hf
export OUTPUT_DIR=/models/llama-7b-trtllm
export TRTLLM_ROOT=/opt/TensorRT-LLM

# 2. 执行权重转换（FP16 精度，张量并行度=2）
python ${TRTLLM_ROOT}/examples/llama/convert_checkpoint.py \
    --model_dir ${MODEL_DIR} \
    --output_dir ${OUTPUT_DIR} \
    --dtype float16 \
    --tp_size 2 \
    --workers 4

# 3. 查看转换产物
ls -lh ${OUTPUT_DIR}
# 预期输出：
# total 13G
# -rw-r--r-- 1 root root 6.5G Jul 10 10:00 rank0.safetensors
# -rw-r--r-- 1 root root 6.5G Jul 10 10:00 rank1.safetensors
# -rw-r--r-- 1 root root 1.2K Jul 10 10:00 config.json

转换完成后，config.json 中记录了模型结构参数（如层数、头数、词表大小）以及量化与并行配置。接下来使用 trtllm-build 命令构建引擎。此阶段需要指定推理时的关键约束，如最大批大小、最大序列长度等，这些参数直接影响显存占用与性能。

# 4. 构建 TensorRT 引擎
trtllm-build \
    --checkpoint_dir ${OUTPUT_DIR} \
    --output_dir ${OUTPUT_DIR}/engine \
    --gemm_plugin float16 \
    --max_batch_size 8 \
    --max_input_len 2048 \
    --max_seq_len 4096 \
    --max_num_tokens 8192 \
    --use_fused_mlp

# 5. 验证引擎文件
ls -lh ${OUTPUT_DIR}/engine
# 预期输出：
# -rw-r--r-- 1 root root 13G Jul 10 10:15 rank0.engine
# -rw-r--r-- 1 root root 13G Jul 10 10:15 rank1.engine

4.3 案例与场景：量化与自定义模型适配

在实际业务中，我们往往需要进一步压缩模型体积或适配非标准结构。以下两个场景最具代表性：

场景一：INT4 权重量化部署。当 GPU 显存有限（如 24GB 单卡部署 7B 模型）时，需要将权重从 FP16 量化为 INT4。TensorRT-LLM 提供了 quantize.py 脚本，支持 AWQ 或 GPTQ 算法。量化过程需要校准数据集，且量化后的精度损失需通过评估集验证。

# 执行 INT4 AWQ 量化
python ${TRTLLM_ROOT}/examples/quantization/quantize.py \
    --model_dir ${MODEL_DIR} \
    --output_dir ${OUTPUT_DIR}_int4_awq \
    --dtype float16 \
    --qformat int4_awq \
    --calib_size 32 \
    --tp_size 1

# 构建 INT4 引擎
trtllm-build \
    --checkpoint_dir ${OUTPUT_DIR}_int4_awq \
    --output_dir ${OUTPUT_DIR}_int4_awq/engine \
    --gemm_plugin float16 \
    --max_batch_size 8 \
    --max_seq_len 4096

场景二：自定义模型结构适配。若模型包含 RoPE 以外的位置编码或自定义 Attention，直接运行 convert_checkpoint.py 会报 KeyError。此时需要继承 BaseModelForCausalLM 并重写 __init__ 中的层映射逻辑。例如，若模型使用 ALiBi 位置编码，需在转换脚本中手动将 model.layers.*.self_attn.rotary_emb.inv_freq 替换为 ALiBi 的斜率张量。

4.4 误区与复盘：转换失败与性能陷阱

在模型转换过程中，工程师常陷入以下误区，导致转换失败或推理性能不达标：

• 误区一：忽略词表对齐。当原始模型词表大小不是 64 的整数倍时（如 LLaMA-2 词表为 32000），TensorRT-LLM 会自动填充至 32000 的倍数。但若在转换后修改了 tokenizer，会导致输入输出维度不匹配。复盘建议：转换前固定 tokenizer 版本，并在 config.json 中记录 vocab_size 与 padding_size。
• 误区二：张量并行度与 GPU 数量不匹配。若使用 TP=2 转换权重，但推理时仅提供 1 张 GPU，引擎加载会直接报错。反之，TP=1 的引擎无法在多卡上利用显存。复盘建议：在转换前明确部署拓扑，并检查 world_size 与引擎 rank 文件数量一致。
• 误区三：过度追求小显存而忽略精度。INT4 量化虽然能将 7B 模型压缩至 4GB 以内，但若校准集选取不当（如仅用 8 条新闻文本），生成质量会显著下降。复盘建议：量化后必须运行困惑度（PPL）评估，若 PPL 上升超过 0.5，应回退至 INT8 或使用混合精度。
• 误区四：混淆 max_seq_len 与 max_num_tokens。前者限制单条序列的最大长度，后者限制所有并发请求的 token 总量。若仅设置前者而忽略后者，在高并发流式输出时会触发隐式截断。复盘建议：使用公式 max_num_tokens = max_batch_size × 平均生成长度 进行估算，并预留 20% 余量。