从Llama Factory到ONNX：跨平台模型导出全攻略-平芜编程栈

从Llama Factory到ONNX：跨平台模型导出全攻略

如果你已经使用Llama Factory完成了大语言模型的微调，接下来可能会面临一个关键问题：如何将微调后的模型部署到不同的运行时环境中？本文将详细介绍如何将Llama Factory的输出转换为ONNX等通用格式，实现真正的跨平台部署。

这类任务通常需要GPU环境支持，目前CSDN算力平台提供了包含相关工具的预置环境，可以快速部署验证。但无论你选择哪种运行环境，下面的技术方案都是通用的。

为什么需要模型格式转换

在模型部署过程中，我们经常会遇到以下痛点：

生产环境可能使用不同的推理框架（如TensorRT、ONNX Runtime等）
边缘设备对模型格式有特殊要求
需要优化模型在不同硬件上的性能表现

ONNX（Open Neural Network Exchange）作为一种开放的模型格式标准，能够很好地解决这些问题。它允许你在PyTorch中训练模型，然后将其导出为通用格式，在其他平台上运行。

准备工作：检查你的微调模型

在开始转换前，请确保你已经完成了以下准备工作：

确认微调后的模型保存完整（通常包含以下文件）：
pytorch_model.bin或adapter_model.bin
config.json
tokenizer.json或相关分词器文件
检查你的运行环境是否满足要求：
Python 3.8+
PyTorch 1.12+（建议使用与训练时相同的版本）
ONNX运行时环境
安装必要的转换工具：bash pip install transformers onnx onnxruntime

从Llama Factory模型到ONNX的转换步骤

1. 加载微调后的模型

首先，我们需要使用transformers库加载微调后的模型：

from transformers import AutoModelForCausalLM, AutoTokenizer model_path = "path_to_your_finetuned_model" tokenizer = AutoTokenizer.from_pretrained(model_path) model = AutoModelForCausalLM.from_pretrained(model_path)

注意：如果遇到显存不足的问题，可以尝试添加device_map="auto"参数让transformers自动分配设备。

2. 准备转换配置

ONNX转换需要指定输入的形状和类型。对于语言模型，典型的配置如下：

import torch dummy_input = { "input_ids": torch.randint(0, 100, (1, 128), dtype=torch.long), "attention_mask": torch.ones((1, 128), dtype=torch.long), }

3. 执行ONNX导出

使用torch.onnx.export函数将模型转换为ONNX格式：

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

关键参数说明： -dynamic_axes：指定哪些维度可以是动态的（如可变长度的输入） -opset_version：ONNX操作集版本（建议使用13或更高）

转换过程中的常见问题与解决方案

1. 显存不足问题

大模型转换可能需要大量显存。如果遇到OOM错误，可以尝试：

减小输入尺寸（如将序列长度从512降到256）
使用CPU进行转换（添加device="cpu"参数）
启用内存优化选项：python torch.onnx.export(..., do_constant_folding=True)

2. 操作不支持问题

某些PyTorch操作可能没有对应的ONNX实现。解决方法包括：

检查并更新PyTorch和ONNX版本
实现自定义操作符
使用替代操作

3. 精度损失问题

转换过程中可能出现精度损失，建议：

比较原始模型和ONNX模型的输出差异
使用FP32精度进行转换
检查量化配置（如果有）

验证转换后的ONNX模型

转换完成后，应该验证ONNX模型的正确性：

import onnxruntime as ort # 创建ONNX运行时会话 sess = ort.InferenceSession("model.onnx") # 准备测试输入 inputs = { "input_ids": dummy_input["input_ids"].numpy(), "attention_mask": dummy_input["attention_mask"].numpy(), } # 运行推理 outputs = sess.run(None, inputs) print(outputs[0].shape) # 应该得到与原始模型相同的输出形状

进阶技巧：优化ONNX模型性能

转换后的ONNX模型还可以进一步优化：

图优化：使用ONNX Runtime的图优化功能python sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
量化：减小模型大小，提高推理速度python from onnxruntime.quantization import quantize_dynamic quantize_dynamic("model.onnx", "model_quant.onnx")
形状推断：为动态形状提供提示python sess_options.add_free_dimension_override_by_name("batch_size", 1) sess_options.add_free_dimension_override_by_name("sequence_length", 128)