《机器学习系统：设计和实现》读后感和一些思考

计算图 利用不同编程接口实现的机器学习程序需要共享一个运行后端。实现这一后端的关键技术是应用无关的计算图。计算图包含计算节点节点之间的边表达计算依赖。计算图可以被同步和异步执行。其实就是一种模型高级的中间表示。
编译器前端 给定一个计算图机器学习框架会对计算图做一系列优化。和硬件无关的优化由编译器前端实现。编译器前端实现包括中间表达自动微分类型推导和静态分析等等。
编译器后端和运行时 机器学习框架利用编译器后端对计算图可以进一步针对硬件的特性例如说L2/L3大小指令流水线长度进行性能优化。最终优化后的计算图通过运行时执行在通用处理器CPU或者是硬件加速器之上。运行时需要实现算子选择和内存分配等技术。

现代机器学习系统需要兼有易用性和高性能因此其一般选择Python作为前端编程语言而使用C和C++作为后端编程语言。前端负责静态分析、类型推导以及自动微分、分布式并行子图拆分等PASS优化后端负责硬件相关的优化如内存优化、图算融合等。

计算图

计算图的作用

对于输入数据、算子和算子执行顺序的统一表达。 机器学习框架用户可以用多种高层次编程语言PythonJulia和C++来编写训练程序。这些高层次程序需要统一的表达成框架底层C和C++算子的执行。因此计算图的第一个核心作用是可以作为一个统一的数据结构来表达用户用不同语言编写的训练程序。这个数据结构可以准确表述用户的输入数据、模型所带有的多个算子以及算子之间的执行顺序。
定义中间状态和模型状态。 在一个用户训练程序中用户会生成中间变量神经网络层之间传递的激活值和梯度来完成复杂的训练过程。而这其中只有模型参数需要最后持久化从而为后续的模型推理做准备。通过计算图机器学习框架可以准确分析出中间状态的生命周期一个中间变量何时生成以及何时销毁从而帮助框架更好的管理内存。
自动化计算梯度。 用户给定的训练程序仅仅包含了一个机器学习模型如何将用户输入一般为训练数据转化为输出一般为损失函数的过程。而为了训练这个模型机器学习框架需要分析任意机器学习模型和其中的算子找出自动化计算梯度的方法。计算图的出现让自动化分析模型定义和自动化计算梯度成为可能。
优化程序执行。 用户给定的模型程序往往是“串行化”地连接起来多个神经网络层。通过利用计算图来分析模型中算子的执行关系机器学习框架可以更好地发现将算子进行异步执行的机会从而以更快的速度完成模型程序的执行。

计算图的组成

计算图由基本数据结构张量(Tensor)和基本运算单元算子(Operator)构成。

静态计算图与动态计算图

1、静态计算图

静态计算图采用先编译后执行的方式该模式将计算图的定义和执行进行分离。

在静态图模式下使用前端语言定义模型形成完整的程序表达后并不使用前端语言解释器进行执行而是将前端描述的完整模型交给计算框架。框架在执行模型计算之前会首先对神经网络模型进行分析获取网络层之间的连接拓扑关系以及参数变量设置、损失函数等信息接着用一种特殊的静态数据结构来描述拓扑结构及其他神经网络模型组件这种特殊的静态数据结构通常被称为静态计算图。静态计算图可以通过优化策略转换成等价的更加高效的结构。当进行模型训练或者推理过程时静态计算图接收数据并通过相应硬件调度执行图中的算子来完成任务。

2、动态计算图

动态计算图采用解析式的执行方式其核心特点是编译与执行同时发生。

动态图采用前端语言自身的解释器对代码进行解析利用计算框架本身的算子分发功能算子会即刻执行并输出结果。动态图模式采用用户友好的命令式编程范式使用前端语言构建神经网络模型更加简洁。

3、两者对比

静态生成和动态生成的过程各有利弊。从使用者的角度可以直观的感受到静态图不能实时获取中间结果、代码调试困难以及控制流编写复杂而动态图可以实时获取结果、调试简单、控制流符合编程习惯。虽然静态图的编写、生成过程复杂但是相应的执行性能却超过动态图。

编译器前端

IR中间表示

中间表示(IR)是编译器用于表示源代码的数据结构或代码是程序编译过程中介于源语言和目标语言之间的程序表示。几乎所有的编译器都需要某种形式的中间表示来对被分析、转换和优化的代码进行建模。在编译过程中中间表示必须具备足够的表达力在不丢失信息的情况下准确表达源代码并且充分考虑从源代码到目标代码编译的完备性、编译优化的易用性和性能。

