部署一个训练完成的神经网络,并不仅仅是把每个算子翻译成GPU内核那么简单。这是一个存在多种合理答案的编译问题。例如,当NVIDIA TensorRT构建一个引擎时,它需要在图中间表示(IR)转换、算子融合、实现方式选择、内存布局以及量化等方面做出相互关联的决策。这一过程被称为调度。不同的调度选择可以在保留模型语义的同时,产生截然不同的推理性能,这使得调度成为深度学习编译器优化的核心问题。

没有任何一种调度策略能在所有模型和部署环境中都表现最佳。对某个网络有加速效果的转换,可能会拖慢另一个网络。某个内核单独运行时速度很快,但与图中其他部分的融合方式、张量布局或内存流量相互作用时,表现却可能很差。激进的调度策略或许能提升性能,却可能违反应用对精度的要求。这些相互耦合、与模型高度相关的影响,往往只有在完整引擎构建并实际测量之后才会显现,这使得选择正确的策略变得越来越难以预测。

本文是系列文章的第一部分(共两部分),旨在构建一种更清晰的性能优化思考方式。文章将探讨以下问题:是什么让一个编译出的引擎优于另一个?为什么单独看起来很有希望的决策,放到整体流程中却可能失败?精度应该如何约束性能优化?文章还将介绍全局性能调优器(Global Performance Tuner)——这是TensorRT 11.1版本中发布的一项自动性能调优功能,并通过一个直观的例子和数学框架来探讨上述问题。

全局性能调优器解决了什么难题

以下这个类比有助于说明TensorRT构建全局性能调优器功能的初衷。

设想一位厨师,他有一道经典菜品,多年来深受三位常客的喜爱。其中一位喜欢辛辣口味,另一位喜欢酸甜口味,第三位则偏好咸口。厨师把辣椒、糖浆、柠檬和盐组合进一道标准配方中,让三位客人都满意。然而某天,一位新客人到店,他不喜欢辣味。继续保留辣椒会让新客人不满;去掉辣椒又可能让第一位老客人失望。

厨师面临这一困境的根本原因在于:这种添加调味料的方式无法规模化扩展——随着顾客数量增加,相互冲突的特殊需求也会增多,而为满足一位顾客所做的调整,可能会损害另一位顾客的用餐体验。

一种可扩展的做法是保留基础菜品,但让调味料变得可配置。厨师为每位常客记录一份合适的调味组合。面对偏好未知的新顾客,厨师可以提供不同调味组合的试吃样品,剔除不合口味的方案,保留顾客最满意的那一种。

将这个类比映射到深度学习编译领域:TensorRT就是那位厨师,其共享的处理流程是基础菜谱,可调的编译器配置是调味料,编译完成的引擎则是成品菜。每个ONNX模型相当于一位顾客。性能调优就是提供试吃样品,而端到端的性能表现则反映了顾客的满意度。

从历史上看,TensorRT的性能优化工作始于内部关键性能指标(KPI)模型。工程师们研究这些工作负载,识别出可优化的模式,并将发现的结果编码为启发式规则,纳入TensorRT的默认调度策略中。每一条启发式规则都凝聚着专家知识,使用户能够自动从那些基于代表性模型发现的优化成果中受益。

然而,随着模型和部署环境日益多样化,这种方法也遇到了类似于上述类比中所描述的扩展性问题。源自某个内部KPI模型的启发式规则被纳入共享策略后,可能有助于处理熟悉的模式,但同时也可能导致某个未曾见过的外部模型出现性能退化。为此再添加一条检查规则或特例处理,或许能保护该模型的性能,但如此反复之后,调度逻辑会变得越来越难以扩展维护。

这个问题通常表现为三种相互关联的形式:

启发式规则的蔓延:每增加一条针对特定模型的启发式规则,就会引入一条新的规则,而这条新规则与现有规则之间的相互作用必须被理解和维护。随着这些规则不断累积,代码会变得难以维护,针对某一类模型所做的改进,也可能无意中影响到另一类模型。

