📚 课程目录第4章模型转换与优化

4. 模型转换与优化：从PyTorch模型到TensorRT引擎的转换流程

4.1 核心原理：为什么需要“转换”而非“直接加载”

PyTorch模型在训练时依赖动态计算图和Python运行时，而TensorRT引擎追求的是静态图上的极致推理优化。两者之间的鸿沟需要通过一个明确的转换流程来跨越。转换的本质是将模型从“训练表达”翻译为“部署表达”，并在此过程中完成三件关键工作：图结构优化（层融合、死分支消除）、精度校准（FP16/INT8量化）以及内核自动调优（选择当前GPU上最快的kernel实现）。

整个转换链路可以概括为：.pt/.pth → .onnx → .engine。其中ONNX（Open Neural Network Exchange）作为中间表示，起到“通用语言”的作用，它剥离了PyTorch的Python依赖，将模型表达为计算图。TensorRT-LLM在此基础上进一步将ONNX图解析为LLM专用的IR（中间表示），并针对Transformer结构进行深度定制优化。

关键认知：TensorRT-LLM并非直接消费PyTorch权重文件，而是通过ONNX作为“翻译官”。跳过ONNX直接尝试加载PyTorch模型是初学者最常见的错误。

4.2 方法流程：标准转换五步法

在TensorRT-LLM框架下，推荐使用官方提供的trtllm-build命令行工具完成转换。整个流程分为五个标准步骤，每一步都有明确的输入输出和验证手段。

导出ONNX：使用PyTorch的torch.onnx.export将模型导出为ONNX格式。需要固定输入维度（或设置动态轴），并确保算子版本与TensorRT-LLM兼容。
准备权重：TensorRT-LLM需要将PyTorch权重转换为自身格式。使用convert_checkpoint.py脚本将HF格式的权重转换为TensorRT-LLM的checkpoint格式。
构建引擎：运行trtllm-build命令，指定checkpoint路径、输出引擎路径、精度（--dtype float16）以及并行方式（如--tp_size 2）。
验证引擎：使用trtllm-run或Python API加载引擎，输入测试样本，对比输出与PyTorch参考结果的误差。
部署集成：将生成的.engine文件与分词器、后处理逻辑打包为推理服务。

# 步骤1：导出ONNX（PyTorch侧）
import torch
from transformers import LlamaForCausalLM

model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
dummy_input = torch.ones(1, 128, dtype=torch.int32)
torch.onnx.export(
    model, dummy_input, "llama2_7b.onnx",
    opset_version=17,
    input_names=["input_ids"],
    output_names=["logits"],
    dynamic_axes={"input_ids": {0: "batch", 1: "seq_len"}}
)

# 步骤2：转换权重格式
python convert_checkpoint.py --model_dir ./llama2_7b_hf \
    --output_dir ./tllm_checkpoint \
    --dtype float16

# 步骤3：构建TensorRT引擎
trtllm-build --checkpoint_dir ./tllm_checkpoint \
    --output_dir ./engine_output \
    --gemm_plugin float16 \
    --max_batch_size 8 \
    --max_input_len 2048 \
    --max_seq_len 4096 \
    --tp_size 1

4.3 案例场景：LLaMA-7B的FP16转换与精度验证

以一个典型的LLaMA-7B部署为例，展示完整的转换参数配置与验证方法。该模型包含32个Transformer层、隐藏维度4096，权重约13GB（FP32）。转换为FP16后，引擎文件约6.5GB，显存占用显著降低。

配置项	推荐值	说明
dtype	float16	默认推荐，精度损失<0.1%
gemm_plugin	float16	启用FP16 GEMM优化kernel
max_batch_size	8~32	影响KV Cache预留空间
max_seq_len	4096	需覆盖最长生成序列
tp_size	1（单卡）	多卡时设为GPU数量

转换完成后，使用Python API进行精度验证。核心指标是最大相对误差和困惑度差异。对于FP16转换，logits的最大相对误差通常应小于1e-2；对于INT8量化，则需关注特定层（如Attention输出层）的误差累积。

# 精度验证脚本（片段）
import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner

runner = ModelRunner.from_dir(engine_dir="./engine_output")
outputs = runner.run(
    input_ids=test_input_ids,
    max_new_tokens=50
)

# 对比PyTorch参考输出
ref_outputs = pytorch_model.generate(test_input_ids, max_new_tokens=50)
max_diff = (outputs["output_ids"] - ref_outputs).abs().max().item()
print(f"最大token ID差异: {max_diff}")  # 期望为0或极小值

实战经验：当转换后的模型出现输出质量下降时，优先检查RMSNorm层的epsilon值是否在ONNX导出时被正确保留。PyTorch默认epsilon为1e-6，但某些ONNX导出工具会错误地将其截断为1e-5，导致推理精度异常。

4.4 常见误区与复盘

在模型转换实践中，以下问题出现频率最高，值得重点规避：

动态轴设置不当：ONNX导出时未声明dynamic_axes，导致引擎只能接受固定长度输入，无法处理变长请求。解决方案：在导出时明确指定batch和seq_len为动态维度。
忽略插件兼容性：某些PyTorch算子（如torch.where、torch.gather）在ONNX中可能被转换为低效的子图。建议在导出前用torch.onnx.export的opset_version参数尝试不同版本，并检查ONNX图中是否有Constant节点异常膨胀。
校准集偏差：INT8量化时，校准集必须覆盖真实推理时的长度分布和内容分布。仅用短文本校准会导致长文本生成时误差放大。
引擎复用陷阱：TensorRT引擎与GPU架构强绑定。在A100上构建的引擎无法在V100上运行。多环境部署时需分别构建，或使用--hardware_compatibility_level参数（但会牺牲部分性能）。

重要提醒：不要在生产环境中直接使用torch.onnx.export的默认设置导出LLM。务必设置do_constant_folding=True（默认开启）并检查_output_metadata中是否包含sequence_length等运行时参数。对于LLaMA等使用旋转位置编码（RoPE）的模型，建议使用TensorRT-LLM提供的export_onnx.py专用脚本，而非通用导出接口。

复盘一次典型的转换失败案例：某团队在转换GPT-NeoX-20B模型时，直接使用HuggingFace的AutoModelForCausalLM导出ONNX，结果生成的引擎在推理时出现CUDA error: misaligned address。排查发现，问题根源在于模型使用了torch.nn.functional.scaled_dot_product_attention，该算子在ONNX导出时被拆分为多个细粒度算子，导致TensorRT在内存对齐优化时出错。最终通过升级TensorRT-LLM至支持SDPA融合的版本，并改用--attention_plugin参数解决。

另一个值得关注的细节是KV Cache的量化。在长序列生成场景中，KV Cache占用显存随序列长度线性增长。TensorRT-LLM支持将KV Cache量化为INT8，但需要在构建引擎时通过--kv_cache_quantization int8显式开启，并在运行时设置--kv_cache_free_gpu_memory_fraction控制预留比例。未开启该选项时，即使模型权重为FP16，KV Cache仍以FP16存储，长序列下容易触发OOM。

最后，建议在CI/CD流程中加入引擎构建回归测试。每次更新模型代码或依赖库后，自动触发一次构建并运行预设的精度测试用例，确保转换流程的稳定性。这能有效避免“本地能跑、上线就崩”的部署事故。