可无限扩展的AI处理器：Tenstorrent虫洞

编者按：

Tenstorrent的可扩展AI架构，有很多特征非常符合我们软硬件融合的理念和思路，大家一起学习：

每个Tensix处理器核是CPU标量处理器+AI计算引擎。既有通用可编程的优势，又有AI张量加速的优势。
条件处理。实现对稀疏矩阵处理的优化。
包处理引擎。负责跟外部的数据交互，以及数据的一些前处理和后处理。实现数据交互的性能优化。
芯片集成类似独立Host CPU的独立的嵌入式CPU，作为系统调度和管理，进一步增强AI芯片的灵活性和图灵完备性。
高可扩展的架构，包括内部总线以及PCIe/ETH，实现近乎无限的扩展性。“无限”可扩展性可以跨处理器核、跨芯片、跨服务器、跨机架。目前基于CPU或GPU的扩展性难以跨越服务器的物理约束（否则整个系统性能会很差），因此把巨服务拆分成微服务的系统解构成为必选项。但是AI的模型发展趋势却是越来越巨型化，这样，提供近乎无限扩展的AI训练系统，可以简化AI算法工程师的工作，让他们从训练集群的软硬件平台的细节约束中抽身出来，把精力专注到算法本身，以此来进一步提高算法设计的效率。

本文内容来自于Tenstorrent在Linley Spring Processor Conference 2021的演讲。

参考文献：

Tenstorrent采用独特的硬件和软件紧密结合的方法。硬件专门用于AI任务，但软件并不复杂。整个软件堆栈只有大约 50,000 行代码。与大多数其他需要定制开发流水线的AI专用ASIC不同，Tenstorrent的芯片具有很强的适应性和灵活性，同时支持所有主流的工具链、框架和运行时。

本次演讲将介绍Tenstorrent的新产品，Wormhole（虫洞），以及支撑Wormhole的大规模定制的微架构。

快速介绍一下Tenstorrent的芯片历史。

Jawbridge是Tenstorrent的第一颗芯片，这是一颗测试芯片，旨在证明架构和性能的可行性。

Jawbridge案例中，6个Tensix处理器内核组成的计算阵列，以及LPDDR4内存，另外有4个支撑乱序的ARC处理器（Synopsys IP）可以运行Linux，承担内部管理和调度。

其他规格如上图所示。

Jawbridge是一款极小的芯片，极低的功耗。凭借有限的预算，将这颗芯片研发并流片，并且证明了他们的设计可行性，以及功耗和性能预期。

GraySkull是Tenstorrent的第一款产品，它在2019年底拿到样片，在过去的一年里，已经得到很多顾客的试用。它目前正在量产阶段。

GraySkull有120个Tensix核组成的计算阵列，8个通道的LPDDR，和PCIe Gen4 x16，以及4个ARC CPU组成的用于调度和管理单元。

该芯片采用 GlobalFoundries 12nm工艺，尺寸为620mm^2。GraySkull 65W，PCIe插卡供电最大为75W，这意味着它可以轻松插入现有服务器，而无需任何额外的辅助电源。

重点介绍Wormhole，Wormhole有一些针对人工智能训练的关键升级。

也包括更快的内存，和用于芯片和芯片连接的以太网接口。

Wormhole在一个月前回到了Tenstorrent实验室（从视频演讲时间算）。

包括DDR、PCIe、以太网和处理器阵列在内的所有功能块都已正常运行。

Wormhole是一个670mm^2的芯片，在GlobalFoundries 12nm工艺节点。

上图是Wormhole的平面图，SOC架构遵循JawBridge、GraySkull设定的演进。

Wormhole有80个Tensix核（少于上一代GraySkull的120个Tensix Core），6个通道的GDDR6，16通道的PCIe Gen4，和16个端口的100Gb以太网，所有这些都是通过NOC总线连接到一起。

通过上述这些明显的变化，Tenstorrent专注于如下一些关键点的改进：

如不断增加的计算性能，如通用PPA改进。

此外，Wormhole也是一个原生的以太网包交换机，这是系统层次的一个关键不同。以太网端口允许将许多芯片链接在一起，以针对大型 AI 网络进行横向扩展。

对芯片的特性做一个高层次的概述。

在Tensix处理器核心方面，已经提高了数学计算单元的吞吐量，改进了SIMD引擎，增加了执行复杂数学和逻辑指令的能力。还增加了每个Tensix核心的SRAM容量。

为了在SOC侧有效地供应数据给计算单元，已经提到过的内存总带宽为384GB/s，以太网端口的总带宽为每秒400GB/s。

重要的是，以太网端口作为扩展，允许芯片间的数据传输，而无需任何软件开销。依赖NOC的包交换能力，Warmhole可以构建非常强大的大规模系统。

尽管矩阵运算性能提高了 2 倍，内存带宽提高了近 3 倍，并且包含 1.6Tbs 的网络交换能力，但 Wormhole 芯片的功率仅翻了一番（达到 150W）并且在相同的 12nm 工艺下，面积增加不到 10%。