性能与精度的纠缠:降低精度的内核以及其他激进的优化手段可能提升性能,但同时也可能引入超出应用可接受容差范围的数值差异。采用更高精度的方案可以降低这种风险,但会牺牲性能。因此,性能与数值可接受性往往无法被独立地分别优化。

性能定制化:对某个模型或某款GPU效果良好的调度选择,未必适用于另一个模型或GPU。把每一个成功的调度选择都编码进一套统一的调度策略中,会使得TensorRT难以针对不断扩展的外部模型、硬件平台和应用需求进行高效的专门化适配。

调优为启发式优化提供了一种具备可扩展性的补充手段。

深度学习编译器性能的数学理论

本节构建了一个关于深度学习编译器性能的抽象数学模型,目的是将两个基本问题区分开来:

编译器能够为一个模型构建出哪些可执行方案

编译器如何从这些可能方案中选出性能最优的一个

概念:模型,模型空间

设\(m\)表示一个模型,\(M\)表示所考虑的所有模型的集合,称为模型空间:

\(m \in M\)

这里,模型是对待编译计算过程的一种抽象描述。它的具体表示形式可以是ONNX图、通过TensorRT API创建的网络、PyTorch模型,或者其他能够完整指定模型计算过程的前端表示形式。

概念:方案,方案空间

给定模型\(m\)作为输入,深度学习编译器会为目标硬件构建出一个可执行的输出,称为方案(solution)。同一个模型通常可以在不同的调度策略下,被编译成不同的方案。设\(s\)表示其中一种方案,\(S\)表示所有可能方案的集合,称为方案空间:

\(s \in S = \{s_0, s_1, s_2, …, s_n\}\)

方案是对目标硬件上可执行计算的一种抽象。具体而言,一个方案可以是一个TensorRT引擎,或是一段经过序列化的GPU内核序列(cubin)。更抽象地说,它也可以是编译器的某种中间表示(IR),只要该表示在剩余的编译流程确定之后,能够唯一确定最终的可执行文件。

每个方案同时也定义了一个从模型输入域到输出值域的函数:

\(s : X \rightarrow Y\)

\(s(x) = y\)

这里,\((x,y)\)是实现该神经网络的GPU内核序列的一组输入-输出对。在这一抽象层面上,模型描述的是预期的计算过程,而方案描述的则是在目标硬件上实现该计算的某一种具体方式。

确定性原理

陈述:当一个确定性程序和它完整的执行环境都保持不变时,该程序就拥有一个固定的底层运行时间。

这一原理作为性能分析的建模假设。严格来说,执行环境包括输入数据、目标设备、软件栈、GPU频率、温度、功耗状态、系统负载等可能影响运行时间的执行条件。但在实践中,并非每一个环境变量都能被完美控制,因此重复的运行时间测量会出现噪声。如果将这种扰动建模为独立的高斯噪声,那么多次测量结果的平均值就是对底层平均运行时间的最大似然估计。

一旦执行环境在抽象所要求的层面上被固定下来,每个方案就可以对应一个特定的运行时间。这种对应关系使得运行时间可以表示为一个数学函数。

概念:运行时间函数

运行时间函数\(T\)将每个方案映射到其在指定目标设备、固定执行环境下的运行时间:

\(T : S \rightarrow \mathbb{R}^+\)

\(T(s_0) = t_0\)

由于环境被视为固定不变,因此在符号表示中予以省略。

概念:构建器算法

构建器算法(也称为性能调度器)是从模型空间到方案空间幂集的一种映射。给定某个特定模型作为输入,构建器算法会在给定精度要求下输出一个方案空间。

\(B: M \rightarrow 2^S\)

\(B(m) = \{s_1, s_2, …, s_n \mid \max_{x \in X} \|s(x) – y\| \leq \epsilon\} = S\)

在上述精度条件中,输入-输出对\((x, y)\)被视为参考值,\(\epsilon\)表示精度容差阈值。参考数据、比较指标和容差均被视为固定的上下文,因此在\(B\)的参数中被省略。

