4、模型导出第一步：使用torch.onnx.export导出LLM模型、动态轴设置、输入输出命名规范

好，咱们进入实战环节了。模型导出，说白了就是把训练好的PyTorch模型，转成ONNX格式。这一步是通往TensorRT-LLM的必经之路。我个人习惯把这一步叫做“炼金术”——把PyTorch的“铁”炼成ONNX的“金”。

你想想看，PyTorch模型在推理时，动态图虽然灵活，但性能上不去。ONNX是静态图，能做的优化就多了。嗯，这里要注意，LLM模型和普通CNN模型不一样，它有几个“坑”需要特别小心。

4.1 torch.onnx.export 的基本用法

先看一个最基础的例子。假设我们有一个简单的LLM解码器层：

import torch
import torch.nn as nn

class SimpleDecoderLayer(nn.Module):
    def __init__(self, hidden_size=768):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(hidden_size, 12)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.ReLU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )
        self.ln1 = nn.LayerNorm(hidden_size)
        self.ln2 = nn.LayerNorm(hidden_size)
    
    def forward(self, x, attn_mask=None):
        # x: [seq_len, batch, hidden]
        attn_out, _ = self.self_attn(x, x, x, attn_mask=attn_mask)
        x = self.ln1(x + attn_out)
        ffn_out = self.ffn(x)
        x = self.ln2(x + ffn_out)
        return x

model = SimpleDecoderLayer()
dummy_input = torch.randn(10, 1, 768)  # [seq_len=10, batch=1, hidden=768]

torch.onnx.export(
    model,
    dummy_input,
    "decoder_layer.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=17
)

这段代码能跑通，但问题很多。我在项目中遇到过，直接用这种写法导出LLM模型，后面做TensorRT推理时，性能惨不忍睹。为什么？因为动态轴没设置。

4.2 动态轴设置——LLM模型的生命线

LLM模型最核心的特点是什么？变长输入。你想想看，用户问的问题有长有短，生成的回答也有长有短。如果ONNX模型把序列长度固定死了，那每次推理都得padding到最大长度，浪费大量算力。

动态轴设置，就是告诉ONNX：这个维度是可变的。具体做法是通过dynamic_axes参数：

dynamic_axes = {
    "input": {0: "seq_len", 1: "batch_size"},
    "output": {0: "seq_len", 1: "batch_size"}
}

torch.onnx.export(
    model,
    dummy_input,
    "decoder_layer_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=dynamic_axes,
    opset_version=17
)

这里我解释一下：{0: "seq_len"}表示第0维（序列长度维度）是动态的，名字叫"seq_len"。第1维（batch维度）也是动态的，名字叫"batch_size"。

核心要点：LLM模型导出时，seq_len和batch_size这两个维度必须设置为动态轴。否则后面做TensorRT推理时，你会遇到各种形状不匹配的错误。

我曾经踩过一个坑：只设置了seq_len为动态，但忘了设置batch_size。结果模型只能以batch_size=1运行，多卡推理时直接报错。嗯，这个教训让我记住了——两个维度都要设。

4.3 输入输出命名规范——别小看这个

命名规范看似简单，但做不好后面会非常痛苦。我个人习惯遵循以下规则：

命名规则 示例 说明
使用小写+下划线 input_ids, attention_mask 避免大小写混淆，TensorRT中统一小写
输入名要有语义 input_ids 而不是 input 方便后续调试和配置
输出名要明确 logits, hidden_states 不要用 output 这种模糊名字
特殊张量加前缀 past_key_values, present_key_values KV cache相关张量要能一眼识别

举个例子，一个标准的LLM模型导出配置应该是这样的：

input_names = ["input_ids", "attention_mask", "position_ids"]
output_names = ["logits", "present_key_values"]

dynamic_axes = {
    "input_ids": {0: "batch_size", 1: "seq_len"},
    "attention_mask": {0: "batch_size", 1: "seq_len"},
    "position_ids": {0: "batch_size", 1: "seq_len"},
    "logits": {0: "batch_size", 1: "seq_len"},
    "present_key_values": {0: "batch_size", 2: "seq_len"}
}

小技巧：命名时尽量和HuggingFace的命名保持一致。这样后面写TensorRT的配置文件时，可以直接复用，减少出错概率。

