基于Dataflow的ML加速架构--Gorgon/WSE

基于Dataflow的ML加速架构--Gorgon/WSE

  • 发布时间:2020-12-12 11:33
  • 访问量:

【概要描述】介绍dataflow处理的基本理念,两种基于dataflow方法的相关架构软硬件实现:Gorgon和Cerebras WSE。

基于Dataflow的ML加速架构--Gorgon/WSE

【概要描述】介绍dataflow处理的基本理念,两种基于dataflow方法的相关架构软硬件实现:Gorgon和Cerebras WSE。

  • 分类:研究院
  • 作者:壁仞科技研究院
  • 来源:
  • 发布时间:2020-12-12 11:33
  • 访问量:
详情

摘要

本文首先介绍dataflow处理的基本理念,接下来介绍两种基于dataflow方法的相关架构软硬件实现,分别为Gorgon和Cerebras WSE。其中,Gorgon是基于2Dmesh CGRA架构,针对DB-ML的混合应用提供了硬件原语支持。Cerebras WSE是基于MIMD的dataflow架构,实现了基于model parallel和data parallel 混合的dataflow策略。

 

Dataflow处理方法

 

狭义的dataflow处理架构是在ISA中的建立一套以dataflow node作为程序的描述方式,一个程序由许多的nodes组成,这些node有无条件的运算,有条件比较判断,barriersync,fork,switch,merge等种类,一个dataflow node的执行,只取决于其操作数是否ready,这里把这种操作数称为token,程序中只有真正的数据依赖。

 

Figure 1: static dataflow machine指令语义[1]

 

上图是静态dataflow的处理单元以及指令,其指令表达中只有opcode,destination,operand几项,没有PC。其中1~5数字标号代表node,LR代表每个node左右输入。

这种dataflow架构非常善于挖掘应用的并行性,但是因为没有PC,操作之间是异步执行的,所以很难确定程序的精确语义,同一笔数据,甚至可能多次运行结果不一致。很难做interrupt和exception处理,debug困难,因此这种架构并不是非常成功。

广义的dataflow处理方法是狭隘dataflow处理架构的扩展,是指在任何软硬件层次,把数据处理流程在空间和时间维度展开,以dataflow的形式执行程序的方法,以此获得最大的并行性,可以在ISA level,可以在micro architecture level,也可以在task level和system level。

例如,在Out of order CPU中使用reservations station和reorder buffer进行的乱序执行的方法,实质上就是dataflow架构的一种实现,虽然CPU语义还是以control flow的形式表达的,但是在一段时间内的处理过程实质上是符合dataflow概念的,即只关注于operands是否ready,而不是根据PC的顺序先后执行。而最后的结果也是和in order执行结果一致的。

又例如,在CNN卷积处理中,算法存在大量的数据复用访问。因此,访问存储器后数据如何在本地register中存储并利用dataflow网络进行数据复用是必须考虑的重要设计目标之一,基于systolic array的dataflow处理方式会有下面几种形式:

·row stationary

weight和input feature存在register中,weight和input feature 以及partial sum都在array中流动。

·weight stationary

input feature和partial sum在array间流动,weight存在计算array中不动。

·output stationary

input feature和weight在array间流动,partial sum存在计算array中不动。

·no local data reuse

数据不在本地register存储复用,通过array算完后写回存储器。

 

Figure 2: Row stationary dataflow[2]

 

这里的array可以指一个乘法器,也可以指一个PE(Processing element)可编程单元。我们也可以把几种策略混合使用,例如同时用weight stationary和output stationary,array间连接的topology,可以是pipeline,也可以是ring或者2D的方式进行流动,也可以从一个node broadcast到所有的array。

 

Gorgon

 

Gorgon是Stanford在CGRA架构的基础上,针对Data-based Machine Learning混合应用,设计的基于dataflow处理的SOC。

CGRA 全称为Coarse Grain Reconfigurable Array, 是将PE(processing elements)或者MEM(Memory tile) 连接到2D mesh的网络上面,其中每个PE都是同构的,可以是单纯的计算器,也可以是可编程的结构,主要进行数学和逻辑操作,PE可以从mem取数,也可以从邻近的PE取数。CGRA的complier将程序的dataflow在空间和时间上展开。

 

Figure 3: CGRA结构图

 

CGRA是固定function加速器和FPGA架构的融合体,既有FPGA的可重配特点,又有固定function加速模块所带来的PPA优势,结合了两者的优点。CGRA具备word level的重配能力,可以根据不同应用算法的特点,快速的实现不同mapping连接配置的PE和MEM的dataflow,具有很强的软硬件一体设计的灵活性,近年来发展迅速。

