领域专用小型语言模型——Profiling 洞察

13 阅读18分钟

本章涵盖:

  • 对移植到 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 算子。

image.png

图 10.1 —— Ort_perf_view 对线性回归 ONNX 模型推理性能统计的渲染结果

Ort_perf_view 可视化会为每个算子显示一个方块,并标注算子名称和该 profiling run 中的调用次数。方框大小反映算子调用所占比例。对于更大模型,这个视图对性能分析就不那么有用,虽然它仍然可以提供一个关于聚合算子调用的高层概览,见图 10.2。

image.png

图 10.2 —— Ort_perf_view 输出,展示大型 ONNX 模型一次推理运行的性能统计

在图 10.2 中,你可以看到,被测试模型的 ONNX 图中每个算子都有成千上万次调用。点击某个算子的方块,会展开该部分,显示每一次调用的详细信息,包括输入、输出值和以微秒为单位的持续时间,如图 10.3 所示。

image.png

图 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。

image.png

图 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:

image.png

图 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:

image.png

图 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:

image.png

图 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 操作类型。

image.png

图 10.8 —— GPT-small 线性回归模型图中 ONNX 算子的总持续时间,单位为毫秒

在图 10.9 中,x 轴显示每个 ONNX 算子被调用了多少次,y 轴列出了 profiling 运行中使用的算子类型。

image.png

图 10.9 —— GPT-2 small 模型图中的 ONNX 算子调用次数

在图 10.10 中,x 轴显示总推理时间百分比,y 轴列出了 profiling 运行中的 ONNX 操作类型。

image.png

图 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 节点。例如,在下面代码中,我们会创建一个自定义节点,将 EyeLikeAdd 算子结合起来,然后看看模型运行是否更高效:

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 图和优化后图的输入层。

image.png

图 10.11 —— 比较 GPT-2 ONNX 转换模型优化前后图的输入层

optimizer 还应用了一些 extended optimization,例如 Skip Layer Normalization Fusion,它会融合全连接层 bias、skip connection 和 layer normalization,见图 10.12。这是因为我们指定了模型架构,也就是 decoder-only Transformer。

image.png

图 10.12 —— 比较 GPT-2 ONNX 转换模型图的一部分,展示应用 Skip Layer Normalization Fusion 前后差异

在 Netron 工具中,你可以检查优化后的图,并看到其他优化也已被应用,例如 Fast GELU Fusion。

通过对两个版本的 ONNX 模型都进行 profiling,我们可以看到性能提升。可以将图 10.13 与原始 ONNX GPT-2 模型的图 10.8 和图 10.9 进行比较。

image.png

图 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 相关。