0%

CSE 599W: Systems for ML

国庆在家闲不住想干活系列……

本篇的内容是陈天奇大佬今年春季在华盛顿大学开的一门课

大佬是上交 ACM 班的本硕,虽然目前还在 UW 读博中,但是在机器学习圈子里面已经有了很高的名望了,他的 MXNet 和 XGBoost 相信很多人就算没用过也肯定听说过(比如我就没用过…)。前段时间他发布的 TVM 也算是开启了深度学习和系统方面探索的一条新道路。

课程介绍上讲的是这门课的目标是填上深度学习算法和系统的实现以及优化之间的 gap,粗略地翻了一下 PPT,后面也有比较多的篇幅是介绍 TVM 的,正是我想了解的!

没找到课程的视频,但是 PPT 可以在上面的课程链接或者这里找到。

下面的内容主要按照每篇 PPT 的整理:


Lecture 1: Introduction to Deep Learning

回顾了一下基本原理和发展历史。

机器学习的过程基本上就是 数学模型+评价指标+参数训练,深度学习则是模型特指各种神经网络。

具体主要涉及到各种不同的模型架构(CNN、RNN、各种变种),目标函数的选择和训练技巧,正则化初始化等等。

这些就不多记了。

Lecture 2 是个实验课,实践怎么搭网络。

Lecture 3: Overview of Deep Learning System

这一节差不多是大纲的性质,每一个小点后面都会分节细讲。

Basic Architecture

基本上所有的深度学习框架都是差不多这个结构,首先来看 User API 层:

这里举了个线性回归的例子来对比手写 Numpy 和框架代码的差别。基本上网络模型都可以比较方便地用一个计算图的结构来表达,节点表示运算操作,边代表数据依赖。

那为了方便用户写代码,一个框架也是一定要有自动求导的功能的(如果反向计算还需要手写那就太瓜皮了)。

然后是 System Components 层:

这里涉及到了首先是计算图的优化,比如一次运行的时候直接过滤掉用不到的图节点(Deadcode Elimination),内存分配方面的优化,图节点和实际计算设备的对应等等。

实际跑图的时候如果有多个设备或者多个工作线程,如何调度以及发挥出计算设备的并行性也是一个需要考虑的问题。

最下面的 Architecture 层:

目前用来支持 DL 的设备也有很多,典型的如 GPU,其他的加速芯片也是越来越多,不同的设备可能要写对应不同的代码,这部分要怎么优化?

现在最常规的做法是每一种不同的计算设备会有开发厂商自己提供支持库,但是这个对框架的开发者来说还是有一个要整合的过程。另外,如果系统中存在多种不同的计算设备,计算任务在多设备上要怎么分配和调度也会是一个很大的麻烦。

为了解决最后的这个问题,目前有一种 Compiler Based Approach,即整个 Architecture 层由一个 High Level Operator Description 加上 Tensor Compiler Stack 来统一解决。这就是之后要提到的 TVM 的设计思路了。

Lecture 4: Backpropagation and Automatic Differentiation

详细解释第三节中的自动求导。

计算机中实现求导这个操作主要有两种方式:基于符号的算术求导和直接用数值进行求导。

算术求导需要构建一棵符号表示树,然后根据各种算术上的求导规则来写公式。缺点在于:如果遇到特别复杂的函数,则需要推导的符号表示树也会很大;然后如果目标只是想要一个导数值,则保存一棵符号表示树就很浪费了;再然后就是这样做容易有误差(?为什么…按公式算不应该误差更小吗)。

数值求导则是按导数的定义做,直接对方程取极限:

$$ \frac{\partial f(x)}{\partial x_i} \approx \lim \limits_{h \to 0} \frac{f(x+he_i)-f(x)}{h}$$

实现起来特别简单,h 取个 1e-6 就差不多了,但是一般只用来检验求导结果用。

然后对于网络中每一层的反向部分,其实求导涉及到的都只是跟本层运算相关的内容:

Backpropagation

上一层传下来的是 $\frac{\partial Error}{\partial z}$,再往下可以通过链式求导法则一直推导下去,而其他需要的则只是与本层运算有关的$\frac{\partial z}{\partial x}$ 和 $\frac{\partial z}{\partial y}$。

更详细的推导可见这里

因此自动求导则是根据以上的规则来创建反向计算图的过程,伪代码以及结果如下:

AutoDiff

BP & AutoDiff