CGRA的flow包括以下5步:

1. 算法描述,编写halide程序。

2. Halide compile,将算法描述转换成directed acyclic graph(DAG)。

3. mapper,将DAG转换成对应CGRA可以实现的指令。

4. Place and route,将map好的DAG放在各个CGRA单元中,并将这些单元的输入输出互联起来。

5. BS builder,将整个dataflow,每个PE的配置,以地址-数据pair表达成16进制编码的形式。这样组成一串bitstream。

这样就实现了多种不同应用算法在同一个CGRA上的实现。

Gorgon 的mesh网格是20x20,网络之间连接了PCU/PMU单元,整个网格有5.1TB/s的双向带宽。其中PCU为pattern compute unit,由16个6级pipeline SIMD lane组成,PMU为pattern memory unit, 主要功能为一块scratch-pad memory,寻址模式可以是FIFO/stride/multi-buffering/coarse-grained pipeline等方式,PMU内部有256KB的单cycle访存SRAM,整个chip有50MB的on-chip memory,和12.5TB/s访存带宽。

Gorgon dataflow是可以重新配置,包括其mesh网络以及PCU和PMU模块。

 

Figure 4: Gorgon整体架构图[3]

 

为了支持database应用的原语,例如partition, merge, join, group-by, filter, Gorgon在PCU中设计了相关逻辑。

 

Figure5: Gorgon PCU中database相关操作处理模块[3]

 

Cerebras WSE

 

Cerebras WSE(wafer scale engine)是世界上第一颗以一整片晶圆为一个芯片的设计。这种做法能极大地提供系统的集成度,理论上,在算力/存储/带宽上都能做到比分离的chip互联起来更好的performance。但是这种设计方式突破了常规芯片功耗/封装/散热的范畴,因此又会有很多额外的penalty。不管怎样,WSE足以载入史册。

 

Table 1 WSE架构参数[4][5]

die size

46225 mm2

晶体管数量

1.2 trillion

core数量

400,000

on-chip memory

18GB

memory bandwidth

9 PB/s

fabric bandwidth

100 Pb/s

工艺

TSMC 16nm

 

 

WSE上一共有400,000块PE单元,用2D fabric的方式连接在一起,从互联结构看来,WSE像个大号的CGRA。每个PE从farbic中接收和输出dataflow,包括数据和control两个部分,PE中的运算单元是一个基于FMAC的datapath,其中FMAC的输入来自registers,SRAM作为local的缓存。另外有一块作为tensor的control的模块,同时控制SRAM/register和FMAC。

 

Figure6: WSE整体结构[4]

 

图7是Cerebras软件栈及流程:

程序员同过各种框架编写各种ML应用,如常见的Tensorflow和Pytorch。接下来是compiler产生LAIR中间表达,将框架的算法mapping到cerebras的kernel上面,生成kernel graph,然后kernel graph通过带入的硬件信息将kernel graph place 到WSE上,例如如何正确地place存储和计算单元,dataflow fabric如何控制等,最后link生成可以执行的文件。

 

Figure7: WSE软件栈[5]

 

WSE中介绍的dataflow策略主要有三种:

§数据并行

§模型并行(layer-pipeline)

§模型并行(within a layer)

数据并行是常见的同构多核处理单元的dataflow策略。例如MLperf中,8块Nvidia A100 GPU做到408 batch size,4096块TPUv4做到65536的batch size。数据的并行特点是每个PE单元比较好编程,缺点是weight update时需要做不同sample之间的reduce操作,对带宽要求比较高。

layer-pipeline的思路是将整个WSE按layer进行mapping,不过引入了virtual device的概念,对于不同结构的layer,能够更细粒度做任务的分配切分,从而避免计算资源的浪费。

 

Figure8:CS-1中的layer-pipeline[4]

 

这里画的是一个kernel在CS-1上的mapping,整个的思路是feature流动,而weight stationary。

大致流程如下:

1)确定kernel处理的基本尺寸,如HW/MNK/element number等。

2)针对kernel的尺寸,确定每个layer所需的硬件存储,算力,带宽。

3)决定kernel是否需要进行data parallel设计,复制因子。

4)得到具体在CS-1上的执行单元所对应的物理尺寸。

5)在CS-1上做place and route,做具体位置的摆放,不同layer之间的互联。

