Behemoth: A Flash-centric Training Accelerator for Extreme-scale DNNs
作者:Shine Kim, Yunho Jin, Gina Sohn, Jonghyun Bae, Tae Jun Ham and Jae W. Lee
@ Seoul National University, Samsung Electronics
本文针对以自然语言处理领域模型为代表的超大规模深度神经网络(DNN)模型的训练问题,构建了一种以闪存为中心存储的训练加速器。
论文为什么要做这个事情?
DNN在自然语言处理、计算机视觉、推荐算法等应用领域有广泛的应用。而增加模型参数的数量对于提升模型的准确度起到至关重要的作用。自然处理领域中各种基于Transformer的模型,其网络中的参数数量达到了超大规模(Extreme-scale)的级别。例如,GPT-3模型的参数数量达到了1,750亿以上。不仅如此,DNN模型的参数规模仍然在继续扩展,如如图1所示:
![NLP领域模型参数数量的扩展趋势](/2021/09/24/FAST21-Behemoth/fig1.png)
模型参数数量的扩展对存储系统提出了以下的两个挑战:
- 内存容量墙(Memory Capacity Wall)
以GPT-3为例,其模型参数大小为700GB。而一个NVIDIA A100的内存容量大概在40GB左右。由此我们可以看出以下几点:
- DNN模型的大小远远超过了一个GPU计算设备的内存容量;
- 用户将不得不采用模型并行的方式,将模型分布在多个GPU上完成训练。但是完成模型并行,需要细致地设计负载均衡机制。
- GPU/TPU上配备有高带宽存储系统(HBM),其价格相当昂贵。
![用模型并行的方法实现超大规模DNN训练](/2021/09/24/FAST21-Behemoth/fig2.png)
- 存储带宽的浪费(Memory B/W Underutilization)
- 随着DNN模型的尺寸扩大,Transformer中各种全连接层矩阵的重用也会变得更加频繁;
- 因此,随着训练计算总量的不断扩张,实际存储的带宽需求并未大幅度提高;而GPU上的高带宽存储系统也随着计算量而不断增加,造成了明显的带宽浪费问题。
![模型训练的需求带宽和实际带宽不匹配](/2021/09/24/FAST21-Behemoth/fig3.png)
本文希望能够将超大规模DNN训练的中心存储由低容量、高花费的HBM替换为高容量、低开销的闪存块,并辅助加速器设计完成单SSD节点上的训练任务。
这个事情之前别人是怎么做的,存在什么问题?
对于超大规模DNN模型训练而带来的内存容量墙问题,有以下两种解决方案:
- 针对张量管理的异构存储系统
由于用多个设备进行分布式训练对训练设备的内存带宽造成了极大的浪费,因此有一些方法考虑建立一种在异构内存系统中进行张量数据迁移的高效存储管理系统*。但是,超大规模的模型甚至在主机的存储器中也放不下,因此并不是解决超大规模模型的训练问题的长久之计。
参考文献: Hotness- and lifetime-aware data placement and migration for high-performance deep learning on heterogeneous memory systems, [TC 2020]
vDNN: Virtualized deep neural networks for scalable, memory-efficient neural network design,[MICRO 2016]
*SuperNeurons: Dynamic GPU memory management for training deep neural networks.[PPoPP 2018]
- 以二级存储为中心的机器学习系统
- BLAS-on-flash:构建一个库,用于针对大型数据集的机器学习算法(例如 ISLE、XML)实现高效的闪存速度。
- Cognitive SSD: 构建一个引擎,用于检索存储在闪存上的 DNN 模型推理所用的非结构化数据。
本工作的主要着眼点在于用闪存实现数千亿参数的超大规模神经网络语言模型的训练。
*: BLAS-on-flash: An efficient alternative for large-scale ML training and inference? [NSDI, 2019]
Cognitive SSD: A deep learning engine for in-storage data retrieval**[USENIX ATC, 2019]**
该论文解决了什么难点,难点存在的原因是什么,作者是如何发现这些难点的?
本论文主要解决了以下难点:
- SSD平台上模型的训练问题:需要设计一个加速器来完成模型的数据并行计算、控制模块间的通信以及充分利用带宽;
- Flash的低带宽问题:SSD有极低的带宽,特别是非顺序访问。除此之外,由于SSD底层有垃圾回收(Garbage Collection)、磨损平衡(Wear Leveling)等机制,其持续写带宽也明显比峰值带宽低很多。
- Flash的耐久性(endurance)问题:SSD只能进行一定数量的P/E操作,这被定义为SSD的耐久性或寿命。用SSD作为DNN训练的存储器明显将缩短其寿命,特别是在随机写时,有写放大(Write Amplification)的现象。
针对以上难点,作者各做了哪些优化,优化背后的思想是什么?
Behemoth计算平台
硬件模块
Behemoth系统设计的整体目的是在单节点上完全容纳超大规模的DNN模型,实现数据并行训练。Behemoth计算平台的整体结构图如下所示:
![Behemoth架构图](/2021/09/24/FAST21-Behemoth/fig4.png)
由于执行过程的底层逻辑是一致的,所以计算加速器的硬件部分可以分为三个部分:计算核心、张量缓存以及平台专用的Flash存储系统。
- 计算核心(Compute Core):计算张量;控制数据转移;指令翻译。
- 张量缓存(Tensor Buffer): 计算核心与NAND Flash之间的一块DDR DRAM区域,将张量暂时保存。
- 平台专用的Flash存储系统(Flash Memory System):代替HBM成为训练过程中的主要存储,用来存储张量信息。相关优化在后面会提及。
整体上来说,Behemoth系统采用了DDR DRAM+NAND Flash的双层存储结构,这一点是与以往训练方案中采用HBM作为单层存储结构本质上不同的。
面向DNN训练的硬件模块组装
利用基本硬件模块,本加速器组装出多个激活计算节点和单个权值计算节点,以实现DNN模型训练的数据并行。如图4所示:
- 权值计算节点(Weight Node):更新并存储权值;
- 激活计算节点(Activation Node): 计算并存储激活值,激活值还将用于反向传播;多个激活计算节点实现了DNN模型训练的数据并行。
- 主机端(Host):负责将训练命令序列(包含计算以及DMA命令)发送到加速器上,在训练结束之后,加速器将训练结果返回给主机端。
训练步骤
DNN模型的训练是一个重复的正向/反向传播的过程。图5(a)展示了单层训练的步骤,其正向传播包含以下几步:
![运行时系统](/2021/09/24/FAST21-Behemoth/fig5.png)
- 从权重节点闪存中读取权重值,到权重节点的张量缓存中;
- 将权重节点张量缓存中的权值通过广播发送到激活节点的张量缓存中;
- 权重张量被加载到片上的SRAM中,由计算核心准备开始计算;
- 激活节点的计算核心完成计算,生成激活张量;
- 激活张量从片上的SRAM中被拷贝到激活节点的张量缓存中,准备写入激活节点的闪存块上;
- 张量缓存中的激活张量写入闪存块中,供反向传播使用;权值张量从张量缓存中被释放。
反向传播与正向传播的步骤类似。当一个迭代结束后,图5(b)展示了权值更新的过程:
- 最终权值的梯度下降张量从激活节点的张量缓存被送至权重节点的张量缓存;
- 收到所有权值的下降梯度后,权值节点将权值下降梯度加载至片上的SRAM;
- 由计算核心来更新权值;
- 更新后的权值张量被写入张量缓存;
- 将更新后的权值由张量缓存写入NAND Flash中以供下一轮迭代使用。
平台专用Flash管理系统——带宽问题
相比于高带宽存储,闪存存在带宽低的问题。根据之前的分析,GPT-3模型训练要求的带宽大概在50GB/s左右,而NAND Flash大概只能提供几个GB/s的带宽,无法满足训练的需求。
事实上,NAND设备的带宽峰值并没有得到很好的发挥。为了优化NAND设备的带宽值,可能的方向有以下两个:
- 尽可能的顺序化对NAND设备的写操作:随机写会造成大量的带宽浪费以及写放大现象;
- 防止NAND设备固件成为性能的瓶颈:固件中的FTL一般都有如垃圾回收、磨损均衡、备份数据管理等机制,这些机制都会影响NAND设备发挥其峰值带宽。
数据分类
![对数据类型进行分类](/2021/09/24/FAST21-Behemoth/fig6.png)
DNN训练的计算过程中涉及的数据读写类型非常固定,可以进行一定程度的分类。
在活跃计算节点上,主要有训练的输入和训练过程中生成的激活值两种张量数据。训练的输入其生命周期长,持续整个训练过程,这样的数据被称为NV数据流(非易失性数据流)。
相比来说,训练过程中生成的激活值是一种中间结果,其生命周期很短,往往只持续几分钟,这样的数据被称为V数据流(易失性数据流)。同样地,在权值节点上也有这两种数据流。上表分析了活跃计算节点和权值节点上数据的不同类型,我们可以发现:
- 在DNN训练中所有数据的写都可以用添加式的顺序写完成。
- 在活跃节点和权重节点上,都各有一个NV数据流和V数据流,两种数据流的数据保留时间相差很大。
由以上分析,NV数据流和V数据流在保存时间、读写权限上有很大的不同,存储在同一段地址空间中并不合适。在存储上,Flash管理系统设计了一种按照流来分层的存储方式,下图展示了激活计算节点的Flash分层管理结构:
![激活计算节点的Flash管理系统的分层结构](/2021/09/24/FAST21-Behemoth/fig7.png)
这样设计的效果是:
- 用物理块地址明确的分离了每种数据,每个数据流在单独的存储空间上运行;
- 每个流有自己的逻辑地址空间、访问权限和基于数据保存时长的P/E周期数。
轻量级FTL
![用轮盘赌算法实现磨损均衡](/2021/09/24/FAST21-Behemoth/fig8.png)
通过对数据的分析可知,对Flash存储的写保证是添加式的顺序写。因此,在之前Flash翻译层(FTL)中复杂的垃圾回收和磨损均衡机制将不再必要。在本方案中移去了垃圾回收机制,并用最简单的轮盘赌(Round-robin)算法来实现充分的磨损均衡,极大程度简化了FTL。如图8中,假如共有4个物理块,一轮迭代需要写入3个逻辑块,采用轮盘赌算法在第一轮迭代写入0/1/2后,第二轮迭代将写入3/0/1块。
写路径的硬件自动化
一般来说,通用的SSD平台实现了读路径的硬件自动化,这样能够最大程度的利用NAND设备的带宽。但是,由于写路径需要去考虑垃圾回收、磨损均衡、数据一致性、冗余数据管理等处理机制,其硬件自动化很难实现。在本平台专用的Flash管理系统中,基于轻量级的FTL,实现写路径的硬件自动化不再是问题。
![写路径的硬件自动化](/2021/09/24/FAST21-Behemoth/fig9.png)
硬件自动的写路径主要由两条流水线组成:
- 写命令流水线:将数据从计算节点的张量缓存转移到Flash存储的SRAM中,其带宽可以达到56GB/s;
- NAND设备写入流水线:从Flash存储的SRAM转移到NAND设备中,其带宽可以达到64GB/s。
注意到,以上设计没有讨论数据恢复的问题。这是因为在DNN训练中的数据多为临时数据;
并且,一旦数据丢失,可以从检查点重新开始执行。
平台专用Flash管理系统——耐久性问题
一般来说,P/E操作会对Flash颗粒造成永久性损伤,Flash颗粒所能承受的最大P/E操作数量被称为Flash的寿命或耐久性。然而,这种损伤的实际体现为颗粒对于数据的保存时间变短,而当数据保持的时间低于预设值的时候,此颗粒就会被认定为已损坏。
因此,当把数据保存时间的预设时间限制降低时,从某种意义上来说,Flash颗粒仍能够发挥数据保存的作用。许多研究已经证明SSD的耐久性会随着对数据保留时间要求的放松而提高。
平台专用的Flash管理系统将NV数据流和V数据流分离开,在V数据流中,如下图所示,数据保留时间的限制仅仅为分钟级别,假设为5分钟。
![张量的寿命](/2021/09/24/FAST21-Behemoth/fig10.png)
图10 张量的寿命
据相关研究表明,当数据保留时间限制为5年时,一般的SSD设备能够承受5万个P/E操作;
当把数据保留时间的限制从5年降低到3天的时候,Flash颗粒所能承受的P/E次数大概能够提升40倍。
而因此保守估计,数据保留时间降低为3天的情况下,V数据流的颗粒承受P/E次数大概为200万次左右,而V数据流的数据量为1.85TB。假定数据带宽为17.6GB/s下,Flash管理系统中V数据流颗粒有5年以上的寿命。[这里写放大系数(WAF)为1,如图11所示,此Flash只支持添加写的顺序写操作,而不涉及垃圾回收等机制。]由此论证了存储系统的耐久性问题。
![DNN训练时数据块的访问特征](/2021/09/24/FAST21-Behemoth/fig11.png)
实验结果
实验设定
验证方法:
- 将Behemoth在存储上的费用开销与传统的基于TPU的DNN训练系统相对比;(验证Flash替换HBM作为主要存储的作用)
- 将Behemoth专用Flash管理系统训练吞吐量与传统Flash管理系统相对比。(验证平台专用的Flash管理系统的作用)
验证平台:
采用MAESTRO作为NPU模拟器平台,采用MQ-Sim作为SSD的模拟器平台。
训练模型:
![评测使用的训练模型](/2021/09/24/FAST21-Behemoth/fig12.png)
图12 评测使用的训练模型
(PxQ:P代表全连接层扩展的宽度倍数,Q代表网络decode或encode层的深度倍数。例如BERT 2x2模型代表在原始的BERT上,每个全连接层的宽度扩展为原来的两倍,网络层数也加深为原来的两倍)
开销对比
开销对比时模拟器的参数设定:
![开销对比中模拟器的参数设定](/2021/09/24/FAST21-Behemoth/fig13.png)
开销对比的结果:
![开销对比的结果](/2021/09/24/FAST21-Behemoth/fig14.png)
以上训练均在10天之内完成,使用TPU v3系列的系统在存储上的开销比Behemoth系统要贵3.65倍左右。
训练吞吐量对比
训练吞吐对比时模拟器的参数设定:
![训练吞吐量对比的参数设定](/2021/09/24/FAST21-Behemoth/fig15.png)
其都采用Behemoth加速器作为计算核心,只有Flash管理上的不同。Baseline采用4个SSD并行达到64个Channel。
对比结果:(下图中的理想情况是指存储访问不花费任何开销)
![训练吞吐量的对比结果](/2021/09/24/FAST21-Behemoth/fig16.png)
- 平台专用的Flash管理系统接近于理想情况;
- 优化后的吞吐量是传统的SSD其训练吞吐量的2.05倍;
- 可以看出,SSD中固件造成的性能瓶颈是导致传统SSD性能差的主要因素。
自己的思考
Behemoth解决了超大规模DNN模型的高效数据并行训练问题。针对存储墙、带宽浪费等问题,创新性地利用Flash作为主要存储完成了单节点上的训练加速器设计。仿真结果显示系统在存储开销和训练吞吐量上都有不错的效果。
- 本文作者: Zhang Xinmiao
- 本文链接: https://recoderchris.github.io/2021/09/24/FAST21-Behemoth/
- 版权声明: 本博客所有文章除特别声明外,均采用 MIT 许可协议。转载请注明出处!