自动求导构建完成反向计算图之后,完整的计算图可以接下来一起用作整体的图优化。

Lecture 5: GPU Programming

在 CPU 上进行数据运算大致有几个过程(按多级流水分):Fetch、Decode、ALU Compute、Write Back。由于 CPU 本来也并不是为了纯运算而设计,因而在 ALU Compute 以外的其他部分会有比较大的计算资源和能耗上的 overhead。

后来增加的向量化指令能够相当程度地改善这种 overhead 的问题,而 GPU 从这个角度来看更像是一种把 ALU 的向量化做的更极致的加速器。

CPU vs GPU

从存储的层次结构上来对比:

Memory Hierarchy

GPU 的大量寄存器就使得它能够以比 CPU 小的多的开销来切换线程,这也就能够支撑起大规模的 SIMT 了。

后面是一些 CUDA 编程的例子,以及如何根据 GPU 的微架构特性高效地发挥出性能来。

Lecture 6: Optimize for Hardware Backends

这一节的内容大概在 System Components 和 Architecture 层之间,一份代码面对不同规模的数据(甚至是不同数据块尺寸的数据)往往不作针对性地调整是达不到最佳性能的。

深入下去需要实际考虑到例如 CPU 的 Cache、GPU 的寄存器等等这些方面,以及 GPU 的多级存储之间的数据搬移开销、数据重用等等,同样是 GPU 也有多种不同的后端。

不同的 Tiling Patterns、Fuse Patterns、Data Layout、Hardware Backend 合起来使得优化工作也变得相当复杂了。

为了解决前面说到的所有这些麻烦的问题,然后这里就引出了 TVM Stack。

Lecture 7: Automatic Code Generation - TVM Stack

各种不同的框架和实际执行运算的各种各样的硬件后端之间其实存在着很大的 gap。

如果从编译器的视角来看待如何解决这个问题,各种框架写的网络可以根据特定的规则转化成某种统一的表示形式,在统一表示的基础上进行一些可重用的图优化,之后再用不同的后端来生成对应不同设备的代码,这就是目前各家都在尝试的设计思路了。

Computational Graph as IR

举例说:TensorFlow 的 XLA 会把高级代码抽象成 XLA HLO 的表示,做目标无关优化之后再用对应后端来生成更深一层的代码。

NVIDIA 的 TensorRT 的优化策略也是在图转化之后的统一表示上做,例如根据设定好的规则来做一些相邻计算单元的合并(Kernel Fusion)等等。

当然这种方式实现的时候会遇到一些同样非常麻烦的问题,一个 operator 需要针对不同的硬件平台、数据格式、精度、线程结构写一堆代码生成规则和优化规则。

到头来是把原本 op 实现的复杂度变成了编译规则的复杂度,绕了个圈以后好像还是很麻烦啊。

TVM 借助了一种叫 Tensor Expression Language 的表示方法,同样采用这种类似表示的还有 Halide(一种图像处理语言)、Loopy(基于 Python 的 kernel 生成器)、TACO(稀疏 Tensor 代码生成器)、Tensor Comprehension(类似 TVM)等等。

这种表示法最初的想法来源于 Halide,核心在于把代码的计算和调度分开

例如一段最原始的 TVM 代码:

1
2
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)

生成得到的 C 代码是:

1
2
3
4
for (int i = 0; i < n; ++i)
{
C[i] = A[i] + B[i];
}

加上额外的调度控制:

1
2
3
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(s[C].axis[0], factor=32)

再生成的代码就变成了:

1
2
3
4
5
6
7
8
9
10
11
for (int xo = 0; xo < ceil(n / 32); ++xo)
{
for (int xi = 0; xi < 32; ++xi)
{
int i = xo * 32 + xi;
if (i < n)
{
C[i] = A[i] + B[i];
}
}
}

甚至于还可以支持绑定中间的 xo 和 xi 到特定的变量上:

