Qwen3模型格式转换与部署：从PyTorch到ONNX再到TensorRT

铭信

108人浏览 · 2026-03-04 01:31:19

铭信 · 2026-03-04 01:31:19 发布

Qwen3模型格式转换与部署：从PyTorch到ONNX再到TensorRT

如果你手头有一个训练好的Qwen3模型，想把它部署到实际应用里，比如做个智能客服或者内容生成工具，可能会发现直接拿PyTorch模型去推理，速度不够快，资源占用也大。这时候，模型格式转换和优化就成了关键一步。

今天，我就来带你走一遍这个流程：怎么把一个PyTorch格式的Qwen3模型，先转成ONNX，做个初步优化，最后再变成TensorRT引擎，让它能在NVIDIA的GPU上跑得飞快。整个过程听起来有点技术性，但别担心，我会尽量用大白话，一步步拆开讲清楚。

1. 准备工作与环境搭建

在开始转换之前，我们得先把“厨房”收拾好，把需要的“厨具”备齐。这里主要就是安装几个关键的Python库。

首先，确保你的Python环境是3.8或以上版本。然后，我们通过pip来安装核心工具。

# 安装PyTorch（请根据你的CUDA版本选择，这里以CUDA 11.8为例）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装Hugging Face Transformers库，用于加载Qwen3模型
pip install transformers

# 安装ONNX和ONNX Runtime，用于模型转换和初步推理
pip install onnx onnxruntime-gpu

# 安装TensorRT相关的Python包
# 注意：TensorRT的安装稍微复杂一些，通常需要从NVIDIA官网下载对应版本的tar包或使用deb/rpm包安装。
# 这里假设你已经安装了TensorRT，并提供了Python绑定的安装方式。
pip install nvidia-pyindex
pip install tensorrt

除了Python包，你还需要一个训练好的Qwen3模型。你可以从Hugging Face Model Hub上下载官方预训练模型，或者使用你自己微调好的模型检查点。我们假设你的模型目录结构是这样的：

/path/to/your/qwen3_model/
├── config.json
├── generation_config.json
├── model.safetensors  # 或 pytorch_model.bin
└── tokenizer.json  # 以及其他tokenizer文件

最后，确认你的机器上有NVIDIA GPU，并且安装了正确版本的CUDA和cuDNN。这是TensorRT能够发挥性能的基础。你可以用nvidia-smi命令来检查。

2. 第一步：从PyTorch到ONNX

ONNX（Open Neural Network Exchange）是一个开放的模型格式标准，它像一个“中间翻译”，能让不同框架训练的模型互相转换和运行。我们把PyTorch模型转成ONNX，是迈向高性能部署的第一步。

2.1 理解转换的核心：追踪与脚本

PyTorch模型是动态的，而ONNX需要的是一个静态的计算图。转换的关键在于，我们要用一些样例输入（叫做“虚拟输入”或“dummy input”）去“跑一遍”模型，ONNX的导出工具会记录下这个计算过程，并把它固化成图。

对于Qwen3这样的大语言模型，我们主要关注其生成文本的核心部分——也就是通常说的model这个对象。转换时，我们需要指定输入和输出的名字、类型以及维度。

2.2 执行转换脚本

下面是一个将Qwen3模型（以Qwen2.5-7B-Instruct为例）转换为ONNX格式的Python脚本。我们主要导出其编码器部分用于序列生成任务。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import onnx

# 1. 加载模型和分词器
model_name = "Qwen/Qwen2.5-7B-Instruct"  # 或者你的本地路径
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # 使用半精度以节省内存
    device_map="auto",
    trust_remote_code=True
)
model.eval()  # 将模型设置为评估模式

# 2. 准备虚拟输入
# 对于文本生成，典型的输入是input_ids和attention_mask
batch_size = 1
seq_length = 32  # 一个初始序列长度
dummy_input_ids = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_length)).to("cuda")
dummy_attention_mask = torch.ones_like(dummy_input_ids).to("cuda")

# 3. 定义输入/输出名和动态轴
# 动态轴允许推理时输入可变长度的序列
input_names = ["input_ids", "attention_mask"]
output_names = ["logits"]
dynamic_axes = {
    'input_ids': {1: 'sequence_length'},  # 第1维（序列长度）是动态的
    'attention_mask': {1: 'sequence_length'},
    'logits': {1: 'sequence_length'}
}

# 4. 导出模型到ONNX
onnx_model_path = "qwen3_7b_instruct.onnx"
torch.onnx.export(
    model,
    (dummy_input_ids, dummy_attention_mask),
    onnx_model_path,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes=dynamic_axes,
    opset_version=14,  # 使用较新的ONNX算子集版本
    do_constant_folding=True,
)

print(f"模型已成功导出到: {onnx_model_path}")

