Press "Enter" to skip to content

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

详细解释第三章中的自动求导。
计算机中实现求导这个操作主要有两种方式:基于符号的算术求导和直接用数值进行求导。
算术求导需要构建一棵符号表示树,然后根据各种算术上的求导规则来写公式。缺点在于:如果遇到特别复杂的函数,则需要推导的符号表示树也会很大;然后如果目标只是想要一个导数值,则保存一棵符号表示树就很浪费了;再然后就是这样做容易有误差(?为什么…按公式算不应该误差更小吗)。
数值求导则是按导数的定义做,直接对方程取极限:
f(x)xilimh0f(x+hei)f(x)h

实现起来特别简单,h 取个 1e-6 就差不多了,但是一般只用来检验求导结果用。
然后对于网络中每一层的反向部分,其实求导涉及到的都只是跟本层运算相关的内容:
Backpropagation
上一层传下来的是 Errorz,再往下可以通过链式求导法则一直推导下去,而其他需要的则只是与本层运算有关的zx 和 zy
更详细的推导可见这里
因此自动求导则是根据以上的规则来创建反向计算图的过程,伪代码以及结果如下:
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 还刚刚开始发展,后面还有一大堆问题留待解决。

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注