理想情况下,构建器算法应生成一个足够丰富的方案集合,使其包含\(s^*\)——即在所建模的方案空间中速度最快且满足要求的最优方案。

\(s^* \in B(m)\)

概念:搜索算法

搜索算法(也称为方案空间求解器)是从方案空间幂集到方案空间本身的一种映射。给定一组候选方案作为输入,搜索算法会从中挑选出一个方案作为输出。

\(A: 2^S \rightarrow S\)

\(A(S) = \arg\min_{s_i \in \underline{S}} T(s_i), \underline{S} \subseteq S\)

需要注意的是,在实际情况中,搜索算法可能会从方案空间中选出一个非最优的方案。在TensorRT中,这种非最优性至少可能来自以下两个方面:

搜索过程中使用的启发式规则:启发式规则可能会在候选方案被纳入完整引擎评估之前就将其剔除。例如,构建器可能依据输入模型某一层的通道维度阈值,过滤掉某些卷积策略。如果某个被过滤掉的策略原本会参与构成端到端性能最优的方案,那么搜索算法就再也无法发现那个最优方案了。

局部性能剖析与端到端执行之间的不匹配:在引擎构建过程中,候选内核可能是在策略选择阶段单独进行性能剖析的,但在推理时,性能是在完整引擎层面进行测量的。单独运行的内核与作为更大计算图一部分运行的同一内核,并不共享完全相同的执行环境(例如,NVIDIA Blackwell GPU中使用的自适应网格调整技术)。因此,各个内核单独测得的运行时间之和,未必等于端到端的整体运行时间。

在理想情况下,构建器算法生成的候选集合包含\(s^*\),而搜索算法能够可靠地选中它:

\(A(S) = s^*\)

因此,一个深度学习编译器可以用算法元组\((A, B)\)来表示,其中\(B\)负责生成候选方案,\(A\)负责从中进行选择。然而在实际实现中,候选方案的生成、剪枝、性能剖析和最终选择往往紧密耦合在一起,因此\(A\)和\(B\)之间的边界在编译器架构中可能并不明确。

设\(B^*\)表示构建器算法家族中的最优构建器算法,\(A^*\)表示搜索算法家族中的最优搜索算法。理想的编译器应当实现元组\((A^*, B^*)\)。对于任意输入模型,这一元组都能生成一个包含全局最优方案的候选集合,并可靠地发现它:

\(\forall A, B, \forall m \in M, T(A^*(B^*(m))) \leq T(A(B(m)))\)

在数学抽象层面,\(A^*\)的一种可能实现方式是暴力搜索算法——评估候选集合中的每一个方案,并选出运行时间最低的那个:

\(A^*(S) = \arg\min_{s_i \in S} T(s_i)\)

尽管在构建时间预算有限的情况下,应用暴力搜索算法并不总是可以接受的——构建时间与推理性能之间可能存在权衡——但\(A^*\)在这里用于表示那类对任意输入模型都一定能返回最优方案的搜索算法。

构建\(B^*\)才是更为根本的挑战。即便是理想的搜索算法,也无法选出一个构建器从未使其可达的方案。与期望一步到位构建出\(B^*\)不同,更现实的做法是通过不断扩展和改进编译器可用的构建路径,逐步逼近\(B^*\)。

全局性能调优器建模

本节将编译过程建模为一系列编译器处理步骤(pass)的组合,进而说明可定制的构建路径如何为具备精度感知能力的性能调优提供基础。

深度学习编译器建模

与传统编译器类似,深度学习编译器也可以被视为一系列处理步骤的组合,每个步骤都将一种中间表示(IR)转换为另一种。如前所述,IR与最终可执行二进制文件之间存在一一对应关系,因此IR也可以被视为一种方案。基于这一抽象,设\(f\)表示一个从方案空间到方案空间的IR自映射:

\(f : S \rightarrow S\)

\(f(s) = s’\)

