Cerebras处理器架构细节

文摘   2024-10-29 08:00   上海  
本文分享基于cs-2处理器,cs-2Cerebras CS-3处理器,大力出奇迹一文中cs-3除了参数区别外,整体架构上差别不大。

大模型的快速发展伴随着模型内存和计算量的指数级增长,同时芯片制程升级、内存迭代、精度变化等也带来了架构上的挑战。虽然在集群scale-out上带来了300x的显著增长,但之前设计上原因导致设备间的内存、计算、通信变得非常复杂。


为了满足ML增长需求及集群scale-out问题,cerebras专门设计了针对NN的细粒度、动态稀疏性处理器,并将处理器核心架构性能提高至少1个数量级,同时在横向扩展方面做数量级的改进和简化。

文章较长,为方便快捷阅读,整体目录如下。

核心架构设计

    全分布式内存架构

    核心datapath

    核心数据流调度

芯片scale-up设计

    PE间fabric网络互联

    PE间通信样例代码

    芯片间互联

    weight -streaming执行流程

    MatMul数组及 GEMM执行

集群scale-out设计





核心架构 - 加速动态稀疏计算

cerebras内核设计,尺寸只有 38760 平方微米,其中一半为48 kB 的内存使用。另一半是 由 110000 个标准单元组成的逻辑处理单元。整个内核以高效的 1.1 GHz 时钟频率运行,峰值功率仅为 30 毫瓦。


全分布式内存架构(Fully Distributed Memory)

在传统GPU内存架构中,通常会有一个共享的中央DRAM来存储数据。但问题在于DRAM速度慢且与处理器相隔较远,即便是HBM也不例外。

一般处理器内部datapath的带宽是内存带宽的100倍,所以可能会存在处理器低效问题。为提高处理器效率,传统解决方法是通过使用本地缓存或者本地寄存器来重用数据。另一种方法就是为datapath提供足够的内存带宽,将内存与处理器核心相邻放置。

在全分布式内存架构中,每个处理器core都有自己的SRAM片上内存,相互之间独立。这种设计允许每个核心都可以在一个时钟周期内访问本地内存,无需缓存,无共享内存方便处理非结构树稀疏等特性。


每个core拥有48kB内存,分成8个32 位的单端口bank,每个bank 6kB,允许并行处理。每个时钟周期可以进行2次完整的内存读取和1次完整的内存写入。所有这些内存都是按内核独立寻址,无共享内存,内存扩展通过fabric完成。这种设计允许处理器核心在每个时钟周期内进行多次内存操作,提高内存吞吐量和处理器的整体性能。

同时拥有256B的本地缓存,一般用于频繁访问的数据。此缓存设计在物理上靠近datapath,频繁的数据访问可以以超低功耗运行。


与GPU相比,WSE-2的内存带宽提高了200倍。



不同BLAS级别计算所需内存带宽对比,除GEMM外的BLAS级别均需非常大的带宽,这对于GPU等架构非常具有挑战。cerebras拥有足够大的内存带宽,可以方便实现向量-标量乘法(AXPY)相关计算,进而加速通用非结构化稀疏性运算。

因为稀疏GEMM实际上只是一系列AXPY操作的集合,每个非零元素对应一个操作。大内存带宽是实现非结构化稀疏性加速的必备条件。


核心datapath

cs-2处理器的核心是一个完全可以编程的处理单元,与通用处理器一样支持一系列通用指令,包括算术、逻辑、加载/存储、比较和分支等指令。这些指令局限于内核本地,与数据一起存储在48 KB本地内存中,每个核可独立运行,支持在整个芯片上进行细粒度动态计算。这些通用指令在 16 个通用寄存器GPR上运行。

同时为处理数据的tensor操作提供硬件支持,这些tensor操作在64b datapath上执行,datapath由四个 FP16 FMAC 单元组成。为了优化性能和灵活性,处理器将tensor作为first-class操作数,类似通用寄存器或内存。

fmac[fpsum]=[fpsum], [fwd_wgt], r_in等式显示了一个 FMAC 指令示例,指令直接将 3D 和 2D  tensor作为操作数进行操作。通过将数据结构寄存器DSRs(共44个)作为指令的操作数来实现,每个DSR都包含了一个指向tensor的描述符及tensor长度、形状和尺寸等关键信息。有了DSR,硬件架构就足够灵活能够原生地支持内存中的4D tensor,fabric tensor、FIFO及循环缓冲区等


核心数据流调度

处理器采用细粒度数据流调度,所有计算均由数据触发。数据和控制信息通过fabric传输,当内核接受到数据后,硬件触发查找任务并执行指令。并在发送时过滤0值,实现原生非结构化稀疏特性。

指令可以终止当前正在运行的任务,然后硬件从构成 PE 代码的任务集中选择一个新任务。它选择一个可运行的任务或者已被激活的任务。

传入的32位消息(wavelets)沿着称为color的虚拟通道传播(共24个),所有color属于同一个物理通道,虚拟通道之间数据传输相互不影响。

硬件可同时支持8个Tensor操作,也称为微线程,在每个周期内硬件可以在不同的Tensor操作之间进行切换。同时调度器根据Tensor输入输出可用性和优先级来选择执行哪个微线程。

与GPU相比,该架构通过细粒度动态计算,稀疏计算的利用率可以提升10倍。




芯片内Scale-up

wse-2整个晶圆构成如下,每个芯片包含10156个core,晶圆包含84个芯片,整体超850000个core。



