Dojo系统详细构成
神经网络
训练节点
训练节点是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
from A to B