2021年09月07日 | Dojo系统详细构成

Tesla在上月的AI Day推出了用于机器学习和神经网络训练的芯片、系统和软件。它们将共同推进用于自动驾驶汽车模型的训练。

Elon Musk和他的芯片与系统设计团队在超过三小时的演讲中提供了大量的技术细节，以下是其中的一些亮点。

神经网络

Tesla设计了一个灵活、可扩展的分布式计算机架构，专门用于神经网络训练。Tesla的架构从D1芯片开始，该芯片有354个训练节点，每个节点有一个强大的CPU。这些训练节点CPU是为高性能的NN和ML任务设计的，32位浮点运算的最大性能为64GFLOPs。

对于拥有354个CPU的D1芯片，32位浮点运算的最大性能是22.6TFLOPs。对于16位浮点计算，D1的最大性能跃升至362TFLOPs。

Tesla推出了两个用于神经网络训练的系统：Training Tile和ExaPOD。一个Training Tile在一个封装中拥有25个连接的D1芯片。一个由25个D1芯片构成的Training Tile有8850个训练节点，每个节点都有上面提及到的高性能CPU。一个Training Tile的32位浮点运算最大性能是565TFLOPs。

ExaPOD将120个Training Tile连接成一个系统，即3000个D1芯片、106.2万个训练节点。一个ExaPOD的32位浮点运算最大性能是67.8PFLOPs。

Tesla对神经网络介绍的细节D1芯片和Dojo神经网络训练系统的推出显示出了Tesla的发展方向。这些产品投入使用的研发投资无疑是非常高的。Tesla很可能与其他公司分享这项技术，创造另一个收入来源，类似于卖给其他OEM的电动车积分。

下表列出了Tesla的神经网络的特点。这些数据是从当天视频中提取的。笔者在一些地方加入了对芯片和系统架构的理解。

Tesla的设计目标是在其芯片和系统中扩展三个系统特性：计算性能、高带宽和计算节点之间的低延迟通信。高带宽和低延迟一直以来都很难扩展到数百或数千的计算节点。看起来Tesla已经成功地将所有三个参数组织在一个连接的二维网格格式中。

训练节点

训练节点是D1芯片上最小的训练单元。它有一个64位处理器，具有4宽标量和4路多线程的程序执行。该CPU还具有8*8向量乘法的2宽向量数据路径。

该CPU的指令集架构（ISA）是为机器学习和神经网络训练任务量身定制的。该CPU支持多种浮点格式，32位、16位和8位。FP32、BFP16，以及一种新的格式CFP8或可配置的FP8。

该处理器有1.25MB的高速SRAM用于程序和数据存储。该存储器使用ECC或纠错代码以提高可靠性。

为了降低训练节点之间的延迟，Tesla挑选了信号在2GHz+时钟频率的一个周期内可传播的最远距离。这定义了训练节点的距离，以及CPU及其支持电子设备的复杂程度。这些参数还允许一个CPU以每秒512G比特的速度与四个相邻的训练节点通信。

训练节点的最大性能取决于使用的算法，常以浮点性能来比较。Training Tile的32位浮点性能（FP32）的最大性能是64GFLOPs。BFP16或CFP8算术的最大性能是1,024GFLOPs。

D1芯片

令人印象深刻的D1芯片是一种专用于神经网络训练的设计。D1采用7nm工艺制造，在一个面积为645平方毫米的裸片中封装了500亿个晶体管。该芯片电线长度超过11英里，功耗为400W左右。

D1芯片有一个带有高速、低功耗SerDes的I/O环，总共有576条通道环绕着芯片。每条通道的传输速率为112Gbps。D1的最大片上传输速率为10Tbps（每秒10兆位）。芯片每侧的最大片外传输速率为4Tbps。

由于D1芯片上的354个CPU中的每一个都有1.25MB的SRAM，加起来就有超过442MB的SRAM。D1芯片的最大性能也是基于354个训练节点的CPU阵列。

D1的32位浮点计算的最大性能达到22.6TFLOPs。16位浮点计算的最大性能为362TFLOPs。

Training Tile

Tesla的Training Tile是扩展AI训练系统的基石。一个Training Tile将25个D1芯片集成到一个晶圆上，并被封装成多芯片模块（MCM）。Tesla认为这可能是芯片行业中最大的MCM。Training Tile被封装成一个大芯片，可以通过一个高带宽连接器与其他Training Tile连接，保留了Training Tile的带宽。

Training Tile的封装包括多层电源和控制、电流分配、计算平面（25个D1芯片）和冷却系统。Training Tile用于IT中心，而不是自动驾驶汽车。

Training Tile提供单个D1芯片的25倍性能，或16位浮点计算的9Peta FLOPs，32位浮点计算则高达565TFLOPs。

12块232配置的Training Tile可以装在一个柜子里，Tesla称它为Training Matrix。

ExaPOD

Tesla描述的最大系统是ExaPOD。它是由120个Training Tile构成的。加起来有3000个D1芯片和106.2万个训练节点。它装在10个机柜中，显然是为IT中心使用的。

ExaPOD的最大性能是16位浮点计算的1.09Exa FLOPs，32位浮点计算的67.8Peta FLOPs。

Dojo软件和DPU

Dojo软件旨在支持大型和小型神经网络的训练。特斯拉有一个编译器来创建软件代码，利用训练节点、D1芯片、Training Tile和ExaPOD系统的结构和能力。它使用的是PyTorch开源机器学习库，并进行了扩展来利用D1芯片和Dojo系统架构。

这些能力允许大型神经网络被分割和映射，以提取不同的并行性、模型、图形、数据的并行度，从而加速大型神经网络的训练。编译器使用多种技术来提取并行性。它可以利用数据模型图并行技术对网络进行转换以实现细粒度并行，并可以优化以减少内存占用。

Dojo接口处理器用于与IT和数据中心的主机通信。它用PCIe 4.0连接到主机，并通过上述的高带宽连接到基于D1的系统。接口处理器还为D1系统提供高带宽的DRAM共享内存。

基于D1的系统可以被细分和划分为称为Dojo Processing Unit的单元。DPU由一个或多个D1芯片、一个接口处理器和一个或多个计算机主机组成。DPU虚拟系统可以根据在其上运行的神经网络的需要，扩大或缩小规模。

总结

Tesla的神经网络训练芯片、系统和软件都非常令人印象深刻。有很多创新，比如保留了巨大的带宽和从芯片到系统的低延迟。Training Tile的电源和冷却的封装看起来很创新。

神经网络训练系统是用于数据中心的，肯定会被用于改进Tesla的AV软件。其他公司很可能也会使用这些Tesla神经网络训练系统。

看起来Tesla希望或依赖这种神经网络训练创新，继续使其基于纯视觉的自动驾驶系统持续改进。这是正确的方向吗？时间会告诉我们，到目前为止，Elon Musk的大部分赌注尽管都有一些时间上的推迟，但都是正确的。

 [参考文章]

Tesla AI Day Perspectives — Egil Juliussen