模型导出与Engine构建：从PyTorch到TensorRT-LLM的最后一公里

说实话，模型训练完只是万里长征走了一半。真正让模型在线上跑起来，还得过「导出与构建」这道坎。我见过太多团队，模型精度刷得漂漂亮亮，一到部署就卡在ONNX导出或者Engine构建上，一卡就是一两周。

这一节，咱们就把这条链路彻底捋清楚。从PyTorch模型导出ONNX开始，到TensorRT-LLM的模型定义，再到Builder配置和Engine序列化。每一步我都会结合踩过的坑来讲。

PyTorch到ONNX：别小看这一步

ONNX导出，说白了就是把PyTorch的动态图变成静态图。你想想看，PyTorch是「边算边建图」，而TensorRT需要的是「图已经定死了」。这个转换，坑不少。

我个人的习惯是，先写一个简单的导出脚本，验证最基本的流程：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(768, 768)
    
    def forward(self, input_ids, attention_mask):
        x = self.linear(input_ids)
        return x

model = MyModel().eval()

dummy_input = torch.randint(0, 50000, (1, 128))
dummy_mask = torch.ones(1, 128, dtype=torch.long)

torch.onnx.export(
    model,
    (dummy_input, dummy_mask),
    "model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "attention_mask": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size", 1: "seq_len"}
    },
    opset_version=17
)

这里有个关键点：dynamic_axes。我刚开始做的时候没加这个，结果导出的ONNX只能跑固定batch size和序列长度，线上稍微变一下就得重新导出，太蠢了。

ONNX导出常见坑与避坑指南

导出过程中，最头疼的是算子不支持。尤其是Transformer里的Flash Attention、自定义激活函数这些。

• 动态控制流：PyTorch里的if-else、for循环，ONNX不一定能处理。我曾经有个模型里用了torch.where做条件判断，导出后TensorRT直接报错。解决方案是用torch.onnx.export的script模式，或者干脆把逻辑改成纯张量运算。
• Reshape操作：ONNX对动态shape的reshape支持有限。我建议用torch.reshape代替view，前者在ONNX里更稳定。
• 数据类型：PyTorch默认float32，但ONNX导出时如果混了float16，可能会出问题。统一用float32导出，后续在TensorRT里再转精度。

TensorRT-LLM的模型定义：ModelConfig是关键

ONNX导出来后，下一步就是告诉TensorRT-LLM：「嘿，我这个模型长什么样」。这就用到ModelConfig了。

ModelConfig定义了模型的架构参数。对于LLaMA、GPT这类Decoder-only模型，核心参数包括：

参数	说明	我的建议
num_layers	Transformer层数	和原始模型保持一致
num_heads	注意力头数	注意分组查询（GQA）的情况
hidden_size	隐藏层维度	通常是768/1024/4096
vocab_size	词表大小	别忘了加上特殊token
max_position_embeddings	最大序列长度	建议设成训练时的2倍

举个例子，一个7B LLaMA模型的配置大概长这样：

from tensorrt_llm.models import LLaMAConfig

config = LLaMAConfig(
    num_layers=32,
    num_heads=32,
    hidden_size=4096,
    vocab_size=32000,
    max_position_embeddings=4096,
    hidden_act="silu",
    intermediate_size=11008,
    num_kv_heads=8,  # GQA
)

嗯，这里要注意num_kv_heads。很多开源模型用了分组查询注意力（GQA），如果你漏了这个参数，构建出来的Engine推理结果全是错的。我当初调这个参数调了一整天，最后发现是KV heads没对齐。

Builder配置：精度、批量大小、序列长度

模型定义好了，接下来是Builder配置。这一步决定了Engine的「脾气」——它能跑多快、吃多少显存、支持多大的输入。

Builder配置的核心有三个维度：

• 精度：FP32、FP16、INT8、FP8。我个人建议，线上推理用FP16起步。INT8需要校准，精度损失要看具体模型。FP8是H100的专属福利，吞吐量能再提30%。
• 批量大小：支持动态batch还是固定batch？我建议用动态batch，max_batch_size设成你线上最大并发数。设太大浪费显存，设太小容易OOM。
• 序列长度：同样建议动态。但注意，max_input_len和max_output_len要分开设。比如输入最长2048，输出最长1024。

from tensorrt_llm import BuildConfig

build_config = BuildConfig(
    max_batch_size=64,
    max_input_len=2048,
    max_output_len=1024,
    precision="float16",
    enable_fp8=False,
    enable_int8=False,
    max_num_tokens=8192,  # 控制显存占用
)

这里有个经验：max_num_tokens这个参数很多人忽略。它控制了KV Cache的显存上限。设得太小，长序列推理会崩；设得太大，显存直接爆。我一般按max_batch_size * max_input_len的1.5倍来设。

Engine构建与序列化：最后一步

配置都好了，终于可以构建Engine了。这一步其实是最简单的，但也是最慢的——一个7B模型的Engine构建，在A100上可能要跑10-20分钟。

from tensorrt_llm import Builder

builder = Builder()
engine = builder.build_engine(
    model_config=config,
    build_config=build_config,
    onnx_file="model.onnx",
    weights_path="model_weights.bin",
)

# 序列化到磁盘
with open("llama7b_fp16.engine", "wb") as f:
    f.write(engine)

构建过程中，你会看到TensorRT在疯狂做图优化：层融合、内存复用、算子替换。这个过程是自动的，但你可以通过日志观察它做了哪些优化。我习惯把日志级别调到INFO，看看有没有warning。

序列化后的Engine文件，就是最终部署的产物。它包含了所有优化后的网络结构和权重。下次加载时，直接反序列化就行，不需要重新构建。

知识体系总览

为了让你对整个流程有个全局视角，我画了一张流程图：

这张图把整个流程串起来了。你从左上角的PyTorch模型出发，经过ONNX导出、模型定义、Builder配置，最后到Engine构建和序列化。每一步都有对应的配置和注意事项。

我个人觉得，整个链路中最容易出问题的是ONNX导出和Builder配置。ONNX导出搞不定，后面全是白搭；Builder配置不合理，Engine性能上不去。所以这两步一定要多花时间验证。

好了，这一节的内容就到这儿。记住：模型导出和Engine构建不是一锤子买卖，线上环境变了，你可能需要重新构建。所以把流程自动化、脚本化，才是企业级的做法。

---

公众号:蓝海资料掘金营，微信deep3321