本章涵盖:
- 对移植到 ONNX 的模型进行 profiling
- 将原始 ONNX profiling 数据转化为洞察
- 为 LLM 优化 ONNX 图
第 5 章介绍了 ONNX 格式和 ONNX Runtime 的能力,先从通用角度介绍,然后介绍其在 LLM 中的应用。第 6 章介绍了几种方法,包括通过 ONNX API 进行 8-bit 量化。本章将探索 ONNX 的其他特性,例如对移植到 ONNX 格式的 LLM 进行性能 profiling,以及可以用来从 profiling 数据中提取有用洞察的工具。
10.1 对移植到 ONNX 的 LLM 进行 Profiling
ONNX Runtime,也就是 ORT,可以在广泛硬件上高性能运行机器学习,也就是 ML,以及深度学习,也就是 DL,模型。不过,为了满足特定用例、模型和设备在延迟、吞吐和内存使用方面的关键性能指标,也就是 KPI,你可能还需要额外的模型优化和 runtime 配置。
ORT 支持在代码中进行性能 profiling。它默认是禁用的,但你可以在调试时通过在 ORT inference session 中把 enable_profiling 选项设置为 True 来启用它:
import onnxruntime as rt
sess_options = rt.SessionOptions()
sess_options.enable_profiling = True
运行时,ORT 会生成一个 JSON 文件,其中包含性能数据,例如线程信息和逐算子延迟。
我们先从一个简单模型开始,看看这个 profiling 文件中包含什么,然后再转向 LLM。我们将复用第 5.2 节中的线性回归示例:
from onnx import TensorProto
from onnx.helper import (
make_model, make_node, make_graph,
make_tensor_value_info)
from onnx.checker import check_model
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
node1 = make_node('MatMul', ['X', 'A'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])
graph = make_graph([node1, node2], 'lr', [X, A, B], [Y])
onnx_model = make_model(graph)
check_model(onnx_model)
为了使用这个模型运行推理,我们需要创建一个 inference session,并指定 execution provider,这里使用 CPU,不需要硬件加速:
from onnxruntime import InferenceSession
sess = InferenceSession(onnx_model.SerializeToString(),
providers=["CPUExecutionProvider"])
我们用三个随机值输入数组测试它:
import numpy as np
res = sess.run(None, {
'X': np.random.randn(2, 2).astype(np.float32),
'A': nup.random.randn(2, 2).astype(np.float32),
'B': np.random.randn(2, 2).astype(np.float32)
}
)
为了测量这个模型的性能,并发现改进方向,我们会在启用 profiling 的情况下运行推理。除此之外,不需要其他代码改动:
import onnxruntime as rt
from onnxruntime import InferenceSession
sess_options = rt.SessionOptions()
sess_options.enable_profiling = True
sess = InferenceSession(onnx_model.SerializeToString(),
sess_options,
providers=["CPUExecutionProvider"])
为了访问生成的 profiling 数据,我们需要使用 end_profiling 方法:
prof = sess.end_profiling()
这个方法会返回生成的 JSON 文件名,该文件遵循 onnxruntime_profile__YYYY-MM-DD_hh-mm-ss.json 命名约定。文件内容看起来如下:
[{"cat" : "Session","pid" :21460,"tid" :14412,"dur" :36164,"ts" :7,"ph" : "X","name" :"model_loading_array","args" : {}},{"cat" : "Session","pid" :21460,"tid" :14412,"dur" :1430,"ts" :36425,"ph" : "X","name" :"session_initialization","args" : {}},{"cat" : "Node","pid" :21460,"tid" :14412,"dur" :2,"ts" :34744604,"ph" : "X","name" :"MatMul_0_fence_before","args" : {"op_name" : "MatMul"}},{"cat" : "Node","pid" :21460,"tid" :14412,"dur" :164,"ts" :34744616,"ph" : "X","name" :"MatMul_0_kernel_time","args" : {"parameter_size" : "0","provider" : "CPUExecutionProvider","op_name" : "MatMul","input_type_shape" : [{"float":[2,2]},{"float":[2,2]}],"node_index" : "0",
"output_type_shape" : [{"float":[2,2]}],"activation_size" :
"32","output_size" : "16","thread_scheduling_stats" : {"main_thread":
{"thread_pool_name": "session-2-intra-op", "thread_id": "14412",
"block_size": [], "core": -1, "Distribution": 0, "DistributionEnqueue": 0,
"Run": 0, "Wait": 0, "WaitRevoke": 0}, "sub_threads": {"20796": {"num_run":
0, "core": -1}}}}},
{"cat" : "Node","pid" :21460,"tid" :14412,"dur" :0,"ts" :34744793,"ph" :
"X","name" :"MatMul_0_fence_after","args" : {"op_name" : "MatMul"}},
{"cat" : "Node","pid" :21460,"tid" :14412,"dur" :0,"ts" :34744806,"ph" :
"X","name" :"Add_1_fence_before","args" : {"op_name" : "Add"}},
{"cat" : "Node","pid" :21460,"tid" :14412,"dur" :56,"ts" :34744808,"ph" :
"X","name" :"Add_1_kernel_time","args" : {"parameter_size" : "0","provider"
: "CPUExecutionProvider","op_name" : "Add","input_type_shape" :
[{"float":[2,2]},{"float":[2,2]}],"node_index" : "1","output_type_shape" :
[{"float":[2,2]}],"activation_size" : "32","output_size" :
"16","thread_scheduling_stats" : {"main_thread": {"thread_pool_name":
"session-2-intra-op", "thread_id": "14412", "block_size": [], "core": -1,
"Distribution": 0, "DistributionEnqueue": 0, "Run": 0, "Wait": 0,
"WaitRevoke": 0}, "sub_threads": {"20796": {"num_run": 0, "core": -1}}}}},
{"cat" : "Node","pid" :21460,"tid" :14412,"dur" :0,"ts" :34744879,"ph" :
"X","name" :"Add_1_fence_after","args" : {"op_name" : "Add"}},
{"cat" : "Session","pid" :21460,"tid" :14412,"dur" :781,"ts" :34744111,"ph"
: "X","name" :"SequentialExecutor::Execute","args" : {}},
{"cat" : "Session","pid" :21460,"tid" :14412,"dur" :828,"ts" :34744075,"ph"
: "X","name" :"model_run","args" : {}}
]
这个 JSON profiling 文件包含一个 ONNX inference session 的原始数据,并按算子和 execution provider 进行拆分。单独看这些数据,它们提供的洞察很有限,因此你需要聚合和转换它们,才能理解模型性能。ONNX Runtime 包含 Ort_perf_view,它可以读取这个 profiling JSON 文件,把汇总统计渲染为 treemap,并计算每个 ONNX 算子的延迟,CPU 和 GPU 都支持。对于小模型,例如本示例中的线性回归器,这个视图很清晰,见图 10.1,因为我们只使用了两个 ONNX 算子。
图 10.1 —— Ort_perf_view 对线性回归 ONNX 模型推理性能统计的渲染结果
Ort_perf_view 可视化会为每个算子显示一个方块,并标注算子名称和该 profiling run 中的调用次数。方框大小反映算子调用所占比例。对于更大模型,这个视图对性能分析就不那么有用,虽然它仍然可以提供一个关于聚合算子调用的高层概览,见图 10.2。
图 10.2 —— Ort_perf_view 输出,展示大型 ONNX 模型一次推理运行的性能统计
在图 10.2 中,你可以看到,被测试模型的 ONNX 图中每个算子都有成千上万次调用。点击某个算子的方块,会展开该部分,显示每一次调用的详细信息,包括输入、输出值和以微秒为单位的持续时间,如图 10.3 所示。
图 10.3 —— Ort_perf_view 展开视图,展示被测模型推理期间所有 Sqrt 算子调用
这种可视化并不是特别有帮助;它只是展示没有洞察的原始数据。接下来,我们换一种更好的方式。
10.2 将原始 ONNX Profiling 数据转化为洞察
在本节中,我们将看一个自定义实现,它可以从原始 ONNX profiling 数据中提取洞察,并以用户友好的方式呈现。为了让示例简单且易于理解,我们会使用一个在 scikit-learn 中训练、然后转换为 ONNX 格式的线性回归模型。相同的 profiling 过程、数据转换和可视化代码,也适用于任何转换为 ONNX 的模型,包括 LLM。
我们会使用 scikit-learn 的 diabetes 数据集作为训练集和测试集。它包含 442 个样本,每个样本有 10 个特征,包括年龄、性别、身体质量指数、平均血压,以及 6 项血清测量指标。目标变量是一年后疾病进展的定量度量。
先下载数据集:
from sklearn.datasets import load_diabetes
data = load_diabetes()
接下来,把数据拆分成训练集和测试集:
from sklearn.model_selection import train_test_split
X, y = data.data, data.target
X_train, X_test, y_train, __ = train_test_split(X, y, random_state=11)
现在可以使用 scikit-learn API 训练模型:
from sklearn.linear_model import LinearRegression
clr = LinearRegression()
clr.fit(X_train, y_train)
可以使用 sklearn-onnx 包将训练好的模型转换为 ONNX 格式:
from skl2onnx import to_onnx
model_def = to_onnx(clr, X_train)
转换后的模型图看起来如下:
ir_version: 8
producer_name: "skl2onnx"
producer_version: "1.14.0"
domain: "ai.onnx"
model_version: 0
doc_string: ""
graph {
node {
input: "X"
input: "coef"
output: "multiplied"
name: "MatMul"
op_type: "MatMul"
domain: ""
}
node {
input: "multiplied"
input: "intercept"
output: "resh"
name: "Add"
op_type: "Add"
domain: ""
}
node {
input: "resh"
input: "shape_tensor"
output: "variable"
name: "Reshape"
op_type: "Reshape"
domain: ""
}
name: "ONNX(LinearRegression)"
initializer {
dims: 10
dims: 1
data_type: 11
name: "coef"
double_data: -60.219814056772776
double_data: -266.4570523220356
double_data: 523.063411259699
double_data: 310.5134699670048
double_data: -336.1614738106952
double_data: 137.33929238424574
double_data: -131.13923283453687
double_data: -1.1492348017316942
double_data: 622.3286855688127
double_data: 60.46645769168251
}
initializer {
dims: 1
data_type: 11
name: "intercept"
double_data: 152.22822762230342
}
initializer {
dims: 2
data_type: 7
int64_data: -1
int64_data: 1
name: "shape_tensor"
}
input {
name: "X"
type {
tensor_type {
elem_type: 11
shape {
dim {
}
dim {
dim_value: 10
}
}
}
}
}
output {
name: "variable"
type {
tensor_type {
elem_type: 11
shape {
dim {
}
dim {
dim_value: 1
}
}
}
}
}
}
opset_import {
domain: ""
version: 13
我们可以通过创建启用 profiling 的 inference session,对这个转换后的模型进行 profiling,如 9.1 节所述:
from onnxruntime import InferenceSession, SessionOptions
so = SessionOptions()
so.enable_profiling = True
sess = InferenceSession(model_def.SerializeToString(), so)
我们会在测试集上运行推理,共 111 个样本,并在最后停止 profiling:
sess.run(None, {'X': X_test})
prof = sess.end_profiling()
现在可以从 JSON 文件中加载原始 profiling 数据:
import json
with open(prof, "r") as f:
js = json.load(f)
接下来,从这些数据创建一个 pandas DataFrame:
from mlprodict.onnxrt.ops_whole.session import OnnxWholeSession
df = pd.DataFrame(OnnxWholeSession.process_profiling(js))
我们使用 mlprodict Python 库中的 OnnxWholeSession 类,将 JSON 内容转换为更清晰的表格格式,见图 10.4。
图 10.4 —— 回归模型示例 profiling 数据的表格格式前几行
从这种更清晰的数据表示开始,构建可视化和提取洞察就容易多了。我们将使用 Plotly Express Python 图表库。为了进行性能分析,我们希望比较模型图中每个 ONNX 算子的总持续时间。为此,我们首先按算子类型分组,并计算每个算子的总持续时间:
gr_dur = df[['dur', "args_op_name"]].groupby("args_op_name")
➥.sum().sort_values('dur')
然后可以把这些数据展示为横向柱状图,其中 x 轴是总持续时间,以毫秒为单位,y 轴是 ONNX 算子类型,见图 10.5:
图 10.5 —— 线性回归模型图中三个 ONNX 算子的总持续时间,单位为毫秒
fig = px.bar(gr_dur, x='dur',
labels={
"dur": "Duration (ms)",
"args_op_name": "Operation type",
},
title='Duration')
fig.show()
这个可视化可以快速显示,在模型推理期间哪些算子耗时最多。另一个有用指标是每个算子被调用的频率。为了可视化这一点,先按算子类型聚合,然后统计每种类型出现次数,并按持续时间排序,使图表与总持续时间视图匹配:
gr_n = df[['dur', "args_op_name"]].groupby("args_op_name")
➥.count().sort_values('dur')
gr_n = gr_n.loc[gr_dur.index, :]
这些信息可以展示为横向柱状图,其中 x 轴是某个算子的出现次数,y 轴是 ONNX 算子类型,见图 10.6:
图 10.6 —— 线性回归模型图中三个 ONNX 算子的调用次数
fig = px.bar(gr_n, x='dur',
labels={
"dur": "Op count",
"args_op_name": "Operation type",
},
title='Occurrences')
fig.show()
为了进行更完整的性能分析,不仅展示每个算子运行多少次,还要展示它占总推理时间的比例。为此,我们先计算每个算子占总推理时间的百分比:
gr_dur /= gr_dur['dur'].sum()
现在基于这些数据创建可视化,见图 10.7:
图 10.7 —— 线性回归模型中每个 ONNX 算子调用占总推理时间的百分比
fig = px.bar(gr_dur, x='dur',
labels={
"dur": "Duration (%)",
"args_op_name": "Operation type",
},
title='Proportion')
fig.show()
这三个可视化可以帮助我们快速发现 ONNX 模型中图优化的机会。
对于 LLM,我们可以不加修改地复用相同的数据清洗、聚合和可视化代码;本章也提供了配套 Colab notebook。该 notebook 覆盖了将 GPT-2 small 模型转换为 ONNX、运行 inline profiling,以及分析和可视化 profiling 数据。由于前面章节已经介绍过加载模型和 tokenizer,以及转换到 ONNX 的代码,这里不再重复。我们只看对 ONNX GPT-2 small 模型 profiling 后生成的图表,见图 10.8 到图 10.10。
在图 10.8 中,x 轴显示持续时间,单位为毫秒,y 轴列出了 profiling 运行中的 ONNX 操作类型。
图 10.8 —— GPT-small 线性回归模型图中 ONNX 算子的总持续时间,单位为毫秒
在图 10.9 中,x 轴显示每个 ONNX 算子被调用了多少次,y 轴列出了 profiling 运行中使用的算子类型。
图 10.9 —— GPT-2 small 模型图中的 ONNX 算子调用次数
在图 10.10 中,x 轴显示总推理时间百分比,y 轴列出了 profiling 运行中的 ONNX 操作类型。
图 10.10 —— GPT-2 small 模型 profiling 过程中每个 ONNX 算子调用占整体推理时间的百分比
10.3 针对 LLM 优化 ONNX 图
Python 版 ORT 可以评估已经转换为 ONNX 的模型——不仅可以评估完整图,也可以评估单个节点,包括自定义节点。它还可以捕获图或节点内的中间结果。转换到新模型版本,或者从 Python ML/DL 框架迁移到 ONNX 图,都可能引入低效之处;评估一个或多个节点可以帮助你发现并修复它们。
先考虑第 10.1 节中的简单线性回归示例模型。为了评估模型的完整 ONNX 图,可以使用 ONNX 的 ReferenceEvaluator 类:
from onnx.reference import ReferenceEvaluator
sess = ReferenceEvaluator(onnx_model)
像往常一样准备模型输入数组:
import numpy as np
x = np.random.randn(4, 2).astype(np.float32)
a = np.random.randn(2, 1).astype(np.float32)
b = np.random.randn(1, 1).astype(np.float32)
feeds = {'X': x, 'A': a, 'B': b}
然后使用前面创建的 ReferenceEvaluator 实例运行模型:
evaluation = sess.run(None, feeds)
我们也可以评估图中的单个节点。对于回归模型中的 MatMul 节点,可以这样做:
node_sess = ReferenceEvaluator(node1)
feeds = {'X': x, 'A': a}
node_evaluation = node_sess.run(None, feeds)
可以通过在 ReferenceEvaluator 实例上把 verbose 参数设置为不同详细程度,启用逐步评估:
sess = ReferenceEvaluator(onnx_model, verbose=verbose)
verbose 参数接受 1 到 4 的整数值。1 是默认值,输出较少;4 是最详细的逐步评估。对于我们的线性回归模型,下面的输出展示了不同详细级别:
------ verbose=1
[array([[0.9837985],
[1.126987 ],
[0.8049259],
[1.6867123]], dtype=float32)]
------ verbose=2
MatMul(X, A) -> XA
Add(XA, B) -> Y
[array([[0.31629467],
[1.5107596 ],
[1.008173 ],
[1.2464874 ]], dtype=float32)]
------ verbose=3
+I X: float32:(4, 2) in [-0.8688125014305115, 0.7680516242980957]
+I A: float32:(2, 1) in [0.6488013863563538, 1.9162827730178833]
+I B: float32:(1, 1) in [2.00530743598938, 2.00530743598938]
MatMul(X, A) -> XA
+ XA: float32:(4, 1) in [-1.2202415466308594, 1.2917602062225342]
Add(XA, B) -> Y
+ Y: float32:(4, 1) in [0.7850658893585205, 3.297067642211914]
[array([[2.835943 ],
[3.2970676],
[1.3345237],
[0.7850659]], dtype=float32)]
------ verbose=4
+I X: float32:(4, 2):1.4935364723205566,-
0.1639600694179535,1.1996209621429443,-0.47337841987609863,-
2.4281582832336426...
+I A: float32:(2, 1):[0.8013423681259155, -0.34330955147743225]
+I B: float32:(1, 1):[-1.6645238399505615]
MatMul(X, A) -> XA
+ XA: float32:(4, 1):[1.253123164176941, 1.1238224506378174,
-2.082425594329834, -0.8871934413909912]
Add(XA, B) -> Y
+ Y: float32:(4, 1):[-0.4114006757736206, -0.5407013893127441,
-3.7469494342803955, -2.5517172813415527]
[array([[-0.41140068],
[-0.5407014 ],
[-3.7469494 ],
[-2.5517173 ]], dtype=float32)]
如前面的代码所示,详细级别 1 只显示模型输出;级别 2 还会显示输入之间的操作序列;级别 3 还包含输入值和中间输出;级别 4 则包含中间步骤的输入。
这个示例中的线性回归遵循如下公式:
y = Add(MatMul(X, A), B)
实现线性回归的另一种方式,是使用类似公式,但向矩阵 A 添加单位矩阵:
y = Add(MatMul(X, Add(A, I)), B)
为了实现它,需要修改原来使用 ONNX 算子构建模型的代码:
node0 = make_node('EyeLike', ['A'], ['Eye'])
node1 = make_node('Add', ['A', 'Eye'], ['A1'])
node2 = make_node('MatMul', ['X', 'A1'], ['XA1'])
node3 = make_node('Add', ['XA1', 'B'], ['Y'])
graph = make_graph([node0, node1, node2, node3], 'lr', [X, A, B], [Y])
在这个示例中,我们使用 EyeLike 算子,它是 ONNX 对单位矩阵的实现。单位矩阵是主对角线上为 1,其余位置为 0 的方阵。然后我们组合这些节点,为更新后的线性回归公式构建最终图。
如本节前面所说,你也可以评估自定义 ONNX 节点。例如,在下面代码中,我们会创建一个自定义节点,将 EyeLike 和 Add 算子结合起来,然后看看模型运行是否更高效:
node01 = make_node('AddEyeLike', ['A'], ['A1'], domain='combined')
然后可以重建图:
node2 = make_node('MatMul', ['X', 'A1'], ['XA1'])
node3 = make_node('Add', ['XA1', 'B'], ['Y'])
graph = make_graph([node01, node2, node3], 'lr', [X, A, B], [Y])
重新构建模型并再次检查:
from onnx.helper import make_opsetid
onnx_optimized_model = make_model(graph, opset_imports=[
make_opsetid('', 18), make_opsetid('combined', 1)
])
check_model(onnx_optimized_model)
这种优化——把两个算子合并为一个优化算子,也就是 AddEyeLike——被称为 fusion。为了评估自定义节点,或者包含该自定义节点的模型图,你必须实现一个自定义类,因为 ReferenceEvaluator 原生只支持内置 ONNX 算子。不过,它也可以接受自定义算子实现。你可以通过扩展 ONNX 的 OpRun 类,并实现其 _run 方法来实现自定义算子:
from onnx.reference.op_run import OpRun
class AddEyeLike(OpRun):
op_domain = "combined"
def _run(self, X, alpha=1.):
assert len(X.shape) == 2
assert X.shape[0] == X.shape[1]
X = X.copy()
ind = numpy.diag_indices(X.shape[0])
X[ind] += alpha
return (X,)
现在可以像往常一样评估模型。唯一不同的是,在创建 ReferenceEvaluator 实例时,还需要传入模型 ONNX 图中包含的自定义节点实现列表:
sess = ReferenceEvaluator(onnx_optimized_model, verbose=2,
➥ new_ops=[AddEyeLike])
对于给定输入,评估过程与不包含自定义算子的模型一样:
x = numpy.random.randn(4, 2).astype(numpy.float32)
a = numpy.random.randn(2, 2).astype(numpy.float32) / 10
b = numpy.random.randn(1, 2).astype(numpy.float32)
feeds = {'X': x, 'A': a, 'B': b}
sess.run(None, feeds)
在测量任何执行加速之前,先确认优化没有降低预测质量。我们会评估两个模型版本:原始版本和融合两个算子之后的版本:
sess_original = ReferenceEvaluator(onnx_model)
sess_optimized = ReferenceEvaluator(onnx_optimized_model,
➥ new_ops=[AddEyeLike])
y_original = sess_original.run(None, feeds)[0]
y_optimized = sess_optimized.run(None, feeds)[0]
print(f"output from original model: {y_original}")
print(f"output from optimized model: {y_optimized}")
print(f"difference: {numpy.abs(y_original - y_optimized).max()}")
两者没有差异:
output from original model: [[-1.3435526 0.11092019]
[ 0.8220897 0.18847765]
[-1.29872 0.882055 ]
[-0.10433175 -0.30559355]]
output from optimized model: [[-1.3435526 0.11092019]
[ 0.8220897 0.18847765]
[-1.29872 0.882055 ]
[-0.10433175 -0.30559355]]
difference: 0.0
现在用更大的矩阵 benchmark 两个模型版本,通过 Python 的 timeit 包计时 10,000 次运行:
x = numpy.random.randn(4, 100).astype(numpy.float32)
a = numpy.random.randn(100, 100).astype(numpy.float32) / 10
b = numpy.random.randn(1, 100).astype(numpy.float32)
feeds = {'X': x, 'A': a, 'B': b}
import timeit
print(f"time with EyeLike+Add: {timeit.timeit(lambda:
➥ sess_original.run(None, feeds), number=10000)}")
print(f"time with AddEyeLike: {timeit.timeit(lambda:
➥ sess_optimized.run(None, feeds), number=10000)}")
优化后的模型有明显改进:
time with EyeLike+Add: 1.1427911999999196
time with AddEyeLike: 0.8636034999999538
你也可以按照 7.1 和 7.2 节中的说明进行评估和 profiling。我在一台配有两颗 Intel Core i5-7300U CPU,2.6 GHz,和 8 GB RAM 的 PC 上运行了 benchmark。
在第 5.6.1 节中,我介绍了 ONNX Runtime,也就是 ORT 的优化级别,以及它们如何应用到 LLM。你也已经看到,图优化可以离线完成,也可以 inline 完成,并且 ORT 提供两组优化:basic,包括本节介绍的 operator fusion;以及 extended。我还介绍了 ORT API 中用于 Transformer 模型的内置 optimizer 类,它可以对 ONNX 模型,或者转换为 ONNX 格式的模型,自动执行离线图优化。现在看看 LLM 的 ONNX 图在优化后如何变化。
在这个示例中,我们会使用 GPT-2 small 模型;完整代码在 10.2 节对应的 Colab notebook 中。不需要硬件加速;所有优化考虑都只针对 CPU。从 Hugging Face Hub 下载 GPT-2 small 权重和 tokenizer,并将模型转换为 ONNX 后,我们可以通过 ORT optimizer 以编程方式离线优化它:
from onnxruntime.transformers import optimizer
num_layer = model.config.n_layer
num_attention_heads = model.config.n_head
hidden_size = model.config.n_embd
onnx_model_path='gpt2onnx.onnx'
onnx_optim_model_path="gpt2onnx-opt.onnx"
optimized_model = optimizer.optimize_model(onnx_model_path,
model_type='gpt2',
num_heads=num_attention_heads,
hidden_size=hidden_size,
use_gpu=False,
opt_level=1,
verbose=True)
optimized_model.convert_float_to_float16()
optimized_model.save_model_to_file(onnx_optim_model_path)
我们向 optimizer 提供了原始 ONNX 模型位置、模型架构类型、attention head 数量、hidden size、没有硬件加速,以及优化级别 1,也就是只应用 basic graph optimizations。优化之后,模型大小约为原始模型的一半,309 MB 对比 622 MB。optimizer 还应用了几种 basic optimization,例如删除 108 个节点,主要是不必要的 Reshape 节点,如图 10.11 所示,该图比较了原始 ONNX 图和优化后图的输入层。
图 10.11 —— 比较 GPT-2 ONNX 转换模型优化前后图的输入层
optimizer 还应用了一些 extended optimization,例如 Skip Layer Normalization Fusion,它会融合全连接层 bias、skip connection 和 layer normalization,见图 10.12。这是因为我们指定了模型架构,也就是 decoder-only Transformer。
图 10.12 —— 比较 GPT-2 ONNX 转换模型图的一部分,展示应用 Skip Layer Normalization Fusion 前后差异
在 Netron 工具中,你可以检查优化后的图,并看到其他优化也已被应用,例如 Fast GELU Fusion。
通过对两个版本的 ONNX 模型都进行 profiling,我们可以看到性能提升。可以将图 10.13 与原始 ONNX GPT-2 模型的图 10.8 和图 10.9 进行比较。
图 10.13 —— 优化后 GPT-2 模型图中每个 ONNX 算子的调用次数和持续时间
分析 LLM 的自动 ONNX 优化,是学习何时以及如何把这些优化手动应用到自定义模型上的好方法。
本章结束了我们对 ONNX 格式和 ORT 的探索。接下来两章会展示如何把这些技术,以及前面章节中的其他技术,与 SLM 在不同环境中的部署结合起来,从而在不过度牺牲质量的前提下,尽可能利用可用计算能力。
总结
- 在
SessionOptions中将enable_profiling设置为True,即可在 ONNX Runtime,也就是 ORT 中启用性能 profiling。 - 原始 profiling 数据会存储在 JSON 文件中,其中包含每个算子执行的计时信息。
OnnxWholeSession.process_profiling可以把混乱的 JSON profiling 输出转换为干净的表格 DataFrame 格式。- Plotly Express 可以为性能分析创建交互式柱状图,展示算子持续时间、调用次数和执行百分比。
ReferenceEvaluator类可以评估整个 ONNX 模型或单个节点,用于调试和优化。- 详细级别 1 到 4 控制调试时可看到多少逐步评估细节。
- Operator fusion 会把连续算子合并为单个优化节点,以加速执行。
- 要评估自定义算子,需要扩展
OpRun类并实现_run。 - ORT optimizer 会自动对 Transformer 模型应用 basic 和 extended optimization。
- 图优化通过删除不必要节点,例如冗余
Reshape操作,减少模型大小。 - Skip Layer Normalization Fusion 会把 bias 层、skip connection 和 layer normalization 融合为单个操作。
- 当你指定模型架构类型时,Fast GELU Fusion 优化会自动应用。
- Netron 工具可以让你可视化检查 ONNX 图,查看应用了哪些优化。
- 模型评估可以确保优化不会损害预测质量或准确率。
- Basic optimization 包括 constant folding、删除冗余节点,以及保持语义的 fusion。
- Extended optimization 会执行复杂节点融合,且与 CPU、CUDA 或 ROCm execution provider 相关。