深度学习编译器的构建器算法可以被建模为一系列将一种方案转换为另一种方案的处理步骤:

\(B(m) = f_{k-1}(f_{k-2}(…(f_1(f_0(m))))) = f_{k-1} \circ f_{k-2} \circ … \circ f_0(m)\)

\(B = f_{k-1} \circ f_{k-2} \circ … \circ f_1 \circ f_0\)

构建器算法\(B\)是一个函数复合,看起来就像模型所经过的一条路径,模型沿着这条路径被转换为某个方案,因此\(B\)也可以被称为构建路径。

需要注意以下几点:

严格来说,至少有一个编译器处理步骤(例如最终的代码生成转换)是从模型空间\(M\)映射到可执行方案空间\(S\)的。然而,编译器处理步骤之间并不总是完全解耦的,一个IR可以被视为唯一确定了由其余固定流程所生成的引擎。因此,将这些处理步骤建模为对一个广义方案状态空间的变换,是一种有用的简化方式。

有些编译器处理步骤在功能上是不可或缺的,但并不参与本文所研究的性能相关选择。这类与性能无关的处理步骤,可以视为恒等映射,或者在构建路径中可调的部分中被省略,因为它们并不能区分出不同的构建路径。

乍看之下,上述函数组合似乎意味着一条构建路径只能生成一个方案(即最终构建出的引擎)。但在实践中,TensorRT的编译流程在生成最终引擎之前,可能会考虑并筛选出许多潜在的策略。由于算法\(A\)和\(B\)紧密耦合,这一候选方案生成过程中的大部分内容在上述函数组合中仍是隐含的。该公式描述的是最终被选中的方案,而不是在这一过程中被考虑或被舍弃的每一个候选方案。

可定制构建路径

性能分析与优化的经验表明,有效的调度策略是与模型相关的。即便运行在同一目标GPU上,不同的模型也可能从各个编译组件中不同的选择中获益。因此,所有模型的最优方案未必都来自同一条通用的构建路径。

因此,构建路径应当是可定制的,从而允许调度决策能够根据模型及其部署环境进行调整。这种可调的构建器算法可以被称为可定制构建路径。

举例来说,考虑两个模型\(m\)和\(m’\),以及一组可用的构建路径家族:

\(\{B_1, B_2, \ldots, B_n\}\)

这两个模型未必共享同一条最优路径。模型\(m\)可能通过\(B_3\)获得最佳结果,而模型\(m’\)可能通过\(B_1\)获得最佳结果:

\(A^*(B(m)) = \arg\min_{B_i(m)} \{T(B_1(m)), T(B_2(m)), \ldots, T(B_n(m))\} = B_3(m)\)

\(A^*(B(m’)) = \arg\min_{B_i(m’)} \{T(B_1(m’)), T(B_2(m’)), \ldots, T(B_n(m’))\} = B_1(m’)\)

全局性能调优器正是基于可定制构建路径这一理念构建的。对于每一个候选方案,它都会依据指定的精度标准对输出结果进行验证,并测量端到端的推理性能。在所有满足精度要求的方案中,搜索算法会从这个扩展后的方案集合里,选出运行时间最短且可行的方案。

该框架由以下三项原则共同定义:

精度感知调优:将端到端精度作为约束条件,用以指导性能调优。

可定制构建路径:使构建路径具备可定制性,从而让调度能够适应每个模型及其部署环境。

庞大的方案空间:构建一个丰富的可行方案集合,然后在数值差异仍处于指定容差范围内的候选方案中,选出速度最快的那一个。

全局性能调优器的数学表示可以写作:

\(\mathcal{B} = \{B_1, B_2, …, B_n\}\)

\(s^* = A(B_1(m) \cup B_2(m) \cup \dots \cup B_n(m))\)

延伸阅读

如需进一步了解,请查阅全局性能调优器的相关文档,并阅读本系列的第二部分,其中将详细介绍全局性能调优器的算法与架构设计,以及使用它的最佳实践。

NVIDIA