这里有个有趣的问题,在现实的芯片中,会有坏的PE出现,不同的chip坏的位置不同,那么mapping的策略应该也有部分差别,原始placer mapping出来的layer pipeline程序可能不同,performance会有不同程度下降。因此在flow上就应该考虑到这一点冗余,比如该算法的算力存储要求是400个PE,placer预先可能配置450,其中50个PE是为了电路冗余考虑。整个流程可能不断迭代,从而得到最优解。CS-1mapping对应资源利用率不一定是100%。

这里的Place原则导致在CS-1上使用layer pipeline策略时很难使用手写kernel,或者效率不高,所以必须高度依赖tool auto-gen的kernel。

 

Figure9:CS-1中的layer-pipeline[5]
 

模型并行(with in a layer),在WSE中,因为一些层weight或者算力需求很大,需要将NN中的单个layer mapping到多个PE单元,以获得更好的performance。

这里以FC(Full connect)层为例子,FC层执行的是向量乘矩阵(GEMV)的操作,其中输入输出的feature都是向量,而weight是一个矩阵。如图10所示,一共有8个不同的块代表8组PE单元,input activation从左到右,代表的是不同的input channel,weight分布在8个PE单元上,每一行对应不同的output channel,输出在左侧, 每次输出所有的output结果,其中的dataflow形式是input activation从上往下传播和不同的output channel相乘,output partial sum从左往右传播并累加得到结果,weight stationary在PE中。

 

Figure 10:WSE中GEMV算法的模型并行实现-withina layer[4]

 

对于上面三种基础的并行dataflow策略,WSE利用单个算法进行了混合mapping优化,以达到最佳的performance

 

Figure 11:WSE dataflow策略选择[4]

 

§Layer运算/存储量大且层数少,WSE尽可能的做模型并行。

§Layer很多,但是运算/存储少,WSE的策略是选择做layer –pipeline

§Layer运算量少且层数很少,WSE的策略是选择模型并行和数据并行。

Cerebras在paper里面主推model parallel,传统的同构多核的架构中,类似GPU/TPU,大都会采用mini-batchSGD,以便将多个batch map在不同的core上面,获得最大的并行性,但是在weight update的时候需要reduce多个sample的weight gradient,在大集群的训练的中,batch越来越大,因此通信互联成为设计优化重点,甚至是NVLINK等互联协议出现间接原因之一。而CS-1采用model parallel的dataflow策略有效的减轻了这个问题,weight update都是local core的,不会有大规模通信消耗,在小的batchsize下面几乎能获得线性的加速比。

 

总结

 

Gorgon是基于CGRA的一个典型的Dataflow处理架构,WSE更像是一个大尺寸的CGRA,微架构上Gorgon加入了Database处理的相关功能支持。Dataflow处理的设计过程是和软硬件协同设计的过程紧密耦合在一起的。在深度学习中,数据流,控制流都比较规整,数据颗粒度较大,具有一定的数据复用性,dataflow设计策略是影响整体系统性能的主要因素之一。

 

参考资料

 

[1]“ComputerArchitecture Dataflow”Prof. Onur Mutlu Carnegie Mellon University

[2]“Eyeriss: AnEnergy-Efficient Reconfigurable Accelertor for Deep ConvolutionalNeural Networks”Yu-Hsin Chen, Tushar Krishna, Joel Emer, VivienneSze

[3]“Gorgon:Accelerating Machine Learning from Relational data”Matthew Vilim,Alexander Rucker, Yaqi Zhang, Sophia Liu, Kunle Olukotun, 2020 ISCA

[4]“Deep Learningat Scale with the Cerebras CS-1”Cerebras Systems NataliaVassilievea HOTCHIP2020

[5]“SoftwareCo-design for the First Wafer-Scale Processor (and Beyond)”Cerebras Systems HOTCHIP2020

近期文章

流体力学与物理导引神经网络

流体力学与物理导引神经网络

现代科学中众多复杂的关键应用依赖着对流体运动的精确预测。然而,由于流体力学方程数值求解的复杂性,基于传统数值方法的高分辨率长周期数值模拟运算量大且难以保证数值稳定性。本文将深入介绍PINN方法在流体计算中的应用并分析其与传统数值方法的差别。
2021-10-25
物理导引神经网络方法分析

物理导引神经网络方法分析

随着GPU能力的提升,支撑深度学习的软硬件生态得到了快速发展。通过深度学习来解决科学计算问题成了一种趋势,其中用深度学习来求解偏微分方程的方法也逐渐兴起。尤其引人注意的是一种称为物理导引神经网络的方法,其为科学计算领域注入了新的活力。
2021-10-18

联系我们

这是描述信息

招贤纳士

公众号二维码
这是描述信息

 版权所有   ©   上海壁仞科技有限公司   |    沪ICP备19047354号