人类对于算力的追求永无止境。但就在昨天,我们似乎前进了一大步:位于美国加州的创业公司 Cerebras 发布了号称全球最大的芯片,而且是专门为 AI 计算打造的。
这款名为 Wafer Scale Engine(WSE)的芯片拥有 1.2 万亿个晶体管,其数量是英伟达最新一代旗舰 GPU Titan V 的 57 倍。它的尺寸已经比一台 iPad 还要大了:
WSE 使用台积电 16 纳米制程工艺打造,面积 46,225 平方毫米,包含 400,000 个核心,片上存储高达 18G,功耗 1.5 万瓦(约等于 6 台电磁炉的功率),内存带宽 9PB/秒,通信结构带宽 100PB/秒。该公司表示,仅用一块这样的芯片即可驱动复杂的人工智能系统,从无人驾驶汽车到监控系统。
Cerebras 的「暴力美学」引来人们惊呼:WSE 的晶体管数量已经超过人类大脑中的神经元了!
在此之前,其实三星也开发过一款 2 万亿个晶体管的闪存芯片(eUFS),但 Cerebras 的这款芯片是用来进行处理工作的。此外,WSE 还实现了 3000 倍的速度提升以及 10000 倍的存储带宽扩展。
这块芯片面积达到了 46,225 平方毫米,而英伟达 GPU Titan V 的核心面积也「仅有」815 平方毫米。
Cerebras 表示,WSE 可以把处理复杂数据所需的时间从几个月缩短到几分钟。
「AI 模型训练时间的缩减将会消除这个行业进步的主要瓶颈,」Cerebras 首席执行官、创始人 Andrew Feldman 表示。他是一个芯片行业资深人士,此前曾将自创的一家公司以 3.34 亿美元的价格出售给 AMD。
目前,Cerebras 已在向少量用户提供这种芯片了,不过价格尚未透露。
今天的人工智能系统主要基于深度神经网络,其高度依赖于强大的计算机算力——大多数公司使用图形计算单元(GPU)处理 AI 模型训练的任务。然而顾名思义,GPU 是为了面向图形计算而设计的,谷歌等公司近年来已经设计出了 TPU 这样的神经网络计算专用芯片。
Cerebras 首席执行官、创始人 Andrew Feldman。
大多数芯片实际上是在一个 12 英寸的硅晶片上集合许多芯片,并在芯片上进行批量加工。但 Cerebras Systems 开发的这款芯片是连接在一个晶片上的单个芯片。互连接的设计是为了让所有的晶体管都能高速运转,这样 1.2 万亿个晶体管就能像一个整体一样运行。
巨型 AI 芯片真的可行?
芯片的大小在 AI 计算中至关重要,因为芯片越大,处理信息的速度越快,得出答案的时间也就越短。如果可以减少训练时间,研究者就可以尝试更多的想法、使用更多的数据去解决新问题。谷歌、Facebook、OpenAI、腾讯、百度以及其他科技巨头都表示,当今 AI 最大的问题是模型训练时间过长。减少训练时间可以克服整个行业的一大瓶颈。
当然,一般芯片厂商不会选择把芯片造那么大,这是因为,在单个晶片上,制造过程中往往会产生一些瑕疵。如果说一个瑕疵会导致芯片出现故障,那么几个瑕疵就会导致芯片报废。如果一个硅晶片上只有一个芯片,那么产生瑕疵的概率几乎是百分之百,而瑕疵会导致芯片报废。为了保证成品率,Cerebras 采用了冗余式的设计:如果有的核心出了问题,内部链接可以跳过坏掉的核心。这种方法大大降低了芯片的制造成本。
Cerebras Systems 的 CEO Feldman 表示,WSE 芯片着眼于 AI 从头进行设计,其中包含众多根本上的创新,通过解决数十年来限制芯片尺寸的问题(如 cross-reticle 连接、成品率、功率输送、封装)来提高性能。「每一个架构决策都是为了提高 AI 的性能。因此,WSE 芯片在有限的功耗和空间范围内,基于工作负载实现了数百甚至数千倍的性能提升。」
这些性能的提升是通过加速神经网络训练的每一部分来实现的。一个神经网络就是一个多级计算反馈回路。输入在回路中传递越快,该回路学习(训练)的速度就越快。加快输入传递速度可以通过加速回路中的计算和通信来实现。
WSE 的超大体积使其可以搭载更多用于计算的核以及更加接近核的内存,从而核可以高效运转。由于这个庞大的核阵列和内存都在一个芯片上,因此所有的通信都是在其上进行的,这意味着它的低延迟通信带宽是巨大的,所以核组的协作效率可以最大化。
这个 46225 平方毫米的硅芯片上有 40 万个 AI 优化的、无缓存、无开销的计算核以及 18GB 的本地、分布式、超高速 SRAM 内存,作为内存层次结构的唯一层次。内存带宽为每秒 9PB。这些核通过一个细粒度的、全硬件的、片上网状连接的通信网络连接在一起,提供每秒 100 PB 的总带宽。更多的核、更多的本地内存和低延迟的高带宽结构一起构成了加速 AI 工作的最佳架构。
Tirias Research 创始人兼首席分析师 Jim McGregor 在一份声明中表示:「虽然 AI 在一般意义上得到应用,但没有两个数据集或两项 AI 任务是相同的。新的 AI 工作负载不断涌现,数据集也持续增大。」
「随着 AI 的发展,芯片和平台解决方案也在不断发展。Cerebras WSE 是半导体和平台设计领域的一项惊人的工程成就,能够在单晶片规模的解决方案中提供媲美超级计算机的计算能力、高性能内存和宽带。」
这种面积的芯片,供电和冷却都是大问题,看起来水冷是必须的了。
Cerebras 表示,如果没有与台积电多年来的紧密合作,公司就不会取得此次破纪录的成就。
台积电高级运营副总裁 J.K. Wang 表示:「我们与 Cerebras Systems 合作制造出了 Cerebras Wafer Scale Engine,这是 wafer scale 开发的一个行业里程碑。」
40 万个 AI 优化的计算核心
WSE 包含 40 万个 AI 优化的计算核心。该计算核心被称为稀疏线性代数核(Sparse Linear Algebra Cores,SLAC),它们灵活可编程,并且为支撑所有神经网络计算的稀疏线性代数进行优化。SLAC 的可编程性确保核心能够在不断变化的机器学习领域运行所有神经网络算法。
由于 SLAC 针对神经网络计算原语(compute primitives)进行优化,因而实现了业界最佳利用率——通常是 GPU 的 3 至 4 倍。此外,WSE 包含 Cerebras 发明的稀疏捕获技术,以加速深度学习等稀疏工作负载(包含 0 的工作负载)上的计算性能。
零在深度学习计算中很常见。一般情况下,相乘向量和矩阵中的大多数元素都为零。但是,乘零运算浪费了硅、功率和时间,因为没有生成任何新信息。
由于 GPU 和 TPU 都是密集型执行引擎——这些引擎在设计时永远不会考虑零的情况——它们会乘以每个元素(即使元素为零)。当 50-98% 的数据为零时,就像深度学习中常出现的情况,大多数乘法运算是浪费行为。由于 Cerebras 的 SLAC 永远不会乘以零,因而所有零数据被过滤掉,并可以在硬件中跳过,使得有用功适时适当地完成。
Cerebras 在发布会上表示,使用这种芯片处理人工智能负载时,单个芯片可以承载整个神经网络的运算。
18GB 的片上内存
内存是计算机架构的关键组件。内存越接近计算,计算速度越快,延迟越低,数据移动的功率效率也越高。高性能深度学习需要大量计算和频繁的数据访问。这就要求计算核心与内存之间极为接近,但在 GPU 中却并非如此,GPU 的绝大多数内存运行速度很慢,并且远离芯片。
WSE 包含了较迄今为止最多的内核和本地内存,并在一个时钟周期内拥有 18GB 的片上内存。WSE 核心本地内存的集合提供了每秒 9PB 的内存宽带——是当前最好 GPU 的片上内存的 3000 倍,内存宽带扩展了 10000 倍。
高带宽、低延迟的通信结构
在芯片面积巨大的情况下,如何保证高速内部通信也是一个巨大挑战。WSE 上所使用的处理器间通信结构被称为 Swarm,与传统通信技术相比,它以非常小的功耗实现了带宽和延迟的突破。
Swarm 提供了一个低延迟、高宽带的 2D 网格,它连接了 WSE 上的 40 万个核心,宽带每秒达 100PB。此外,Swarm 支持单字节活跃消息,这些消息可以通过接收核心来处理,并且不需要任何软件开销。
路由、可靠的消息传递以及同步均在硬件中处理。消息自动激活每个到达消息的应用处理程序。Swarm 为每个神经网络提供独特和优化的通信路径。基于运行中特定用户定义的神经网络,软件将通过 40 万个核心来配置最优通信路径,以连接处理器。
通常情况下,消息会遍历一个具有纳秒延迟的硬件链接。一个 Cerebras WSE 的聚合宽带为每秒 100PB。TCP/IP 和 MPI 等通信软件则不再需要,因而可以避免性能损失。同时,这种架构的通信能量成本远低于 1 皮焦耳/比特,几乎比 GPU 低了 2 个数量级。Swarm 通信架构融合了极大的宽带和极低的延迟,使得 Cerebras WSE 能够以较当前任何可用解决方案更快的速度学习。