Tenstorrent特别关注Wormhole的系统级架构。

AI 训练工作负载的复杂性正在飙升。OpenAI 声称，训练最强大的网络所需的计算量每 3.5 个月翻一番。

这样，就导致人工智能训练的工作量持续增长，它们已经跨越了多个服务器和机架。扩展性成为AI训练的一个关键能力。而近乎无限地扩展到大规模网络的能力是 Tenstorrent Wormhole 的关键优势。

Tenstorrent Wormhole提供两种形态：

一种是作为一个PCIe卡，可以轻松插入服务器；
但更重要的是，作为一个Wormhole模块，该模块由一个芯片和6个GDDR6存储器组成。它还将所有16个以太网端口接到板上，模块的每一边有四个端口。

使用Wormhole模块，Tenstorrent设计了nebula（星云），一个4U服务器包含32个Wormhole芯片。

整个芯片连接在一个完整的网络中，每一侧的以太网带宽为每秒100GB。

很明显，在边缘的芯片，并没有使用所有的以太网端口，它们可以连接到交换机或其他服务器，以便进一步扩展。

在一个nebula服务器中，有384GB的DRAM，以及超过TB/s的以太网，可以以对软件完全透明的方式，扩展连接到其他外部服务器系统。

接下来是Galaxy，一个48U的机架，有8个完全连接的Nebula服务器。

该机架还包含 4 个 AMD Epyc 服务器和一个共享内存池。该机架提供 >3TB 的 GDDR6 和 256Gb 的外部以太网链路。通用 AMD Epyc 服务器和内存池连接到以太网网格。

通用的服务器和内存池都通过以太网连接到Nebula服务器。

类似地，Nebula服务器之间也通过以太网直接连接，图片展示了机架内的服务器到服务器的连接。

Tenstorrent 认识到并非所有 AI 工作负载都是同质的。48U机架可以有不同配置。例如，计算能力更低，在这种情况下只有四个Nebula服务器。但是需要更大的内存池。这种类型的配置更适合内存密集型模型，例如深度学习推荐系统。

这是一个完整的48U的机架，称之为超级计算机。

正如你所看到的，它像一个2D网格一样，每个Wormhole服务器连接在另一个服务器的对等端。

因此对整个机架进行软件处理，就像一个大而均匀的Mesh网络。

Wormhole模块，重要的是里面有Tensix核心，这很重要，稍后会讲到。

当然，也可以支持其他拓扑，不仅仅是Mesh。

最后，横向扩展功能并不止于此，Wormhole在机架级连接方面也是非常灵活的。

上图是一个完全连接的机架到机架系统的超级计算机拓扑结构，它只是简单地以标准的方式扩展网格。

与此拓扑中单机架级别上的连接类似，一个机架边缘的Wormhole模块直接连接到下一个机架上的对应模块。

可以类比成网格的x或y维度上的进一步扩展。

对于软件来说，它看起来像是一个由 Tensix 内核组成的大型同质网络。片上网络透明地扩展到许多服务器机架，无需任何痛苦的软件重写。理论上，它们的网状网络可以以全带宽和均匀带宽扩展到无穷大。

实际上，这为阵列中的每个处理器提供了充分的带宽，并在网格上具有统一的带宽，因为Wormhole有一个内置的交换机。

超级计算机拓扑不需要昂贵的交换机。事实上，机架式开关的顶部仅用于与集群外网的连接。

另一种选择是，采用更传统的数据中心拓扑。这样，机架之间的带宽则更少，但仍然是统一的网格和完全相同的连通性。

所有Nebula服务器都通过TOR连接。

这样，就会有一个连接带宽的分层，芯片到芯片是全带宽，但服务器到服务器和机架到机架的带宽是有约束的。

然而，这种拓扑的好处是它的弹性非常好。它更容易支持多租户和冗余。

Wormhole被设计成完全灵活的横向扩展连接，而不会影响系统的软件视图。

因此，芯片到芯片、服务器到服务器和机架的直接拓扑和带宽可以优化到工作负载的最优规模。

接下来，将谈谈Wormhole如何使能深度学习弹性扩展。

弹性扩展，意味着可以将人工智能模型的大小扩展到数万亿个参数，并有效地利用这些参数。

通过把很多服务器和很多机架放在一起，可以实现数以万计的处理器以成本高效的方式连接成一个大的集群系统。

通常，这样就会有多层的分层结构，从处理器核内部到芯片间，再到服务器间的带宽，其编程模型各不相同。从芯片到服务器再到机架，会有不同的带宽。

这是一种简单粗暴的方式，增加了很多的计算量，易用性不好。

编译器设计师和模型设计者需要考虑物理硬件和层次结构，他们调优模型反映了硬件而不是实际想要的模型。

所以，分层结构越复杂，编译器的设计就越复杂，很多工作就难以自动化。

