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

4. 模型转换与优化

在完成环境准备与模型选型后，我们面临的核心任务是将 HuggingFace 格式的权重转换为 TensorRT-LLM 可执行的引擎文件。这一过程并非简单的格式改写，而是涉及图结构解析、算子映射、权重重排与精度量化的深度优化。只有正确完成转换，才能让 TensorRT-LLM 的 Kernel 自动调优、显存池化、KV Cache 管理等能力真正生效。

4.1 核心原理：从 PyTorch 权重到引擎的三层递进

理解转换流程，需要先厘清三个关键产物：原始权重目录 → 转换后的 checkpoint 目录 → 最终 engine 文件。HuggingFace 模型通常以 PyTorch 的 state_dict 形式存储，权重命名遵循 model.layers.0.self_attn.q_proj.weight 这类层级结构。而 TensorRT-LLM 需要的是按 层（Layer）、张量并行（TP）、流水线并行（PP） 重新切分后的权重布局，同时算子融合方式（如 QKV 融合）也要求权重按特定维度拼接。

转换脚本 convert_checkpoint.py 正是承担这一“翻译”职责：它读取 HuggingFace 配置（config.json）与权重，根据目标精度（FP16/INT8/INT4）执行量化，并输出 TensorRT-LLM 自定义的 checkpoint 格式。随后，trtllm-build 工具读取该 checkpoint，结合网络定义、插件选择、显存预算等参数，编译生成最终的 engine 文件。整个链路中，checkpoint 是中间桥梁，engine 是运行时唯一依赖的产物。

关键认知：不要试图直接加载 HuggingFace 权重到 TensorRT-LLM 运行时。引擎文件内部包含了经过算子融合、权重重排、量化缩放因子打包后的二进制数据，其内存布局与 PyTorch 完全不同。

4.2 方法流程：三步完成转换与精度配置

下面以 Llama-3.1-8B-Instruct 为例，演示从 HuggingFace 权重到 FP16/INT8/INT4 引擎的完整命令流程。假设模型已下载至 /models/llama-3.1-8b-hf，且 TensorRT-LLM 已正确安装。

步骤 1：转换 checkpoint（含量化）

FP16 精度转换命令如下：

# 进入 TensorRT-LLM 示例目录
cd /opt/TensorRT-LLM/examples/llama

# FP16 转换（无量化）
python convert_checkpoint.py \
    --model_dir /models/llama-3.1-8b-hf \
    --output_dir /models/llama-3.1-8b-trllm-fp16 \
    --dtype float16 \
    --tp_size 2 \
    --pp_size 1

若需 INT8 或 INT4 量化，则增加 --quantize_ckpt 参数。TensorRT-LLM 支持多种量化方案，常用选项如下表：

精度方案	命令行参数	权重类型	激活类型	典型显存节省
FP16	`--dtype float16`	FP16	FP16	约 50%（对比 FP32）
INT8 (W8A8)	`--quantize_ckpt=w8a8`	INT8	INT8	约 75%
INT8 (W8A16)	`--quantize_ckpt=8`	INT8	FP16	约 50%
INT4 (W4A16)	`--quantize_ckpt=4`	INT4	FP16	约 62.5%

INT4 量化示例（AWQ 或风格）：

# INT4 AWQ 量化（需要先有校准数据集）
python convert_checkpoint.py \
    --model_dir /models/llama-3.1-8b-hf \
    --output_dir /models/llama-3.1-8b-trllm-int4 \
    --dtype float16 \
    --quantize_ckpt=4 \
    --calib_dataset /data/c4_sample.jsonl \
    --calib_batch_size 8 \
    --tp_size 2

步骤 2：构建引擎

checkpoint 转换完成后，使用 trtllm-build 生成引擎。关键参数包括最大批大小、序列长度、KV Cache 策略等：

trtllm-build \
    --checkpoint_dir /models/llama-3.1-8b-trllm-fp16 \
    --output_dir /models/llama-3.1-8b-engine-fp16 \
    --gemm_plugin float16 \
    --max_batch_size 32 \
    --max_input_len 2048 \
    --max_seq_len 4096 \
    --max_num_tokens 8192 \
    --use_fused_mlp

注意：--max_num_tokens 控制 inflight batching 的预算，建议设置为 max_batch_size × max_seq_len 的 1/4 到 1/2，以平衡显存占用与调度灵活性。若设置过小，高并发下会触发隐式刷新，降低吞吐。

步骤 3：验证引擎完整性

构建完成后，可用 trtllm-infer 快速验证输出是否正常：

trtllm-infer \
    --engine_dir /models/llama-3.1-8b-engine-fp16 \
    --tokenizer_dir /models/llama-3.1-8b-hf \
    --prompt "Explain the concept of KV cache in LLM inference." \
    --max_output_len 128

4.3 案例：FP16 与 INT4 引擎的显存占用对比

以 Llama-3.1-8B 为例，FP16 权重占用约 16 GB（80 亿参数 × 2 字节）。若使用 INT4 量化，权重降至约 4 GB，但需额外存储缩放因子（per-group 或 per-channel），实际约 4.5 GB。在 2 张 A100-80G 上以 TP=2 部署时，FP16 引擎每卡权重约 8 GB，INT4 引擎每卡仅约 2.3 GB，节省的显存可用于增大 KV Cache 或批大小，从而提升吞吐。

实际测试中，INT4 相比 FP16 的延迟通常增加 5%～15%（取决于算子融合与显存带宽），但吞吐（tokens/s）可提升 1.5～2 倍，因为更大的 KV Cache 允许更高的并发批处理。若你的场景是延迟敏感型（如在线对话），FP16 可能更合适；若是吞吐敏感型（如离线批量处理），INT4 优势明显。

4.4 常见误区与复盘

误区 1：跳过 checkpoint 直接构建引擎。部分用户误以为 trtllm-build 可直接读取 HuggingFace 目录。实际上，构建器只接受 TensorRT-LLM 格式的 checkpoint，必须先运行转换脚本。
误区 2：量化后精度下降不可接受。INT4 量化对 8B 以上模型通常能保持 PPL 损失 < 0.5，但需注意：--quantize_ckpt=4 默认使用 RTN（round-to-nearest）算法，若模型对量化敏感，应改用 AWQ 或并准备校准集。
误区 3：忽略 TP 与 PP 的一致性。转换时的 --tp_size 必须与构建引擎时的预期一致。若转换时 TP=2，构建时也需对应 2 卡环境，否则引擎加载会报错。
误区 4：FP16 引擎不需要量化配置。即使不量化，也需显式指定 --dtype float16，否则默认 FP32 会导致显存翻倍且速度下降。

  复盘建议：在正式基准测试前，建议先构建 FP16 引擎作为基线，验证精度与性能符合预期后，再尝试 INT8/INT4。每次量化后，用相同的 prompt 集对比输出 logits 的余弦相似度（应 > 0.99），确保量化损失在可接受范围。

完成引擎构建后，下一步即可进入性能基准测试环节。但在此之前，请务必确认引擎文件路径、tokenizer 路径以及精度配置均已记录在案，这些参数将直接影响后续测试脚本的编写与结果解读。