在此基础上编译流程就可以在前后端直接增加更多的优化流程这些优化流程以现有IR为输入又以新生成的IR为输出被称为优化器。优化器负责分析并改进中间表示极大程度的提高了编译流程的可拓展性也降低了优化流程对前端和后端的破坏。

机器学习框架的中间表示

在设计机器学习框架的中间表示时需要充分考虑以下因素

1) 张量表达。机器学习框架主要处理张量数据因此正确处理张量数据类型是机器学习框架中间表示的基本要求。

2) 自动微分。自动微分是指对网络模型的自动求导通过梯度指导对网络权重的优化。主流机器学习框架都提供了自动微分的功能在设计中间表示时需要考虑自动微分实现的简洁性、性能以及高阶微分的扩展能力。

3) 计算图模式。主流机器学习框架如TensorFlow、PyTorch、MindSpore等都提供了静态图和动态图两种计算图模式静态计算图模式先创建定义计算图再显式执行有利于对计算图进行优化高效但不灵活。动态计算图模式则是每使用一个算子后该算子会在计算图中立即执行得到结果使用灵活、便于调试但运行速度较低。机器学习框架的中间表示设计同时支持静态图和动态图可以针对待解决的任务需求选择合适的模式构建算法模型。

4) 支持高阶函数和闭包 。高阶函数和闭包是函数式编程的重要特性高阶函数是指使用其它函数作为参数、或者返回一个函数作为结果的函数闭包是指代码块和作用域环境的结合可以在另一个作用域中调用一个函数的内部函数并访问到该函数作用域中的成员。支持高阶函数和闭包可以抽象通用问题、减少重复代码、提升框架表达的灵活性和简洁性。

5) 编译优化。机器学习框架的编译优化主要包括硬件无关的优化、硬件相关的优化、部署推理相关的优化等这些优化都依赖于中间表示的实现。

6) JIT(Just In Time)能力。机器学习框架进行编译执行加速时经常用到JIT即时编译。JIT编译优化将会对中间表示中的数据流图的可优化部分实施优化包括循环展开、融合、内联等。中间表示设计是否合理将会影响机器学习框架的JIT编译性能和程序的运行能力。

常见编译器前端优化方法

1、无用与不可达代码消除

如图6.5.2所示。无用代码是指输出结果没有被任何其他代码所使用的代码。不可达代码是指没有有效的控制流路径包含该代码。删除无用或不可达的代码可以使得中间表示更小提高程序的编译与执行速度。无用与不可达代码一方面有可能来自于程序编写者的编写失误也有可能是其他编译优化所产生的结果。

2、常量传播、常量折叠

常量传播如图6.5.3所示如果某些量为已知值的常量那么可以在编译时刻将使用这些量的地方进行替换。

常量折叠如图6.5.3所示多个量进行计算时如果能够在编译时刻直接计算出其结果那么变量将由常量替换。

3、公共子表达式消除

如图6.5.4所示如果一个表达式E已经计算过了并且从先前的计算到现在E中所有变量的值都没有发生变化那么E就成为了公共子表达式。对于这种表达式没有必要花时间再对它进行计算只需要直接用前面计算过的表达式结果代替E就可以了。

编译器后端

概述

如图7.1.1所示编译器后端处于前端和硬件驱动层中间主要负责计算图优化、算子选择和内存分配的任务。首先需要根据硬件设备的特性将IR图进行等价图变换以便在硬件上能够找到对应的执行算子该过程是计算图优化的重要步骤之一。前端IR生成是解析用户代码属于一个较高的抽象层次隐藏一些底层运行的细节信息此时无法直接对应硬件上的算子算子是设备上的基本计算序列例如MatMul、Convolution和ReLU等需要将细节信息进行展开后才能映射到目标硬件上的算子。对于某些前端IR的子集来说一个算子便能够执行对应的功能此时可以将这些IR节点合并成为一个计算节点该过程称之为算子融合对于一些复杂计算后端并没有直接与之对应的算子但是可以通过几个基本运算的算子组合达到同样的计算效果此时可以将前端IR节点拆分成多个小算子。然后我们需要进行算子选择。算子选择是在得到优化的IR图后需要选取最合适的目标设备算子。针对用户代码所产生的IR往往可以映射成多种不同的硬件算子但是生成不同的算子执行效率往往有很大的差别如何根据前端IR选择出最高效的算子是算子选择的核心问题。算子选择本质上是一个模式匹配问题。其最简单的方法就是每一个IR节点对应一个目标硬件的算子但是这种方法往往对目标硬件的资源利用比较差。目前来说对于现有的编译器一般都对每一个IR节点提供了多个候选的算子算子选择目标就是从中选择最优的一个算子作为最终执行在设备上的算子。总的来说在机器学习系统中对前端生成的IR图上的各个节点进行拆分和融合让前端所表示的高层次IR逐步转换为可以在硬件设备上执行的低层次IR。得到了这种更加贴合硬件的IR后对于每个单节点的IR可能仍然有很多种不同的选择例如可以选择不同的输入输出格式和数据类型我们需要对IR图上每个节点选择出最为合适的算子算子选择过程可以认为是针对IR图的细粒度优化过程最终生成完整的算子序列。最后遍历算子序列为每个算子分配相应的输入输出内存然后将算子加载到设备上执行计算。

