这个系列的开山之作:【FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness】
主要是通过计算换访存的思路,减少 cuda kernel 对 global memory 的访问量,在 sequence length 长,整体计算过程偏 memory bound 的情况下有很好的效果。
先来看下常规的 Self Attention:
$$O = Softmax(Q * K^T) * V$$
假设 Sequence Length 是 L,Head dim 是 D,这个过程在常规实现中需要放在 3 个 cuda kernel 中,伪代码:
1 | // [L, Dim] * [Dim, L] -> [L, L] |
对 global memory 的访存量为 $4L * Dim + 3L * L = O(LDim + L^2)$
重新看一下 softmax 是在做什么:
$$Softmax({x_1, x_2, …, x_N}) = \left \lbrace \frac{e^{x_i}}{\sum^N_{j=1}e^{x_j}} \right \rbrace^N_{i=1}$$
值得注意的是,目前通用的 softmax 实现中为了防止数值溢出还需要再额外减掉一个 max:
$$\begin{aligned}
m &= Max^N_{n=0}x_n\\
Softmax({x_1, x_2, …, x_N}) &= \left \lbrace \frac{e^{x_i - m}}{\sum^N_{j=1}e^{x_j - m}} \right \rbrace^N_{i=1}
\end{aligned}
$$
因此这里计算 max 需要一次独立的全局 reduce,计算分母的 sum 再需要一次独立的全局 reduce,最后分别计算每一个元素的 softmax 值。三个步骤之间存在数据依赖。
【Online normalizer calculation for softmax】 这篇 paper 提出了一种能将上面的 3 步 softmax 合并成 2 步完成的思路。
考虑原始计算步骤中分母求最大值以及求和的部分,这里需要 3 个独立的循环:
$$\begin{align}
m_i &= \max(m_{i-1}, x_i) \\
d_i &= d_{i-1} + e^{x_i - m_N} = \sum^i_{j=1}e^{x_j - m_N} \\
softmax_i &= \frac{e^{x_i - m_N}}{d_N} \\
\end{align}$$
公式中产生数据依赖的原因是 $(2)$ 需要依赖 $m_N$,而 $(3)$ 需要依赖 $m_N$ 和 $d_N$。
如果能有这样一个 $d_i’$:
$$d_i’ = \sum^i_{j=1}e^{x_j - m_i} = d_{i-1}’ + e^{x_i - m_i}$$
则 $(2)$ 对 $m_N$ 的数据依赖就解除了,虽然序列的中间部分的值不相等但 $d_N$ 与 $d_N’$ 是等价的。
而 $d_i’$ 这个序列存在递推性质:
$$
\begin{aligned}
d_i’ &= \sum^i_{j=1}e^{x_j - m_i} \\
&= \sum^{i-1}_{j=1}e^{x_j - m_i} + e^{x_i-m_i} \\
&= \left ({\sum^{i-1}_{j=1}e^{x_j - m_{i-1}}} \right ) * e^{m_{i-1} - m_i} + e^{x_i-m_i} \\
&= d_{i-1}’ * e^{m_{i-1} - m_i} + e^{x_i-m_i} \\
\end{aligned}
$$
这样 $m_i$ 和 $d_i’$ 就可以在一个 kernel 中计算完成了,kernel 可以减少到两个:
$$\begin{aligned}
m_i &= \max(m_{i-1}, x_i) \\
d_i’ &= d_{i-1}’ * e^{m_{i-1} - m_i} + e^{x_i-m_i} \\
softmax_i &= \frac{e^{x_i - m_N}}{d_N’}
\end{aligned}$$
但是这样对于加速来说还是不够,重新推几层:
Index | 0 | 1 | 2 | 3 |
---|---|---|---|---|
$m_i$ | $x_0$ | $\max(m_0, x_1)$ | $\max(m_1, x_2)$ | $\max(m_2, x_3)$ |
$d_i’$ | 1 | $e^{x_0-m_1} + e^{x_1-m_1}$ | $e^{x_0-m_2} + e^{x_1-m_2} + e^{x_2-m_2}$ | $e^{x_0-m_3} + e^{x_1-m_3} + e^{x_2-m_3} + e^{x_3-m_3}$ |
可以发现由于 exp/log 计算的特性,$d_3’$ 除了正常从 $d_0’$、$d_1’$、$d_2’$ 按顺序推出以外,乱序也是可以得到的,例如 2->1->0->3:
$$\begin{aligned}
m_{21} &= \max(x_2, x_1) \\
d_{21}’ &= e^{x_2 - m_{21}} + e^{x_1 - m_{21}}\\
m_{210} &= \max(x_0, m_{21}) = \max(x_0, x_2, x_1)\\
d_{210}’ &= d_{21}’ * e^{m_{21} - m_{210}} + e^{x_0 - m_{210}}\\
&= (e^{x_2 - m_{21}} + e^{x_1 - m_{21}}) * e^{m_{21} - m_{210}} + e^{x_0 - m_{210}}\\
&= e^{x_2 - m_{210}} + e^{x_1 - m_{210}} + e^{x_0 - m_{210}}\\
m_3 = m_{2103} &= \max(x_3, m_{210}) = \max(x_3, x_0, x_2, x_1) \\
d_3’=d_{2103}’&=d_{210}’ * e^{m_{210} - m_3} + e^{x_3 - m_3}\\
&=(e^{x_2 - m_{210}} + e^{x_1 - m_{210}} + e^{x_0 - m_{210}}) * e^{m_{210} - m_3} + e^{x_3 - m_3}\\
&=e^{x_2 - m_3} + e^{x_1 - m_3} + e^{x_0 - m_3} + e^{x_3 - m_3}\\
\end{aligned}$$
如果定义分块计算时 $d_{xy}’ = d_x’ * e^{m_x - m_{xy}} + d_y’ * e^{m_y - m_{xy}}$,则按照 (2->1)->(0->3) 的迭代顺序:
$$\begin{aligned}
m_{21} &= \max(x_2, x_1) \\
d_{21}’ &= e^{x_2 - m_{21}} + e^{x_1 - m_{21}}\\
m_{03} &= \max(x_0, x_3) \\
d_{03}’ &= e^{x_0 - m_{03}} + e^{x_3 - m_{03}}\\
m_3=m_{2103} &= \max(m_{21}, m_{03}) = \max(x_2, x_1, x_0, x_3)\\
d_3’=d_{2103}’ &= d_{21}’ * e^{m_{21} - m_3} + d_{03}’ * e^{m_{03} - m_3} \\
&= (e^{x_2 - m_{21}} + e^{x_1 - m_{21}}) * e^{m_{21} - m_3} + e^{x_0 - m_3} + (e^{x_0 - m_{03}} + e^{x_3 - m_{03}}) * e^{m_{03} - m_3} \\
&= e^{x_2 - m_3} + e^{x_1 - m_3} + e^{x_0 - m_3} + e^{x_3 - m_3}\\
\end{aligned}$$
同样可以得到相同的 $d_N’$。这样,我们就可以得到这种方式最大的一个特性:$m$ 和 $d’$ 的迭代计算操作同时满足交换律和结合律,任意分块分别计算 $m$ 和 $d’$ 之后,将所有子块结果重新聚合在数学上完全等价,即序列中 max 值带来的影响可以延迟到最后一步再被修正。
这样 Online softmax 就可以通过分块并行得到进一步的加速了。
继续回到一开始的 FlashAttention,我们把 Online Softmax 放进去,这里的 Softmax 是对 [L, L] 结果中的每一行做一维的 softmax:
$$\begin{aligned}
Softmax^L_{r=1} &= Softmax(X_{r, 1}, X_{r, 2}, …, X_{r, L})^L_{r=1}\\
&=\left \lbrace \left \lbrace
\frac{e^{X_{r, i} - M_{r, L}}}{\sum^L_{j=1}e^{X_{r, j} - m_{r, L}}}
\right \rbrace^L_{i=1} \right \rbrace^L_{r=1} \\
\end{aligned}$$
原始实现中:
$$\begin{aligned}
X_{r, i} &= \sum^{Dim}_{j=1}Q[r, j]K[j, i]\\
M_{r, i} &= \max(M_{r, i-1}, X_{r, i}) \\
D_{r, i} &= D_{r, i-1} + e^{X_{r, i} - M_{r, L}} = \sum^i_{j=1}e^{X_{r, j} - M_{r, L}} \\
Softmax_{r, i} &= \frac{e^{X_{r, i} - M_{r, L}}}{D_{r, L}} \\
\end{aligned}$$
同样可以替换成 $D_{r, i}’$:
$$\begin{aligned}
D_{r, i}’ &= D_{r, i-1}’ * e^{M_{r, i-1} - M_{r, i}} + e^{X_{r, i}-M_{r, i}} \\
Softmax_{r, i} &= \frac{e^{X_{r, i} - M_{r, L}}}{D_{r, L}’} \\
O_{r, c} &= \sum^L_{i=1}(Softmax_{r, i} * V[i, c]) \\
\end{aligned}$$
把 $O_{r, c}$ 的累加过程拆开看:
$$\begin{aligned}
SubSum_{r, c, i} &= SubSum_{r, c, i-1} + Softmax_{r, i} * V[i, c]\\
&=SubSum_{r, c, i-1} + \frac{e^{X_{r, i} - M_{r, L}}}{D_{r, L}’} * V[i, c]\\
&=\sum^i_{j=1}\frac{e^{X_{r, j} - M_{r, L}}}{D_{r, L}’}V[j, c]
\end{aligned}$$
可以发现 $SubSum_{r,c,i}$ 也是依赖于 $M_{r,L}$ 和 ${D_{r,L}’}$,运用与 online softmax 相似的方式,可以在这里增加一个 $SubSum_{r,c,i}’$:
$$\begin{aligned}
SubSum_{r,c,i}’ &= \sum^i_{j=1}\frac{e^{X_{r, j} - M_{r, i}}}{D_{r, i}’}V[j, c]\\
&=\sum^{i-1}_{j=1}\frac{e^{X_{r, j} - M_{r, i}}}{D_{r, i}’}V[j, c] + \frac{e^{X_{r, i} - M_{r, i}}}{D_{r, i}’}V[i, c]\\
&=\left (\sum^{i-1}_{j=1}\frac{e^{X_{r, j} - M_{r, i-1}}}{D_{r, i-1}’}V[j, c] \right) * \frac{e^{M_{r, i-1} - M_{r, i}}D_{r,i-1}’}{D_{r,i}’} + \frac{e^{X_{r, i} - M_{r, i}}}{D_{r, i}’}V[i, c]\\
&=SubSum_{r,c,i-1}’*\frac{e^{M_{r, i-1} - M_{r, i}}D_{r,i-1}’}{D_{r,i}’} + \frac{e^{X_{r, i} - M_{r, i}}}{D_{r, i}’}V[i, c]\\
\end{aligned}$$
最终整理一下,在一个 $i = (1, L)$ 的循环中可以完成:
$$\begin{aligned}
X_{r, i} &= \sum^{Dim}_{j=1}Q[r, j]K[j, i]\\
M_{r, i} &= \max(M_{r, i-1}, X_{r, i})\\
D_{r, i}’ &= D_{r, i-1}’ * e^{M_{r, i-1} - M_{r, i}} + e^{X_{r, i}-M_{r, i}}\\
SubSum_{r,c,i}’ &=SubSum_{r,c,i-1}’*\frac{e^{M_{r, i-1} - M_{r, i}}D_{r,i-1}’}{D_{r,i}’} + \frac{e^{X_{r, i} - M_{r, i}}}{D_{r, i}’}V[i, c]\\
\end{aligned}$$
最终的输出结果为:
$$O_{r, c} = SubSum_{r,c,L}$$
写成伪代码是:
1 | for (r = 1 to L) |
当然,与 Online softmax 一样,$D_{r, i}’$、$SubSum_{r,c,i}’$ 也具有分块满足交换律和结合律的特性:
$$\begin{aligned}
D_{r, xy}’ &= D_{r, x}’ * e^{M_{r, x} - M_{r, xy}} + D_{r, y}’ * e^{M_{r, y} - M_{r, xy}}\\
SubSum_{r,c,xy}’ &= SubSum_{r,c,x}’ * \frac{e^{M_{r, x}-M_{r, xy}}D_{r, x}}{D_{r, xy}’} + SubSum_{r,c,y}’ * \frac{e^{M_{r, y}-M_{r, xy}}D_{r, y}}{D_{r, xy}’}\\
\end{aligned}$$
则在常规矩阵乘法计算中可以用到的 tiling 分块策略也同样可以用在这个算法上得到终极加速了。
【FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning】
待续…
]]>Talk 的视频链接在:这里
在过去的快两年时间里,我们依靠 TVM 打造了 NIO 全栈自研的 AI 引擎,支持各种感知、规控等等自动驾驶算法在 NT2.0 平台上的高效运行。我们几乎是看着 ET7 交付启动时只有 AEB,到后来逐步版本迭代可以支持 NOP+,再到今年 NAD 也即将要在我们的引擎支撑下开放测试了。
由于 talk 的时间限制,PPT 上对很多部分都是简单带过了,这里想稍微细节展开讲一下,也算是对来 NIO 以后工作的一个总结。
不知道算不算目前业界部署 TVM 范围最广的团队了(笑)。NT2 这代的车已经不止跑在全国各地的路上了,去年也已经出口到欧洲。
去年最早我们的部署目标只有 ET7,到后来 ES7、ET5 大规模交付,今年剩下的 ES8、EC7、ES6、EC6 也全面交付之后,路上能看到的 TVM inside vehicle 就到处都是了 ~(撒花)
这里想说以下我们为什么选择了 TVM。
虽然说实话得承认有一个很大的原因是我们团队初期的几个成员在加入之前都是主要做 TVM 的。认真地回答则是我们其实也考虑过各种可选择的方案,那时候应该就是 TVM 和 MLIR 两个开源项目相对比较有活力了。我们认为 TVM 的当时的现状更加成熟一些:前端上,tf、torch、onnx 等等各种框架的支持都有了;后端有 AutoTVM 以及 Ansor 这样可用的工具,并且我们在 Ansor 中已经证明了通过搜索 codegen 是能够达到 SOTA 的性能的;另外像 TVM 中 BYOC 这样的 feature 也让我们在实际部署中有了更大的灵活性。
当然,MLIR 的架构更加灵活,想要把这些工作在 MLIR 中做到相同的水平当然也是可以期待的。只是凭借我们组的这么点人力,想要在短时间内做到可用基本上就比较难了。要知道从我们入职开始搭建 AI 引擎到 ET7 直接开启首批用户交付差不多只有半年时间。
值得一提的是,TVM 有 tianqi 来控制发展大方向,相比各家百花齐放的 MLIR 而言,感觉各个项目能更踏实地落到能真正用起来的实处上。我们也很期待能把后续工作切到 TensorIR、MetaSchedule 上,更多地把 TVM Unity 生态中的东西用起来。
目前,我们实际部署过在车上跑的模型应该也有几十个了,模型包含 CV 类的、transformer 等等各种类型的模型,甚至还见过在上家公司做推荐业务时见过的 DeepFM 类模型,可以说自动驾驶真实是各种 AI 算法大综合的一个场景了。我们早期模型中存在比较多的控制流逻辑和动态 shape 情况,后来几个重点模型即使是上车前已经做了很大压缩和简化了,也还是会包含一部分动态 shape(例如 Lidar 点云数据天生就是稀疏 + 动态的)以及不少需要处理业务逻辑的 custom op 算子,模型本身还是相当复杂的。
也是受益于 TVM 中 BYOC 以及 VM 的灵活性,我们可以在模型中集 TRT + Cutlass + Custom Op + Ansor 等各种优化引擎的能力于一体,将每个部分的性能和精度优化到极致。
早期我们还会需要靠 TRT 来处理模型中比较规整的静态 DNN 部分;而随着我们自身的优化能力逐步起来之后,目前上车的业务模型已经全面去掉了 TRT,性能与量化后精度全面超越 TRT 的版本,在某几个重点的感知模型中至多能达到 TRT 的 1.5 倍性能。
一些内部的数据没法公开,PPT 里面就简单贴了下公开 benchmark 的结果。
去掉 TRT 之后,我们可以对模型本身做更多灵活有意思的事情,也才终于真正可以有底气地说是全链路全栈自研,全链路自主可控了,车上跑的 AI 模块里面的每一个地方我们都能拿得出来 source code。
量化算法以 PTQ int8 + uint8 为主,对于一些精度敏感性较好的层会再继续尝试 int4 + uint4。如果遇到 scale 没办法通过进一步的误差分析和其他修正算法拉回来的情况就只能 fallback 回浮点了。量化模块处理完之后再将图过一遍 MixedPrecision pass,尽可能多的将剩下的浮点部分转成 fp16 处理。
这里有个也许是我们业务场景带来的特殊量化需求:部分量化。
车上环境对车规验证要求比较严格,每一次自动驾驶算法的模型更新都需要大量的仿真和实车路测验证才能通过验收。在目前比较普遍的一个 Backbone + 多个不同功能的检测 head 的模型结构中,有时候一些小版本更新只会更新部分检测 head 中的结构和参数。
我们通过控制量化算法的更新范围,可以做到一个模型中未改动部分的输出数值与之前的版本完全一致(误差在浮点 1e-10 以下严格一致),则小版本更新时所需的测试只需要重点关注改动部分即可。从引擎这一侧出发,我们给出了严格控制量化算法下 x86 云端 GPU 与车端 Orin 芯片结果一致性的保证,也是帮助业务和测试团队节省了很大的测试负担。
图优化部分没什么太多可说的了,继续找到一些可优化的 pattern 优化掉,都是各家遇到了之后会做的工作。
由于后续我们没有继续在 Ansor 中做 TensorCore 支持的开发了,目前我们重度依赖 Cutlass 来处理 Conv/Matmul 这样的计算密集算子。
论坛上也有好几个帖子问过相关的问题,其实之前我有个实验性的 custom sketch 是可以支持 TensorCore 跑起来的:
https://github.com/jcf94/tvm/tree/tensorcore_sketch_support/python/tvm/auto_scheduler/test_sketch
另外基于我们以前 AutoTensorCore 工作配合 Ansor 也能有不错的效果,只是后续没有精力继续完善了
后续更加合适的方向当然还是依靠 TensorIR 和 MetaSchedule 再来把 TensorCore 的性能推到一个新的高度了
拿到量化 + MixedPrecision 之后的计算图后,图上尽可能多的部分会通过 BYOC 圈给 Cutlass。相比社区版本的 Cutlass,我们扩展了更多 epilogue pattern 的支持,将 Conv/Matmul 前后的 q/dq、bias/residue add 以及连接着的 activation 等等边角的算子都尽可能圈进去。
剩下一些零碎的访存算子以及其他不太适合使用 TensorCore 来计算的算子则由 Ansor 来 cover 了。举例来说,大多数 CNN 模型的输入图像 channel 会是 3,也就导致第一层 conv 通常会有一个对 TensorCore 来说非常不友好的计算 shape,我们在 Ansor 中扩展了 CUDA dp2a/dp4a 指令的支持,这样第一层 conv 可以做到相比强行使用 TensorCore 更好的性能。其他像优化 Warp Shuffle 实现 cross thread reduction 等等也是社区版本上原本就有基础的实现,我们的版本就是在基础上继续扩充了更多各种场景下的 Ansor Sketch,让 Ansor 的搜索空间能够包含更强的性能,搜索得更有效率。
当然,我们也一直在跟进和在 MetaSchedule 上进行着性能探索工作,期待未来的有一天我们能将对 Cutlass 的使用逐步迁移到 MetaSchedule 上。
这个其实也是个比较有意思工作,大思路上与 TVM Unity 的想法是完全一致的。我们的初衷是为了能在 TVM 中方便地引入其他系统能力。比如 TensorFlow、Torch 中就有像 NNI 这样成熟的量化工具。
我们通过 MSC 可以很方便地完成 Relay graph 到 Torch、TensorFlow 等等各种第三方框架代码的转换,通过其他工具完成处理,最后重新转回 TVM 中完成最后的后端 codegen 部分。我们目前的整套量化算法就完全是在 torch 中实现的,借由这样一套灵活的架构,其他包括 Gym、Nas、稀疏化、QAT 等等工作也可以更方便地得到实现。
最后这部分算是我近半年来觉得做的最有趣的一块工作了。TVM 中主要依靠 VM runtime 实现动态 shape 的支持。由于业务模型中的部分动态性没有办法避免,VM 也是我们唯一的选择。
如果实际跑过的同学应该也会有体会,虽然 VM 已经是个相当轻量级的 runtime 了,但是相比静态的 GraphRuntime 而言,kernel 之外的 overhead 还是会更大一些。尤其是像在 Orin 这种 ARM 核性能并不是太强的硬件架构上,CPU 部分的负担有时候会超过 GPU,当 workload 太多的时候整个系统可能会 bound 在 CPU 上,最严重的极端情况下 GPU kernel 甚至会被卡住发不出去。
VM 的指令设计非常精巧,但在实际的业务优化过程中,我们逐渐发现它还是有点复杂了。(话说我看新版的 Relax VM 里面已经把 ByteCode 指令的种类减少了很多了,不知道是不是也有这方面的思考)
我们的想法是,能不能从现代超标量处理器对于 CPU 指令的优化经验中借鉴一些来,对 VM 目前的执行方式做一些好玩的工作呢。于是就有了我们现在的魔改版 VM。针对 VM 指令,我们加了一些基础调整策略:
做完重排和合并之后,我们还想过可以做一个动态多发射(多线程 issue 最后顺序 commit)的超标量虚拟机(可惜简单试了下发现没有太大的收益…后面倒是有机会也希望能更多地探索一下)。不过,至少把 memory allocation 相关的指令合在一起之后就可以基于这个做类似 GraphRuntime 那样的静态初始化 memory planing 了。
当所有执行 kernel 的 memory 有机会变成静态之后,更大的收获是 CUDA Graph 也可以上了!所以可以看到我们最终做完指令级优化的模型相比刚转出来的时候能减少 95% 的 ByteCode 指令数量,实际部署环境下对 CPU 的负载也能大大减少了。
这里另外可以提的一个工作是基于 VM 可以往一个模型里打包多个 ByteCode Function 的能力,我们增加了可以每一次 inference 都能够动态选择执行模型某一部分的能力。这样车端的某些检测 head 在不需要这么高的帧率时可以动态关闭以节省算力和功耗。
Runtime 部分其他的内容诸如重写了全新的内存池等等都是我们在业务摸索遇到了实打实的坑之后慢慢加上的。
我们几乎是把整个 TVM 里面对 CUDA Stream 的使用都重新理了一遍,基于内存池把 memory allocate、free 等等各种管理操作也都绑定到了 stream 上,实现了一个几乎是全异步的 runtime。
非常幸运能在几年前从 TVM 开始 AI Compiler 的探索之旅,也非常开心能看到 TVM 每年都还这么有活力,朝着更强大的方向继续发展着。
至少我们现在已经证明了 TVM 这套编译栈架构是足以撑得起几十万辆车上车规级业务的稳定部署的。后续我们也还将基于 TVM 继续在自研芯片上支持 NT3.0 的下一代车型和自动驾驶功能。
]]>后面每一年再有新的更新也会继续补充到后面。
Compute Capability: 9.0
SM 结构:
每个 Process Block:
比较显著的改动是 FP32/FP64 Core 的数量翻倍了,加上 H100 主频相比 A100 更高,以及 SM 数量的增加,体现到最终 H100 单卡可以达到相比 A100 做到 3 倍的浮点性能。
这代 TensorCore 的单 cycle 吞吐相比 Ampere 再翻了一倍,反映到单卡上也是 3 倍多 A100 的同类型运算性能。
更进一步能让数据上更好看的是 TensorCore 增加了 FP8 数据类型支持(可选 4 位指数 + 3 位底数或者 5 位指数 + 2 位底数),相比自己的 FP16 是两倍性能,与 A100 的 FP16 比就是 6 倍性能了。
软件层面加了个 Transformer Engine,专为 transformer 结构做的优化,估计也是通过加速库 api 的形式提供。
SM 里面多出来的这个 TMA 相当于是把之前针对 TensorCore 上做的一些软件层面的 memory 搬运优化固化到了硬件上,进一步优化了 TensorCore 的数据访存效率。
感觉 TensorCore 上真是能搞的都搞了:计算提升主要靠工艺堆料和新的数据类型,Sparse 和这次这个 TMA 都是针对访存的提升。
有点好奇 TensorCore 现在在 GPU 上的芯片面积能占到多少了?
感觉会不会已经像是一小部分通用计算部件挂一个巨大 NPU 的结构了,与一众新兴 AI 芯片公司的设计殊途同归了属于是。
看 NV 明年还能再往 TensorCore 上整点啥。
编程模型方面,Hopper 新增了一种 DPX 指令,主要用于动态规划算法的优化。白皮书中的举例是生物 DNA 序列匹配上的 Smith-Waterman 算法(算不算蹭最近几年因为新冠火起来的生物信息方面研究的热点,AI 上搞得差不多了,NV 继续往其他有可能突破的方向上发力)。
直接看到这个我还真想象不到动态规划和 GPU 是怎么能联系在一起,这里也没有给指令的 api 和更多实现细节。
]]>1 年间感觉世界发生了很多变化,技术方面,感觉我们这边是每个月都在 rush 新东西。
前段时间 TESLA 说他们要比之前说的纯视觉方案要走的更加极致了,连近距离的毫米波雷达都要取消了。惊叹于他们对自己技术实力的自信(和他们的胆子是真的大啊)的时候,也确实挺让人好奇这两年间 TESLA 的 FSD 有了什么样的改进。
YouTube 的原始直播资源在:这里。
同样,今年 B 站也是有不少人搬运的,顺便也贴一个在这里:
Join us to build the future of AI → https://www.tesla.com/ai
0:00:00 - Pre-show
0:13:56 - Tesla Bot Demo
0:29:15 - Tesla Bot Hardware | Hardware Architecture
0:34:22 - Tesla Bot Hardware | Hardware Simulation
0:39:40 - Tesla Bot Hardware | Actuators
0:45:12 - Tesla Bot Hardware | Hands
0:47:24 - Tesla Bot Software | Autonomy Overview
0:49:55 - Tesla Bot Software | Locomotion Planning
0:52:20 - Tesla Bot Software | Motion Control and State Estimation
0:55:00 - Tesla Bot Software | Manipulation
0:56:44 - Tesla Bot Software | What’s Next?
0:58:00 - FSD Intro
1:04:32 - FSD | Planning
1:12:11 - FSD | Occupancy Network
1:19:17 - FSD | Training Infra
1:25:48 - FSD | Lanes and Objects
1:34:22 - FSD | AI Compiler & Inference
1:40:34 - FSD | Auto Labeling
1:47:45 - FSD | Simulation
1:53:33 - FSD | Data Engine
1:56:50 - Dojo Intro
2:02:30 - Dojo Hardware
2:13:47 - Dojo Software
2:26:25 - Q&A
去年 Tesla Bot 只是相当于结尾菜单的感觉,今年的 AI Day 已经是把前半部分的大头都给了它。(TESLA 是真想把这玩意搞出来呀…)
我们还是跳过继续关注 FSD 部分吧。
开场是个比较经典的无保护左转场景的例子,这是一个没有单独做左转转向红绿灯的十字路口,对面有车辆、行人过马路,对侧车道也不断有车辆从右侧向左侧驶过:
假设目前感知得到的场景中的对象有 20 个,他们相互之间可能存在超过 100 中相对交互组合,考虑对场景中所有对象做多智能体联合轨迹规划可能会出现组合爆炸的情况。而往往在这种场景中,其他对象都是动态的,留给 planning 的时间是很短,多迟疑一会,可能之前可行的行动方案就被错过了。
FSD 采用了一套 Interaction Search 的框架去解决这个规划问题:
上面的决策过程相比去年的版本已经出现了一个新的概念:Occupancy。
从演示图上,我们可以发现这个图上的概念其实看着都非常像雷达点云图了,TESLA 的思路有点像是用 8 个摄像头的图像结果中直接得到类似全向雷达的扫描效果(个人想法)了。
Occupancy network 直接从视频输入中得到一个反映了每个位置是否有物体占用的向量空间,也因为有视频时序的关系,对于被遮挡的位置也可以预测是否存在障碍物等等。同时,输出结果中也会包含物体类别的语义信息(直接也做了物体识别了?)。
跑一次 10 ms 左右,相当于最佳情况下能跑到 100 帧了。
接下来,终于可以看到整个网络的全貌了:
对比去年的网络结构,可以看到网络的主体部分基本上是差不多的。
去年我看的时候有个未解之谜是输入图片的 channel 是多少终于破案了!小哥今年终于提到说 12bit 的原始输入图像是 4 channel,相比常规处理后的 3 channel RGB 图像可以多一些信息,并且用原始图像的好处是不需要 ISP 了,少了这里的一个步骤之后整个流程可以更快一些。
原始图像去畸变,过 RegNets 和 BiFPNs,再到 Spatial Attention 和时序处理,这里这些都是跟去年网络中的思路完全一致。
后面的部分是今年 Occupancy Network 的改变,通过多次反卷积,把时序模块得到的 [C, T, X, Y, Z] 结果通过一系列反卷积展开成 [C, 16X, 16Y, 16Z] 的最终空间信息,即前面所说的 Occupancy 向量空间了。
为了得到更高的分辨率,最终还有两个 MLP 用于从这个向量空间中把需要的信息提取出来。输入 [x, y, z] 的空间坐标,分别可以得到每个位置的占用概率以及分类语义。可以看做是通过反卷积和 MLP 的 query 实现了去年网络中不同 task 的 head 识别的任务,网络的复杂度被大大简化了。
下面也给了个例子:
场景中右侧是一辆前后两节弯折的大型巴士,模型一开始标记是红色物体的前方有另外一个蓝色物体,随着车辆继续前行,两个物体被标记成了相同颜色,即网络识别出来了这是属于同一辆完整的车,并且巴士弯折的部分也可以在 occupancy 空间中很精确地标识出来。(想了下,如果用去年的网络处理这种场景可能确实比较麻烦,大概率还是会把这俩识别成两个不同的物体框了)
确实让人忍不住想要称赞 666 了。
本来以为这部分要结束了,没想到上面这张图的更完整版是这样的:
Occupancy Network 不仅是能得到物体的空间占用以及类别语义,更重要的是把地面的语义也同样带上了(事实上前面的 Volume Output 空间本身是包含地面的),这个向量空间是完整的,连地面的上坡下坡起伏的信息都能精确地识别出来,这个对于后续自动驾驶的决策而言也有非常大的实用价值。也不用再在 head 中加一个 ground 的目标或者专门用一个后处理阶段来搞这些东西了。
What’s more! NeRF States! 惊了 … Occupancy Network 的输出信息甚至可以应用到光场识别上。后续可能能进一步用作 3D 重建等 CV 任务。
这一节主要是介绍针对视频训练部分做的优化,做的工作也是相当硬核,虽然看起来都是比较零碎的,但是也真的可以称得上是对全链路上的每个步骤都做完充足的优化了。
都是一些工程部署上的细节,比如优化文件存储,减少 copy,甚至设置环境变量等等,这里不展开了,贴一张总结图:
另外想吐槽的一个点是,其实这里面大部分的零碎优化点通常每个优化组都会做,但是只能算是一些配置调整,单独写出来作为工作是不够 solid 的。
接下来细节展开介绍车道线和目标识别。
回顾一下,最早的传统做法是把车道识别看做一个实时的图像分割任务,直接从输入图片的 2D 空间中识别出不同的车道线信息,然后在后融合中把这些信息带上。换到 BEV 的方案后其实本质上对于 lane feature 的获取也还是从 2D 图片中来的,只是相当于把融合位置提前了。
这样做的缺陷在于对于一些车道线非常复杂的十字路口或者根本没有车道线的情况下,这种 2D 识别的方案效果会非常不好,几乎没办法得到一个可用的结果。
以下是 TESLA 目前的 lane 模型方案:
左侧是 Occupancy Network 的视觉感知输出,加上一些地图信息之后得到了一组包含更丰富的 lane 信息的识别空间。这里的地图信息可以不需要是高精地度,就是常规的导航地图提供的那一点点信息就足以。
然后是 transformer again!
用常规的方式去标记出复杂场景中的每一条可行的车道线还是过于复杂,因此 FSD 最终采用了一种看起来可能更复杂,但是对于这个问题来说其实是一种很好的抽象的“语言学”的方式去解决。他们定义了一套词法语言来描述每一条车道线的生成规则(灵感来源可能是人类例如说问路的时候可能会有的对路径点的描述方式?),然后用一个 decoder 从上一步的 dense world tensor 中把这些信息解出来。
每条车道线上的每一个点都可以视作是 decoder 句子中的一个“词”,正如一条车道线由线上的多个点构成,decoder 工作中处理的一个“句子”表征了一条车道线。
这里的“车道线”甚至其实已经不是识别出来的了,可以相当于是预测出来的可以让车辆开的路线,即使对于实际没有车道线的场景,在模型的输出空间中也能够得到对应的结果。
又到了我最想看的部分,然而发现内容还是跟去年一样少,对于这个部分 TESLA 藏得还是太深了。
出大问题,这小哥的英语是真的难听懂啊…
针对前面提到的 Lane network 的 attention 部分,他们做了很多 op 上的优化以及针对 int8 准确率的优化。
而在全车部署框架的层面,有用的信息其实只有接下来两张图:
从这里基本上能比较确定地看出来是编译期把图都切分好,然后在车上做确定性调度了。编译期 graph partitioner 需要知道的是每个网络中的每个部分的依赖关系、运行时间等等信息。然后做全局的图切分和放置排布,NN linker 把切分好的所有子图都打包在一起,最终车上根据前面定制好的 schdule 去混合跑所有模型即可。
Graph partitioner 对于每个模型在什么时候用到了片上的什么资源是确切知道的,因此就能通过对模型中每个运行部分的排布让片上的硬件资源得到全局的高效利用。
图上的 partitioner 排布逻辑也比较好猜,就是尽可能地想办法把 FSD chip 上的两个 NPU 用满。
(话说对于 FSD chip 好像好多年都没见升级了?当然我猜测不升级的原因有可能是目前的硬件出货量已经太大了,即使升级了新硬件,老的一样要继续维护升级,对于 TESLA 来说维护这样很多代可能不一定划得来,成本也会增加很多)
后面展示网络可视化的部分有点搞笑了…
后面标注和数据处理部分我兴趣不大,跳过去直接看 Dojo 超算部分的介绍。
去年的 AI Day 上主要介绍了 Dojo 的逻辑架构,有很多工程上的实现细节是没有披露的,今年直接开场就抛出来了,以证明他们不止是理论设计可行,实际做出来之后也真的能用起来。
Dojo 自研了 VRM(Voltage Regulator Module)来解决高度集成之后芯片模块的供电和散热问题,迭代了 14 个版本之后才达到了他们预期的 CTE 要求。
然后设计了一种数据传输专用的处理器(Dojo Interface Processors)来负责板间互联和高速数据传输,给 Tile 上的 640GB 内存提供超过 18TB 每秒的运算带宽,和超过 1TB 每秒的网络交换带宽。
最终交付使用的 Dojo 会是下面这个样子(预期 2023Q1):
性能参数方面跟去年公布的应该差不多。
软件部分主要是基于 PyTorch 开发了 Dojo 的插件,也是通过 AI 编译器来 offload 到 Dojo 上。
分布式方面的优化也是比较常规的,这里给了个基于 2D-Mesh 做 AllReduce 的例子:
最后的性能展示部分主要体现出硬件和 compiler 两方面的提升带来的加速效果:
有了高度集成的 Dojo 之后,之前 6 个大 GPU 机柜可以直接用一个 75mm 高的 tile 直接代替了。真是令人惊叹…
照例跟去年一样稍微更新下我们这边的进展。(2021 回顾里面重新更新进展的时间其实已经是到 2022 年中了,这里更新的进展是 2023 年初时候的状态)
首先是 BEV 终于已经完成上车了,而且看样子效果非常不错。NT2 平台的车主这时候应该已经都用上了 NOP+ 了吧,B 站可以找到的很多测评视频中都给出了相当好的评价。至于 NAD,我也不太清楚确切的计划,不过大概率今年是也会上的(至少测试肯定会放出来了)。
我们基于 tvm 做的自研引擎已经在车上跑了 1 年多啦!可以公开的信息是在 Orin 这块 NV 的板子上我们已经全面干掉了 TRT,从量化精度到性能等各方面全面超越。23 年 tvm conference 上分享完了之后如果有空我也会另外写篇帖子详细展开一下这段我们都做了啥。
]]>回顾一下今年特斯拉的 AI day。
直播的原始视频发在了 YouTube 上,这里。
发现 B 站刚好有人把完整的 3 小时都搬过来了,就拿过来贴在这里:
顺便 YouTube 评论区置顶的这条贴了视频的几个关键时间点目录:
Join us to build the future of AI → https://www.tesla.com/ai
0:00 - Pre-event
46:54 - AI Day Begins
48:44 - Tesla Vision
1:13:12 - Planning and Control
1:24:35 - Manual Labeling
1:28:11 - Auto Labeling
1:35:15 - Simulation
1:42:10 - Hardware Integration
1:45:40 - Dojo
2:05:14 - Tesla Bot
2:12:59 - Q&A
正片差不多 38 分钟左右开始,一开始上来就放了一段自动驾驶片段。应该是辆 Model 3 或者 Y,可以在中间的中控上看到实时的环境识别和规划路径的可视化效果。左边红色的线是道路边缘,白色和橙色的应该是识别出来的车道线,绿色的是规划出来的目标路径。不得不说视频里面的效果还是很流畅的,中间也经过了一些复杂的交汇车和等待行人的十字路口等等场景。
46 分钟马斯克上来招了一波人 …
之前演示的视频就是目前特斯拉在车上采用的纯视觉方案,8 个摄像头通过神经网络处理后构建出一个 3 维的向量空间,然后所有的识别和规划都在这个空间中完成。
他们用了 ResNet 作为 Backbone,输入是摄像头采集过来的 1280 * 960 的 12 位原始图像。这里用的直接是摄像头的 RAW 数据,可以减少一些视觉处理算法的预处理等等,理论上能最大限度地保留信息?1280 * 960 这个尺寸不知道是不是他们摄像头的原始尺寸,感觉分辨率偏低,不过输入尺寸太大了确实也影响推理性能。
接下来过一个 BiFPN 去对 Backbone 中的多层数据做一个特征融合。融合后的特征最后再被拿过来做具体的识别和分类。视频里有一页给了 cls 和 reg 的输出尺寸分别是 640 * 480 * 1 和 640 * 480 * 4,猜测 1 大概是描述某个像素点是不是有物体,如果有物体的话输出的 4 则是物体的分类属性。
这里有个重要的点是前面到 BiFPN 为止的数据和特征都是被接下来不同的分类任务共享的,后面接上不同的检测头就可以从原始的图片数据里面分别检测出物体(可以是车、人、其他障碍物等等)、车道线、红绿灯等等。每个检测头都是独立的,比如这里给出的细节是 lane 只用了个简单的全连接,其他的目标识别则是用了更复杂一些的网络。这种 one model + multi task heads 的方式:
后面出结果的时候也是一帧图片进去,跑完整个网络直接就能同时出来所有不同类型对象的检测结果。
单摄像头拍到的视频帧会少了个很重要的深度信息,因此每个摄像头得到的结果是没办法简单地直接放在一起用的,且单独一个摄像头检测出来的结果也很容易会产生偏差,也会有一个对象可能同时出现在多个摄像头的画面里面的情况。所以接下来是一个叫 occupancy tracker 的部分,需要把所有摄像头识别出来的车道线、物体等等根据它们的空间关系融合到一起去(就是前面视频里面提到的 3 维向量空间)。
“It’s very easily said much more difficult to actually achieve.” Emm … 这句话好像在哪里听过的样子,不过确实这个想想就很麻烦。
模型先放到后面说,首先第一个技术点是不同车上摄像头的位置会有偏差,每个摄像头拍到的画面也都会有一定的畸变,每个画面到最后提取完 feature 再去做调整就不太好操作了,效果可能也不好。所以他们直接采用的是前融合的方式,在摄像头的 raw 数据后面加了一个 Rectify 层,在获取到图片数据的时候就用上相机的标定数据等等信息,对图像做去畸变矫正以及变换到一个标准的 virtual common camera 的视角上去。
这样是不是不同车上的数据在最后训练和识别的时候都能够完全统一到一起去了?666
然后回到这个融合模型上,说实话细节部分我还没有完全理解(惭愧…),他们用了个 transformer 来做多目 image 到 BEV(Birds’s-Eye-View?) 空间的转换和 fusion。不同摄像头图片上提取到的 feature 通过 context summary 和 positional encoder 合到一个网络里面去,得到这个 transformer 里面的 Q。每一张图片和它们的 feature 分别变成了 K 和 V。
一开始我想了好久也没理解他这里说的 Query 和 Key/Value 是啥,后来一看这不就是 transformer 结构里面的 QKV 吗,希望应该不是我理解错了。
如果 do all the engineering correctly(听起来确实挺麻烦),Transformer 的最终结果就可以拿来做融合后的最后识别了。输出的 multi-camera features 的尺寸是 20 * 80 * 256。
后面的演示结果里面 Multi-Cam 前融合得到的车道线和物体检测结果都确实要比做后融合的效果要清晰和稳定的多。
下一步是对所有视频帧的数据进行时序融合,只有这样才能进一步确定下来每个检测对象的速度、方向、被遮挡等等的信息。
每一帧的 feature(20 * 80 * 256 * 60)加上 IMU 中得到的车本身的姿态信息(1 * 1 * 4 * 60)再加上位置编码(1 * 1 * 40 * 60)下采样后得到的 video feature(20 * 80 * 300 * 12)会存放到一个 feature 队列中。Feature 队列中数据的出入队需要同时基于时间和距离信息决定,时间很好理解,太久之前的信息就没有用了。距离信息这里举了个例子是如果在一个路口之前提前看到了车道上对这个路口的车道标识,那后面在路口等待红绿灯时之前的车道标识信息也是可以用上的。
在视频信息融合上他们尝试过 3D Conv、transformer、RNN 等等,最后选了个空间 LSTM。车在每一个时间点上的感知能力是有限的,如果把整个感知空间定义成一个很大的二维平面(W * H * C),其实每一次只需要更新车周围的平面位置的数据就好了,车周围平面的 feature 尺寸是 20 * 80 * 256,用 IMU 参数修正之后对应到大二维平面上,每个点都是一个 RNN(???)。(RNN 的详细细节也没完全看懂,后面再细看…)
回顾一下整个网络的结构(假设是所有都放在一个 batch 里面的实现,大胆推测一下每个部分的输入输出):
规控这部分咖喱味口音略重,有点不太容易听清楚。
规划上主要的矛盾点在于:
那折中可行的方案就是把这两种结合起来。首先用粗粒度的搜索从整个决策空间中搜出最优解所在的一个小范围空间(Convex Corridor),然后在这个局部空间里面用优化方法求最优解。
Vector Space -> Coarse Search -> Convex Corridor -> Continuous Optimization -> Smooth Trajectory
这里举了个例子,导航给的路径是过了这个路口之后在下个路口左转:
方案一:过路口后减速,提早左转,可能会需要插到别的车队中去,too uncomfortable(需要过路口后迅速刹车,break pretty harshly)
方案二:过路口后加速冲过去,在左侧的车之前左转,这样会有充足的距离和时间,但是也有错过左转窗口的风险
对这两种方案之间做搜索优化后就可以找到一条特别平顺的路径来完成这个左变道的动作。目前他们的算法可以提前规划 10 秒的决策,实际执行的效果基本上也是非常贴合规划路径的。
下一个例子用来说明的是规划算法除了自身规划以外,还需要给环境中的其他对象的行为也做一个预测。窄道交汇车:
发现前方车辆时,预测出对方车辆的行进路线有大概率会卡死整条道,因此在这里的规划上选择我方向右侧避让(这里是右侧刚好有空间,有点好奇如果没有避让空间会怎么样)。等到后面发现对向车辆停住时,修正对对方的预测,此时中间的路可以通行,因此看到视频里面我方停顿了一下之后马上流畅地继续开过去了。
这个例子感觉做的非常漂亮,已经很像人类老司机会做的判断了。
下面这个才终于是前面提到的粗粒度搜索 + 细粒度优化的例子。视频里面首先搜出一个灰色的凸空间作为接下来路径优化的硬约束,之后再在这个小空间里面去做连续平滑路径的规划。
再往下是搜索算法的选择,给的是一个停车场泊车的例子,蓝框是目前的位置,绿框是最终目标要停到的位置:
启发式搜索(A*)是比较常规的选择,首先对比了两种 A* 的方案:
他们发现 Heuristic 信息在搜索过程中非常重要,需要一种比导航路径更加有效的全局信息来指导搜索过程,于是他们想到了神经网络。
如果把整个规划抽象成全局信息、决策、更新全局信息后再继续决策的过程,本质上就跟 alpha go 要解的是一样的问题了,因此也可以用类似的成熟方案去实现,直接上 NN + 蒙特卡洛搜索。(666666…)其他比如人工干预信息、距离、时间以及某些决策是否会让乘客舒服之类的都可以作为额外的 cost function 加到蒙卡树里面去。
甚至都还没把地图信息加上,就只需要 288 次搜索就能够找到非常好的结果了。(换句话说如果再在这个基础上加上高精地图是不是就更有效了?)
最终的系统架构如下:
本身感知模块产出的环境信息(基于摄像头得到的鸟瞰图)以及识别出的特定目标就是给规控做路径规划和决策的,现在就是相当于在下面多加了一个 NN Planner 部分,输入是感知 NN 里面某个中间层的 feature 用作 MCTS 的全局信息。
接下来这一段是我关注比较少的部分,主要是介绍 TESLA 如何标注数据以及通过模拟器来生成大量的训练数据。其中涉及到的自动化标注以及自动化场景重建等也都是相当硬核的工作。
这里就先略过了…
接下来跟硬件部署相关的终于到了我自己的工作领域了。
模型和算法部署上车的目标是最小化延迟以及最大化帧率,车上部署的是一块叫 FSD Computer 的板子,集成了两块自研的 SOC。
每块 SOC 看起来应该就是他们在 19 年发布的 FSD Chip,片上 3 个 ARM v8 A72 cluster 一共 12 个 CPU 核,一个 Mali G71 GPU 以及两个 NPU 单元。
这里面 Mali GPU 的算力比较低(FP32/FP64 600 GFLOPS),基本上只是用来跑一些基础的预处理用,主要的网络计算还是要放到单块 (INT8 36.86 TOPS) 的 NPU 上跑。
整体部署链路看起来比较常规,也是采用了 AI Compiler 的方式去做图优化以及 kernel fusion,再 schedule 到 NPU 上。关于这里的 schedule,视频里面介绍得偏少,看左下角的动图很像是编译期就决定好的静态调度?
比较可惜的是我想了解的更多的还是他们车上的部署部分,但是 AI day 上很快就跳到下一个 HPC 训练部分的环节了。
TESLA 之前使用的也是比较传统的 NVIDIA GPU 集群的方案,他们目前已有的规模是:
然后…That’s not enough!
常规 HPC 集群能用的最灵活的分布式架构就是下面这种 2d mesh 的方式了:
每个计算单元可以是一台多卡 GPU/NPU 的服务器,然后通过 Infiniband 的高速网络互联。但是问题在于这样 scale up 计算的规模是很容易的,但往往受限于节点间的数据传输带宽,全局同步等等,就比较难 scale up 全局的 throughput 以及 latency 了。
Dojo 的基本架构还是这种 2d mesh,但是达到了 extremely high bandwidth & low laytencies,然后通过软件层面的 AI Compiler 做到各种灵活的模型配置以及数据的 locality。
对于 mesh 中每个计算单元尺度的选择上:如果太小则每个单元的计算能力会很小,大量计算单元之间的 sync 代价会很高;如果太大(例如比较常规的 8 GPU 卡甚至更大尺度的服务器),计算单元之间的传输带宽又会受到限制(比如在这种尺度上就只能通过网络来通信了)。
他们定位出的核心矛盾是 bandwidth 和 latency,因此计算单元的设计也是极致的 laytency 导向的:
Dojo 最终架构中的 training node 单元是一个单核 4 路超线程的超标量 CPU,CPU 和 CPU 之间通过 1 cycle latency 的 mesh 互联。
单 CPU 的算力数据是峰值 1024 GFlops 的 BF16/CFP8 (2Ghz * 4 个 Matmul 单元 * 8 * 8 * 2FMA)和 64 GFlops 的 FP32。
最终集成的 D1 chip:
362 TFlops(前一页 PPT 上说一块芯片上是 354 个 nodes,如果照单 CPU 的算力算应该是需要 362 个 nodes,不知道是不是哪里没对上…),巨大的访存带宽。但从一块 D1 的算力性能来说基本上是 GPU/NPU 的级别了,但是他们这里还是强调这玩意是 CPU 的架构,应该可以算是个众核 CPU 芯片了。
恐怖点在于在这么高的片上、片间互联带宽下,这玩意的扩展性可以说是吊打 NV GPU 啊。
D1 system 图中的每一个小黄块都是一块 D1,在这种架构下,可以非常轻松地扩展到 50W 个 CPU(50W TFlops?)
再以上面这个图中的 D1 system 作为一个基础 unit,进一步可以构建出下面这样的 trining tile:
可以达到恐怖的 9 PFlops BF16/CFP8 以及 36 TB 访存带宽。
由 120 个 Training Tile(3000 块 D1 芯片,超过 1 M 个 CPU nodes)组成的 ExaPOD 可以达到 1.1 EFlops 的算力。
可以 … 这么简单粗暴就达到超算 top 500 第一名的水平了。(更正一下,想起来超算 top 500 的算力应该是 linpack 的 fp32 实测结果,这里 Dojo 给的 1.1 EFlops 是 BF16/CFP8 的理论峰值性能,那还是要差不少的)
NIO 从 2021 年开始针对 NT2.0 这一代的自动驾驶平台全栈自研。截止本次更新(2022 年),目前我们自研的感知算法 + orin 硬件平台已经在 ET7、ET5、ES7 上交付。
目前主要体现的地方就是方向盘后仪表盘上实时显示的感知结果:
包括车道线、车辆、行人等等。这些是综合了车上激光雷达、各个视觉摄像头等多种传感器的融合感知结果,其实显示在仪表盘上的只是目前开放的一小部分,完整的结果会被用于紧急刹车等安全功能以及未来开放的 NAD 这些辅助/自动驾驶功能中。
据我所知目前国产自动驾驶参赛选手好像只有理想官宣已经用上了 BEV,目前 NIO 这边的 BEV 进度应该是在 2022 年内可以完成,等达到足够的路测里程后就可以 OTA 更新给目前 NT2.0 的所有用户了,可以预见之后感知的效果也会比目前的软件版本更上一层楼。
]]>最近在比较多的帮团队面试,因为我们这边主要是 focus 在 AI 编译栈以及更偏底层一些的算子实现方面的工作,在面一些编译方向的候选人时,我这边最后一关通常是让他们写一个高效的矩阵乘的实现。
然后结果一直都没有遇到能写出来让我觉得还可以的人…唉,是我要求的太高了吗?
根据候选人的背景和平时的擅长领域,对这个问题我会做一些不同的调整:
并不要求一定要把 code 写的很好,写不出来也没关系,毕竟还真不一定是每个搞这块的人都有自己手撸这些的经验,但是我想看到候选人对这个问题里面一些细节的思考。
面了挺多在这方面有好多年工作经验的候选人了,得到的答案都让我不太满意,里面也不乏简历上写了有比较多的算子调优、优化经验的人,我甚至都开始怀疑是不是我自己的认识出了问题…所以打算自己好好写一写这个问题,这篇先从简单的 SIMD 开始吧,也先不考虑多线程这些问题。
测试设备是我自己的 MBP,一块 Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
的 CPU,虽然架构号还挺高的,不过挺可惜没有 AVX512。
首先的问题是一个达到性能极限的程序大概能跑多快?这里引用一个分析:
里面提到的 cpufp 这个小工具写的挺好的,我自己也经常在机器上跑这个测。不过我现在用的这个 commit 在 mac 上有一点点小问题,稍微 fix 一下以后在我的 mac 上可以跑出来:
1 | Thread(s): 1 |
这里的 fma 是 avx2 下的乘累加指令,已经是我这块 CPU 上能支持的最高效的计算指令了。暂时先不考虑睿频这些的影响(2.60GHz 加上 AVX2 FMA 的指令特性是能算出一个理论上限的),我们可以把这里测出来的作为能够跑出来的峰值性能的上限。
先找个计算库看看性能情况,OneDNN 就不错,拉下来 build 一下,test 目录中有个 benchdnn 可以直接跑,这里就把目标设的小一点,跑一个 [128, 128] x [128, 128] 的矩阵乘吧:
1 | $ OMP_NUM_THREADS=1 ./benchdnn --matmul --verbose=99 --mode=p 128x128:128x128 |
平均的 Gflops 是 106.513,差不多 78% 左右的峰值性能,毕竟是开源的 OneDNN 版本,可能再多做一些参数调整或者换上 blas 以后性能能再好一些,基本上可以认为跑到 80% 左右的性能可以算很不错了。
然后,嗯…做完 Ansor 之后一直没有好好写点跟它相关的东西(主要是因为我懒),其实还挺惭愧的。
不过也是最近在知乎上看到个帖子,有同学质疑 Anosr 测下来只能跑到一半不到的性能(虽然里面有一些其他方面的原因,这个结论还是让人挺沮丧),就顺便也贴一个随手写的测试结果:
1 | import numpy as np |
大概跑个 500 ~ 1000 步回来看结果,这个算子不算大,搜起来非常快,多跑个几分钟就跑完了:
1 | $ TVM_NUM_THREADS=1 python gemm.py |
差不多 82% 左右的峰值性能,可能这里面会有一些测量误差,已经差不多算比较满意的成绩了。谦虚一些可以看成跟前面 OneDNN 的结果基本相当吧,在实际的模型里面应用的时候结合点别的优化手段还可以再提升一点点。
用 Ansor/TVM 的好处是所有的一切动作对我们来说都是可控的,想知道它生成出来的 schedule 长啥样,加点代码打出来就好了:
1 | inp, _ = load_best_record(log_file, task.workload_key) |
首先是这个 schedule 的 state:
1 | Placeholder: placeholder, placeholder |
这里面用上的优化策略其实很简单,就是基本的一些常规操作 Tiling、Unroll、Vectorize。
虽然我前面把核数限制到了单核上,但是结果还是带上了 Parallel 的 attribute,这个在我看来好像是个代码里面的 bug …
试了下这个 kernel 在 TVM_NUM_THREADS=4 下面居然还能够跑到 400 左右的 Gflops,多核扩展性看起来也还可以哦。
对应的 TVM ir:
1 | primfn(placeholder_2: handle, placeholder_3: handle, c_1: handle) -> () |
这个 schedule 里面最核心的计算块其实是在做一个 [4, 8] x [8, 16] = [4, 16] 子矩阵运算,TVM ir 层面对最内层的 j_c.3 做了向量化,并且对 i_c.3 做了循环展开。
1 | for k.1 (0,8) |
考虑 A、B、C 三个矩阵各自的访存顺序,C 和 B 都是在 j 方向上连续访问,因为我们 Ansor 默认的策略也是直接在这个维度上做向量化。AVX2 的指令长度是 256 位,对应到 float32 上是 8 个,所以虽然这里搜出来的 vectorize 的长度是 16,其实翻译成指令会是两条 fma。
在这个子矩阵的计算里面,每次是从 A 矩阵读出一个值,broadcast 成 8 份,然后跟 B 矩阵里面连续读出来的 8 个数做对应位置相乘,再累加到输出上。
存放 C 矩阵的输出需要用上 4 * 2 * 8 / 8 一共 8 个 256 位的向量寄存器,B 矩阵的数据可以重复复用,只需要两个寄存器,这里把 A 的 4 个值分别 broadcast 需要 4 个寄存器,一共 14 个寄存器可以解决:
1 | vector_b_0 = vector_load(b[0-7]) |
嗯 … 看起来 A 这里的寄存器其实还可以再复用到同一个上,可以进一步把寄存器数量压到 11 个,不过这么做也就额外引入了一个数据依赖?况且 16 个寄存器还没有全部用满,这段代码的效率还有再进一步提高的空间。
看下最后生成的汇编指令是什么样子的吧:
1 | print(func.get_source("s")) |
从完整的汇编代码里面截取了其中的一小段:
1 | vbroadcastss -1052(%rcx,%r12), %ymm9 |
上面这个是 8 段类似代码里面的两段,可以发现 LLVM 在往下 lower 成汇编指令的时候进一步对计算逻辑里面的 k.1 也做了循环展开。
以第一段举例,如前面所分析的,输出用了 ymm0 ~ ymm7 这 8 个向量寄存器。每一段里面把 A 的数据反复 broadcast 到 ymm9 这一个向量寄存器上,ymm10 和 ymm11 用来存 B 矩阵的数据(好吧…还是只用了 11 个)。
如果把 M、N、K 都改成 144 再试一下,发现很容易就能搜出一个性能能跑的更高的 kernel:
1 | Execution time of this operator: 0.047 ms |
从汇编里面可以找到这段 schedule 能够跑出 93% 峰值性能的 schedule 的秘密:
1 | vbroadcastss -2324(%rdx,%rbx), %ymm14 |
除去最后一个往 ymm13 里面做 broadcast 的指令,这段代码生成得是真的是工整(当然其实太工整也不一定是好事,有时候为了充分利用流水线还得做指令重排把它们弄的“乱”一些,这是另外一个问题了)。跟之前类似,ymm12 和 ymm13 存放 B 矩阵的 16 个连续元素,每次取一个数 broadcast 到 ymm14 中,所有结果累加到 ymm0 ~ ymm11 这 12 个向量寄存器中,整个过程中一共用满了 15 个向量寄存器。
那么问题来了,寄存器用的更多性能就一定好吗?
当然不是……严格来说这得去分析指令流水线,能够打满流水,把所有的计算部件都用上才能发挥出最大的计算性能。
不过如果把问题简化一下,从过往经验上来看,如果我们假定计算的 pattern 已经排布的非常高效了,通常计算优化到最后都很容易最终会 bound 到 memory 上(这一点在 CPU 和 GPU 上都是成立的,为什么 A100 的 TensorCore 要出个 2 比 4 的稀疏?也是因为算力已经压榨到极限了,最后跟不上的反而是访存,这样其实相当于压缩了一倍的访存。Emm…这又是另外一个问题了)。至少如果已知还有资源没用上,当然还是想办法把所有能用的东西都用起来。
这里还有个用满 16 个寄存器跑到 99% 的峰值性能的例子,基本上可以说是把这个游戏玩到底了(虽然不用全部用上其实也有办法达到接近性能峰值,不过这个也是个很好的例子):
Ansor 在这一点上的局限性其实在于 tiling 的每个 split factor 的选择都是跟算子的原始尺寸相关的,为了避免引入一些对 index 的 if else 判断,默认的 split factor 采用的都是每个 axis 的约数。因此在 M == N == K == 128
的 case 中 Ansor 永远不可能搜出来 144 中最后那种 micro kernel 的尺寸(6 * 16 * 6)。这个局限性在最严重时候的体现是我们在业务中曾经遇到过两个尺寸非常接近只在某一个维度上有略微差别的 op,axis 是 59 的矩阵乘会与 60 的有很大的性能差距,原因就是 59 这个质数长度的 axis 在我们的策略中没办法 split,因此始终搜不出一个效率比较高的 micro kernel。
曾经想过的几种策略:
希望后面有机会的时候能把这些想法都尝试一下吧。
关于 2 这一条,之前看过一个也很有意思的工作做的更激进:为了提高计算密度,对某些利用率不高的计算过程做有损变换,最后再加一个修正的 stage 把结果调回来。类似把 dilation 卷积变成常规卷积做,最后把结果重新修正回来这种方式,即使浪费了计算量也引入了额外的修正操作,可能在某些情况下还是能有性能收益的。
可惜忘了论文题目了…后面找到再回来补上吧。
回到一开始的 SIMD 本身上,向量化矩阵乘还有别的实现方式吗?
当然。
对于一个 A、B 矩阵均为 NN layout(非转置)的矩阵乘运算,上面的做法其实是把这个基础的三重循环:
1 | for (i = 0 to n) |
变成了:
1 | for (i = 0 to n) |
的过程。
如果对 A、B、C 矩阵加上转置,则可以让这个计算过程根据需要做到各种方向上的数据连续存储,也可以做到更多种实现方式。
可以对 k 方向上做向量化吗?
当然。
很多候选人一看到这个 i,j,k 三重循环的表达式,上来就直接给我把 k 拆了,然后看了一会感觉好像哪里不太对,就进行不下去了。
k 在这里是个 reduce 的方向,常见的向量化部件设计上确实比较少有能够直接处理这种操作的,也有像 arm v8 的 neon 上就有提供了一条叫 sdot 的指令,可以对 int8 的数据做向量点积,最后累加到一个 int32 的寄存器上。这个过程的示例是从一个转置的 B 矩阵开始,做完转置之后 A、B 就都在 k 方向上连续了:
1 | for (i = 0 to n) |
然后再往下:
1 | for (i = 0 to n) |
额,上面这种方式其实我没有实际写过…
如果非要在不支持 dot 指令的硬件上实现 k 方向上的向量化,只是相当于手动实现一下这个 vector_dot 的过程:
1 | for (i = 0 to n) |
咦,好像看起来也是有可能达到比较高的效率的。注意,虽然最后这里对 v_c 的 v 个元素做 reduce 求和的操作我不确定是不是各硬件平台上就有现成的指令可以直接做,不过可以看到中间最核心的代码块一直是在做高效的向量访存和向量乘加,如果 k 方向上的长度是一个比较大的值(即 n / v
是个比较大的值),则即使最后只是挨个把这 v 个元素做标量加法应该也是能有一定的性能收益的。
似乎 MKL 里面的实现就是比较多的采用了把 NN 矩阵转置成 NT 以后再算的,不知道是不是主要用的这种实现方式。
]]>另:引申一下 tvm 中提供了一条叫 rfactor 的 schedule primitive,是针对 parallel 做的,不过实现方式上跟这个思路也有一些相似点。Reduce axis 在 tvm 中本身并不支持直接做 parallel 或者 vectorize,就把这个过程拆成两步做,可以在其中的一步做上 parallel / vectorize,只要 k 足够大就是能够有性能收益的。
硬/软件接口那篇有介绍过 GPU 的结构,当时也是以 Fermi 架构为例的,NV 很有意思的是会用一些历史上杰出的科学家的名字来命名自己的硬件架构。
总体上,NV GPU 用到的 SIMT 基本编程模型都是一致的,每一代相对前代基本都会在 SM 数量、SM 内部各个处理单元的流水线结构等等方面有一些升级和改动。这篇暂时不涉及到渲染管线相关的部分,其他诸如多少 nm 工艺、内存频率提升等等也都先略过,只关注计算相关的硬件架构演进。
关于初代 GPU 的架构,找到的资料不太多,基本上都是从 Fermi 开始的。
Compute Capability: 2.0, 2.1
每个 SM 中包含:
我的理解是做一个双精 FMA 需要用到两个 CUDA Core?所以是 32 / 2 = 16
Compute Capability: 3.0, 3.2, 3.5, 3.7
这一代 SM 整体结构上跟之前是一致的,只不过升级完了以后又往里面塞进去了更多的运算单元,其他部分也没有做太大的改动。
每个 SM(这里叫 SMX 了)中包含:
Kepler 是附近几代在硬件上直接有双精运算单元的架构,不用通过单精单元去做双精运算了,所以对比前后几代的双精浮点的性能话会发现 Kepler 要高出一截。
Compute Capability: 5.0, 5.2, 5.3
可能是觉得 Kepler 往一个 SM 里面塞了太多东西,其实最终效率也并没有那么高,这一代的 SM 开始做减法了,每个 SM(SMM)中包含:
Kepler 里面 192 这个数字也被诟病了(不是 2 的倍数)。
这些硬件单元的流水线分布也不再是像 Kepler 那样大锅炖了,而是有点像是把 4 个差不多像是 Fermi 的 SM 拼在一起组成一个 SM:
每个 Process Block 里面是:
图上没有看到之前 lane 的标记,不过我猜应该也还是 4 条,两条 CUDA Core 的 lane,1 条 SFU,1 条 LD/ST Unit。
应该是工艺和频率的提升,Maxwell 每个 CUDA Core 的性能相比 Kepler 提升了 1.4 倍,每瓦性能提升了 2 倍。对 CUDA Core 的详细结构没有再介绍,姑且认为从 Fermi 开始一直到以后 CUDA Core 内部的结构都没有什么改变。
另外一点是,前面说到的双精单元在这一代上也移除了。
也许是觉得认为只有少数 HPC 科学计算才用的上的双精单元在这代上不太有必要吧。
Compute Capability: 6.0, 6.1, 6.2
这一代可以说是有了质的飞跃,还是先从 SM 开始:
可以看到一个 SM 内的部分作了进一步的精简,整体思路是 SM 内部包含的东西越来越少,但是总体的片上 SM 数量每一代都在不断增加,每个 SM 中包含:
一个 SM 里面包含的 Process Block 数量减少到了 2 个,每个 Process Block 内部的结构倒是 Maxwell 差不多:
单个 Process Block 的流水线增加到 6 条 lane 了?
其他质变的升级包括:
然后 NV 现在已经把 Infiniband 行业的龙头 Mellanox 给收购了……说不定那时候就已经有这个想法了呢
CUDA Core 在这一代也终于有了升级,现在硬件上直接支持 FP16 的半精计算了,半精性能是单精的 2 倍,猜测应该是一个单精单元用来算两个半精的计算。
Compute Capability: 7.0, 7.2
又一个针对深度学习的质变 Feature,Tensor Core!
看到 SM 的时候我们会发现这一代除了多出了一个额外的 Tensor Core 的单元以外,怎么 SM 的体积看起来好像又加回去了,每个 SM 中包含:
事实上相比 Pascal 而言,单个 SM 中的单精运算单元数量是一致的,相当于把 Pascal 中的每个 Process Block 进一步地又拆成了 2 个,每个 Process Block 中包含:
这里把原本的 CUDA Core 给拆开了,FP32 和 INT32 的两组运算单元现在是独立出现在流水线 lane 里面了,这一设计的好处是在前几代架构中 CUDA Core 同时只能处理一种类型的运算,而现在每个 cycle 都可以同时有 FP32 和 INT32 的指令在一起跑了。Pascal 中需要 6 个 cycles 来做一组 FMA,现在在 Volta 中只需要 4 个 cycles。
另外每个 Warp Scheduler 还有了自己的 L0 指令 cache。
这一代还改进了一下 MPS,现在从硬件上直接支持对资源的隔离,方便多任务共享 GPU。
其他一些比较重要的改进:
最重大的改动不用说也知道是 Tensor Core 了。
Tensor Core 的思路从系统设计上还是相当直接的,目前深度学习的 workload 中最主要的计算量都在矩阵的乘加上,因此为了专门去高效地支持这些 workload,就增加一些专用于矩阵运算的专用部件进去。
这个也是常见的 AI ASIC(比如 Google 的 TPU、其他厂商的各种 xPU 等等)通常采用的思路,只不过 ASIC 可以从一开始就是针对特定的 workload 去的,因此设计上可以更直接更激进一些,直接上大量的 MMU(Matrix Multiply Unit),然后采用例如脉冲阵列这种设计去最大化它的 throughput。
而 NV 的 GPU 毕竟还要用作其他一些通用的运算,所以只能往原本的 SM 流水线里面插进去一些额外的专用部件 lane 了。开个脑洞,要是哪一天发现除了 FMA 以外还有其他另外一种形式的运算有大量的需求,未来的 GPU 设计里面说不定也会出现其他 xx Core。好在 FMA 除了深度学习以外在 HPC 的 workload 里面也是挺常见的,这个设计以后还是比较有用的。
Tensor Core 这个部件直接从 SM 的寄存器里面取两个 FP16 的矩阵作为输入,进行全精度的矩阵乘之后得到的结果可以是 FP16 或者 FP32 的,然后累加到 FP16/FP32 的 accumulator 里面去。数据类型选择 FP16 作为输入然后输出 FP32 猜测可能是为了保证结果不溢出,然后在加速部件设计等等方面做了一些 trade off。
所以 FP16 in -> FP16 out 和 FP16 in -> FP32 out 哪一个性能更好呢…
我没有测过,但是猜测可能默认结果是 FP32 out 更快?反而是输出 FP16 需要从 FP32 再转一次?
接下来道理我们都懂了,那 Tensor Core 要怎么用呢?这个部件的编程模型在一开始接触的时候可能会有一些坑。
我们知道常规的 CUDA 代码需要制定 grid 的结构、block 的结构,然后其实我们写的 kernel 代码都是针对每一个单独的 thread 的,可以认为是 thread level 的编程。对一个子矩阵的 FMA 运算存在比较多的数据重用机会,这时候如果只是一个 thread 算一个矩阵块的 FMA 就比较浪费了,因此 Tensor Core 的设计是用一整个 warp 去共同完成一个 FMA 运算,一个 warp 中的 32 个 thread 可以复用寄存器里面的数据。CUDA 对 Tensor Core 的指南里面把这个叫做 “WMMA warp-wide macro-instructions”。所以 Tensor Core 的编程模型直接就是针对一整个 warp 写的。
事实上,Tensor Core 的代码写起来还是有相当多的限制的,CUDA 给 Tensor Core 提供了以下这些 c 的 API:
1 | template<typename Use, int m, int n, int k, typename T, typename Layout=void> class fragment; |
PTX 的指令应该更多一些,不过我没有详细看过。
首先用来做乘加的矩阵都需要放在这个叫 wmma::fragment
的变量里面,这个本质上就是定义了一个要放在 SM 寄存器上的存储空间,但是需要提供详细的 FMA 参数:
Use
是这个 fragment
在 FMA 运算里面的角色,可选项有:matrix_a
、matrix_b
和 accumulator
,含义就是字面意思,也没什么需要再解释的了。__half
,accumulator 是 float
,然后 m、n、k 都是 16。Layout
可选项有两个 row_major
和 col_major
,代表这个 fragment
在内存里面实际存储的行列主序情况。load_matrix_sync
和 store_matrix_sync
分别是把数据写到 fragment
空间里面和从这里面取出来写到别的地方去。fill_fragment
对 fragment
初始化。mma_sync
就是对整个 warp 调用 Tensor Core 去跑完这一个 FMA 运算了。
常规的写法也是先把矩阵 A、B 都 load 到 shared_memory 上,然后再从 shared_memory 里面取对应 FMA 块大小的数据到 fragment
里面,mma_sync
跑完,最后从 fragment
里面把结果写到外面去。
这里的注意点是上面这些代码(包括 fragment
定义以及下面几个函数的调用)都是针对 warp 的,即我们在写代码的一开始就需要考虑到每个 block 里面的 thread 结构,保证一个 warp 的 32 个 thread 执行的代码是完全相同的。相应地,对矩阵的分块也是需要在写代码的时候就考虑清楚,我们要保证每个 warp 处理的 a、b 矩阵的大小刚好是这个地方设定好的 m、n、k。
看起来确实相当麻烦,不过想想可能好像也还好,本来如果要写出性能很好的 CUDA 代码来,每个 warp 要算多少东西也是需要精细考虑清楚的。
Volta 这一代对 SIMT 的编程模型也做了改变。
在之前的 SIMT 流水线中,如果一个 warp 的指令里面出现了分支,这些分支块是不能被同时执行的。所以一直以来写 CUDA 代码都会要有一个原则是不要在一个 warp 里面出现不同的分支,要不需要花费两倍的时间去处理。
这一代开始把 PC 和调用栈做成了每个线程独立的:
现在呢,每个分支里面的指令可以在更细粒度的层面上进行混合调度了,也可以手动插入一些在 warp 层面同步的指令进去:
白皮书后面给了一个可以从这个改动上得到收益的 Starvation-Free Algorithms 的示例,修改带锁的双向链表的时候,不同 thread 可能会被 block 在锁上,以前的架构应该基本上不太可能能处理得了这种 case,新架构就保证了即使有些 thread 还在等待锁,另外的 thread 也有可能先拉出来跑。
可能也是因为这样所以 1 个 Dispatch Unit 配 1 个 Warp Scheduler 了?因为线程指令的实现事实上更加复杂了。
所以其实最后还是同时只能执行一个分支里面的一部分,这个 upgrade 我暂时还没有想到具体的应用场景会有多常出现(上面这个带锁双向链表我觉得写在 CUDA 里面就很不常见啊…),以及会具体有多少性能收益,说不定还是原本的那种简单的设计更直接更高效一些呢。(期待一下未来的硬件里面会不会把这个恢复回去……)
以前 CUDA 编程原则里面不要写分支的那条在新架构下我觉得还是适用的,不写分支就不会有这么多额外的麻烦要考虑了。
另外有一个 Cooperative Group 的新设计倒是看起来感觉更有用一些。原本的 __syncthreads( )
是针对一个 block 里面的所有 thread 做同步的,现在可以对不同 block 的不同 thread 单独定义同步组了,CUDA launch 的时候会把同一个组的一起 launch 上去,同步可以在一个更加细粒度的层面上完成。
Compute Capability: 7.5
这一代的改进看起来只是对 Volta 的小加强,Compute Capability 大版本号都仍然是 7。Tesla 计算卡系列只更新了 T4 这种推理卡,增加了 Int8/Int4 的推理能力,其他主要的设计更新重心都放在了像实时光追这种渲染能力上,让 Gforce 系列游戏卡的能力有了质的飞跃。
SM:
每个 SM 中一共有:
每个 SM 中有 4 个 Process Block,对于每一个 Process Block 来说有:
相比原本的 Volta SM,可以看到 FP64 又被去掉了,LD/ST 单元砍半。(是不是因为 Turing 系列没有打算出 T100 这种大卡,而推理和渲染对 FP64 以及访存的需求没有那么大?)
Tensor Core 部分相比 Volta 的初代设计增加了 Int8 和 Int4 支持。且毕竟占了很大一块芯片面积,Tensor Core 现在除了 AI 场景以外在这一代开始也能够在游戏和渲染中用上了,Turing 支持 Deep Learning Super Sampling(DLSS)深度学习超采样技术,看名字应该是通过神经网络进行视频插帧。
其他的提升主要集中在渲染管线和新加的这个 RT Core 光线追踪能力上,这里就跳过了。
Compute Capability: 8.0, 8.6, 8.7(Orin)
算是又一个大版本更新:
还是先来看下 SM 的结构:
每个 Process Block:
从单元数量上来看,Ampere 跟 Volta 是基本一致的,也有 FP64 的核心。把 Volta、Turing、Ampere 三代放在一起比就有种 Tik-Tok 的感觉出来了 … 每一代大卡通过工艺升级压进去更多的计算单元,FP64 和 LD/ST 单元的数量都是拉满的,然后会有一个小改款出不带 FP64 的推理卡以及专注提升图像性能。
这一代比较大的改进在 Tensor Core 以及各个部分的 memory 工作效率上。
唯一的区别是 Tensor Core 的数量相比上代减少了一半,但可不要觉得 Ampere 的性能因此就下降了。
Volta/Turing 中的每个 Tensor Core 每个 cycle 可以执行 64 个 FP16/FP32 的混合精度 FMA 计算,而 Ampere 中的 Tensor Core 虽然数量减少了,但是每个的内部都堆了更多的料,单 cycle 的吞吐量提升了 4 倍到 256 个 FMA。一来一回,单个 SM 的 Tensor Core 计算能力还是提升了两倍。
Tensor Core 可以支持的数据类型扩展到了 FP16、BF16、TF32、FP64、INT8、INT4 甚至是 Binary。(TF32 的数据类型是为了能够更好地支持 FP32 的 Tensor Core 计算)
另外一个最大的改进在于 NV 给 Tensor Core 增加了一种特殊的结构化稀疏运算能力:
网络中训练出来的 Weight 通过预处理以 2:4 的压缩比处理成一个稀疏矩阵的结构,每 4 个数据块精简掉 2 个块。可以看到这个结构跟通常的 BSR、CSR 等稀疏存储结构是完全不同的,压缩好的系数矩阵的数据量变成原来的一半,然后再加一个标记了 2:4 稀疏位置的 indices 矩阵,估计是需要用 CUDA 自己的稀疏处理 API 来做到。
上图中也比较清楚地解释了 Tensor Core 的稀疏计算原理,indices 矩阵通过两个选择器直接把输入数据的对应位过滤出来,然后做点乘。
很有意思的是白皮书在这里给了一张细化到每个点上的性能对比图:
一连串的翻倍提升看起来确实挺吓人的。首先 Tensor Core 在本身吞吐量以及 A100 的总 SM 数量上综合有 2.5 倍的硬计算能力提升,如果用上稀疏计算,则吞吐量可以再涨一倍到 5 倍。各级 Memory 之间的访存带宽以及容量也有了大幅提升。
接下来,之前 Volta 上 Tensor Core 只能在一个 warp 的 8 个线程间共享数据,现在这个共享范围扩展到了整个 warp 的 32 个线程,计算时对寄存器的访存可以减少 2.9 倍。Ampere 上还增加了一个新的异步访存指令,之前从 Global Memory 拷贝数据到 Shared Memory 其实是要先经过寄存器中转的,现在可以直接 bypass 寄存器了进行异步数据拷贝了。
剩下的几个看上去倒是并没有太 promising 的样子:
TVM 是在 Halide 的基础上发展而来的一套编译架构。早一点版本的 TVM 其实还能看到 Halide 是作为一个 git 的 submodule 放在 TVM 目录里的,当时 TVM 自身的代码中还有一个 HalideIR
的 namespace,也有不少的结构是直接继承了 Halide IR 里的内容:
1 | ... |
[INFA][IR] Build and Evolve Low-level IR. Remove HalideIR dep. #3533 这个 PR 之后,Halide 相关的部分逐渐从 TVM 中删掉了,之后可以说从代码层面已经跟 Halide 没有关系了。
目前 TVM IR 的结构基本上还是从 Halide IR 一脉相承,可能以后 TVM IR 进一步演化以后跟 Halide 的差别就更大了。
expr.h
里面定义了 TVM IR 的两个基础结构:Expr 和 Stmt,分别是语法表达式和语法树节点的基类。
Expr 的派生类有加减乘除、IntImm、FloatImm 等等,从文法上可以做一些 symbolic 的处理。
Stmt 的派生类有 AttrStmt(语法树属性节点)、Store(数据存储节点)、Allocate(数据 Buffer 分配节点)等等。
每个 Stmt 结构本身表示一个独立的语法树节点,但是语法树节点之间相互嵌套,通过 Stmt 的 body(Stmt 的通常结构)等成员继续向下查看就能够看到一颗完整的抽象语法树(AST)了。
例如 IfThenElse 这个 Stmt 的结构:
1 | class IfThenElse : public StmtNode { |
判断条件是个 Expr,then_case 和 else_case 都是另外两个 Stmt,再展开又是两棵子树。
TVM 中定义了一些 ir_pass 来处理语法树,通过对语法树的修改和调整来完成编译优化的过程。
各种 ir_pass 的核心结构是 IRVisitor 和 IRMutator,从名称上也可以很容易看出来,IRVisitor 的功能是遍历语法树收集信息,本身对语法树的访问是只读的,然后通过 IRMutator 完成 ir_pass 需要的语法树修改需求。
看下 IRVisitor 的结构:
1 | class TVM_DLL IRVisitor { |
这里面最重要的实现其实是自己构造了一个虚函数表:
1 | IRVisitor::FVisit& IRVisitor::vtable() { // NOLINT(*) |
之后 IRVisitor 的派生类只需要去重载针对不同 Stmt 类型的 Visit_()
函数就好了。
看 TVM 代码包括 debug 的时候一开始会觉得两眼一抹黑连语法树长啥样都没有个概念,也不知道 build 过程中每个 ir_pass 具体做了什么事情,就很想要有个工具能把语法树打成图看下。
最开始 TVM 里面是找不到这种现成的工具的,后来在论坛里有看到别人提的这方面相关 RFC,不过我后续也没有去关注最后到底有没有收进 repo 里了。其实了解完 IRVisitor 的实现之后,自己写一个语法树的 Dumper 还是挺简单的。
ir_visitor.cc
里面有一个 PostOrderVisit 的示例:
1 | class IRApplyVisit : public IRVisitor { |
按照对语法树后续遍历的顺序对每个语法树节点应用 f_()
函数,不过要实现语法树输出的目标,光靠这个还不够,需要再稍微进行一点点的扩充。
我们通过一个栈来记录下语法树上每一个 stmt 的从属关系,首先扩展一下上面那个 Visitor 来做到在访问 stmt 节点的前后分别调一个外部函数:
1 | class IRPrePostOrderVisitor : public IRVisitor { |
PrePostOrderVisit()
相应的在 ir_visitor.h
里也要添加一下。
接下来再往 api_pass.cc
里添加一下_PrePostOrderVisit
的函数注册:
1 | TVM_REGISTER_API("_PrePostOrderVisit") |
这样在 C++ 部分的工作就完成了,之后是 Python 这一层,我们的 dump 目标是语法树,所以找个能直接拿到 stmt 结构的地方,build_module.py
就很不错。
lower()
是从 TVM IR 往能够运行的代码编译的第一步,涉及到多种 ir_pass 的使用,不同 ir_pass 的前后插入 dump ast 的代码可以帮助我们快速搞清楚每个 ir_pass 到底实际做了什么事情。
1 | def lower(...): |
接下来看一下我们需要在 dump_ast()
里面写上什么:
1 | def dump_ast(stmt): |
思路也是很简单的,我们只要在 PrePostOrderVisit 访问 stmt 节点前将节点入栈,然后访问节点结束后将节点退栈就完事了。
这个地方也体现出上一篇中 TVM 特别搞出来的这套跨 Python 和 C++ 混合运行的 PackedFunc 机制的便利性,这里事实上我们是从 Python 层开始,调了一个 C++ 的函数,然后在这个 C++ 的函数里面又回调了两个 Python 的函数,并且这个过程中数据还是存在我们在 Python 层创建的结构上的。
试一下下面这段示例代码:
1 | n = tvm.var("n") |
通过 dump_ast 打出来的 dot 代码:
1 | digraph { |
通过 GraphViz 这种 dot 可视化工具处理一下:
当然这个实现还是相当简单了,根据每个节点的类型等等还可以再加一些更复杂的判断逻辑,控制一下输出的内容量等等,以及其实可以看到里面还有很多重复的节点也都可以被筛掉。
Relay IR 的整体处理结构跟 TVM IR 一致,用类似的方法也可以把 Relay 那一层的 AST 打出来。
]]>之后的工作有很多是要跟神经网络的编译执行相关,准备继 TF 拆包之后再来记一下我对 TVM 的探索过程。这个系列的坑会开多久也不一定,毕竟之前 TF 的其实也只做了一点点微小的工作,还有很多方面的内容没有看,另外 TF 要更到 2.0 之后可能有不少地方已经有一些变动了。
(也许有空也会看一下 XLA 吧,不过那个应该是要合到 TF 的拆包系列里面去了)
一些基础的背景可以见前面一篇对 TVM 作者一门课的课程记录:【CSE 599W: Systems for ML】
第一篇先从代码中一些基础的结构开始,最早读代码的时候没看明白一些东西是怎么实现的,为了看懂细节吃了不少的苦头。
参考文档:【TVM Runtime System】
Node 和 NodeRef 这两个类在 TVM 中几乎是所有对象的基类了,Node 是功能本体,NodeRef 可以看成是对 Node 的一个指针引用。举例来说 TVM IR 语法树中的两个结构 Statement 和 Expression 的实际存储对象是 StmtNode
和 ExprNode
,从 Node 继承而来,但是在被其他结构用到的时候用的却是 Stmt
和 Expr
两个结构,从 NodeRef 继承而来。
Node 这个结构本身没什么好看的,看一下 NodeRef 的实现:
1 | /*! \brief Base class of all node reference object */ |
最关键的点在于 NodeRef 对它的 ->
运算符做了一下重载,即返回自己实际代表的实体对象,所以 NodeRef 以及其派生结构虽然实际是个对象,但是在代码里面可以当成指针来用,->
运算符之后直接跟的就是本体 Node 的成员。
1 | /*! \brief Container of all statements */ |
以上面的 Stmt
类为例,其他很多 NodeRef 的派生类结构也会用一个 TVM_DEFINE_NODE_REF
的宏来扩展出这部分的关键代码。
那么问题来了,为什么要设计成这个样子呢。
这里谈一下我自己的理解,如果有问题的话也欢迎看到的同学指正一下。
一方面可能是为了方便下面 C++ runtime 和 Python 部分所有结构的无缝衔接;另一方面从上面的代码也可以看到,NodeRef 重载了 ->
运算符之后返回的是个 const 的指针对象,所以通过 ->
访问到的实际的成员结构都是只读的了,这就保证了 TVM 这套复杂系统里面各种数据结构的安全性。
当然严格的只读访问在某些情况下是不够的,所以 TVM 提供了 CopyOnWrite()
的机制,如果某个 NodeRef 类的定义中包含了 TVM_DEFINE_NODE_REF_COW
这个宏的话,可以通过 NodeRef.CopyOnWrite()
获得一个可修改的 Node 指针,之后对成员内容的修改均通过这个指针来做就可以了。
1 | /*! |
话说 CopyOnWrite()
这个函数名称我感觉可能不是特别确切,也许改成 GetMutablePtr()
之类的会更好点?因为这个实际上并不 Copy,直接调用这个函数返回的是对这个 NodeRef 自己所指代对象的指针,之后的改动也都是对这个 Node 自身做的。
如果确切希望实现 Copy 的语义,则需要像前面注释里面示例的那样,先用另一个 ref2
复制一份 ref
,之后再在 ref2
上进行修改。
TVM 的整个软件栈涉及到很多高层脚本语言(Python、JavaScript)和 C++ 运行时的交互,因此这里提供了一套 PackedFunc 的基础用来把整个过程方便地串接起来。
第一次看到这种实现时真的是被惊到了,感觉非常神奇。
在 C++ 层面创建一个函数,可以直接进行本地调用:
1 |
|
也可以通过 API 注册之后(要注册到 TVM 的 C++ 运行时库里面去),从 Python 层进行调用:
1 | // register a global packed function in c++ |
1 | import tvm |
反过来 Python 层写好的函数也可以直接从 C++ 层调用:
1 | TVM_REGISTER_GLOBAL("callhello") |
1 | import tvm |
Python 层定义的 callback(msg)
通过 callhello
传递给 C++ 层,C++ 层执行时直接从输入参数中得到了 Python 的函数对象并调用执行。
实现方面,C++ 层的 PackedFunc 是一个对 std::function 对象的封装结构,而 Python 层面tvm.convert
实际上是把 Python 的函数用 ctype 做了一下封装:
1 | TVMPackedCFunc = ctypes.CFUNCTYPE( |
PackedFunc 对函数的输入参数、返回值的解析处理做了比较精巧的处理,最终达到了从 API 层面看上去非常好的使用体验。
]]>Amazing!
有空的话可以试一下把 TVM 里面的这部分内容单独扒出来,这个实现思路真的非常有意思。
生活总是不知道会在什么时候给你点“惊喜”。
过去的一年间发生了挺多事情的,对我自己来说有好事也有坏事,上上下下大起大落,最开始的计划被一次次地打乱,每一次都算是艰难的抉择吧。
看了下博客也有半年没更了,更让我觉得惊讶的是上一次挂了“随笔” tag 的居然已经是2年前了(可见这两年读研生涯中我是多么正经…),准备来理一理过去这大半年的流水账。
大概是去年6月多投出去的第一份简历,到10月中旬签下 offer 为止大概经历过 NVIDIA、大疆、华为、网易游戏、网易云音乐、腾讯、阿里的面试,然后聊过几家独角兽(后来有一天接到厦门航空招聘组的电话真的是吓到我了!招副机长?还专业不限?入职后再做飞行培训?……也许我应该应下来跟他们多聊聊的)。
NV 的笔试题很简单,然后6月份就迎来了我人生中的第一次电话面试。虽然当时完全没怎么准备,却也只能硬着头皮上了。1 小时的电面里面大概问了我快40道题,幸而对 GPU 和 C++ 这块还算熟悉,有惊无险。之后是 4 位面试官一共 4 小时的车轮面(一人一小时,连着一个一个上的那种)。
虽然面对 NV 时很匆忙,但还是很感谢这段面试给我的经验,经历过车轮面之后后面再来啥都不怕啦。NV的校招流程整体给我的体验真的非常好,谈薪资和 offer 之前都不管你未来是不是有可能拒,一律先免费给做一次入职体检!以至于后来抉择完选了别家的时候真的是觉得很不好意思,觉得有点愧对一直帮我处理了整个校招过程的HR小姐姐。我投的岗位是 Deep Learning Performance Architect,其实工作内容都是我很想做的方面,只可惜后来有了另外的考虑。
网易游戏是线下笔试,当时运气不好的是笔试的日子刚好遇上实验室有另外的事情,只好放弃。网易云去了杭州现场,还找在网易的上届师兄蹭了一顿他们的食堂(不得不说!猪厂的食堂是真的赞啊!什么都有,而且员工吃饭免费!要啥自行车!)。可惜投的岗不太对,面我的小哥说他们部门是做推荐算法的,要的是纯算法方向,看完我简历说来他们这里不太合适,可以帮我推荐到百度的部门。额,知道是这个情况,那后来就干脆开始跟小哥聊天了(原来小哥自己就是从熊厂跳过来的)。最后还是婉拒了他的推荐,真的都是很好的人啊。
华为没啥好说的,聊得还行,奇怪的是我投的明明是杭州的岗,最后说录的却是深圳。体验就没前两个好了。
阿里的 PAI 算是我之前很早就看上的部门了,一年前就在知乎上看到了他们发的研究分享以及招聘公告,然后在他们校招开放的第一天开始了内推的过程。一路下来,3面和交叉面都给了我很大的压力,完全是被按在地上摩擦的那种。虽然觉得自己当时没有表现好,但是经历过面试就能感觉到他们是真的对我感兴趣的方向非常了解,我也相信如果我去了这里也一定能够学到我真正想要的东西。他们家跟 NV 岗位的区别可能是工作重心会更偏系统一点,NV应该是跟偏向加速卡这块。纠结了很久。
P.S.可能最后杭州这个工作地点对我来说也比其他城市有更高的吸引力吧。
然后跳出来半路截胡的是腾讯和大疆,之前做梦都没想到过能够拿到这么高的 package。遥想当年本科毕业的时候,想着拿个十几二十W就顶天了(当时没定下来要不要保研的时候身为一个魅粉还差点就去了珠海…),面对好几倍的数字我是真的犹豫了。腾讯的方向是安全平台,要处理巨量的流量数据所以需要高性能高并发的支持,总的来说还算是系统方面的工作。大疆的岗位是个新部门,目标倒是跟我想做的比较契合,只是经历过阿里的面试之后,我觉得自己在这方面缺的东西不是一点半点,本来研究生阶段的工作很大程度上都已经是我们自己在一个新方向上摸着石头过河了,工作后再去到一个需要一边学一边慢慢建设的部门,我确实没有太大的信心了。
后来跟很多人商量了我自己的想法:在一开始工作的时候还是更想去一个成熟一点有足够积累的部门,我希望在离开学校的一两年间能尽可能地先提升自己。过来人都是劝我说没必要跟钱过不去(捂脸)……当然,薪资的吸引力还是很大啊,只是也许真的是我太年轻?
至少我到现在还是怀有这样的想法,也希望未来真的开始工作之后自己不会为当初做的这个选择而后悔吧。
本以为已经结束了所有校招日程,可以老老实实忙手头上最后的一些事情准备毕业了,结果10月底收到的一封邮件又改变了我原来的计划。
Apple?做 iPhone 和 Mac 的那个 Apple?又是一件做梦都没有想到过的事情发生了。
很难形容当时刚看到邮件那一刻是什么样的心情,二十多年来我跟苹果唯一的交集就是用过一个 iPod Nano,结果突然有一天天上掉下来一扇门,打开以后就能去到地球的另一边,去到一个拥有全球最高市值的公司。
还有人会犹豫吗?
于是得到导师的同意之后开始了又一轮的面试,之后是等HC、背调、准备签证材料等等等等一堆麻烦的事情。
等到真正弄完所有的手续,等签证材料寄到已经是1月底了,能够约到的最近一次面签时间是1号的北京使馆。于是31号跑去上海FedEx的仓库人肉取了签证要用的表,连夜赶到北京。
然后,果然还是被 check 了。我记得这时候贸易战还没开打,但是也许计算机和 AI 相关的内容早就已经成了敏感词了吧,签证官根本就没问几个问题,要了个人简历扫了一眼之后直接就给了条子。
可是我原本的打算是2月开始实习的,这样可以赶在学校毕业答辩之前回来。从大使馆出来之后也想过要不就这样算了吧。
现在想来如果当初没有继续坚持下去,也许就不会有再后来的这些麻烦事了吧。嗯…当然我人生中也会少掉一段珍贵的经历。
这都是后话了。
接下来这一等就是2个月。
真的非常感谢我在 Apple 的主管 Jay(要说这也是一种缘分呐,我喜欢的第一偶像周董叫 Jay,之前也经常用这个作为自己的英文名),最初是他给了我这次实习的机会,在我签证被 check 之后还愿意帮我推迟实习的日期,并且在以后的日子中给了我很多帮助。
4月6号,终于登上了飞往西雅图的航班。
更幸运的是,去 Cupertino 做入职培训时正好赶上了公司内部的一个学术交流会,让我有机会走进 Apple Park 以及 Steve Jobs Theater(应该是果粉圣地吧)!
最惊讶的发现大概是原来国际巨头里面也会有重复造轮子这样的事情了,以前一直以为可能只有国内的公司才会有不同部门“很巧地”刚好都做了差不多的东西的情况……
西雅图从纬度上来看应该算很北方了,却觉得一点也不像国内北方的气候。也许是因为靠海?这边下雨很多。总体上对一个来自江南地区的人来说,这里还算挺舒服的。
神奇的是虽然说下雨多,但是即使雨天也很少看到有带伞的人,有时候走在路上发现只有我一个人打伞的时候就很尴尬。所以后来我也不带了。
在 Apple 每周都是 955 的生活,做的是自己很感兴趣的方向,每周大大小小的会也能听到学到很多新鲜的东西。这样的日子真的很赞啊。
Office 里的同事们都很友善。偶尔下午开完会之后,大家也会约着下楼,走过几条街,然后买杯咖啡带回来继续干活。
但是某天突然意识到从过道前走过的人中可能 CMU 的一抓一大把,坐在隔壁的小哥是 StandFord 毕业的高材生,前几天跟我聊过天的另外一个实习生小姑娘是 MIT 来的。于是会觉得应该要更脚踏实地地对待每一项任务,我差的还很远啊,我要学的东西还很多啊!
一开始的日子现在想起来都觉得好笑。因为不熟悉美国超市的情况,想买双拖鞋都花了我一两个星期才找到……
然后一个人去了 Pike Market、Space Needle、附近的一些博物馆。某个周日去 Pioneer Square 参加他们的 Underground Tour 时导游都笑话我这么好的日子跑来听历史课真是可惜了。
收到的对我英语的最高评价是 Underground Tour 中一对老夫妇惊讶地问我:“Where did you pick up your English?” 额……惭愧惭愧,我觉得我听力是挺不错的,不过口语是真的不太行哈哈。
主管说问我愿不愿意留下来做 Full-time?
以前是真的从没想过还有出国工作这个选项。但是如果能拿着美元的工资,做的还是自己最感兴趣的方向。Why not?
遗憾的是我过来实习用的 J-1 签证会有一个回国工作的限制,而豁免手续似乎也不是我当时的情况下能顺利办下来的。
更麻烦的是学校没有远程答辩的选项。由于最早的签证被 check,实习开始的晚了,一系列连锁因素导致中间还卡着一个关口必须回国参加毕业答辩。而我拿到的 J-1 签证只给了单次入境的机会,这意味着回国后需要重新再办一次签证,还存在没有办法回来做完后半程实习的危险。
嗯……真的居然又被 check 了!🙄
大概是要凉了。
卡塞尔学院对你来说是一扇门,打开这扇门你就会进入新的世界,但那样你就再也回不去了……你每做出一个新的选择,其他选项就消失了。自始至终,你都只有一条路走。
——《龙族IV》
对于一个毕业生来说,临走前夕最惨的事莫过于告诉你还不能走了。
没错,匆忙从大洋彼岸飞回来,然后签证被卡,毕业也遇到了麻烦。还能更惨一点吗?
嗯,可以的,人一旦开始倒霉,好像所有不好的事情都会接踵而来。当然另外发生的一些事情也是后话了,不想说了。
得知答辩日期的时候已经只剩下不到一周时间了,然后赶紧买了回国的机票,预约新的面签,在路上修改大论文以及准备答辩的材料。
Anyway,中间的事情一言难尽,最后的结果就是按学校的流程,一旦错过了6月的毕业窗口就只能等11月才能拿到双证了。学院说一切按流程办事,而且当初请假出去实习时“发生任何事情后果自负”的字是我自己签了的,好吧,确实没毛病。
只能说这也是我自己的原因吧,一切都处理的太着急了,没有把该准备的东西都提前准备好。不作死就不会死,自己作死的苦果只有自己吃了。
惊完之后的几天后开始后悔:
会想如果有时光机,那要回到哪一刻更好呢?回国前、5月初、或者更早到2月……不,这样一想过去后悔的事情多了去了,有太多“如果当时不这么做会怎么样”了。
世界上有好多事情是可以重来的,但更多的只会有一次机会。
时间就是这么公正。
最后开始接受现实并且思考解决办法,然后理了理这段日子我收获了什么,付出的代价是不是真的有这么大。
理性思维战胜了感性思维。
理工头脑:画个表格列一下?
不不不,那还是算了吧。
一年前可能不会有这种感触:除了签证以外,出个国真的是很平常的一件事情。之前有时实验室的一些项目中也有要去别的省出差的时候,也会面对完全陌生的环境,会有很多新鲜的事物,感受到不一样的文化和氛围。可能区别就是路上时间会更久,当地文化的差异更大吧。
这 1 个月的海外实习经历对我即将起步的工作生涯来说是真的很珍贵了。体会到了大公司内部轻松而又非常有序的节奏,长了很多见识,也体验到了他们是怎么看待 AI 领域的未来的。
另外……至少从生产力工具层面,软粉表示真香!Mac 对程序员来说是真的很友好啊。Pad + Pencil 的体验感觉也要比 Surface 更好一点。
几天前看到以前的一篇博文收到了评论:【October】。还是有不少感慨的。
可能是性格原因吧,我常常会自己反思自己以前做过的各种事情。更重要的是,如果我写的东西能给到看到这些文字的人任何帮助或者鼓励的话,对我来说也会是很开心的事情。
一直觉得自己是个对自己没什么信心的人。以至于有时候发生了一些好事的时候,第一反应不是高兴,而是会迟疑自己是不是真的 afford(想了半天没想到怎么表达比较好,还是这个单词比较顺)?然后小心翼翼地珍惜。相反遇到麻烦的时候却能比较快地接受现实,可能习惯了对自己没有那么高的预期吧,然后能很快恢复过来:既然事情已经这样了,有空自怨自艾不如想想怎么解决,生活还是要继续啊。
做过很多未来规划,但是也经常被一些未知的因素打断。幸而这大半年来,大多数的“意外”都是向着更好的方向改变,嗯,最后这当头一棒挺重的,不过也还算受得了吧。
也想感谢一路上所有帮助过我和给我肯定的人们。
我不知道自己之前做的这些选择是不是都是对的,可能也没有人能说自己的评判是绝对正确的。唯有希望未来做出每个抉择的时候也都能有自己清晰的想法,不要让自己后悔!
]]>2019 第一篇!
也是去年年底的时候心血来潮写的这个小东西:
中间断断续续往里面补了一些功能,修修 bug 什么的,然后就很久没有再管过了。
虽说一开始想要的本来不是个能写 GraphViz 的文本编辑器,结果写到最后还是写成了个文本编辑器。
话说 vscode 里面的 GraphViz 插件其实也挺不错的,
结果我自己平时还是觉得用这个更顺手哈哈
有天忘了在找什么的时候路过发现这个叫 CS Academy 的 OJ 上提供了一个简易的图编辑工具:
左边是文字形式记录的图关系,右边呢可以通过点击和拖动等等直观对图进行编辑,然后修改一边另外一边也会跟着同步更新。
刚看到的时候觉得惊了!这不就是我最早想要的玩意嘛!
然后果断 F12 开始研究起来。
嗯,这个网站本身应该是用某种框架做的,所有的动作都被包含在一个大概 10 万行左右的超大 js 文件里面。虽然从这里面直接扒出来基本上应该是不太可能的,不过看了下几个位置的 click、鼠标拖动事件等等差不多也可以知道实现的原理了,又在 d3.js 里找到个类似的东西,照着改改很快也山寨了一个功能差不多的出来。
然后我就发现更纠结了。
最初的想法是通过右边的图形界面操作生成左边的中间结果,然后手动细调之后再进一步生成 GraphViz 的 dot 代码,最后通过 GraphViz 生成 svg 的图。
结果等到这个 Graph Editor 写出来之后,却发现这样瞎搞还不如从一开始就老老实实直接写 dot 代码来的顺畅。
这大概就是花了好大的精力实现了最初的想法之后发现自己走了歪路的感觉了。
说起来这算是我第一个真正走到了实用阶段还一直在继续改的脑洞了,歪过一次之后现在决定这玩意还是老老实实就写成个编辑器吧。
目前是加了个多模版的启动页,然后用 PPT 花几分钟做了个图标。
对!你没看错,PPT!
一直没有真正好好学过前端相关的东西,所有这些都是有想法之后现场在网上找的教程和资料。
等以后有空了再改大概得想办法照顾一下这个界面的美观性(目前负分……),然后把逻辑部分的代码也好好重构一下。然后虽说现在把 Ace Editor 的自动补全功能打开了,但是这里面似乎不自带补全规则,之后还得看看怎么加。再然后还想把 GraphViz 语法的文档功能加进去,我自己经常用的时候还得开 GraphViz 的网站查有的东西怎么写,生产力啊生产力!
Keep coding!
2019 大家继续加油呀!
]]>不知不觉居然写了十篇了……写这个的初衷是觉得自己是个容易忘事的人,不找个地方做点笔记可能回过头就忘了自己看过什么了。
这篇要分析的是 TensorFlow 自带的 Allreduce 实现。
TensorFlow 中常规的使用操作是:
1 | import tensorflow as tf |
但是实际上 TensorFlow 目录下的 __init__.py
里面并没有把全部的内容都包含进去,另外的内容需要通过:
1 | from tensorflow.xxx.xxx import xxx |
这样的方式直接导入单独的包。
然后更奇怪的是有的 API 明明不在 contrib
中,按理说应该算是官方库中正式内容了,但是在官网的文档中却找不到。也可能是正要从 contrib
往官方库中移?
举例来说前面提过的 StagingArea 就是这样,然后这里的 Allreduce 操作的 API 也是,这就让人很好奇是不是还有别的什么不常用到的 API 可能在某些场景下会有奇效?
扯远了,
首先是 tensorflow.python.ops
下的 collective_ops.py
这个文件,一共包含了 3 个 API:
1 | def all_reduce(t, group_size, group_key, instance_key, merge_op, final_op, |
基本的用法示例如下:
1 | with tf.device('cpu:0'): |
参数中的 t 是每个 device 上要进行 reduce 的本地 Tensor,group_size 是一共需要参与的 Tensor 数量,merge_op 和 final_op 分别是 reduce 聚合时和聚合结束之后要做的事情。
然后比较坑的是 group_key 和 instance_key 这两个参数,注释写的太不明确了(还是我理解能力有问题?)。instance_key 应该是用于标识完整的一次 reduce 操作,而 group_key 具体是怎么来的还不是很清楚。测试的时候整体换成另外的数字都是可以正常工作的。
另外 all_reduce 和 broadcast 都似乎不能在 GPU 和 CPU 间进行工作?选择两块 GPU 卡之前是没有问题的,而再加上一个 CPU 就报错了。
其中的具体实现来自于 gen_collective_ops
这个包,那我们就知道这又是个用 C++ 写的然后封装到 python 下面用的操作了。
接口定义在 tensorflow/core/ops/collective_ops.cc
中,C++ 部分的实际实现在 tensorflow/core/kernels/collective_ops.cc
中,但是这里的 CollectiveReduceOpKernel
、CollectiveBcastSendOpKernel
、CollectiveBcastRecvOpKernel
三个类也主要还是用于设置输入输出的参数信息、检查结果等等,更进一步的实现要到 tensorflow/core/common_runtime/base_collective_executor.cc
中。
继续深挖下去之后,在 collective_param_resolver_local.cc
中看到了这样一段:
1 | cp->instance.impl_details.collective_name = |
……
嗯,所以这里的实现又是注册、又是各种封装的,结果到最后只有两种算法。
目前这里的 all_reduce 只有环状通信的算法实现,broadcast 则只有二叉树广播方式的算法实现。
要想搞清楚这个 group_key 是怎么回事,还得从 tensorflow/core/kernels/collective_ops.cc
这个接口开始。
CollectiveReduceOpKernel
、CollectiveBcastSendOpKernel
、CollectiveBcastRecvOpKernel
三个类的 ComputeAsync()
方法的基本结构都差不多:
CanProceedWithCompute()
方法检查各项参数,对 group_key 的检查也在这里完成单机环境下的 CanProceedWithCompute()
最后会调用 CollectiveParamResolverLocal::CompleteGroupLocal()
,group_key 在这里只是完全作为一个 map 的关键字来使用:
1 | { |
若某个 group_key 是第一次被使用,则与之关联的 GroupRec 会用当前创建该 op 的默认设备类型来作为整个 GroupRec 的设备类型。
当下一个 collective_op 创建时,再对目标的设备类型和根据 group_key 找出来的 GroupRec 作对比,不一致则报错,因此这里确实是限制了参与 reduce 或者 broadcast 的所有 op 都要是同一种设备类型的。
另外需要注意的是,同一个 group 中参与 all_reduce 和 broadcast 的 op 必须要和设备独立一一对应,所以也不可以在一块卡上同时发起 broadcast 的 send 和 recv,或者在同一块卡上的两个变量间进行 allreduce。
这个设备限制我觉得还是比较奇怪的,从 CPU 向当前节点下的所有 GPU 设备广播我觉得是很常规的一种逻辑啊……
instance_key 的作用也与 group_key 类似,InstanceRec 在这里主要是用来维护几个 mutex,主要负责多个 op 之间的同步,同时发生的多个 collective_op 操作只要 instance 不一样相互之间是不影响的。
To be continued.
]]>去年毕业开题大概开的也是这个方向,然后前段时间面试的时候:
面试官:你最近做了什么?
我:我 follow 了 FB 和 UCB 他们在 ImageNet 上做训练的工作,大概试了下他们论文里面的想法。
面试官:哦,那你最多用到多少卡?
我:e……最多跑到 16 块 V100。
说完这句自己突然感到一阵尴尬,人家工作多少块卡我多少块卡?真好意思说复现……(捂脸)
然而我也没办法啊,实验室资源已经算不少了,我还想拿一块卡假装模拟两块卡用呢。0.0
Anyway,趁着正好要读大法的论文,把之前其他人的东西都理一理。
前面这几段笔记还是去年的时候写的,看了下过了这么长时间,有的已经正式发了论文了,有的是 arXiv 上有了新版,那就下新的重新看一遍吧。
还是留下 OneDrive 的链接:
要说其他更多的材料的话,其实这几年在 ImageNet 上面冲分的有很多,Stanford 做过一个 DAWNBench,上面也记录了好多研究者跑 ImageNet 的最短时间、最小开销等等:
Facebook 做的大规模 GPU 集群训练,顺便推一波 Caffe2。
256 块 GPU,1 小时内跑完整个 ImageNet 的数据集训练,想想都可怕。
这篇文章主要是从一个“炼丹师”的角度来写的,主要关注在学习率、minibatch 等等之间的关系,以及它们对训练带来的影响。文章的主要工作是提出了一种根据 minibatch 大小调整学习率的扩放规则,以及一种在训练时预跑的机制。
这些可能大多都是训练中总结出来的一些 tricks 吧。
随着网络模型和数据规模的不断增长,训练时间也跟着增长了,为了保证训练时间能够保持在一个有意义的限度之内,需要一些额外的工作。
文章的立足点是通过增大训练时的 batch size 来加快训练速度。例如 ResNet-50 这个网络通常大家采用的 batch 大小是 256,在 8 块 P100 上大约需要 29 个小时才能跑完整个 ImageNet 数据集(应该指的是原始的 ResNet 论文)。
但是如果加大 batch size 呢?改变 batch size 会影响到最终的收敛和训练结果的正确率。
文章接下来的工作就是通过调学习率等等其他的一些方法,把加大 batch size 带来的影响控制在一个可接受的范围内,最终做到 8192 个 batch 训出来的结果跟原本的效果持平。当然最重要的,batch size 增大了 32 倍,总的训练时间也能够降下来,用 256 块 GPU 在 1 小时内训完了整个 ImageNet!
为了维持训练结果的准确率,他们在过往积累的大量经验上得出了一个令人难以置信地简单却又有效的规则:
其他超参数维持不变就好。
多次测试之后,可以发现应用了上面这个规则时,不同 minibatch 的 SGD 不仅最终准确率非常接近,而且训练的收敛曲线都是非常相似的。
但是在两种情况下,前面讨论的一些假定会有问题:
训练初始的几个 epoch 可能会引起网络参数巨变
在某些网络中,如果一开始设的学习率太大,初始的几个 epoch 会产生过大的梯度,上面这条规则就可能不太适用了。为了解决这个问题,文章于是提出了预热训练的阶段。
minibatch 的大小不能无限增长
就拿前面 ResNet-50 为例,这个规则只够保证 minibatch 在增长到 8192 的时候能维持训练效果,再大就保证不了了。
接下来是训练预热的机制:
经过实践调整之后,他们的实现里面是经过 5 个 epoch,学习率逐渐从$\eta$增长到$k\eta$。
分布式 SGD 在具体实现上还有很多细节的地方,可能稍微变动一下就会影响到整个模型的超参数,这里提了几个注意点吧:
关于我最关心的梯度聚合的实现,他们没有采用 Parameter Server 的形式,而是所有 worker 直接做 allreduce。
Allreduce 分成 3 个阶段:
1、3 两个阶段的操作,如果数据量大于 256 KB 则用 NVIDIA 的 NCCL 库来完成,否则就简单地各自把数据传回 CPU 再做汇总。
阶段 2 用到了递归二分倍增算法以及令牌环算法……嗯,基本上也没有什么特别的,就是想办法高效实现 all_reduce 和 all_gather。
软件栈用的是 Facebook 自家开源的 Gloo,具体的通信方面都是靠这个通信库来完成,就不涉及到具体的实现了(……)。以及 Caffe2。
硬件是 Facebook 的 Big Basin GPU 集群。每个节点是 8 块 NVLINK 的 P100(DGX-1),50 Gbit 的 InfiniBand。
预估一下 ResNet-50 的数据量:大约参数量是 100MB 左右,单 P100 跑完一次大约 120ms,也就是峰值带宽大概要在 100MB * 2 / 0.125s = 12.8 Gbit/s 的程度,加上一些额外的 overhead 应该不会超过 15 Gbit/s,对这套 IB 来说完全不是问题。
后面的训练设置描述得很详细,包括很多超参数要怎么取值,卷积层和全连接层参数的初始化分别采用不同的方法等等。
单节点的 baseline 是 8 卡(k=8),每块卡的 BatchSize 是 32(n=32),相当于单机一轮的 BatchSize 是 256,90 个 epoch 之后得到了 23.6% 的错误率,表现可以说是非常不错了。
多机训练的时候上了 32 个节点,一共 256 块卡,总的 BatchSize 达到了 8k,相较单机扩大了 32 倍。如前面所描述的,单机训的时候学习率选在 0.1,这里则是从 0.1 开始经过 5 个 epoch 逐渐增长到 3.2,即原来的 32 倍。用了这套 Gradual warmup 的预热机制之后,最终网络收敛在了 23.74%,几乎与单机一致了,并且从测试结果中可以看到后期甚至是连 error 的下降曲线都是基本一致的。
效率方面,他们使用的这套 Allreduce 算法得到了非常好的效果,基本上已经可以看到线性的性能表现。
做大规模训练的话,这篇论文基本上可以说是教科书级别的示范了。
上面 Facebook 用 256 块 P100 做了大 minibatch 的 ImageNet 训练之后,感觉军备竞赛就开始了。
这是日本的一个研究机构,用 1024 块 P100,更大的 BatchSize!!15分钟跑完了 ImageNet。
完成这个任务主要有两方面的挑战:
这里给了个表,对比了一下目前用 ResNet 跑 ImageNet 的几个工作:
Team | Hardware | Software | Minibatch size | Time | Accuracy |
---|---|---|---|---|---|
MSRA(ResNet作者) | Tesla P100 * 8 | Caffe | 256 | 29 h | 75.3% |
Tesla P100 * 256 | Caffe2 | 8192 | 1 h | 76.3% | |
SURFsara | KNL 7250 * 720 | Intel Caffe | 11520 | 62 m | 75.0% |
UC Berkeley | Xeon 8160 * 1600 | Intel Caffe | 16000 | 31 m | 75.3% |
This work | Tesla P100 * 1024 | Chainer | 32768 | 15 m | 74.9% |
嗯…话说,本来以为 UCB 最发的文章是在 KNL 上跑的,刚刚才发现他们当时最新的工作是用居然是 8160(这可是通用处理器啊???害怕,下面那篇真得好好看看了)
具体的实现,这里说是参照 Facebook 的工作做的,比较额外的就是加大了 BatchSize,然后软件方面换了他们自己写的 Chainer 框架,因此这块写的比较简略。可能还有个差别是这里用了 RMSprop 吧,另外也提供了他们所使用的 RMSprop 的 Warm-up 方案。
NCCL 和 OpenMPI 意味着他们写的这个框架是基于 MPI 做的通信,而且还做了卡间数据传输。主要还是用的同步模型,然后做 Allreduce。
训练用的单精,然后划重点,为了减少通信的数据量,数据传输的时候用的**半精浮点**!!测试之后发现对结果影响不大(666666….)。
硬件方面,双路 E5-2677,8 块 P100 的配置,应该是没有 NVLINK 的,IB 用的 FDR(56 Gbps)。
2018 年更新,这篇已经发在了 ICPP 上,也增加了很多内容,基本上已经是一篇新的文章了
接下来上台的是 UC Berkeley。
最早看到他们 arXiv 上的论文版本写的是 24 分钟跑完,可能被同期其他同行的工作刺激了一下,后来更新把标题里面的 24 去掉了,最终实现里面确实也是把这个时间缩短到了 10 分钟的量级。
他们家应该是跟 Intel 合作,用的是 Intel Caffe 和 Intel 的计算设备。
17 年的论文中使用了两种硬件方案:
18 年的正式版本还额外增加了 P100 的测试内容。
基本的结果是:2048 颗 8160 用 11 分钟跑完 100 个 epoch 的 AlexNet,准确率收敛在 58.6%;2048 块 KNL 用 20 分钟跑完 90 个 epoch 的 ResNet-50,然后 14 分钟跑到 64 个 epoch 的时候就已经收敛到 74.9% 了。
他们不出意外地也用了同步的更新策略,这里给的理由是同步可以保证整个计算过程的确定性,并行实现与串行实现的行为是一致的,这一点在设计和调试以及优化网络的时候尤其重要。有一些前人的工作也证明了异步的方式是不稳定的(例如 Revisiting Distributed Synchronous SGD 以及前面 Facebook 的工作等等)。
异步策略会引入更多的影响因素,而且由于各个 worker 上的参数存在延迟,对整体的算法收敛也会有影响。话说前面几篇论文里面提出的优化方案也是,在异步训练的时候完全是不同的表现,不能直接适用过去。
Introduction 最后点了一下这篇论文的一些要点:
嗯……Related Work 里面把前面两篇文章都提了一下,还是很有趣的。
第三章分析了一下大 Batch 训练的好处。理论上改变 BatchSize 是不会影响到一个 epoch 的总浮点计算量的,但是随着单个 step 的 BatchSize 的增长,一个 epoch 需要的迭代次数就减少了,整体通信的次数也减少了。
总之整体的运行时间可以减少,最理想情况当然是可以做到线性。表中给的总时间应该是以 GPU 数量取 log 来作为一次通信的 overhead 影响因子,因为这里考虑的是 Allreduce 同步的通信方式,log 级的估算基本上还是可以接受的。
话说我觉得这个问题还得换个角度更深入考虑一下,一个 epoch 中通信的次数减少了,但是如果增加节点数就意味着单次通信的总数据量要增加,最后还是反映在带宽上。
另一方面 BatchSize 变大之后单个 step 的计算时间也增大了,说起来通信这部分的 overhead 可能也更容易被 overlap 在计算时间之后。这里还只是一些比较模糊的定性讨论,数据并行能不能达到线性增长还是得看通信带来的 overhead 有多大和实际跑的结果。
另外一个原因是说大 Batch 可以提高 GPU 的利用率。
我觉得本来跑满一块 GPU 应该是基本吧。
其实像前面 Facebook 的工作,单块 P100 能够支持的 ResNet-50 的 BatchSize 应该远不止 32,不知道他们最后训的时候 P100 的使用率大概是个什么情况。当然他们也有实际收敛效果方面的考虑,以他们当时的优化策略还并不能够保证 BatchSize 再大以后的收敛性能。
在模型的选择上,他们提出了一个扩放率的概念,即计算和通信的比率,对比 AlexNet 和 ResNet-50 这两个网络:
Model | Communication # parameters | Computiation # flops per image | Comp/Comm Scaling Ratio |
---|---|---|---|
AlexNet | 61 Million | 1.5 Billion | 24.6 |
ResNet-50 | 25 Million | 7.7 Billion | 308 |
ResNet-50 的扩放率大概是 AlexNet 的 12.5 倍,所以 ResNet-50 的可扩放性更好,能达到更高的弱扩放效率。
但同时大 Batch 也会带来一定的问题,比如 BatchSize 不能无限往上加,大到一定程度之后,准确率就降下来了。所以后面的重点就在于怎么调整学习率这些超参数来保证增大 BatchSize 之后的训练效果(LARS 登场预告)。
基本的原则跟 Facebook 提出的观点一致:
但是比 Facebook 的工作有了更深入的研究:
他们发现不同层参数和梯度的二范数比值($\frac{||w||_2}{||\nabla w||_2}$)差距可能会很大,例如 AlexNet 中第一个卷积层和第一个全连接层的比值可能会差上千倍,猜测这个会成为影响大 BatchSize 训练精度的一个重要原因。
于是他们提出了 LARS(Layer-wise Adaptive Rate Scaling)算法来分别控制每一层的学习率。
666666
整体算法如下:
把一个算法扩大规模到更多的处理器上时,通常通信都会成为最大的 overhead。
所以这里从两个角度分析了通信比计算慢这一点,首先是硬件运算力(每个 flop 需要花费的时间,例如 P100 的峰值性能是 11TFlops 左右,$\frac{1}{11TFlops}$大约是$0.9*10^{-13}s$左右)会远小于理论传输速度(这里用的是$\frac{1}{Bandwidth}$,例如 56Gb/s 的 FDR,大约是$0.2*10^{-9}s$的水平),远小于延迟:
$$Time-Per-Flop << \frac{1}{Bandwidth}<<Latency$$
去年的论文中还分析了在 45nm 工艺的 CMOS 处理器上,通信消耗的能量也会比计算更大:
Operation | Type | Energy(pJ) |
---|---|---|
32-bit Int Add | Computation | 0.1 |
32-bit Float Add | Computation | 0.9 |
32-bit Int Multiply | Computation | 3.1 |
32-bit Float Multiply | Computation | 3.7 |
32-bit Register Access | Communication | 1.0 |
32-bit SRAM Access | Communication | 5.0 |
32-bit DRAM Access | Communication | 640 |
能量这个角度比较神奇,感觉应该体现的不是很明确,正式论文中已经删掉了这一段。
设总的 epoch 数为 E,图片数为 n,BatchSize 为 B,网络参数为 |W|。
当 E 和 n 保持不变的时候,总的计算量是固定的,总的迭代次数是 $E* \frac{n}{B}$,通信量是$|W|*E*\frac{n}{B}$。当 BatchSize 增大之后,总迭代次数减少,那么通信次数也减少了,通信量减少。
感觉这里再加一个 worker 数量等等这样的参数会更好一点。
更大的 BatchSize 有利于增大计算通信比,因此更大的 BatchSize 有利于网络向大规模进行扩展。
后面的测试都得到了非常好的结果。
这篇是腾讯和香港浸会大学合作的成果,话说 HKBU 的这几个作者之前还发过一个叫 DLBench 的各个深度学习框架性能测试的文章,没想到又看到了他们的名字。
这篇文章主要做了三件事:
最好的成果是 2048 块 P40 在 6.6 分钟把 ResNet-50 跑到了 75.8% 的准确率,以及 4 分钟跑完 95 个 epoch 的 AlexNet。
他们搭了一套名叫 Jizhi 的训练系统(机智?):
他们整体采用 fp16 来完成整个训练以提高整体运算的吞吐量,在额外应用 LARS 进行优化时遇到了半精范围溢出的问题。于是最终做了这么一件事:
这……
使用这套方案之后,最终在 64k 的超大 BatchSize 下 ResNet-50 能够收敛到 76.2% 的准确率。
通信方面,首先是 Tensor 聚合,由于某些卷积层的参数量特别小,直接应用 Ring-Allreduce 等等算法去发送大量的小参数可能很难把带宽性能发挥出来,所以采用一个 buffer 池来收集完成前向运算的参数,当池中的参数量达到某个阈值时,再一次性发送出去。
我倒是比较好奇这块是怎么在 TensorFlow 中实现的。
根据参数量的大小,也分别采用不同的 Allreduce,并且在多级硬件框架中采用分级的 Allreduce 方法。
基本上就是这些。
再是索尼的这篇文章。
他们的成果是 2176 块 V100 在 224 秒中达到 75.03% 的准确率,以及 1088 块 V100 下达到 91.62% 的扩展效率。
为什么是这么奇怪的两个数字…
基本的思路还是跟前面的工作大同小异,他们采用了动态调整 Batch Size 配合 LARS 来作为超参数的优化策略,并且同样配合使用了 fp16/fp32 混合精度,并且在通信方面采用了 2D 环网的 Ring-Allreduce 模式(日本的研究者真是喜欢在网络拓扑上做文章,之前他们那台超算 京 也是在网络方面有独特的设计)。
NIPS 2017 上有一篇论文从数学层面分析过神经网络收敛方面的问题,得出的基本结论就是训练时间更久才能填上收敛的 gap(…嗯,虽然这个从经验上早就能得出来了)。前面的各位都是在考虑怎么把随机梯度下降的 BatchSize 提上去并且同时又把训练精度控制住,比如说各种花式处理学习率,花式改 BatchNorm,动态调整 BatchSize 等等。
这篇 11 月底的新工作则是换了个思路:大 Batch 带来的问题会不会就是 SGD 本身的限制,无论是 Momentum 还是 RMSprop 本质上都只是对 SGD 的优化修正算法,那如果干脆就不用 SGD 呢?
于是这里改用了 K-FAC(Kronecker-factored approximate curvature) 这个二阶的优化方法,配合 Allreduce 和半精来做分布式训练,测试结果是要比一阶的 SGD 收敛更快,仅仅 35 个 epoch 就可以把准确率收敛到 75%。同时,使用这种方法在 1024 块 V100 上也能够在 10 分钟就把 ResNet-50 收敛到 74.9%。更进一步的优化策略(针对 K-FAC)还能再减少计算量和内存的占用量。相关工作中也说这里更多的是对 ICLR 2017 上的一个 K-FAC 工作的改进。
K-FAC 是利用各层参数得到的费希尔信息矩阵(Fisher Information Matrix)计算一个系数,再乘上参数的梯度来进行更新。数学部分先跳过了,大致的过程跟 SGD 差不太多,只是 Kronecker 系数的计算涉及到了更多的东西,相对 SGD 来说,这部分计算是额外的 overhead。
软件方面,他们的工作是在 Chainer 上完成的,这个框架平时完全没接触过,但是前面 2017 年那篇日本的工作用的也是这个,猜测应该是作者是日本的研究者或者在日本这个框架更流行吧。
以 3 层网络为例,基本的流程如下:
Stage 1、2 分别是网络的前向和后向计算,之后的 3、4、5 是 Allreduce 计算 Kronecker 系数,最后 Stage 6 用每层计算出来的系数进行梯度更新。
其他诸如学习率的调整和预热的方法这里也有用。
Facebook 的工作是我看到的最早在大规模训练实现这块引路的了,之后 UC Berkeley 提出的优化方案又把整个效果往前推进了一步,其他人后来的文章都是集合了当时各种高效方案的实现,最近看到的索尼的那篇应该算是这个方向的集大成了。
总结下来,基本的参考点大概有几个:
想办法增大计算吞吐量、减少通信的 overhead
应用各种调超参的 trick 加速收敛
这篇来研究一下 TF 中的一些高级 API。
TensorFlow 由于一直是在开源的社区环境中发展起来的,早期的一些 API 都比较简单粗暴(更直白地说就是不那么好用),以至于在它之上封装的更友好的 Keras 可能在大部分的使用者群体中会有更高的出现率。后来的 TensorFlow 中也有吸收 Keras 里面一些比较好的结构,有出现像 tf.layers
这样的更高层封装,可以期待一下 2.0 以后会不会大幅优化上层的编码 API 吧。
那这里说的高级 API 是什么呢?
官网的 guide 里面列了几个:
tf.keras
,官方提供的从 TensorFlow 本身向 Keras API 兼容的实现(感觉怪怪的,底层库中包含了一个对高层封装的兼容???)其他的还有包括像 StagingArea(构建软件流水,神器!!!)等等目前还在 tf.contrib
中处于实验阶段的很多东西,开源的力量太强大了,每隔一段时间就有很多新功能被社区添加进库中。
像 tf.keras
和 Estimator 的设计都是为了让用户能够更方便地编写网络,话说简单看了下 Estimator 的用法,API 的设计方式应该大概率是从 Keras 里面借鉴的。
具体的使用这里就不多记了,这里 写了个很小的例子,直接开始拆 Estimator 的实现吧。
核心是 tf.estimator.Estimator
这个类,先看初始化参数:
1 | __init__( |
第一个是用来构建网络的模型函数,具体后面详细分析;model_dir
用于指定保存的参数、checkpoint、log 信息等等存放的目录;再后面的几个都是一些额外的配置选项。
让我觉得非常怪的是官网的介绍页面说 Estimator 的优势是不需要用户建图……我真是一脸懵逼。或许对于 TF 内置的一些事先建好的 Estimator 是这样吧,但是如果想自定义呢?……写文档的人吹的有点过了吧。
创建 Estimator 时,首先初始化各项配置参数信息(值得一提的是 model_dir
是允许被 config
中的选项覆盖的),设置训练或者验证时用的数据分布策略(DistributionStrategy,后面再详细分析),设置参数和图节点分布的 device_fn
;之后简单检查 model_fn
的参数是否符合规范,然后完处理完 warm_start
的一些设置就结束了。
传入的 model_fn
是用于构建 Estimator 代表的模型网络的核心函数,它能够接受的参数名有严格的规定:
tf.estimator.ModeKeys.PREDICT
、tf.estimator.ModeKeys.TRAIN
和 tf.estimator.ModeKeys.EVALUATE
;接下来是 Estimator 类的三个调用方法 evaluate、predict 和 train,从字面上就能够看出来各自对应的是什么功能了(Keras 里面对应的 API 应该是 evaluate、predict 和 fit)。
1 | evaluate( |
三个方法的共同参数是这个 input_fn
,这是类似前面 model_fn
一样,也需要 Estimator 的创建者写好的输出数据产生函数。这个函数的返回值是一个二元组 (features, labels)
对应了 model_fn
的前两个输入参数。
train 中的 steps 表示从哪里开始训练,Estimator 将首先从保存的 checkpoint 中找到最接近的保存点,然后开始这次的训练,max_steps 则简单地就是训练的 batch 数了。
1 | def _train_model(self, input_fn, hooks, saving_listeners): |
如果在初始化时没有配置 _train_distribution
项,则会使用默认的方式来执行 train 操作,最终把 model_fn
也绑定出来:
1 | estimator_spec = self._call_model_fn( |
向 model_fn
中传入输入数据以及 ModeKeys.TRAIN
,接下来实际的执行函数是:
1 | def _train_with_estimator_spec(self, estimator_spec, worker_hooks, hooks, |
添加 TensorBoard 中的 Summary、创建参数保存点、如果有 saving_listeners
则额外添加到运行的 hooks 中,之后:
1 | with training.MonitoredTrainingSession( |
嗯……这段代码是不是很熟悉,没错,官方建议的常规 TensorFlow 训练代码就是要写成这个格式。
至此,train 部分基本上分析完了(带 DistributionStrategy 的版本后面再说),整个过程就是把一套常规的 TensorFlow 代码的各个部分做了几级封装,要说有什么特别的就是它把 Summary 和 Saver 都默认包含在内了。
如果按照这个格式解开成普通的 TensorFlow 代码的话,可以说是非常好的官方范例了。
然后再注意到 model_fn
的返回值,前面也提到了 evaluate、predict 和 train 这三个实际执行的方法其实最终都是把 input_fn
中产生的数据传给 model_fn
来跑,这里的控制差别就需要配合对不同的 mode 选项的分支判断来做,所以一个 model_fn
函数写出来大概是这个样子的:
1 | def model_fn(features, labels, mode): |
不管是哪种模式,model_fn
最终的返回值都需要通过 EstimatorSpec
这个结构来传出去,其属性有:
1 |
|
[1]
的 Tensor;minimize()
方法返回的那个;Metric
类,或者一个 (metric_tensor, update_op)
的元组;要求 ModeKeys.TRAIN
模式返回的必须包含 loss 和 train_op,ModeKeys.EVAL
模式返回的必须包含 loss
,ModeKeys.PREDICT
模式返回的必须包含 predictions。
前面说了,Estimator 这套 API 的出现是因为开发者希望能够方便用户快速搭网络,并且易于扩展到各种不同的计算结构上。
那么 Estimator 本体上面已经拆过了,没什么神秘的,就是个简单封装,距离实现单机到多卡/多机的这种扩展其实还差挺多的,而 DistributionStrategy 就是用来补上中间那个 Gap 的。
官方提供的资料是个 Github 上的 README:Distribution Strategy
里面给的使用的例子都非常简明易懂:
1 | distribution = tf.contrib.distribute.MirroredStrategy() |
下面的三行是个普通的 Estimator 的使用,跟前面一样,唯一区别的就是在 tf.estimator.RunConfig
中创建一个 DistributionStrategy,然后作为 config 选项传递给 Estimator 即可。
确实很方便啊。
如果有看过官方的 benchmarks中对多卡/多机的写法的话,可以发现那个写法大致上跟 DistributionStrategy 的设计非常像。
_train_model_distributed
的大致结构是这个样子的:
1 | with self._train_distribution.scope(): |
_train_distribution.scope
中封装的是一些 variable_scope
和 custom_getter
,用于在用tf.get_variable()
创建的变量之上再套上一层额外的封装控制。
对普通的
tf.get_variable()
套上variable_scope
之后可以控制这个变量创建的设备位置等等很多东西,第一次看到这种写法的时候还觉得这像是一种 hack 的方法,但是不用改变原本里面的代码,还是非常方便的。
以单机多卡为例,call_for_each_tower
是在每块 GPU 卡上跑一遍 Estimator 中给定的 model_fn
,即在每一块卡上独立创建一份数据并行的网络。基类里面这个函数是留空的,要求具体的实现类来完成这个部分。
话说这里我有个疑问……model_fn
中的正常写法应该是对训练模式返回一个包含了 Optimizer.minimize()
的 EstimatorSpec,但是多卡并行的过程中不是应该需要做梯度的聚合平均之后再更新到每个变量上吗?而且不同的并行模式下,这部分的处理方式应该也是不一样的,不知道这套 API 要怎么把这些全都统一起来。
看一下 MirroredStrategy 的这个实例里面是怎么实现这个函数的吧,中间部分的代码是这样的:
1 | # TODO(isaprykin): Create these threads once instead of during every run() |
distribution.worker_devices
中保存了单机中的多块 GPU 卡设备,这里直接对每块卡挂一个 shared_variable_creator
并且开一些的一对一的线程去处理。
shared_variable_creator
用于处理多卡之间的参数共享,在 device_id 为 0 的设备上调用 get_variable()
函数是创建新变量,并且保存到给定的 shared_variable_store
dict 中;在 device_id 大于 0 的设备上调用 get_variable()
则会尝试共用前面创建好的变量。
接下来看一下创建线程这部分的逻辑。
主线程和多个子线程之间的控制这里用了 should_run
和 has_paused
两个 threading.Event()
来控制。开始的时候,每个线程都调用 should_run.wait()
来等待,等待主线程调用对应的 should_run.set()
来唤醒它们。主线程随后阻塞在 has_paused.wait()
上,等到每个线程完成自己那部分图的构建之后再用 has_paused.set()
唤醒。
话说为啥一定要用多线程来实现这个部分呢……感觉就用普通的单线程循环一样可以做到这里想要的事情。
那么对梯度的聚合和最终的 apply 呢?似乎这部分代码里面根本没看到啊,每个线程的 run()
函数基本上就是跑完各自的网络部分就没了。唯一看上去非常让人介意的是 run()
中在执行 main_fn()
前套了一堆 Python 的 Context,难道又是用有点 hack 的方法完成的?
其中 MirroredTowerContext 这个结构继承了 distribute_lib.TowerContext,只用于call_for_each_tower
中用于处理多块卡之间相同代表数据的同步。
问题是我还是没有看到处理 reduce 等等这些的代码。
然后……抱着一种怀疑的想法,我重新打开了 Optimizer 中关于梯度计算部分的代码!!发现这里已经跟当时拆包第二篇(Optimizer in TF)里看到的不一样了。
例如新的 apply_gradients
中增加了这样的一段:
1 | if distribution_strategy_context.has_distribution_strategy(): |
那么 DistributionStrategy 是如何配合 Estimator 把原本单机的代码直接扩展开来就很明白了。
当时最早拆包开始时大概是 TensorFlow 的 1.6 版左右。
查了下 Optimizer 中增加部分的 git 记录,差不多是在今年 3 月底的时候加上的,应该是在 TensorFlow 的 1.7 版左右,然后后来又有过一次较大的改动。
再看一下官方文档中对 DistributionStrategy 的设计思想。
首先是一些底层的概念:
call_for_each_tower(fn, w)
,fn 是模型函数,w 代表一些 Wrapped values。这个函数的调用过程中就包含了变量的 unwrapping 和 merging,假定在设备 d0
上 fn(w0)
得到的结果是 (x, a, v0)
,在设备 d1
上 fn(w1)
得到的结果是 (x, b, v1)
。首先在调用函数之前,w 需要被解包变成 w0 和 w1 然后分别调用 fn 函数。返回的结果有三种情况,第一个值都返回了一个相同的对象 x
,则最终 merge 之后还是对象 x;第二个值是每个设备不一样的,则 merge 之后是一个 PerDevice 对象(其实就是个设备和对应值的 map);第三个值是每个设备返回的分别是一组 Mirrored 对象的成员,则 merge 之后是一个 Mirrored 对象。所以 call_for_each_tower(fn, w)
在这里返回得到的就是一组 (x, PerDevice{...}, Mirrored{...})
reduce()
这种所有设备共同参与的一个操作就需要这种封装更新一个变量的常规操作如下:
d.distribute_dataset()
中,然后创建一个 iteratord.call_for_each_tower()
来分别创建网络模型,并且最终各自得到一组梯度/变量对:d0
上有 {(g0, v0), (g1, v1), ...}
,d1
上有 {(g'0, v0), (g'1, v1), ...}
等等这样d.reduce(VariableAggregation.SUM, t, v)
或者 d.batch_reduce()
来对梯度求和,并且对应到各自的变量上:{(Sum(g0, g'0), v0), (Sum(g1, g'1), v1), ...}
d.update(v)
来对每一个变量进行更新3、4 两步如果用 Optimizer 中的 apply_gradients()
方法可以自动完成(……这就是 Optimizer 后来加进去那部分代码的作用),或者在一个 Cross-tower context 中调用 _distributed_apply()
方法也可以。常规的网络层都应该在 Tower context 中被调用。
话说这个
_distributed_apply()
为什么前面带下划线啊喂,这个方法本来不打算直接给人调的吧???大概是 API 还没最终设计好。
嗯,所以 Estimator 本身一点都不神奇,真正这套机制麻烦的地方在 DistributionStrategy 里面,手写一个 DistributionStrategy 应该会是一件很麻烦的事情。
不知道未来这套机制会如何改进得更好用一些。
2018 12月更新
最近正在试图手动 DIY 一个 DistributionStrategy,发现到处都找不到相关的资料,官方的文档方面对这部分也是写的不明不白。
试了下自己继承一个 DistributionStrategy 类,但是发现这个基类的几乎所有功能都是交给另外一个 DistributionStrategyExtend 类来做的,而且直接从空类开始写缺的东西也有点太多了,还没下手成功。准备之后再试试直接继承一个现有的类比如 MirroredStrategy 然后重载掉里面的功能函数试试看。
Estimator 这套机制想要达到的目标是非常好的,但是似乎……由于 TensorFlow 本身过于庞大和复杂,不知道什么时候这两个东西才能真正成为方便用户使用的好接口。
从名字上直译过来应该是用于暂存的区域,这套 API 用于跨 step 地把数据保存到网络 data path 之外的地方,然后可以在另外的 step 中把保存下来的数据取出来。
解释上看起来挺绕的,而实际上用这套 API 实现出来的效果就是——软件流水。
例如以下面这个由三个阶段组成的计算过程为例:
在 a/b 和 b/c 之间分别加入 StagingArea 即:
更重要的是加入了暂存结构之后,事实上 a、b、c 三个计算阶段的依赖就被解耦了:
对执行的流程稍微进行一些修改:
1 | step1: A1 |
原本必须按顺序执行的三个计算阶段现在就可以互不相关地并行执行了,在某些计算与 I/O、数据通信共同存在的环境中,原本可能存在的数据延迟、等待等等就有可能通过流水线的方式隐藏掉!(例如多机分布式训练的情况,实测效果非常好)
关于这套 API 如何使用的介绍这里就不多记了,直接来看 TensorFlow 是怎么实现它的。
StagingArea 这个类在 tensorflow/python/ops/data_flow_ops.py
中,很早以前应该是在 tf.contrib
里面的,大概试验成熟之后移到正式的包部分了。
put()
和 get()
这两个方法的实现分别调用了 gen_data_flow_ops.stage()
和 gen_data_flow_ops.unstage()
,然后你会发现虽然一开始是从 from tensorflow.python.ops import gen_data_flow_ops
中引入了这个包,但是源代码里面是找不到这个包的。
原因在于这里面的东西都是在 C++ 层代码中定义然后在编译过程中生成的,追到 tensorflow/core/ops/data_flow_ops.cc
中可以看到大量用 REGISTER_OP
宏注册的 op,其中就有 StagingArea 用到的 stage()
、unstage()
等等函数。
当然到这里为止还是没办法找到它的实现,因为 REGISTER_OP
宏只是负责 Python 与 C++ 的接口部分的处理,具体 C++ 层调用的实际内容还要再往 Kernel 里面找:tensorflow/core/kernels/stage_op.cc
。这里才是真正最底层的实现内容了,然后还能看到很多 REGISTER_KERNEL_BUILDER
宏,用于把 C++ 部分编译成的库与上面的接口绑定起来。
在拆包第三篇简单记过 TensorFlow 中 Op 的创建方式,嗯,在这里用上了。
然后就发现这玩意的实现就是个双向队列的封装,没啥神奇的……╮(╯_╰)╭
1 | std::deque<Tuple> buf_; |
国庆在家闲不住想干活系列……
本篇的内容是陈天奇大佬今年春季在华盛顿大学开的一门课。
大佬是上交 ACM 班的本硕,虽然目前还在 UW 读博中,但是在机器学习圈子里面已经有了很高的名望了,他的 MXNet 和 XGBoost 相信很多人就算没用过也肯定听说过(比如我就没用过…)。前段时间他发布的 TVM 也算是开启了深度学习和系统方面探索的一条新道路。
课程介绍上讲的是这门课的目标是填上深度学习算法和系统的实现以及优化之间的 gap,粗略地翻了一下 PPT,后面也有比较多的篇幅是介绍 TVM 的,正是我想了解的!
没找到课程的视频,但是 PPT 可以在上面的课程链接或者这里找到。
下面的内容主要按照每篇 PPT 的整理:
回顾了一下基本原理和发展历史。
机器学习的过程基本上就是 数学模型+评价指标+参数训练,深度学习则是模型特指各种神经网络。
具体主要涉及到各种不同的模型架构(CNN、RNN、各种变种),目标函数的选择和训练技巧,正则化初始化等等。
这些就不多记了。
Lecture 2 是个实验课,实践怎么搭网络。
这一节差不多是大纲的性质,每一个小点后面都会分节细讲。
基本上所有的深度学习框架都是差不多这个结构,首先来看 User API 层:
这里举了个线性回归的例子来对比手写 Numpy 和框架代码的差别。基本上网络模型都可以比较方便地用一个计算图的结构来表达,节点表示运算操作,边代表数据依赖。
那为了方便用户写代码,一个框架也是一定要有自动求导的功能的(如果反向计算还需要手写那就太瓜皮了)。
然后是 System Components 层:
这里涉及到了首先是计算图的优化,比如一次运行的时候直接过滤掉用不到的图节点(Deadcode Elimination),内存分配方面的优化,图节点和实际计算设备的对应等等。
实际跑图的时候如果有多个设备或者多个工作线程,如何调度以及发挥出计算设备的并行性也是一个需要考虑的问题。
最下面的 Architecture 层:
目前用来支持 DL 的设备也有很多,典型的如 GPU,其他的加速芯片也是越来越多,不同的设备可能要写对应不同的代码,这部分要怎么优化?
现在最常规的做法是每一种不同的计算设备会有开发厂商自己提供支持库,但是这个对框架的开发者来说还是有一个要整合的过程。另外,如果系统中存在多种不同的计算设备,计算任务在多设备上要怎么分配和调度也会是一个很大的麻烦。
为了解决最后的这个问题,目前有一种 Compiler Based Approach,即整个 Architecture 层由一个 High Level Operator Description 加上 Tensor Compiler Stack 来统一解决。这就是之后要提到的 TVM 的设计思路了。
详细解释第三节中的自动求导。
计算机中实现求导这个操作主要有两种方式:基于符号的算术求导和直接用数值进行求导。
算术求导需要构建一棵符号表示树,然后根据各种算术上的求导规则来写公式。缺点在于:如果遇到特别复杂的函数,则需要推导的符号表示树也会很大;然后如果目标只是想要一个导数值,则保存一棵符号表示树就很浪费了;再然后就是这样做容易有误差(?为什么…按公式算不应该误差更小吗)。
数值求导则是按导数的定义做,直接对方程取极限:
$$ \frac{\partial f(x)}{\partial x_i} \approx \lim \limits_{h \to 0} \frac{f(x+he_i)-f(x)}{h}$$
实现起来特别简单,h 取个 1e-6 就差不多了,但是一般只用来检验求导结果用。
然后对于网络中每一层的反向部分,其实求导涉及到的都只是跟本层运算相关的内容:
上一层传下来的是 $\frac{\partial Error}{\partial z}$,再往下可以通过链式求导法则一直推导下去,而其他需要的则只是与本层运算有关的$\frac{\partial z}{\partial x}$ 和 $\frac{\partial z}{\partial y}$。
更详细的推导可见这里。
因此自动求导则是根据以上的规则来创建反向计算图的过程,伪代码以及结果如下:
自动求导构建完成反向计算图之后,完整的计算图可以接下来一起用作整体的图优化。
在 CPU 上进行数据运算大致有几个过程(按多级流水分):Fetch、Decode、ALU Compute、Write Back。由于 CPU 本来也并不是为了纯运算而设计,因而在 ALU Compute 以外的其他部分会有比较大的计算资源和能耗上的 overhead。
后来增加的向量化指令能够相当程度地改善这种 overhead 的问题,而 GPU 从这个角度来看更像是一种把 ALU 的向量化做的更极致的加速器。
从存储的层次结构上来对比:
GPU 的大量寄存器就使得它能够以比 CPU 小的多的开销来切换线程,这也就能够支撑起大规模的 SIMT 了。
后面是一些 CUDA 编程的例子,以及如何根据 GPU 的微架构特性高效地发挥出性能来。
这一节的内容大概在 System Components 和 Architecture 层之间,一份代码面对不同规模的数据(甚至是不同数据块尺寸的数据)往往不作针对性地调整是达不到最佳性能的。
深入下去需要实际考虑到例如 CPU 的 Cache、GPU 的寄存器等等这些方面,以及 GPU 的多级存储之间的数据搬移开销、数据重用等等,同样是 GPU 也有多种不同的后端。
不同的 Tiling Patterns、Fuse Patterns、Data Layout、Hardware Backend 合起来使得优化工作也变得相当复杂了。
为了解决前面说到的所有这些麻烦的问题,然后这里就引出了 TVM Stack。
各种不同的框架和实际执行运算的各种各样的硬件后端之间其实存在着很大的 gap。
如果从编译器的视角来看待如何解决这个问题,各种框架写的网络可以根据特定的规则转化成某种统一的表示形式,在统一表示的基础上进行一些可重用的图优化,之后再用不同的后端来生成对应不同设备的代码,这就是目前各家都在尝试的设计思路了。
举例说: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 | C = tvm.compute((n,), lambda i: A[i] + B[i]) |
生成得到的 C 代码是:
1 | for (int i = 0; i < n; ++i) |
加上额外的调度控制:
1 | C = tvm.compute((n,), lambda i: A[i] + B[i]) |
再生成的代码就变成了:
1 | for (int xo = 0; xo < ceil(n / 32); ++xo) |
甚至于还可以支持绑定中间的 xo 和 xi 到特定的变量上:
1 | C = tvm.compute((n,), lambda i: A[i] + B[i]) |
话说这样出来的代码就可以用在 CUDA kernel 里面了:
1 | int i = threadIdx.x * 32 + blockIdx.x; |
具体后续的调度部分的设计,首先需要保证生成的代码在逻辑上要能跑出正确的结果,常见的手工优化代码的方法也都要包含在内,并且要能够方便引入其他额外的新技术。
目前 TVM 的调度部分还在继续开发中,已经从像 Halide、Loopy 这种成熟的语言中吸取过来的方法有例如 Loop Transformations、Thread Bindings、Cache Locality 等,针对 GPU 还开发了一些方法例如 Thread Cooperation、Tensorization、Latency Hiding 等。
再额外的就是 TVM 还用了 Auto-tuning,由于 TVM 的论文还没看,不确定我的理解对不对。Schedule Space 模型的自动调优就是尝试不同的优化方法组合,然后在整个策略空间里面搜索哪一种优化效果最好最终就采用哪一种吗?
末尾给的一些测试中,TVM 表现出了相当不错的性能结果。
当然,TVM 还刚刚开始发展,后面还有一大堆问题留待解决。
上一节的 TVM 是一个纯软件栈,这一节就来探索一下用于深度学习的专用硬件。
DL 的疯狂发展对计算硬件也有了越来越高的需求,而且不同应用场景的需求还可能会差很多,例如数据中心和移动终端上面的 AI 设备就完全要往两个极端去考虑。
上面这张图讲的是在 DL 的发展过程中,对数据的尺寸以及存储精度的需求也在不断变化,低精度可以节省空间以及加速运算,但是这也要在硬件本身可以直接支持低精运算的前提下才能有效果(硬件是64位双精的,你要那它跑 int8?那就呵呵了…emm,此处并不是针对某 sw 哈哈)。另外,一些出现的新算法是不是能够用硬件高效实现也很关键,实现不了的话可能还是要选择老算法更好。
不断发展的 DL 算法在实验室里面可以任意瞎搞,效果好就好了,但是如果要应用到实际的生产环境中,那能不能实现/怎么高效实现就非常重要了。
再再另一方面,摩尔定律也逐渐受限,更低纳米制程的工艺难度越来越大,所有这些问题最后都会导向一个终极的解决方案,那就是 DL 专用的计算芯片/硬件了。
下面用 TPU 来举了个栗子:2015 年流片的 ASIC,92 TOPS 的峰值性能,相比 K80 有 30~80倍的性能功耗比。这些数据看着都吓人。
那为什么这么强呢?原因在于它直接硬件支持 8 位的整数 Inference(相比 16 位半精要节能 6~30 倍),大量的乘加运算部件(MACs)以及大量的片上存储。
TPU 主要的峰值运算性能都来自于右边的一大堆矩阵乘单元和累加器。光片上的 Unified Buffer 和 MMU 就占到了整块 TPU 超过 50% 的芯片面积。
这种设计也在于尽可能地提高数据的重用程度,提高计算密集度。
总结一下,像 TPU 这样的 DL 专用的加速器相对 CPU 和 GPU 主要有三个方面的特点:
下面是举了 3 种 Hardware/Software Co-design 的方法:
。。。上面的这三个感觉理解的比较模糊。
再往下才是本节的重点内容——VTA。
TVM 构建的是软件栈,硬件加速器方面,他们也提出了一套开源的 FPGA 加速器设计方案,即 VTA(Versatile Tensor Accelerator)。
VTA 针对不同的带宽、存储和精度需求可以自定义 Tensor Core 的形状、数据类型、内存子系统分配、支持的运算操作等等;对不同类型的代码提供 CISC 或者 RISC 的指令集支持;并且还做了一些 Latency Hiding 的工作。
大致的框架设计如下:
IF 单元从 DRAM 中获取到下一条指令后,会根据类型将其发送到目标部件对应的队列中。
Load 单元负责准备激活函数和计算核的存储资源、提供Micro-Op 的缓存,取出来的数据放在 Load Buffer 中。
计算单元负责根据前面的 Micro-Op 以及准备好的数据执行 ALU 运算或者 GEMM 运算,更新寄存器的内容。
Store 单元负责把前面 Store Buffer 中的寄存器值写回 DRAM。
整体的运行依靠多个任务队列来维护数据依赖关系,基本上是个数据流的设计。(。。。)
VTA 的控制代码部分则依靠 TVM 来生成。
加上 VTA 之后,整个 TVM 的完整架构显得更复杂了:
顶层是各种成熟的深度学习框架,TVM 充当编译器的角色,底层的硬件执行部分则由 VTA 来实现。
还是在 System Components 这层,继续深入分析 DL 训练过程中存在的问题。
回到前面的自动求导部分,这里抛出来一个问题是为什么自动求导是采用往计算图中扩展反向计算的数据计算通路的方式,而不是直接在原来的图上进行反向计算(Backprop in Graph)?
这个 Backprop in Graph 这里也没有更详细的说明,我大概理解成递归返回的那种样子,正向计算是不断递归向下,然后每层递归退出的时候执行反向,完美!
其实说起来,本质上递归的这个顺序也就是数据依赖关系的拓扑序。
原因呢,则是在于内存上。
State-of-art 的很多模型都可能会有资源受限的问题,现在确实很多效果好的新模型都越来越大了,一方面计算量在增长,另一方面内存会成为一个很大的麻烦:CPU 的内存还好一点,如果是用在 GPU 上,目前单块卡的显存最多也只有几十 G 的量级。
先来看一下前向部分的内存使用情况,以下面这几个简单的运算组成的计算图为例:
朴素做法是为每个计算节点都开一块内存,则图越大、计算节点越多,需要的内存量越大。
然后我们发现这个过程中有很多内存是用过之后就不会再用了的,因此更高级一些的方案是配合内存池进行动态内存分配,例如上面 mul
算完之后所占有的内存就不再需要了,因此这块内存可以被内存池回收,然后用在 exp
的计算上。
再有另外一种方式则是静态内存规划,即拿到计算图之后,就按照尽可能重用内存的方案事先分配内存,基本上达到的效果应该要跟上一种内存池的动态分配方案差不多。这种做法有点类似编译器的寄存器重命名的过程。
基本的分配原则也非常简单,只是应用的时候另外要注意如果分配不好是有可能要影响计算的并发性的:
下面的几个内存规划的例子也都比较简单:根据拓扑顺序依次分配和回收;或者先从起点到终点找一条最长路径,把路径上的内存全部设为 Inplace 重用,然后再找别的路径等等。
这里举的内存重用的例子都特别简单,实际上每个计算节点需要的数据尺寸和内存大小都不一定一样,不可能这么简单地就分配好了。
回到前面那个自动求导的两种方案的问题,可以很容易地体会出来往计算图中扩展反向通路的方案非常容易做内存优化:
只需要把 Inplace 和 Reuse 用好即可,在 MXNet 上测试出来的效果也非常好。
深度学习的 BP 算法中存在的最大问题在于反向运算时需要用到前向的一些结果,这事实上就大大地限制了 Reuse 策略的发挥,因为前向算出来的结果总需要找个地方暂存着。
针对一些内存需求特别大的场景,可以采用计算换空间的折中方案:
只保存前向结果的一部分,当反向运算中需要时,再重新从前面开始把缺失的部分重新算一遍,用 25% 的额外计算量可以把整体的内存使用降到原来的开方级别,在某些场景下还是有非常不错的收益的。
实验室的一位师姐之前在 RNN 的优化里面用到过这种方法,节省下内存之后可以跑更大的 BatchSize,最后得到的效果非常好。
再回到前面两种反向方案的讨论中,这一节的最后给了一个特别有趣的点:内存重用的优化方案从某种程度上来看有点像递归中的内存分布!!
666666!
System Components 这层的另外一个方面是并行调度的问题。
用户写好一个计算图,如果框架没有能力把机器上所有的硬件资源全都调动起来那就太浪费了。
关于 DL 中模型并行数据并行这块就不再多提了,在常规的数据并行中,计算和通信之间存在着一个 Gap:
这样一张有着复杂的计算、通信需求的计算图要如何才能比较好地并行起来呢,答案就是我们需要一个自动的调度系统。
首先计算图本身可以很好地描述计算之间的数据依赖关系,那在这个基础上的 Scheduler 设计感觉其实也没啥好说的,基本上都是很常规能够想到的解决方案。TensorFlow、MXNet 等等的基础设计原则都是这样,只是实现上可能有所不同。
同样是队列的调度方案,这里的这种是以每个变量为单位有一个自己的队列,TensorFlow 中是线程池中的每个线程会有个自己的队列。
这一节把目标放在上一节调度图中的 Synchronization 部分。
大量的篇幅是对 Allreduce 的讨论,也没啥好说的,跳过跳过…
Parameter Server - Worker 架构的同步异步,也没啥好说的…
Emm…这两节不是我偷懒,主要是内容比较基础,跟着 PPT 就好了,没什么特别值得注意的。
除去后面没有资料的几个 Guest Talk 以外,这节算是课程内容的最后一部分了,主要讲的是实验室的成果上线进行实用的过程中可能会有的问题,例如:
在下面这个视频应用中:
终端设备的采集、处理、数据传输等各个不同有码率、功耗等限制,云端提供服务的部分也有带宽和花费开销的限制,事实上从 Workload 到 Budget 之间也存在一个巨大的 Gap。
下面的内容主要从模型压缩和服务系统两个方面来介绍。
这部分是介绍如何对一个网络模型进行压缩。
首先是矩阵/向量的低秩分解,可以应用在全连接层和卷积层中,能够有效地降低整体的计算量和存储量。
这块还是自己的数学知识比较缺乏,暂时不往下细看了。
然后是网络剪枝:训好一个网络之后,通过一个 01 的 Mask 把参数中的某些部分置为 0,再重新训练达到之前相同的预测精度,不断重复以上过程并且逐渐增加 Mask 中 0 的比例,最后就可以得到一个想要的剪枝结果。
权值共享:对参数矩阵进行重新采样,把实际值存在一张表中,然后参数矩阵改成存储对应实际值在表中的索引。这也是尽可能地减少存储冗余。
低比特量化:这个也比较容易理解,就是降低数据类型的存储精度,32 位单精降到半精、int8 甚至是二进制的 01 值。预测结果可能会有一定的精度损失,但是在可以接受的精度损失范围内可以大大节省参数的存储量,并且配合上低精度的硬件预算部件也能大大加快运算速度。
知识蒸馏:用一个训练好的大模型来训练一个小模型。
Knowledge Distillation 这块感觉挺神奇,但是还没细看,不是很理解。
这里还给了一些参考的论文资料:
一个比较好的服务系统需要达到几个目标:
然后举了个叫 Nexus 的系统为例,后面就不细看了。
]]>主要是整理几个多线程小例子,读写锁和定时器队列是面试的时候考到的,线程池是自己很早就想写着玩的(挖了好久的坑一直没填,内存池也是……)。
整理出来的代码会放在这里:
参考了一下别人的实现方案:
总体上还是比较容易理解的,每个工作线程队里。
当任务队列为空时,每一条工作线程在一个条件变量上等待;当任务队列非空时,则空闲线程直接从任务队列中取出封装好的 std::function
来执行。
1 | void ThreadRunner::operator()() |
定时器即让一个注册好的回调函数在某个设定时间到达之后执行。
话说现在应该有很多系统级的 API 都能够提供这种延迟执行或者阻塞当前线程一段时间的功能,但是当定时的任务比较多的时候,就有可能还是自己写一个定时器队列来维护性能更好点了(主要是不确定到时候所用的定时 API 是怎么实现的,例如如果所有的这些定时任务都用 epoll_wait 的超时触发做,那 epoll 的底层实际上是用红黑树来维护的,猜想效率应该不会太差)。
另外还有一个定时精度的问题,不同的 API 可能能够提供的精度支持也是不一样的。
当时面试的时候被问到这个是远程共享写代码,由于之前没接触过这种 case,一开始不是很明白面试官到底想考我什么,就觉得定时器这种东西不是现有的 API 一大堆嘛,然后一脸懵逼。
最后写了一个线程每隔一个最小的周期(Tick)处理一次队列中的任务,把到达时间的任务取出来执行,任务列表就用优先队列或者说小根堆来实现,保证每次都能够高效地把设定时间最靠前的任务取出来。
网上查了下,基本上比较不错的思路有两种,一种就是小根堆的实现,另外一种是 Linux 内核中采用的时间轮的定时器方法。
由于小根堆的堆顶元素是最近触发的一个任务,因此还可以动态改变下一次处理需要等待的 Tick 时间来节省 CPU 资源的消耗。但是感觉这种方式会不会影响定时精度?以及这样似乎没有办法处理“在堆顶元素之前又插入了一个更近的任务”这种情况,所以最后想想还是觉得固定 Tick 比较稳妥。
然后就是重点要来学习一下时间轮(Timing-Wheel)了。
这个思路其实很简单,想象有一个钟表盘,上面有一圈时间刻度(从 1 到 N),有一个指针每隔 Tick 时间转动一个刻度,而所有的任务也是按照等待时间除以 Tick 分布在每个刻度上面,当指针扫到某个刻度时,即遍历一遍这个刻度上的任务列表,把到时间的拿出来执行即可。
当 N 足够大时,显然这个定时器的各个操作的复杂度都是 O(1)
的。
但是如果 N 过大,则这个时间轮的内存消耗将会比较大,因此又有了多级时间轮的优化思路:以钟表上的时针、分针、秒针为例,时针分为 12 格,分针、秒针分别为 60 格,这样就有一共 132 个格子,每个格子都是一条任务链表。一分钟以内的任务被加到秒针的格子中,一小时以内的任务加到分针的格子中,更久的任务加到时针的格子中。秒针每个 Tick 移动一次,每次移动后直接把当前格子中的所有任务取出来执行;分针则是每 60 个 Tick 移动一次,每次移动把当前格子中的任务下放到秒针轮对应的格子中;时针也是类似。
更厉害的优化思路可见下面这篇:
最后这个读写锁是面试完之后让半小时内写完,然后邮件回面试官交作业……
结果读写锁的逻辑很快写完之后 main 函数的测试 case 却调 bug 调了好久……狂汗
读写锁从功能上来看也可以叫做共享锁,相对普通的互斥锁来说,它会有三种状态:未锁定、读锁定和写锁定。
读写锁的读操作是可以共享的,即处于读锁定的时候另外一个线程还可以继续对其施加读锁定,而写锁定跟一般的锁一样是独占的。
因此基本的思路就是用一个读计数来记录读操作,当读写锁处于解锁或者读锁定状态时,获取读锁即增加读的锁定计数,而此时想要获取写锁就需要等到所有的读操作都释放之后才行。
写操作可以用一个布尔变量来记录,当读写锁处于写锁定状态时,继续获取读锁或者写锁都需要等前一次的写操作释放。
下面分别是简单地直接用一个锁的版本:
1 | void rwlock::getRead() |
以及用条件变量实现的版本:
1 | void ReadWriteLock::getRead() |
协程。
严格地说这个跟多线程关系不大,其实这个东西本身大概类似一种用户态的线程,然后关键在于这个线程的运行和调度都是要靠用户自己来管理的。
很多语言里面都自带协程的支持,例如,举一个别人的例子,一个 js 写的生成斐波那契数列的函数:
1 | function fibonacii() { |
用协程的写法则是:
1 | function *fibonaciiUsingYield() { |
相当于每一次这个函数都断点在 yield 对象上,当调用 next 函数时,函数向前执行一轮。
那么如何在本身不支持协程的 C/C++ 里面把这种特性用上呢?
答案是手动实现一个状态机……
详见轮子哥的考不上三本系列:
下面是按照这种方式写的一个简单例子:
1 | /* *********************************************** |
Result:
1 | 0 |
一开始是分了 Basic 和 Project 两个大标题,主要还是想写三个面试遇到的多线程例子的实现,然后发现前面 Basic 部分的东西写的有点多了……篇幅有点长,然后想了想还是把 Project 部分另外开一篇吧。
这大概就是计划赶不上变化?……就是这么任性。
这里来理一下 C 和 C++ 里面与多线程相关的 API。
在 C11 之前,C 标准库里面本身不带线程支持,通常需要用 UNIX/LINUX 系统库里面的 clone、fork 之类或者用 POSIX 标准的 pthread 库来实现。虽然这些东西都包含在 glibc 的库里面了,但是事实上不算 C 语言本身的标准,所以在 cppreference 里面是搜不到的(当时奇怪了好久)。
pthread 底层用来创建线程的实现应该也是调的 clone、fork 这些系统调用来完成的。
至于 mutex 锁、信号量这些原语,从 Linux 2.6.x 版本内核之后都是通过 FUTEX 系统调用来实现的,具体以后有空再看。
另外需要对线程和进程这两个概念作一下强调,Linux 中本身不分进程和线程统称为 task,后来用于区分这两个概念主要是看一个 task 所拥有的资源情况。进程有自己的内存空间,线程共享父进程的内存空间。
先来看一下三个系统接口。
系统调用 | 描述 |
---|---|
fork() | 创建父进程的完整副本,复制父进程的资源,包括所有内存的内容。写时复制。 |
vfork() | 创建的子进程与父进程共享数据段,并且由 vfork 创建的子进程先运行,父进程在子进程返回前保持阻塞。 |
clone() | 可以对创建的子进程做更多控制,启动的时候指定需要执行的函数内容。 |
需要注意的是这几个 API 在 Mingw 里面是用不了的。
Bash on Windows 大法好~~~
fork() 的使用方式特别简单:
1 | /* *********************************************** |
结果:
1 | # jcf @ J-CF-MSI in /mnt/c/Users/jcf/Desktop/multithread [15:45:20] |
fork() 创建产生的子进程除了这个函数返回的 pid 值以外,与父进程完全一致,父子进程也都会从 fork() 函数返回之后继续向下执行相同的内容。另外,如果把本地的变量值的地址打出来,会发现它们的虚拟地址也都是一样的。
这里用了一个写时复制的技术,fork 出来的子进程一开始直接用的是父进程的内存空间,所有内容包括虚拟地址都是跟父进程完全一致的,直到发生了数据更改,才会在物理地址空间中作一个新的映射。这也就提高了创建子进程的效率,因为分配新的页表也会是一件比较耗时的工作。
然后是 vfork(),把上面这段代码中的 fork 直接改成 vfork 之后会得到这样的运行结果:
1 | # jcf @ J-CF-MSI in /mnt/c/Users/jcf/Desktop/multithread [15:46:37] C:134 |
vfork 调用之后,父进程会被阻塞,所以可以看到不同于之前的情况,这里永远都是 son 这句先输出,执行完毕之后才会轮到父进程执行。
那为什么会炸了呢……
在 fpid 等于 0 的分支末尾加上 exit(0);
之后程序就能够正常执行了:
1 | # jcf @ J-CF-MSI in /mnt/c/Users/jcf/Desktop/multithread [15:47:10] C:134 |
原因是 vfork 不同于 fork 的一点是,创建出来的子进程直接共享父进程的数据段,当子进程跑完之后,他会像一个正常的进程一样对自己的栈空间等等做回收,则之后当父进程开始执行的时候自身的内存数据就被子进程破坏掉了一部分,这也是为什么前面第一次父进程的 aaa 输出来的结果是不对的,而加上 exit(0);
之后,父进程可以正常输出 2。
话说网上说 vfork 出来的子进程如果用 return 来返回的话会出现很奇怪的 bug,不过我这里测试的时候没有见到,可能跟 gcc 和系统库的版本有关系。
那么为什么会有 vfork 这个看上去有点问题的实现呢?
这就需要提到另外一个系统接口 exec 了。exec 的作用是拿另外一个程序的代码来替换到当前进程的正文、数据和堆栈,简单地说就是用来启动一个新程序。
exec 的接口实际上是一套,一共 6 个函数,具体的这里先不展开了。
vfork 自身设计的目标是为 exec 服务的,当需要创建一个新进程来执行一段完全不同的代码时,vfork 直接共享父进程地址空间的做法是开销最小的,即保证先有一个子进程,然后调用 exec 来载入一段新的代码并且创建自己的独立地址空间,在子进程开始新程序或者退出之前,内核保证父进程一直处于阻塞状态。
1 | /* *********************************************** |
结果大概是这样:
1 | # jcf @ J-CF-MSI in /mnt/c/Users/jcf/Desktop/multithread [16:07:41] C:1 |
最后是 clone,先看下 man 里面的定义:
1 | /* Prototype for the glibc wrapper function */ |
fn 是需要执行的函数指针,即 clone 出来的子进程需要执行的函数内容。
child_stack 就明显是给子进程分配的系统堆栈空间的位置。
flags 用于描述子进程需要从父进程中继承哪些部分的内容,因此通过这个值可以控制产生进程、线程、甚至非父子关系而是兄弟关系的进程等等,功能强大。
后面的就是传入新进程中的参数了
测试代码:
1 | /* *********************************************** |
上面这份代码中有两个地方需要额外注意一下:
#define _GNU_SOURCE
的宏,表明下文代码不可移植,可能会用到一些非 GNU 标准的内容(例如 clone)。结果:
1 | # jcf @ J-CF-MSI in /mnt/c/Users/jcf/Desktop/multithread [16:46:24] |
这三个接口的最底层涉及到的都是 do_fork()
这个调用,只是传入的参数不同,clone 可以认为就是个 do_fork()
的 API 外衣。
pthreads 的全称应该是 POSIX Threads,是 POSIX 的线程标准,它定义了一套 C 语言标准的线程控制 API,由一个 <pthread.h> 的头文件和一个线程库来实现,主要包含了:线程管理、互斥锁、条件变量、线程同步等等这些线程操作的 API。
1 | /* *********************************************** |
结果:
1 | [Running] cd "c:\Users\jcf\Desktop\multithread\" && gcc test.c -o test && "c:\Users\jcf\Desktop\multithread\"test |
话说从使用方式上来看,pthread_create()
的接口跟 clone 就特别像,大概率底层实现就是用 clone 做的,不过传入的线程函数的格式不太一样。
pthread 也提供了互斥锁和条件变量这些结构: pthread_mutex_t
、pthread_cond_t
。具体的使用方式跟后面的差别不大,下文再整理。
头文件 <semaphore.h> 中也提供了信号量的支持。
C11 之后,标准库里面提供了线程支持,包含在头文件 <threads.h> 中(这下可以在 cppreference 里面查到啦)。
基本也就是线程创建、等待、互斥、条件变量等等的支持,感觉看上去跟 pthread 基本一致。
而且不知道为什么,虽然在标准库里面查到了这个库,但是似乎用的人特别少。
那 C 的部分还是用 pthread 吧。
从知乎讨论上面来看,大家对 C++11 的 thread 意见还是比较大的。
C++ 这部分……其实涉及到的东西非常多,需要一堆不同的库联合起来一起用:
线程支持库 <thread> 提供了线程创建、调度、等待等等一系列管理操作;
互斥库 <mutex> 提供了基本的互斥量 mutex,RAII 的锁控制方式 lock_guard 和 unique_lock 等等;
条件变量库
异步支持库 <future> 提供了像 promise、future、async 等等这种异步语义(在 Nodejs 里面用过,之前还真没听说 C++ 里面还带这种玩意);
原子操作库 <atomic> 提供了一系列与原子操作相关的支持;
另外 C++ 中任何可以被调用的东西都是函数对象,前面用来创建线程用的目标函数也需要由函数对象库 <functional> 来管理,std::bind、std::invoke 等等这些管理参数调用,也可以用 lambda 表达式等等。
话说 <thread> 库是不是也还是 pthread 的封装???
有关 std::function
和 Lambda 表达式,很早之前稍微有记过一些:
C++11 中的基础互斥锁结构是 std::mutex
,用法应该基本跟 pthread 的一样。<mutex> 中额外还提供了两个符合 RAII 标准的锁控制封装,以更加异常安全的方式来管理互斥锁。
std::lock_guard
就是个简单的互斥封装容器,构造时锁定给定的锁,然后析构的时候自动释放。事实上它能操作的锁不一定只限于 std::mutex
,任何有 lock()
和 un_lock()
两个成员函数的对象都可以。
std::unique_lock
功能更多一点。构造时可选地对传入的锁上锁(也可以选择不锁),析构时自动释放。并且同时它还提供了 lock()
、try_lock()
、unlock()
等等这些成员函数,使用起来就更灵活了,除了离开作用域自动析构释放这一点之外,在作用域中还可以手动控制加锁解锁。
条件变量需要结合互斥锁一起使用,这里的 std::condition_variable
尤其在 wait 的时候必须配合 std::unique_lock
来用。
条件变量的核心操作是等待(wait)和唤醒(notify),通常情况下,需要在条件变量上等待的线程需要:
std::unique_lock<std::mutex>
锁(重要!!);wait()
、wait_for()
或者 wait_until()
,这三个函数需要把前面的 unique_lock 作为参数传入,执行时将原子地释放传入的 unique_lock,然后挂起当前线程进入等待状态;在条件变量上执行唤醒操作的线程需要:
std::unique_lock<std::mutex>
,其他方式管理也行;notify_one()
或者 notify_all()
,则其他处于等待状态的线程会被唤醒。文档中对 notify_one
的描述是会唤醒一个等待的线程,但是并不一定是哪一个,跟进入 wait 状态的线程的先后顺序无关,notify_all
则是唤醒当前正处于等待状态中的所有线程。
由于 wait 操作本身自带对 unique_lock 的加锁解锁操作,因此 notify 这边也需要注意前面这个 mutex 锁的状况,必须保证 wait 在调用的时候 mutex 是锁上的,然后 wait 被唤醒的时候锁是开着的。如果 wait 唤醒时试图获取锁失败则会被阻塞在等互斥锁的状态,这个下面有测试。
还有一个重要的注意点是多个线程对某个条件变量的 notify 和 wait 操作可以看成是对一个原子变量的顺序操作,这就意味着如果先调用 notify,再调用 wait,则 wait 是不会从唤醒中恢复的。
看上去这个注意点很正常啊,正常就应该是这样的啊。但是实际多线程操作中非常容易出现:自认为 notify 会发生在 wait 以后,实际执行却不是,然后导致死锁的 bug。
Talk is cheap, show me the code!
1 | /* *********************************************** |
做了个小测试,对于上面打上标记的 4 块代码:
cv_m.unlock()
删掉,则结果就会死锁!!!原因恰恰是在于 wait 被唤醒的时候要首先试图获得锁,由于 cv_m 这时候是锁着的,然后 wait 线程就被阻塞在获取互斥锁的状态了。所以看上去最好的方式是 notify 的时候根本就别管锁?如果不上锁,不就没这么多麻烦了吗……事实上,更好的写法应该是这样的:
1 | /* *********************************************** |
三种 wait 函数均有一种附加条件的多参数调用方式,等同于在一个 while 循环中调用单参数版的 wait 函数。另外用一个 ready 变量标识等待情况,在 notify 时,cv_m 这个锁实际上是用于保护这个 ready 变量用的。用这种方式写则无论 notify 代码块发生在 wait 线程前还是发生在之后,wait 线程均会正常返回了。
这个头文件里面的内容很有意思,核心的类主要是 std::promise
、std::packaged_task
、std::future
这几个。
std::future
是一个对异步操作结果的封装类,一般需要配合 std::async
、std::packaged_task
和 std::promise
的异步操作使用。
1 | /* *********************************************** |
std::async
的作用是在另外一个线程中异步(默认操作,也可以设置在调用线程中同步执行)地执行给定的函数,函数执行的返回值则是由一个 std::future
对象来接收。
std::packaged_task
则是一个可调用目标(函数、Lambda表达式或者其他函数对象,即 std::function
的对象)的类模板封装,这个封装主要也就是把函数对象的执行和返回值给分开。package_task 对象可以直接加参数调用,或者放在另外一个线程中调用,结果会在函数体执行完毕之后存到对应的 future 结构中。
最后是 std::promise
,这个对象感觉有点像 placeholder 占位符的作用。promise 通过 set_value()
来提供数据,在 future 绑定的 promise 准备完成之前,future 的 get()
会阻塞所在的线程。
其中大概尤其与 Cache 有关的内容比较有代表性,于是准备根据几次回忆出来的自己和小伙伴遇到的面试题好好理一理 Cache 这块的内容。
上图是一张最基础也最常见的存储器层次结构图,表达的是计算机系统中各级存储器的速度、容量、价格等等的金字塔关系。
Cache 这个思想本身特别简单,利用的核心原理就是数据的局部性,即把最常用到的东西放在最容易拿到的地方,这种局部性即包含了数据的空间局部性也包含了使用数据的时间局部性,但实际用起来效果却是非常地好,并且广泛应用在各种不同的场合中,例如:CPU 中的 Cache 用来加速对主存数据的访问,TLB 可以看成是对虚拟地址-物理地址转换的页表的 Cache,分布式环境中有的时候也会做个本地数据缓存,也是 Cache 的思想。
其实《硬软件接口》里面已经有介绍过 Cache 了,前面也有记过:
编程的时候 Cache 对程序员透明,写代码的人直接看到的是内存(……更准确地说看到的是虚拟的地址空间),但是 CPU 实际执行的时候访问的数据全部都是从 Cache 里面来的。每次访问一块新的内存数据时,首先检查 Cache 中是否存在,有就返回,没有就触发一次 Cache Miss,从主存把数据 load 到 Cache 中以后再返回。
然而 Cache 这个简单的想法在实际实现的时候可能会复杂的多,例如 Cache 分级,然后 Cache 还有多种数据和地址的映射关系,什么时候更新写回,地址冲突的时候怎么替换等等很多麻烦的问题都需要在实际实现的时候考虑。
这里还有两篇知乎上的小白科普文:
面试的时候我和我的小伙伴们在 Cache 上踩了很多坑,虽然平时工作中很少有遇到需要考虑的这么深的,也就一般很少想到 Cache 这一层的工作内容了,但是终究说起来还是自己的基础没打扎实。
L1 Cache 标记数据块用的是虚拟地址还是物理地址?
往后的 Cache 里面肯定用的是物理地址,这个没有任何悬念,关键在这个最靠近 CPU 的 L1 Cache 上。
这个问题一开始我脑子里冒出来的是当时看《硬软件接口》虚拟内存那章时画的那张图:CPU 先拿虚拟地址去查 TLB,然后再找 Cache,那妥妥的就是物理地址了。
后来一想这本书上是以一个假想的简单数据通路结构为例进行介绍,可能细节上后来有变动,然后去翻了下《量化》,里面没有详细讲,但是看图感觉像是用的虚拟地址。
再后来去网上搜这个问题的时候,说物理地址的也有,说虚拟地址的也有,还各有各的理由,于是就更迷了。
直到我在这个专栏文章中看到了这张图:
MMU 是 CPU 中用于虚拟地址-物理地址转换的工作单元(主要是页表、TLB 甚至多级 TLB 等等组成),这张图中根据芯片实际实现时 MMU 和 Cache 的位置关系,则很明显 L1 Cache 可能用虚拟地址也可能用物理地址了。
那我们来考虑一下这两种实现的区别:
Logical Cache/ Virtual Cache
感觉这种实现访问速度应该更快一点,因为不需要经过地址翻译就能够直接在 Cache 中搜索数据了。
问题是每个进程的虚拟地址空间都是独立的,如果有多个进程在同一个核上切换,则每次切换的时候都需要把 Cache 刷掉(话说我其实不太清楚这个代价会有多大,也不知道这个会不会成为制约性能的一个原因),当然,也有可能有多个虚拟地址对应到同一个物理地址上的这种情况,那这种情况要怎么处理就需要考虑更多东西了。
……真麻烦啊。
Physical Cache
查 Cache 前不管怎么样先把地址转换做了,这样前面 Logical Cache 存在的问题也就不存在了,而且也不需要一切换进程就刷一遍整个 L1 Cache,如果有多个虚拟地址对应到同一个物理地址上也无所谓。
问题大概就在于 MMU 的地址转换上,如果 TLB Miss 了,查页表还是要访问到主存上,那再查 Cache 就需要花更多的时间了。
CPU_cache 的维基百科页对地址转换这块也有大段篇幅的介绍。Cache 受地址转换影响在 Latency、Aliasing、Granularity 等几个方面都需要考虑到,根据用虚拟地址还是物理地址进行查找和标记,常见的也有下面四种 Cache 实现方式:
Physically Indexed, Physically Tagged(PIPT)
对应前面的 Physical Cache,没啥大毛病,就是慢
Virtually Indexed, Virtually Tagged(VIVT)
对应前面的 Logical Cache/ Virtual Cache,带来的麻烦问题很多
Virtually Indexed, Physically Tagged(VIPT)
可能是最好的一种实现方式了,目前市场上的一些现代处理器应该大多都是基于这种方式或者在这个基础上优化的。
由于查虚拟地址的索引和查 TLB 可以同时进行,这种方式的延迟会比 PIPT 低很多,但是实际数据还是要等到 MMU 把物理地址算出来之后对比 tag 才能确定。
另一方面由于用了物理地址作为 tag,VIVT 中可能会有的虚拟地址冲突的问题也解决了。
Physically Indexed, Virtually Tagged(PIVT)
这个…大概只会出现在文献中,实际实现的时候会集合 PIPT 和 VIVT 所有的缺点。
另外,对 Cache 数据的查找和标记还需要考虑到全相联、组相联等不同的映射关系的实现。
对 Context Switch 和 Cache 访存需要大概花费多少个 Cycle 有概念吗?
并没有……卒。
看一下网上找到的答案:
同样都是用 SRAM 做的,为什么会有速度差异呢?原因大概有以下几个:
容量大小
显而易见的是,Cache 的容量会随着级数的增加而增大。由于 Cache 需要做到随机访存,即能够直接访问到存储器的任意一个位置,在制程和设计完全一致的情况下,容量越大就需要花费更多的时间来做到随机访存(延迟跟容量的开方大致成正比)。
芯片上与 CPU 的距离
在芯片面积有限的情况下,L1 Cache 会被放置在离 CPU 核心非常近的地方,而 L2 Cache 就只能放到边缘位置了。
L1 中的指令 Cache 会在 Fetch 单元附近,数据 Cache 会在 Load/ Store 单元附近,L2 Cache 就要在 CPU 流水线外面了,L3 更远要在核外了。
具体实现的差异
L1 和 L2 在设计时的侧重点会有所区别,L1 更注重速度,而 L2 要在 L1 Miss 之后才发挥作用,因此更注重节能和容量。
查询一个地址时,L1 Cache 会把多个 Cache Line 的 tag 和数据全部取出来,然后再比较 tag 看哪一个命中或者都没命中。
而 L2 Cache 虽然也是 N 路组相联,但是比较时会先取 tag,当找到命中的之后再去把对应的数据取出来。
L3 做在核外,通常是多个核共享,因此还需要额外考虑一致性等等更多的东西。
关于 Context Switch 这点,进程和线程的切换其实都要涉及到上下文的切换。
进程切换时由于虚拟地址空间不同,因此需要切换页面映射、刷新 TLB 等等,如果是线程切换则开销会小很多。
这个问题直接有别人解答了,感觉说的也还算清楚:
当时被问到这个问题的时候有点懵逼,不确定面试官想表达的是什么意思,我回答 “Index、tag、cache line data” 的时候被否决了……然后后来就没有答上来了。
如果指的不是全相联、组相联这种实现的话,可能想问的是硬件实现?
我对 Cache 硬件结构的了解就只是知道它是用 SRAM 做的,其他的就不懂了,卒。
这一块能够找到的资料也比较模糊,最后是从王齐的《浅谈 Cache Memory》中找到了比较靠谱的答案:
组相联方式组织的 Cache 会分成两个部分,Tag 部分和数据部分分开存放,例如一个 8 路组相联的 Cache 结构是这样的:
左边是地址 Tag 以及当前 Cache Line 的状态,右边是实际存放的数据。
这两类字段由于功能和特性不同,会使用两种不同类型的存储器来存放。Tag 阵列多使用 CAM(Content Addressable Memory)来存放,以利于并行查找,数据字段用的才是多端口多 Bank 的 SRAM。一般说的 Cache 大小也都指的是 SRAM 数据块的大小,CAM 这部分不包含在内。
CAM 对应的应该是前面结构图中一个 Set 中的多个 Way 的结构。首先根据需要访问的虚拟地址确定 Index 找到在哪个 Set 中,然后对该 Set 中的多条记录并行进行 Tag 的比对。
CAM 的基本结构如下:
上图的 CAM 有 3 个 Word,分别对应一条横向的 ML(Match Line),每一个 Word 由 4 个 Bits (CAM Cell)组成。在每一列中,Bits 分别与两个 SL(Search Line)对应。
使用 CAM 进行查找时,首先需要把需要搜索的目标(Tag)放入 Search Data Register/Drivers 中,分解成多个 Bits 之后,通过 SL 发送到所有的 CAM Cell 中,每个 Cell 的 Hit/Miss 信息会向右传递给各自的 ML,最终 ML 汇总得到自己的 Hit/Miss 情况,这样就能够确定下来在当前的多个 Way 里面有没有命中的数据了。
后面更细节的就跳过了。
假如你是一个 Intel 的工程师,有一天你的竞争对手 AMD 推出了一款新的 CPU,然后你想要知道关于其中的 Cache 的信息,要怎么做?没有其他任何的资料,只能通过实际测试的方式。
这个问题是紧接着上一个的,由于我被问到上一个问题的时候已经懵逼了,这题基本上完全没答上来,其实后来想想应该至少能把 Cache 的大小测出来,当时回答的时候真的是表现得太差了。
参考:
根据第一篇的程序稍微改了一下:
1 | /* *********************************************** |
结果:
1 | Stride: 1, Line Size: 4 Bytes, Average Cost: 0.002586 us |
我笔记本是 i7 7700HQ,在网上可以查到详细的 Cache 信息
Cache: | L1 data | L1 instruction | L2 | L3 |
---|---|---|---|---|
Size: | 4 x 32 KB | 4 x 32 KB | 4 x 256 KB | 6 MB |
Associativity: | 8-way set associative | 8-way set associative | 4-way set associative | 12-way set associative |
Line size: | 64 bytes | 64 bytes | 64 bytes | 64 bytes |
Comments: | Direct-mapped | Direct-mapped | Non-inclusive Direct-mapped | Inclusive Shared between all cores |
Cache Line 比较好测,上面的代码中可以明显看到从 64 Bytes 开始,平均访问时间出现跳变。
L1、L2、L3 各自的大小……说实话我觉得这样测效果并不好,变化倒是确实有,就是其中可能还有很多误差成分在。
CSAPP 上面讲 Cache 的那章用的是类似的方法,根据这样的测试结果可以画出一张“存储器山”的图。
这个问的主要应该就是 Cache 的一致性、写回、写直达等等这些方面的内容。
另外还有一个叫做伪共享(False Sharing)的问题:
多核 CPU 通常都是 L1、L2 每个核独立,共享 L3。如上图这种情况,两个核实际操作的数据是独立的,但是它们恰好在一个 Cache Line 里面,则其中一个作了修改之后,另一个的 Cache Line 也会跟着失效,引起了本来不必要的效率问题。
]]>顺手来理一下 LCA 的板子。
前面在 【Tarjan 大佬的算法们】 中提到过他的离线 LCA 算法,就从这里开始。
原始算法具体见前文吧,还引用了别人的一个链接,里面有个动画演示的挺清楚的。
离线 LCA 的关键在于 dfs 遍历整个树的过程中,对于被询问的点对 (v, w)
,要求其中一个点要先被访问过,然后再遍历到另一个点则可以通过查询前一个点所属的并查集来找到它们的 LCA,例如 v 先被访问,接下来遍历到 w,则它们的 LCA 是 getfather(v)
,反正若 w 先被访问,接下来遍历到 v,则它们的 LCA 是 getfather(w)
。查询结束再将 v 和 w 合并到同一个集合中(向树上父节点的方向合并)。
Tarjan 论文中的原始算法的树节点应该是保证有序的,然后通过先处理每一对 (vi, wi)
,使得 vi < wi
来保证每次访问到 wi 时,它对应的 vi 都被访问过。
在通用的树中就只能在每一层 dfs 中询问跟当前层搜的父节点相关的边来查找这种对应关系了。
记两道模版题:
【题意】
给出一棵树,询问树上的两个点,要求回答两个点之间的最近距离。
【分析】
树首先是任意给的,假定我们建出来的树根是 R,则对于一次询问的点对 (a, b)
,可以在这棵树上找到它们的 LCA,记为 c。则 a 和 b 之间的距离就是:
$$Dis(R, a) - Dis(R, c) + Dis(R, b) - Dis(R, c) = Dis(R, a) + Dis(R, b) - 2*Dis(R, c)$$
在找 LCA 的 dfs 过程中顺手把每个节点到根节点的距离记出来即可。
【题意】
跟前面类似,加了个条件是可能有多棵树。
【分析】
每对点对是否在同一棵树上可以在读取数据的时候直接用并查集判断联通性记下来,之后对所有未访问过的点做 Tarjan LCA 即可。
【模版】
1 | /* *********************************************** |
dfs 过程中是先询问边还是先进行下一层的 dfs 不影响结果。
话说虽然前向星写习惯了,不过看人家写的数组邻接链表也挺好看的。
前面 【HDU 3966】 用过一次 LCT 和树链剖分来维护树上一条路径上的点值,翻回去看的时候……妈呀,以前题解怎么都没好好写,自己都看不懂了,顺便拿到这里重新理一下思路。
LCT 的核心思路是用 splay 森林来维护 Preferred Path,核心操作是调整 Preferred Path 的 Access 操作,对于 LCT 来说,找 a、b 两点的 LCA 只需要两步:
Access(a)
,此时 a 所在的 splay 树即为从根节点到 a 点所构成的 Preferred Path;Access(b)
,找到从根节点到 b 点所构成的 Preferred Path 之后,此时 Splay 的根即为 a 和 b 的 LCA。当然也有可能两次 Access 操作之后的 Preferred Path 的根不是同一个,那就意味着这两个点是分属于两棵不同的树,不存在 LCA,特判一下就好了。
所以上面那题的 LCT 模版是:
1 | /* *********************************************** |
从最后的求解过程上来看,应该比离线的要更顺一点,毕竟在线算法不用多考虑记录输出顺序什么的。
还是看前面 【HDU 3966】 的时候……发现树链剖分忘得差不多了,顺手再补一个树剖找 LCA 的模版吧。
树链剖分的原理是把整棵树按照一条一条树边组成的链划开,每条链相当于一个区间,那对树上的某条路径的操作就成了对树上的一条或者多条树链的操作了,具体维护区间的部分可以用树状数组啊、线段树啊什么的来做。
划分树链的基本思路是对整棵树进行轻重边划分,定义 size(x)
是以 x 为根的子树的节点个数,这里的重边就是 x 和它节点数更多的一个子节点(size 更大的一个子节点)组成的边。
这里会有两个性质:
(father, son)
是一条轻边,则 size(son) <= size(father)/2
O(logn)
划分好轻重边之后,从根节点开始把所有连着的重边连起来,就成了一条重链,重链就是前面说的需要从树上剖出来的树链啦。
这个过程可以用一次 bfs 加两次队列遍历来完成:
那么 LCA 要怎么找呢?
在树链结构上找 a、b 两点的 LCA 即走完各自所在的树链,沿着它的轻边向上找,直到找到各自路径上的两个点在同一条树链上,那么深度较浅的那个点就是 a、b 两点的 LCA 了。
由于从树根到某一个点的路径上的轻边的个数不会超过 O(logn)
这个性质,每次找 LCA 的复杂度是 O(logn)
。
模版如下:
1 | /* *********************************************** |
话说从 HDU 提交的结果上来看,树链剖分是最快的,LCT 次之,Tarjan 离线最慢。
]]>严格上来说本篇不应该算在拆包里面,因为记的是 TF 团队最近发的一篇论文里面的东西。
前面拆包的第二篇记过关于 TensorFlow 中的数据流模型实现,实际上这套数据流模型已经是非常完备的,只是目前大家用 Python 搭出来的简单网络形式还很难把它的真正潜力发挥出来。
正当我们往这个方向做的时候,得,Google 发论文了。
这篇 Dynamic Control Flow in Large-Scale Machine Learning 发表在 EuroSys 18 上,系统结构方向的 B 类会议。
其实文中所提到的几乎所有内容都是 TensorFlow 原有的,或者说 TensorFlow 当初设计架构的时候就已经考虑到了未来这种使用方式的需求,这篇文章只是整理了一下这部分的设计思路(内容大部分跟以前发的控制流白皮书是一致的,见 TF 拆包第二篇),然后做了一定的测试,从实践上证明给大家看这样做是有效的。
首先提了一下深度学习中对控制流的需求,主要是像 RNN、MoEs 这样的任务中会明确地需要一些控制流的支持。但从更宏观的角度来看,使用动态控制流对任何应用都是有用的,理论上可以比较好地地把计算和通信部分给 overlap 开,尤其对提高异构系统(CPU、GPU、TPU等等)的计算效率是有很大的好处的。
目前常见的一些机器学习框架基本上都是用数据流图的方式来组织计算。
关于如何实现数据流的控制部分,主要有两种方式:
除了 TensorFlow 以外,别的框架似乎都很少用数据流这个词来指代自己的设计,可能原因就在这里?其他框架虽然整体计算还是以数据流图的形式做的,但并不是真正用一套数据流的运行时去支撑的。
用 TensorFlow 来举例,方式 2 的写法通常是:
1 | python for i in range(xxx): |
多轮控制是写在 Python 层的代码中,每一次循环只跑训练的一步。
恐怕我们见到的大部分 TF 代码都是这个样子的吧。
方式 1 则是控制部分已经是图的一部分了,那最后我们只要 sess.run(total_train_step)
一次,就能够达到跟前面一样的训练效果。
单一的计算图更便于进行全图的优化,且这种实现能保证整个计算过程都停留在运行时里面(而不是像原先那样,跑一轮进退一次运行时,再跑一轮再进退一次运行时),减少很多不必要的开销。
数据流运行时的特性是一旦某个 op 的依赖都满足了,它就马上可以被调度执行了,在 out-of-graph 方式中,这种数据流的调度粒度只限定在一次 step 中,而 in-graph 方式甚至能把并行性扩展到多次 step 间,这样就能够最大程度地挖掘数据流异步、并行的能力了。
最初的 TensorFlow 白皮书中也有介绍过关于数据流部分的实现,但是并没有给出详细的设计方案以及测试结果,这篇文章就是把这部分补上。
总的来说,本文的内容包括:
话说前两部分都是 TensorFlow 原有的。。。。。。
2、3、4、5 章的大部分内容与 TensorFlow 拆包(二):TF 的数据流模型实现 中记录的类似,就不多重复了。
需要额外提一下的是,由于跨 step 的 op 有可能被并行执行,这也就意味着可能要用上更多的内存。TensorFlow 的控制流中也考虑了内存的问题,建立在 GPU 上的 frame 如果使用的显存超过某个上限则会自动做与 CPU 的内存切换的动作,把不用的部分数据换出去,把接下来要用的数据换进来。
例如 tf.while_loop()
的函数接口中就有个 swap_memory
的参数。
6666666666….
说实话,这篇文章的测试结果部分我觉得写的有点乱。
前面都是搬以前原有的内容,然后在本文的重点部分又写的这么乱,Google 的大佬们你们是认真的吗?
测试的系统配置是 Intel 服务器配上 K40 和以太网,每个节点一块卡,某些例子中用到了 8 卡的 DGX-1 V100。
一开始的两个测试用的是构造出来的模拟算例。
图 11 的结果感觉有点迷。
图 12 是模拟 RNN 的结构,把一个类似 8 层的 RNN 计算分布在 8 块卡上,把 tf.while_loop()
支持的并行 iteration 数从 1 调到 32,可以发现并行性发挥出来之后效果确实是挺好的,最高大约有 5 倍左右的性能提升。并行 iteration 数为 1 的时候其实就相当于跟 out-of-graph 一样。
后面模型并行的测试是把一个实际的 8 层 LSTM 分布在 8 块卡上,具体的并行方式与图 12 的测试类似。在 1~8 块卡上分别测试,加速比也还可以。
接下来对一个单层 LSTM 的测试是对比是否开启内存交换。不开内存交换时,序列长度加到 600 就出现超内存的现象了,而开启内存交换则可以在保证能跑的前提下还不会损失性能。
从追踪出来的 profiling 结果中也能看到,在这种计算模式下内存拷贝和 GPU 计算 overlap 得比较好,这也是性能不受影响的重要原因。
再下一个测试是固定 LSTM 的序列长度为 200,调整 Batch Size 的大小来对比动态 RNN 和手动循环展开的效果。动态 RNN 稍微损失了一点点性能,但是差距不大。
另一方面动态 RNN 比手动做循环展开在内存方面有更大的优势,类似上一个测试,动态 RNN 开了内存交换之后可以跑更大的 Batch Size。
最后是对 DQN 强化学习网络的测试,尽管 DQN 现在已经不用了被其他更好的方法替代了(???),还是希望能从它的测试中展现一下动态控制流的效果。
DQN 中包含了多个网络,根据不同的情况需要做出许多不同的操作。使用动态控制流的方法把所有的操作都包含在一个计算图中之后,最终能够比原始情况得到 21% 的性能提升。
总结来看,这篇文章重新整理了有关 TensorFlow 中控制流部分的实现思路,证明 in-graph 方式的纯数据流实现是有意义的。但是我对它的测试部分并不太满意,用到的是模拟的 workload,说服力不够,并且感觉测试的内容还是偏少。
]]>