Tenstorrent的设计目标是设计一个完全可（无限）扩展的微架构，这样前面提到的那些严格的层次结构就会被移除，这样就可以非常方便的添加更多的处理器核心和服务器。

当然，呈现给模型设计者或程序员的视图只是一个无限网格。

一方面，这将消除模型设计者的约束，他们只关注模型的看法以及模型应该什么样，而不需要关心硬件。这是一个统一的硬件。

这个横向扩展问题非常困难，特别是对于定制 AI 芯片。即使是在横向扩展硬件领域处于领先地位的NVIDIA，也需要模型开发人员处理这些严格的带宽、延迟和编程层次结构。

如果Tenstorrent声称将这一痛苦任务自动化的说法是属实的，那么他们已经颠覆了整个行业。

要了解Tenstorrent是如何做到这一点的，我们需要回顾一下横向扩展训练的历史。

在CPU集群上基于分块的数据并行，将其拆分到整个集群中。在CPU上有中央参数服务器。

由于带宽限制，这并没有很好地实现扩展。

使用GPU集群改善了这一点，GPU集群之间的all-reduce，使得ResNet和BERT的突破，但这些方法也有局限性。

DistBelief框架，能够结合其他类型的并行，不仅仅增加批量的大小，还可以分割模型和数据，然后分散到集群中。

简单回顾一下。

在某种程度上，基于批量的并行性是有局限性的。

最大的约束在于批量大小，如果继续增加批量大小，模型将不再收敛并难以达到高精度。

对于较大的模型，DRAM 容量成为一个限制因素，因为整个模型都在所有节点上进行复制。

中间计算甚至需要存储存储在片内SRAM。

更大的DRAM容量和带宽所产生的影响是每个节点的成本飙升。这种缩放方法适用于小型模型，但不再适用于大型模型。

随着更大的模型开始发挥作用，我们开始看到用于有效扩展的库的出现。

用户现在可以跨服务器集群指定和组合模型、流水线和分片数据并行性。

本质上，用户将跨节点拆分网络的模型和层。这允许模型继续扩展，因为它们不需要在每个节点内复制整个模型。

这种缩放形式的最大问题是必须手动完成。开发人员必须选择将哪些层映射到哪些硬件并控制数据流。

上图是两个机架，或者说16台nebula服务器，或者说512颗Wormhole所组成的集群系统。

尽管包含许多芯片、服务器和虫洞芯片机架，但软件本身只能看到这个核心网。严格的层次结构被删除，模型开发人员不受约束。

编译器会根据网络拓扑自动在网络中有效地放置和路由迷你张量，而不是让模型开发人员担心。

添加更多服务器可以扩展网格并允许模型扩展而无需担心。

这样，就允许大量模型流水线并行。网络中的层可以使用任意数量的资源来满足计算需求。不需要大量计算的层可以使用一半服务器或一半芯片，但需要大量计算的另一层可以跨机架延伸。

第 4个示例表明，如果您的网格足够大，甚至可以跨服务器机架拉伸单层。编译器会看到网格的大小并根据它们的大小映射层，无论是内核网格是1颗芯片还是内核网格是跨许多服务器机架的10,000颗芯片。

在扩展模型时，您需要在各个节点和硬件单元上对模型层进行切片和映射。

除了流水线并行之外，单个张量操作也可能变得足够大，以至于单个张量核心硬件单元无法单独执行。一层上的一个节点内的这些张量操作也需要由研究人员手动切成小张量。

Tenstorrent 的目标是创建一个可以本地放置、路由和执行迷你张量操作图的架构。

迷你张量被实现为 Tenstorrent 架构的本机数据类型。这意味着开发人员不必担心张量切片（slicing）。每个迷你张量都被视为单个数据包。这些数据包有一个数据payload和一个header，用于在核心网格内识别和路由数据包。Tensix 核心直接在这些迷你张量数据包上完成计算，每个核心都包括一个路由器和数据包管理器以及大量的SRAM。路由器和数据包管理器处理同步并发送计算数据包以沿着网状互连流动，无论它是在芯片上还是通过以太网在芯片外。

从软件的角度来看，数据包的传输在整个内核网格中看起来是一致的。在同一芯片上的内核之间发送数据包看起来与在不同芯片上的内核之间发送数据包相同。因为每个芯片和 NOC 本身就像一个交换机，迷你张量数据包可以沿着核心网路由到它需要去的下一个核心。

与其他网格架构相比，Tenstorrent 网格更大且更具可扩展性。FPGA 处于非常精细的水平，需要大量时间来手动调整。CGRA的运行标量图，但它们仍然有许多限制因素。Tenstorrent在其矩阵引擎中有多个teraflops和更大的内存大小。NOC、数据包管理器和路由器智能地处理芯片内和芯片间通信，让模型开发人员专注于拼图的其他部分。这使得它可以更有效地扩展 AI 工作负载，同时也更容易开发。