PE间fabric网络互联

晶圆内部高速通信通过2D mesh网络拓扑的fabric实现,每个PE包含1个处理器、1个路由器及本地内存。PE 的路由器通过双向链路直接连接到其自己的 CE 和网状中四个最近相邻 PE 的路由器。指向其自身 CE 的链接称为 RAMP,指向四个相邻 PE 的链接由其基本方向表示。路由器是 PE 用于发送和接收数据的唯一通信设备。


wavelet信息共32位,包含优化过的FP16数据元素 + 16位index控制信息,使用静态路由,即上文中提到的color虚拟通道进行数据传输。cerebras专门提供了csl软件语言编写代码,编译后在处理器上运行。


PE间通信样例代码

官网提供了通信、任务激活和控制流程及阻塞等代码样例,以下就通信流程做个简单介绍。
当wavelet数据达缓冲区到后,main-task被激活,然后选择关联任务。然后数据被移动到寄存器中,方便task访问该数据。

// 7 is the ID of a color with some defined data routingconst red: color = @get_color(7);
// On WSE-3, 2 is the ID of an input queue which will// be bound to our data task.const iq: input_queue = @get_input_queue(2);
// For WSE-2, the ID for this task is created from a color.// For WSE-3, the ID for this task is created from an input queue.const red_task_id: data_task_id = if (@is_arch("wse3")) @get_data_task_id(iq) else @get_data_task_id(red);
var result: f32 = 0.0;
task main_task(wavelet_data: f32) { result = wavelet_data;}
comptime { @bind_data_task(main_task, red_task_id);
// For WSE-3, input queue to which our data task is bound // must be bound to color red. if (@is_arch("wse3")) @initialize_queue(iq, .{ .color = red });}

芯片间互联

芯片间的刻线上做桥接实现互联通信,并通过源同步source-synchronous并行接口减少时钟偏差,提高数据传输准确性。而且还内置冗余机制,实现自动纠错,提高稳定性。同时利用液冷解决数据中心散热问题。

互连网络设计用于在晶圆级别上实现统一的带宽和数据传输,使得整个晶圆就像一个单一巨大的芯片一样工作。这种设计允许数据在晶圆上的不同芯片之间以全带宽进行传输,同时保持了低延迟和高可靠性。通过在晶圆级别集成,可以实现比传统芯片间连接更短、更快的数据通路,这对于高性能计算和大规模并行处理应用来说是非常有利的。

下图可以看出,wse-2带宽提升近1个数量级,每比特功率效率提高近2个数量级。





weight-streaming执行流程

wse-2性能和容量足以运行当时的GPT3(wse-3则可以覆盖当前最大模型),无需分区等处理。系统做存算分离,模型weight存在memoryx中,计算过程中weight流式传输到cs-2系统中,基于上文提到的数据流机制执行计算。weight-streaming执行流程可见下面动‍图‍。‍



MatMul 数组及GEMM执行

每个单独的weight都触发一个独立的AXPY操作,weight计算完成就会被丢弃,硬件随即转向下一个元素。由于这些weight不会存储在芯片上,模型的大小就不会受到芯片上内存容量的限制。在反向传播阶段,梯度会逆向传输回MemoryX单元进行weight更新。

神经网络层本质上是进行矩阵乘法运算。由于CS-2的大规模能够利用晶圆上所有850,000个核心,将它们作为一个单独的、巨大的矩阵乘法器。通过在核心网格上分割tensor维度,CS-2可以并行地执行大量的计算任务,从而加快了神经网络的训练和推理速度。

计算执行的工作原理如下:
对于像GPT这样的变换器模型,activitions有三个逻辑维度:批量(Batch)、序列(Sequence)和隐藏维度(Hidden dimension,即B、S和H)。把这些tensor的维度分布在晶圆上核心的二维网格中。隐藏维度沿着x轴方向分割,而批量和序列维度则沿着y轴方向分割。这种布局允许在序列和隐藏维度上有效地进行weight broadast和reduce操作。


fabric包含32位的数据及控制命令,数据包触发FMAC数据计算任务,控制包触发同步等控制任务。


按列进行broadcast,通过硬件数据流机制,weight数据触发FMAC计算。然后再通过其他命令做reduce等操作。下图展示了weight广播流程,一行weight数据中的非0数据在不同列中传输,触发FMAC计算,进而产生稀疏加速。


以下两图则展示了weight 与 ativation相乘,然后累加及reduce的完整GEMM流程。


由于cs-2处理器独特架构及互联,使得支持的矩阵可以非常大,即便100000 * 100000的matmul层也可直接运行。芯片FP16稀疏性能达75 petaflops,FP16密集性能达7.5 petaflops。



集群Scale-out

针对参数量大的模型,一般采用TP/PP/DP等联合方式做优化,但会带来大量的通信开销,扩展性方面也会由于计算/内存/通信间的相互约束,而受到限制。


Cerebras 架构支持在单个芯片上运行所有模型,只需数据并行复制即可进行扩展,无需任何复杂的模型并行分区。通过swarmx进行互联,其采用spine-leaf网络拓扑结构,方便扩展。







更多资料:
hotchips cerebras架构分享:
https://www.bilibili.com/video/BV145seeZEAQ/
cerebras编程模型:
https://sdk.cerebras.net/computing-with-cerebras
cerebras白皮书


IT奶爸
实践是检验“专家”的唯一标准。一群认真执着的IT奶爸的学习和分享。
 最新文章