1
2
3
4
5
6
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(s[C].axis[0], factor=32)
s[C].recorder(xi, xo)
s[C].bind(xo, tvm.thread_axis(“blockIdx.x”)
s[C].bind(xi, tvm.thread_axis(“threadIdx.x”)

话说这样出来的代码就可以用在 CUDA kernel 里面了:

1
2
3
4
5
int i = threadIdx.x * 32 + blockIdx.x;
if (i < n)
{
C[i] = A[i] + B[i];
}

具体后续的调度部分的设计,首先需要保证生成的代码在逻辑上要能跑出正确的结果,常见的手工优化代码的方法也都要包含在内,并且要能够方便引入其他额外的新技术。

目前 TVM 的调度部分还在继续开发中,已经从像 Halide、Loopy 这种成熟的语言中吸取过来的方法有例如 Loop Transformations、Thread Bindings、Cache Locality 等,针对 GPU 还开发了一些方法例如 Thread Cooperation、Tensorization、Latency Hiding 等。

再额外的就是 TVM 还用了 Auto-tuning,由于 TVM 的论文还没看,不确定我的理解对不对。Schedule Space 模型的自动调优就是尝试不同的优化方法组合,然后在整个策略空间里面搜索哪一种优化效果最好最终就采用哪一种吗?

末尾给的一些测试中,TVM 表现出了相当不错的性能结果。

当然,TVM 还刚刚开始发展,后面还有一大堆问题留待解决。

Lecture 8: Hardware Specialization in Deep Learning

上一节的 TVM 是一个纯软件栈,这一节就来探索一下用于深度学习的专用硬件。

DL 的疯狂发展对计算硬件也有了越来越高的需求,而且不同应用场景的需求还可能会差很多,例如数据中心和移动终端上面的 AI 设备就完全要往两个极端去考虑。

Evolution of Deep Learning

上面这张图讲的是在 DL 的发展过程中,对数据的尺寸以及存储精度的需求也在不断变化,低精度可以节省空间以及加速运算,但是这也要在硬件本身可以直接支持低精运算的前提下才能有效果(硬件是64位双精的,你要那它跑 int8?那就呵呵了…emm,此处并不是针对某 sw 哈哈)。另外,一些出现的新算法是不是能够用硬件高效实现也很关键,实现不了的话可能还是要选择老算法更好。

不断发展的 DL 算法在实验室里面可以任意瞎搞,效果好就好了,但是如果要应用到实际的生产环境中,那能不能实现/怎么高效实现就非常重要了。

再再另一方面,摩尔定律也逐渐受限,更低纳米制程的工艺难度越来越大,所有这些问题最后都会导向一个终极的解决方案,那就是 DL 专用的计算芯片/硬件了。

下面用 TPU 来举了个栗子:2015 年流片的 ASIC,92 TOPS 的峰值性能,相比 K80 有 30~80倍的性能功耗比。这些数据看着都吓人。

那为什么这么强呢?原因在于它直接硬件支持 8 位的整数 Inference(相比 16 位半精要节能 6~30 倍),大量的乘加运算部件(MACs)以及大量的片上存储。

TPU Block Diagram

TPU 主要的峰值运算性能都来自于右边的一大堆矩阵乘单元和累加器。光片上的 Unified Buffer 和 MMU 就占到了整块 TPU 超过 50% 的芯片面积。

这种设计也在于尽可能地提高数据的重用程度,提高计算密集度。

总结一下,像 TPU 这样的 DL 专用的加速器相对 CPU 和 GPU 主要有三个方面的特点:

  1. 通常需要显示管理片内的存储子系统,而 CPU 的 Cache 是隐式的,对程序员透明,GPU 则是可以有自己的 L1 Cache 同时也可以手动维护(Shared Memory);
  2. 计算主要以 Tensor (矩阵或者向量)为单元,CPU 主要就是标量运算了,配合 SIMD 的话则跟 GPU 一样主要是向量计算;
  3. 如果为 Inference 设计则不需要太高的精度,低比特量化之后可能更适合。

下面是举了 3 种 Hardware/Software Co-design 的方法:

  1. Tensor 化(Tensorization):把矩阵-向量运算变成矩阵矩阵运算,4x8 的矩阵乘上 8x1 的向量变成两个 4x4 和 4x1 的乘积(…提高计算密集度?)
  2. 存储结构针对计算内容做优化:面向卷积优化则需要一个较大的激活函数 Buffer(空间重用),较小的 FIFO 队列(存储参数);面向 GEMM 优化则需要分配较多的空间用以累加器的块存储。
  3. 低比特量化:可以线性提高存储带宽。

。。。上面的这三个感觉理解的比较模糊。


再往下才是本节的重点内容——VTA

TVM 构建的是软件栈,硬件加速器方面,他们也提出了一套开源的 FPGA 加速器设计方案,即 VTA(Versatile Tensor Accelerator)。

VTA 针对不同的带宽、存储和精度需求可以自定义 Tensor Core 的形状、数据类型、内存子系统分配、支持的运算操作等等;对不同类型的代码提供 CISC 或者 RISC 的指令集支持;并且还做了一些 Latency Hiding 的工作。

大致的框架设计如下:

VTA Design

IF 单元从 DRAM 中获取到下一条指令后,会根据类型将其发送到目标部件对应的队列中。

Load 单元负责准备激活函数和计算核的存储资源、提供Micro-Op 的缓存,取出来的数据放在 Load Buffer 中。

计算单元负责根据前面的 Micro-Op 以及准备好的数据执行 ALU 运算或者 GEMM 运算,更新寄存器的内容。

Store 单元负责把前面 Store Buffer 中的寄存器值写回 DRAM。

整体的运行依靠多个任务队列来维护数据依赖关系,基本上是个数据流的设计。(。。。)

VTA 的控制代码部分则依靠 TVM 来生成。

加上 VTA 之后,整个 TVM 的完整架构显得更复杂了:

Programmability Challenge

顶层是各种成熟的深度学习框架,TVM 充当编译器的角色,底层的硬件执行部分则由 VTA 来实现。

Lecture 9: Memory Optimization

还是在 System Components 这层,继续深入分析 DL 训练过程中存在的问题。

回到前面的自动求导部分,这里抛出来一个问题是为什么自动求导是采用往计算图中扩展反向计算的数据计算通路的方式,而不是直接在原来的图上进行反向计算(Backprop in Graph)

这个 Backprop in Graph 这里也没有更详细的说明,我大概理解成递归返回的那种样子,正向计算是不断递归向下,然后每层递归退出的时候执行反向,完美!

其实说起来,本质上递归的这个顺序也就是数据依赖关系的拓扑序。

原因呢,则是在于内存上。

State-of-art 的很多模型都可能会有资源受限的问题,现在确实很多效果好的新模型都越来越大了,一方面计算量在增长,另一方面内存会成为一个很大的麻烦:CPU 的内存还好一点,如果是用在 GPU 上,目前单块卡的显存最多也只有几十 G 的量级。

先来看一下前向部分的内存使用情况,以下面这几个简单的运算组成的计算图为例:

朴素做法是为每个计算节点都开一块内存,则图越大、计算节点越多,需要的内存量越大。

然后我们发现这个过程中有很多内存是用过之后就不会再用了的,因此更高级一些的方案是配合内存池进行动态内存分配,例如上面 mul 算完之后所占有的内存就不再需要了,因此这块内存可以被内存池回收,然后用在 exp 的计算上。

再有另外一种方式则是静态内存规划,即拿到计算图之后,就按照尽可能重用内存的方案事先分配内存,基本上达到的效果应该要跟上一种内存池的动态分配方案差不多。这种做法有点类似编译器的寄存器重命名的过程。

基本的分配原则也非常简单,只是应用的时候另外要注意如果分配不好是有可能要影响计算的并发性的:

  • Inplace:可行的情况下尽可能地原地存储,即把输出存到输入的内存里(前提是这份输入只被一个计算节点依赖)
  • Reuse:不再用到的内存尽可能地重用。

下面的几个内存规划的例子也都比较简单:根据拓扑顺序依次分配和回收;或者先从起点到终点找一条最长路径,把路径上的内存全部设为 Inplace 重用,然后再找别的路径等等。

这里举的内存重用的例子都特别简单,实际上每个计算节点需要的数据尺寸和内存大小都不一定一样,不可能这么简单地就分配好了。

回到前面那个自动求导的两种方案的问题,可以很容易地体会出来往计算图中扩展反向通路的方案非常容易做内存优化:

只需要把 Inplace 和 Reuse 用好即可,在 MXNet 上测试出来的效果也非常好。

深度学习的 BP 算法中存在的最大问题在于反向运算时需要用到前向的一些结果,这事实上就大大地限制了 Reuse 策略的发挥,因为前向算出来的结果总需要找个地方暂存着。

针对一些内存需求特别大的场景,可以采用计算换空间的折中方案:

只保存前向结果的一部分,当反向运算中需要时,再重新从前面开始把缺失的部分重新算一遍,用 25% 的额外计算量可以把整体的内存使用降到原来的开方级别,在某些场景下还是有非常不错的收益的。

实验室的一位师姐之前在 RNN 的优化里面用到过这种方法,节省下内存之后可以跑更大的 BatchSize,最后得到的效果非常好。

再回到前面两种反向方案的讨论中,这一节的最后给了一个特别有趣的点:内存重用的优化方案从某种程度上来看有点像递归中的内存分布!!

666666!

Lecture 10: Parallel Scheduling

System Components 这层的另外一个方面是并行调度的问题。

用户写好一个计算图,如果框架没有能力把机器上所有的硬件资源全都调动起来那就太浪费了。

关于 DL 中模型并行数据并行这块就不再多提了,在常规的数据并行中,计算和通信之间存在着一个 Gap:

这样一张有着复杂的计算、通信需求的计算图要如何才能比较好地并行起来呢,答案就是我们需要一个自动的调度系统。

首先计算图本身可以很好地描述计算之间的数据依赖关系,那在这个基础上的 Scheduler 设计感觉其实也没啥好说的,基本上都是很常规能够想到的解决方案。TensorFlow、MXNet 等等的基础设计原则都是这样,只是实现上可能有所不同。

同样是队列的调度方案,这里的这种是以每个变量为单位有一个自己的队列,TensorFlow 中是线程池中的每个线程会有个自己的队列。

Lecture 11: Distributed Training and Communication Protocols

这一节把目标放在上一节调度图中的 Synchronization 部分。

大量的篇幅是对 Allreduce 的讨论,也没啥好说的,跳过跳过…

Parameter Server - Worker 架构的同步异步,也没啥好说的…

Emm…这两节不是我偷懒,主要是内容比较基础,跟着 PPT 就好了,没什么特别值得注意的。

Lecture 12: Model Serving

除去后面没有资料的几个 Guest Talk 以外,这节算是课程内容的最后一部分了,主要讲的是实验室的成果上线进行实用的过程中可能会有的问题,例如:

  • 延迟限制:在云上提供服务的时候,Batch Size 没办法做到越大越好;如果是在终端设备上提供服务,则能够支持运行的模型本身也有轻量级的要求;
  • 资源限制:设备有功耗、内存限制;云服务则有花费开销的限制;
  • 精度限制:提供多级 QoS 等等。

在下面这个视频应用中:

终端设备的采集、处理、数据传输等各个不同有码率、功耗等限制,云端提供服务的部分也有带宽和花费开销的限制,事实上从 Workload 到 Budget 之间也存在一个巨大的 Gap。

下面的内容主要从模型压缩和服务系统两个方面来介绍。

Model Compression

这部分是介绍如何对一个网络模型进行压缩。

首先是矩阵/向量的低秩分解,可以应用在全连接层和卷积层中,能够有效地降低整体的计算量和存储量。

这块还是自己的数学知识比较缺乏,暂时不往下细看了。

然后是网络剪枝:训好一个网络之后,通过一个 01 的 Mask 把参数中的某些部分置为 0,再重新训练达到之前相同的预测精度,不断重复以上过程并且逐渐增加 Mask 中 0 的比例,最后就可以得到一个想要的剪枝结果。

权值共享:对参数矩阵进行重新采样,把实际值存在一张表中,然后参数矩阵改成存储对应实际值在表中的索引。这也是尽可能地减少存储冗余。

低比特量化:这个也比较容易理解,就是降低数据类型的存储精度,32 位单精降到半精、int8 甚至是二进制的 01 值。预测结果可能会有一定的精度损失,但是在可以接受的精度损失范围内可以大大节省参数的存储量,并且配合上低精度的硬件预算部件也能大大加快运算速度。

知识蒸馏:用一个训练好的大模型来训练一个小模型。

Knowledge Distillation 这块感觉挺神奇,但是还没细看,不是很理解。

这里还给了一些参考的论文资料:

  • Compression of deep convolutional neural networks for fast and low power mobile applications. ICLR 2016
  • Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. ICLR 2016
  • “XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks. ECCV 2016
  • EIE: Efficient Inference Engine on Compressed Deep Neural Network. ISCA 2016
  • MCDNN: An Approximation-Based Execution Framework for Deep Stream Processing Under Resource Constraints. MobiSys 2016

Serving System

一个比较好的服务系统需要达到几个目标:

  • 写应用的时候要有很高的灵活性
  • 应用跑在 GPU 上要有很高的效率
  • 满足各种不同的延迟 SLA 需求

然后举了个叫 Nexus 的系统为例,后面就不细看了。