通用硬件优化算子拆分和算子融合

深度学习算子按其对资源的需求可以分为两类计算密集型算子这些算子的时间绝大部分花在计算上如卷积、全连接等访存密集型算子这些算子的时间绝大部分花在访存上他们大部分是Element-Wise算子例如 ReLU、Element-Wise Sum等。在典型的深度学习模型中一般计算密集型和访存密集型算子是相伴出现的最简单的例子是“Conv + ReLU”。Conv卷积算子是计算密集型ReLU算子是访存密集型算子ReLU算子可以直接取Conv算子的计算结果进行计算因此我们可以将二者融合成一个算子来进行计算从而减少内存访问延时和带宽压力提高执行效率。

例如“Conv + Conv + Sum + ReLU”的融合从图7.2.1中我们可以看到融合后的算子减少了两个内存的读和写的操作优化了Conv的输出和Sum的输出的读和写的操作。

除了上述针对特定算子类型结构的融合优化外基于自动算子生成技术还可以实现更灵活、更极致的通用优化。以 MindSpore 的图算融合技术为例图算融合通过“算子拆解、算子聚合、算子重建”三个主要阶段如图让计算图中的计算更密集并进一步减少低效的内存访问。

图7.2.2中算子拆解阶段Expander将计算图中一些复杂算子composite op图中Op1、Op3、Op4展开为计算等价的基本算子组合图中虚线正方形框包围着的部分在算子聚合阶段Aggregation将计算图中将基本算子basic op如图中Op2、拆解后的算子expanded op组合融合形成一个更大范围的算子组合在算子重建阶段Reconstruction中按照输入tensor到输出tensor的仿射关系将基本算子进行分类elemwise、 broadcast、reduce、transform等并在这基础上归纳出不同的通用计算规则如 elemwise + reduce 规则elemwise + reduce在满足一定条件后可以高效执行根据这些计算规则不断地从这个大的算子组合上进行分析、筛选最终重新构建成新的算子如图中虚线正方形包围的两个算子 New Op1 和 New Op2。图算融合通过对计算图结构的拆解和聚合可以实现跨算子边界的联合优化并在算子重建中通过通用的计算规则以必要的访存作为代价生成对硬件更友好、执行更高效的新算子。

算子信息

针对不同特点的计算平台和不同的算子为了追求最好的性能一般都需要选择不同的数据排布格式。机器学习系统常见的数据排布格式有NCHW和NHWC等。
对于不同的硬件支持不同的计算精度例如float32、float16和int32等。算子选择需要在所支持各种数据类型的算子中选择出用户所设定的数据类型最为相符的算子。

数据精度和存储方法

通常深度学习的系统一般使用的是单精度floatSingle Precision浮点表示。这种数据类型占用32位内存。还有一种精度较低的数据类型float16其内部占用了16位的内存。由于很多硬件会对float16数据类型进行优化float16半精度的计算吞吐量可以是float32的2∼8倍且float16可以占用的数据更小这样可以输入更大的BatchSize进而减少总体训练时间。接下来我们详细看一下半精度浮点数与单精度浮点数的区别。

如图7.3.5其中sign代表符号位占1位表示了机器数的正负exponent表示指数位Mantissa为尾数位。其中float16类型的数据采用二进制的科学计数法转换为十进制的计算方式如下

算子选择的过程

其中算子信息主要包括了支持设备类型、数据类型和数据排布格式三个方面。经过编译器前端类型推导与静态分析的阶段后IR图中已经推导出了用户代码侧的数据类型。下面介绍算子选择的基本过程。

首先选择算子执行的硬件设备。不同的硬件设备上算子的实现、支持数据类型、执行效率通常会有所差别。这一步往往是用户自己指定的若用户未指定则编译器后端会为用户匹配一个默认的设备。然后后端会根据IR图中推导出的数据类型和内存排布格式选择对应的算子。