# 5. (可选) 简单验证导出的ONNX模型
onnx_model = onnx.load(onnx_model_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型格式检查通过。")

运行这个脚本后，你会得到一个.onnx文件。这个文件包含了模型的结构和权重，可以被ONNX Runtime或其他支持ONNX的后端加载。

3. 第二步：用ONNX Runtime进行验证与初步优化

拿到ONNX文件后，先别急着往TensorRT转。我们应该先用ONNX Runtime跑一下，验证转换是否正确，同时它本身也提供了一些优化选项，可以提升推理速度。

3.1 验证转换正确性

最基本的验证就是确保ONNX模型和原来的PyTorch模型，在相同输入下，输出结果差不多。

import onnxruntime as ort
import numpy as np

# 1. 创建ONNX Runtime会话
providers = ['CUDAExecutionProvider']  # 使用GPU
session = ort.InferenceSession(onnx_model_path, providers=providers)

# 2. 准备与之前相同的输入（转移到CPU并转为numpy数组）
input_ids_np = dummy_input_ids.cpu().numpy()
attention_mask_np = dummy_attention_mask.cpu().numpy()

# 3. 运行ONNX推理
ort_inputs = {
    session.get_inputs()[0].name: input_ids_np,
    session.get_inputs()[1].name: attention_mask_np,
}
ort_outputs = session.run(None, ort_inputs)
onnx_logits = ort_outputs[0]

# 4. 运行PyTorch推理（用于对比）
with torch.no_grad():
    torch_outputs = model(dummy_input_ids, attention_mask=dummy_attention_mask)
    torch_logits = torch_outputs.logits.cpu().numpy()

# 5. 比较结果（允许微小的数值误差）
print(f"PyTorch输出形状: {torch_logits.shape}")
print(f"ONNX Runtime输出形状: {onnx_logits.shape}")
if np.allclose(torch_logits, onnx_logits, rtol=1e-03, atol=1e-05):
    print("✓ 转换验证通过，输出结果基本一致。")
else:
    print("⚠ 输出结果存在较大差异，需要检查转换过程。")

3.2 使用ONNX Runtime的性能优化

ONNX Runtime提供了图优化功能，可以自动进行算子融合、常量折叠等优化。我们可以直接使用它来生成一个优化后的模型，这个过程很快，而且优化后的模型依然可以在ONNX Runtime上运行，通常能获得即时的速度提升。

from onnxruntime.transformers import optimizer

# 使用ONNX Runtime的优化器
optimized_model = optimizer.optimize_model(
    onnx_model_path,
    model_type='gpt2',  # Qwen基于类似GPT的架构，可选'gpt2'
    num_heads=32,       # 根据你的模型配置修改，例如Qwen2.5-7B是32头
    hidden_size=4096,   # 根据你的模型配置修改
)
optimized_onnx_path = "qwen3_7b_instruct_optimized.onnx"
optimized_model.save_model_to_file(optimized_onnx_path)
print(f"优化后的模型已保存至: {optimized_onnx_path}")

经过优化后，你可以再次用上面的验证脚本跑一下优化后的模型，确保功能正常。现在，我们手里就有了一个经过验证和初步优化的ONNX模型，为转换到TensorRT打好了基础。

4. 第三步：转换为TensorRT引擎

这是追求极致性能的关键一步。TensorRT是NVIDIA推出的高性能深度学习推理SDK，它会对模型进行更深层次的优化，包括层融合、精度校准（如FP16/INT8）、内核自动调优等，并为特定的GPU生成高度优化的引擎。

4.1 使用trtexec工具进行转换

NVIDIA提供了一个非常方便的命令行工具trtexec，它通常随TensorRT一起安装。我们可以用它来将ONNX模型转换为TensorRT引擎（.plan或.engine文件）。

# 基础转换命令，使用FP16精度以平衡速度和精度
trtexec --onnx=./qwen3_7b_instruct_optimized.onnx \
        --saveEngine=./qwen3_7b_fp16.engine \
        --fp16 \
        --workspace=4096 \ # 指定最大工作空间大小(MiB)，大模型需要更多
        --minShapes=input_ids:1x1,attention_mask:1x1 \ # 最小输入形状
        --optShapes=input_ids:1x32,attention_mask:1x32 \ # 优化输入形状（常用尺寸）
        --maxShapes=input_ids:1x512,attention_mask:1x512 # 最大输入形状

# 如果你想尝试INT8量化以获得更快速度（可能需要校准数据）
# trtexec --onnx=./qwen3_7b_instruct_optimized.onnx \
#         --saveEngine=./qwen3_7b_int8.engine \
#         --int8 \
#         --workspace=4096 \
#         --calib=<校准数据缓存文件> \
#         ...（动态形状参数同上）

参数解释一下：

--onnx: 指定输入的ONNX模型路径。
--saveEngine: 指定输出的TensorRT引擎文件路径。
--fp16: 启用FP16（半精度）模式，能显著减少内存占用并提升速度，大多数情况下精度损失可接受。
--workspace: GPU工作空间大小。复杂的融合操作可能需要更多临时内存，如果转换失败提示内存不足，可以适当增大这个值。
--minShapes/optShapes/maxShapes: 这是动态形状配置，非常重要！它告诉TensorRT引擎需要支持的输入尺寸范围。optShapes是预期最常出现的尺寸，TensorRT会针对这个尺寸进行深度优化。

4.2 在Python中加载并运行TensorRT引擎

转换成功后，我们就可以在Python代码中加载这个.engine文件进行推理了。这里需要用到TensorRT的Python API。

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

class Qwen3TensorRTRunner:
    def __init__(self, engine_path):
        # 1. 加载TensorRT引擎
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f, trt.Runtime(self.logger) as runtime:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        
        # 2. 创建执行上下文
        self.context = self.engine.create_execution_context()
        
        # 3. 分配输入输出缓冲区（Host和Device）
        self.inputs, self.outputs, self.bindings = [], [], []
        self.stream = cuda.Stream()
        
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding)) * self.engine.max_batch_size
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            # 在GPU上分配内存
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            
            self.bindings.append(int(device_mem))
            if self.engine.binding_is_input(binding):
                self.inputs.append({'host': host_mem, 'device': device_mem})
            else:
                self.outputs.append({'host': host_mem, 'device': device_mem})
                
    def infer(self, input_ids_np, attention_mask_np):
        # 4. 将输入数据复制到GPU
        np.copyto(self.inputs[0]['host'], input_ids_np.ravel())
        np.copyto(self.inputs[1]['host'], attention_mask_np.ravel())
        
        for inp in self.inputs:
            cuda.memcpy_htod_async(inp['device'], inp['host'], self.stream)
        
        # 5. 设置动态输入形状（如果转换时指定了动态形状）
        # 假设我们只动态了序列长度维度
        self.context.set_binding_shape(0, input_ids_np.shape) # input_ids
        self.context.set_binding_shape(1, attention_mask_np.shape) # attention_mask
        
        # 6. 执行推理
        self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
        
        # 7. 将输出数据从GPU拷贝回CPU
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
        
        self.stream.synchronize()
        
        # 8. 处理输出（这里假设第一个输出是logits）
        output = self.outputs[0]['host']
        # 根据实际的输出形状进行reshape，这里需要根据模型定义来调整
        # 例如，可能需要reshape为 (batch_size, seq_length, vocab_size)
        return output
    
    def __del__(self):
        # 清理资源
        with self.engine, self.context:
            pass

