基于Dataflow的ML加速架构--Gorgon/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

上一个: 在壁仞科技研究院,我们思考的问题

下一个: 加入图卷积的多智能体强化学习

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多