理想情况下算子选择所选择出的算子类型应该与用户预期的类型保持一致。但是由于软硬件的限制很可能算子的数据类型不能满足用户所期待的数据类型此时需要对该节点进行升精度或者降精度处理才能匹配到合适的算子。

算子的数据排布格式转换是一个比较耗时的操作为了避免频繁的格式转换所带来的内存搬运开销数据应该尽可能地以同样的格式在算子之间传递算子和算子的衔接要尽可能少的出现数据排布格式不一致的现象。另外数据类型不同导致的降精度可能会使得误差变大收敛速度变慢甚至不收敛所以数据类型的选择也要结合具体算子分析。

In-Place算子

在内存分配流程中我们会为每个算子的输入和输出都分配不同的内存。然而对很多算子而言为其分配不同的输入和输出地址会浪费内存并且影响计算性能。例如优化器算子其计算的目的就是更新神经网络的权重例如Python语法中的’+=‘和’*=‘操作符将计算结果更新到符号左边的变量中例如’a[0]=b’语法将’a[0]’的值更新为’b’。诸如此类计算有一个特点都是为了更新输入的值。下面以Tensor的’a[0]=b’操作为例介绍In-Place的优点。图7.4.6左边是非In-Place操作的实现step1将Tensor a拷贝到Tensor a’step2将Tensor b赋值给Tensor a’step3将Tensor a’拷贝到Tensor a。图7.4.6右边是算子In-Place操作的实现仅用一个步骤将Tensor b拷贝到Tensor a对于的位置上。对比两种实现可以发现In-Place操作节省了两次拷贝的耗时并且省去了Tensor a’内存的申请。

模型推理

汇编语言优化

对于已知功能的汇编语言程序来说计算类指令通常是固定的性能的瓶颈就在非计算指令上。计算机各存储设备类似于一个金字塔结构最顶层空间最小但是速度最快最底层速度最慢但是空间最大。L1-L3统称为cache(高速缓冲存储器)CPU访问数据时会首先访问位于CPU内部的cache没找到再访问CPU之外的主存此时引入了缓存命中率的概念来描述在cache中完成数据存取的占比。要想提升程序的性能缓存命中率要尽可能的高。

下面简单列举一些提升缓存命中率、优化汇编性能的手段

1循环展开尽可能使用更多的寄存器以代码体积换性能

2指令重排打乱不同执行单元的指令以提高流水线的利用率提前有延迟的指令以减轻延迟减少指令前后的数据依赖等

3寄存器分块合理分块NEON寄存器减少寄存器空闲增加寄存器复用

4计算数据重排尽量保证读写指令内存连续提高缓存命中率

5使用预取指令将要使用到的数据从主存提前载入缓存减少访问延迟。

寄存器与NEON指令

ARMv8系列的CPU上有32个NEON寄存器v0-v31如图10.4.2所示NEON寄存器v0可存放128bit的数据即4个float328个float1616个int8等。

针对该处理器可以采用SIMD(Single InstructionMultiple Data单指令、多数据)提升数据存取计算的速度。相比于单数据操作指令NEON指令可以一次性操作NEON寄存器的多个数据。例如对于浮点数的fmla指令用法为fmla v0.4s, v1.4s, v2.4s如图10.4.3所示用于将v1和v2两个寄存器中相对应的float值相乘累加到v0的值上。

PTQ训练后量化

量化公式

假设r表示量化前的浮点数量化后的整数q可以表示为

round(⋅)和clip(⋅)分别表示取整和截断操作 $q_{min}$ 和 $q_{max}$ 是量化后的最小值和最大值。s是数据量化的间隔z是表示数据偏移的偏置z为0的量化被称为对称Symmetric量化不为0的量化称为非对称Asymmetric量化。对称量化可以避免量化算子在推理中计算z相关的部分降低推理时的计算复杂度非对称量化可以根据实际数据的分布确定最小值和最小值可以更加充分的利用量化数据信息使得计算精度更高。

具体流程

使用直方图统计的方式得到原始FP32数据的统计分布 $P_f$
在给定的搜索空间中选取若干个 $q_{min}$ 和 $q_{max}$ 分别对激活值量化得到量化后的数据 $Q_q$
使用直方图统计得到 $Q_q$ 的统计分布;
计算每个 $P_f$ 与 $Q_q$ 的统计分布差异并找到差异性最低的一个对应的 $q_{min}$ 和 $q_{max}$ 来计算相应的量化参数常见的用于度量分布差异的指标包括KL散度(Kullback-Leibler Divergence)、对称KL散度(Symmetric Kullback-Leibler Divergence)和JS散度(Jenson-Shannon Divergence)。