# 使用示例
if __name__ == "__main__":
    runner = Qwen3TensorRTRunner("./qwen3_7b_fp16.engine")
    
    # 准备一个简单的输入
    test_input_ids = np.array([[100, 200, 300]], dtype=np.int32)
    test_attention_mask = np.array([[1, 1, 1]], dtype=np.int32)
    
    # 进行推理
    trt_output = runner.infer(test_input_ids, test_attention_mask)
    print(f"TensorRT推理完成，输出logits形状（展平后）: {trt_output.shape}")
    # 注意：实际使用时需要将展平的输出重塑为合适的形状，并接上后续的解码（如beam search）逻辑。

这段代码提供了一个基本的TensorRT引擎加载和推理框架。在实际的文本生成场景中，你还需要在它外面包裹一个循环，来实现自回归的token生成过程。

5. 总结与后续步骤

走完这三步——从PyTorch到ONNX，再到TensorRT——我们相当于给Qwen3模型做了一次“深度改装”，让它从实验室环境走向了高性能的生产部署。

用ONNX Runtime验证和初步优化，好比是出厂前的质检和基础调校，确保模型转换无误且运行顺畅。而转换到TensorRT，则是针对NVIDIA GPU这个特定“赛道”进行的专业级改装，通过极致的算子融合和内核优化，把推理速度推到一个新的高度。

在实际操作中，你可能会遇到一些挑战，比如动态形状的支持、更复杂的模型结构（如多模态）、或者INT8量化的精度校准。对于这些，TensorRT提供了更高级的Python API（builder, network, parser）让你能够进行更精细的控制。此外，整个流水线也可以和Triton Inference Server这样的推理服务框架结合，实现模型的高效托管和并发服务。

刚开始接触时，建议从FP16精度的动态形状转换开始，这是性价比最高的优化方式。等你熟悉了整个流程，再逐步尝试INT8量化等更进阶的优化手段。希望这篇教程能帮你顺利跨出模型高性能部署的第一步。