4.4 完整导出流程与常见问题

下面是一个完整的LLM模型导出流程。我以GPT-2为例，但思路适用于所有decoder-only模型：

import torch
from transformers import GPT2Model, GPT2Config

# 1. 加载模型
config = GPT2Config()
model = GPT2Model(config)
model.eval()

# 2. 准备dummy输入
batch_size = 1
seq_len = 10
dummy_input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
dummy_attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)
dummy_position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)

# 3. 定义动态轴
dynamic_axes = {
    "input_ids": {0: "batch_size", 1: "seq_len"},
    "attention_mask": {0: "batch_size", 1: "seq_len"},
    "position_ids": {0: "batch_size", 1: "seq_len"},
    "last_hidden_state": {0: "batch_size", 1: "seq_len"},
    "past_key_values": {0: "batch_size", 2: "seq_len"}
}

# 4. 导出
torch.onnx.export(
    model,
    (dummy_input_ids, dummy_attention_mask, dummy_position_ids),
    "gpt2_model.onnx",
    input_names=["input_ids", "attention_mask", "position_ids"],
    output_names=["last_hidden_state", "past_key_values"],
    dynamic_axes=dynamic_axes,
    opset_version=17,
    do_constant_folding=True,
    verbose=True  # 建议打开，可以看到导出过程中的详细信息
)

print("导出成功！")

警告：导出时如果遇到"Unsupported operator"错误，通常是因为opset版本太低。建议使用opset_version=17或更高。如果还不行，可能需要自定义ONNX算子。

我在项目中遇到过最头疼的问题，是导出时模型内部有if条件分支。ONNX不支持动态控制流，所以导出前一定要确保模型的前向传播路径是固定的。解决办法是：用torch.jit.script先转成TorchScript，再导出ONNX。但这是另一个话题了，咱们后面再细聊。

4.5 验证导出的ONNX模型

导出完别急着走，先验证一下。用onnxruntime跑一遍，确保结果和PyTorch一致：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
ort_session = ort.InferenceSession("gpt2_model.onnx")

# 准备输入
inputs = {
    "input_ids": dummy_input_ids.numpy(),
    "attention_mask": dummy_attention_mask.numpy(),
    "position_ids": dummy_position_ids.numpy()
}

# 推理
outputs = ort_session.run(None, inputs)

# 和PyTorch结果对比
with torch.no_grad():
    pt_outputs = model(dummy_input_ids, dummy_attention_mask, dummy_position_ids)

print(f"ONNX输出形状: {outputs[0].shape}")
print(f"PyTorch输出形状: {pt_outputs[0].shape}")
print(f"最大误差: {np.max(np.abs(outputs[0] - pt_outputs[0].numpy()))}")

如果最大误差在1e-5以内，恭喜你，导出成功了。如果误差很大，检查一下模型是否在eval模式，或者有没有dropout层没关掉。

总结一下今天的内容：

  • 动态轴设置是LLM模型导出的关键，seq_len和batch_size必须设为动态
  • 命名规范要统一，建议和HuggingFace保持一致
  • 导出后一定要用onnxruntime验证，确保精度无损
  • 遇到控制流问题，考虑先用TorchScript过渡

嗯，模型导出这一步看似简单，但做不好后面全白费。我见过太多人在这步翻车，然后花大量时间排查问题。所以，耐心点，把每一步都做扎实了。

ONNX导出核心流程 PyTorch模型 .eval() + 关闭dropout 准备dummy输入 batch=1, seq=10 设置动态轴 seq_len + batch_size 输入输出命名 小写+下划线+语义化 torch.onnx.export opset_version=17 ONNX验证 onnxruntime推理对比 常见问题 ⚠ 控制流不支持 ⚠ opset版本过低 ⚠ 动态轴遗漏 ⚠ 命名冲突 ⚠ 精度误差过大 核心：动态轴 + 命名规范 + 验证 = 成功的ONNX导出

好了，这一章的内容就到这里。记住，导出ONNX只是第一步，后面还有更多挑战等着我们。但只要你把这一步做扎实了，后面的路就好走多了。

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

← 上一章 📖 返回目录 下一章 →