软硬件融合 · 3月22日

Tesla Dojo:极致可扩展的AI训练芯片和平台

编者按

Tesla Dojo芯片和相应的整个集群系统,跟传统的设计理念有很大的不同。其基于整个POD级的超强的扩展性和全系统栈协同设计能力,值得我们大家学习和借鉴。

Tesla Dojo是Jim Keller主导的项目,整个架构跟Tenstorrent的设计理念有很多相似之处。后面我们会花时间一起学习Tenstorrent的架构和系统。

这里,让我们先学习Dojo的芯片设计和系统架构吧!


参考文献:

  1. Tesla's AI chip REVEALED! (Project Dojo)
  2. https://zhuanlan.zhihu.com/p/...,青031,浅谈Tesla Dojo集群——大人,时代变了
  3. https://semianalysis.com/tesl...

Tesla DOJO综述

image.png

Tesla Dojo的最小单位是处理器Node。处理器的基本特征:向量扩展、4路超线程、顺序流水线、针对ML定制指令的ISA。

image.png

354个Node组成一个Chip。

image.png

然后25个Chip Die封装在一个基板上,组成一个Tile。
image.png

然后,12个Tile组成一个Cabinet。
image.png

然后10个Cabinet组成一个ExaPOD。

1 ExaPOD = 10 Cabinet = 120 Tile = 3000 D1 Chip = 1, 062, 000 Node.
image.png

一个POD,即一个集群。Dojo系统的每个Node都是完全对称的,是一个POD级完全UMA的架构。或者说,Dojo的扩展性,跨过了芯片、Tile、Cabinet,达到了POD级别。

Dojo是目前扩展性最高的系统平台。


Tesla DOJO详情

image.png

DOJO项目是特斯拉目前最新一代的AI训练平台。

image.png

AI训练对性能和神经网络训练能力的需求是无止境的。正因为如此,特斯拉针对自己的自动驾驶AI训练场景,设计了Dojo,一个超快的训练超级计算机系统。

随着AI训练算法模型复杂度飞速增长,对集群的规模需求也越来越大。如上图,短短的两年时间,对用于训练的GPU数量就增长了几乎8倍。
image.png

Dojo的目标是:

  • 实现最佳的AI训练性能;
  • 支持所有的更大更复杂的模型;
  • 与此同时,能够兼顾功耗和成本效益。

在这些目标的驱动下,特斯拉设计了Dojo。

image.png

考虑到能够容纳更大的模型,必须增强集群的可扩展能力。于是,特斯拉构建了一个分布式计算架构。

这个架构的计算单元可以是CPU、GPU、加速器等,他们有很小的内存,然后通过一个二维网络把这些单元连接起来。

这个架构的最大特点是,它有很好的扩展性,可以很容易按比例缩放计算资源。

但扩大带宽非常困难,减少延迟也非常困难,特斯拉在设计的一些办法突破了这些限制。
image.png

特斯拉设想了一个大型的计算平面,它充满了非常健壮的计算单元,支持大内存池,并以非常高的带宽和低延迟的以2D网格结构互连。

在此基础上,对超大规模的神经网络进行划分和映射,提取不同的数据并行度模型图。然后神经编译器将利用空间和时间局部性,这样它可以减少局部区域的通信足迹,并因此减少全局通信。这样带宽利用率就可以随着想要的计算平面不断缩放。
image.png

要从栈顶一直到栈底整个系统的优化协同,来解决这个问题,并消除任何层级上的所有瓶颈。

让我们以一种由内而外的方式开始这段旅程。首先,从最底层的芯片开始。

image.png

前面说过,芯片由计算单元组成。

最小规模的实体被称为训练节点,这个节点的选择对于确保无缝伸缩非常重要。如果选择地太小,它会运行得很快,但是同步的开销会很大,软件会占主导地位。

如果选择地太大,在实际硬件中实现时会很复杂,最终会遇到内存瓶颈问题。特斯拉把延迟和带宽作为主要优化点。接下来看看特斯拉是怎么做到的。
image.png

需要考虑的是选择信号在一个非常高的时钟周期内所能穿越的最远距离。在Dojo中,频率是2GHz。在这种情况下,在2GHz周围画了一个方框,然后用金属线把他们填满,这就是整个芯片可以提供的最高带宽。
image.png

然后在网格上添加机器学习的计算节点,节点有一个很大的SRAM池,最后是一个可编程的CPU核来控制。

image.png

高性能的训练节点是一个64位的超标量CPU,优化了矩阵乘法单位和向量SIMD。支持FP32、FP16和一种新的格式CFP8(可配置的FP8)。

有1.25MB的支持快速ECC的SRAM,以及针对性设计的低延迟高带宽的网络交换。

image.png

这是整个设计中的最小实体,计算能力可以超过万亿次浮点运算。

image.png

我们来看看这个计算节点的微架构。这是一个非常强大的架构。

有一个顺序的超标量CPU,4路标量和2路向量的流水线。虽然向量和标量流水线可能会乱序但对于最纯粹的流水线,我们仍然称其为有序的。

它还具有四路超线程,这提高了利用率。

因为我们可以同时进行计算和数据传输,我们针对机器学习工作负载专门优化定制了ISA。它有一些特性,如转置、收集、链接遍历、广播等等。

image.png
即使在物理后端领域,我们也将其模块化,这样我们就可以开始在任何方向上连接这些训练节点,并开始形成我们设想的计算平面。
image.png

当点击354个这样的训练节点时,得到的计算阵列,它能够提供362TFLOPS的机器学习计算。

image.png

还有连接这些节点的的高带宽结构

image.png

在这个计算阵列周围,我们为它提供了高速并且低功耗的Serdes,576个Serdes使芯片能够有最大的I/O带宽。这里需要强调的是一个比较点是,这是目前最先进的网络交换芯片带宽的两倍多。网络交换芯片被认为是I/O带宽的黄金标准。

image.png

如果我们把所有这些集合到一起,我们训练优化的芯片D1,是七个纳米技术制造,包含500亿个晶体管,面积645毫米平方。

大家会注意到,百分百的面积是用于机器学习训练和带宽,没有资源浪费,也没有向前兼容的包袱。这是一个纯机器学习的机器。

image.png

这是Flip BGA封装的D1芯片。完全由特斯拉团队设计,从架构到GDS再到封装。

D1芯片拥有GPU级别的性能,也具有CPU级别的灵活性,以及两倍的网络芯片级的I/O带宽。

image.png

如果要画出纵坐标上的I/O带宽,以及水平坐标上的万亿次浮点运算的性能。这里给出最先进的机器学习芯片的坐标,也包括一些初创公司。你很容易就能看出为什么Tesla的设计如此出类拔萃。

image.png

已经有了这个基本的物理构建模块,那么如何围绕它设计系统呢?

image.png

既然D1芯片可以无缝连接,无需任何胶水,我们就可以轻松地把它们放在一起。

image.png

把50万个训练节点放在一起,形成计算平面。这是1500个D1芯片无缝连接在一起。

image.png

然后在两端添加Dojo接口处理处理器,这是连接到数据中心中标准主机的主桥。

它的一边连接PCIe Gen4,另一端通过一个高带宽的互联接口连接到计算平面。

接口处理器不仅提供主桥,还为计算平面提供高带宽的DRAM共享内存。

image.png

此外,接口处理器还可以允许我们有一个更高的根网络连接,以实现计算平面。

我们必须想出一种将这些芯片整合在一起的新方法。

image.png

这就是所说的训练块,这是系统的基本规模单位。
image.png

这是一个突破性的集成,把25个D1 DIE紧密集成到一起,这样能够保持尽可能最大的带宽。

image.png

此外,还创建了一个连接器,一个高带宽、高密度的连接器,它保持了从训练Tile出来的高带宽。

然后是9个1TB带宽的连接器,确保能够满足大量的IO带宽。这可能是芯片行业中最大的MCM多芯片模块。设计非常不容易,当时还没有工具。Tesla的工程师创造了新的方法使之成为现实。

有了高带宽I/O的计算平面Tile,接下来需要考虑如何供电了。

image.png

特斯拉设计了一种垂直供电的新方法。定制了一个自定义的电压调节模块,可以直接回流到这个扇出晶圆上。

在芯片封装里,将PCB级别的回流技术引入到这个晶圆上,集成了很多的东西。

但特斯拉并没有就此止步。

image.png

整合了整个电气、散热和机械部件。为了适配训练Tile的52V DC输入的完全集成接口。

这是前所未有的。这是一个惊人的工程,计算平面完全正交于电源和散热,这使得高带宽的计算平面成为可能。
image.png

它是一个9千万亿次的训练Tile,这成为特斯拉系统规模扩展的基本单位。

image.png

这是实物,真不敢相信,能以每秒九千兆次速度运行的系统居然可以捧在手里。

image.png

这是在dojo上运行的min gpd2的测试结果。

image.png

下一步,该如何形成一个计算集群。目前,大家应该明白了为什么模块的部分要讲这么仔细。

image.png

一个2 × 3共6个Tile组成一个Tray,2个Tay组成训练矩阵,放在一个Cabinet里。性能可以达到每秒100千万亿次计算。

就停止在这吗?没有。Cabinet

image.png
特斯拉进行了无缝整合,打破了Cabinet墙。通过保持高带宽,无缝地集成了这些Tile。

没有带宽分流,没有带宽剪辑,所有的Tile都用相同的带宽无损耗地连接在一起。

image.png

这样,就有了一个百亿亿次的POD。由10个Cabinet组成。

超过一百万个训练节点。

对这个训练节点进行了细致的分析,这里的一百万个节点,他们的带宽都是一致的。
image.png

不仅仅是硬件,软件对于确保可伸缩性非常重要。并不是所有的工作都需要一个巨大的集群,

所以从一开始就规划好。
image.png

计算平面可以被细分、可以被划分成单元,称为dojo处理单元(Dojo Processing Unit,DPU)。

DPU由一个或多个D1芯片组成,它也有接口处理器和一个或多个主机。这可以根据任何算法的需要放大或缩小,任何网络都可以运行。

image.png

用户需要做什么?

他们必须最少限度的更改脚本,这是因为我们拥有强大的编译器套件。

编译器考虑了细粒度的并行性,并将神经网络映射到计算平面上。

image.png

编译器使用多种技术来提取并行性。利用数据模型图并行技术,不仅可以将网络转换为细粒度并行,还可以进行优化以减少内存占用。因为Dojo的高带宽特性的Fabric,模型的并行性可以扩展到非常高的水平,业界目前无法达到的水平。

其他一个平台的扩展性可能被限制在芯片的范围,而我们由于高带宽,可以扩展到整个Tile甚至超过Tile的范围。

因此,大型网络可以在低批处理的情况下有效地映射到Dojo,并提取利用率和新的性能水平。

编译器能够处理高层次的动态控制流,如if, else等。

image.png

编译器引擎只是我们整个软件套件的一部分。该堆栈由Pytorch的Dojo扩展组成,该扩展确保了ML科学家所习惯的相同的用户级接口。

编译器会动态地生成代码,以便在后续执行时重用。它有一个LLVM后端编译器,可以为硬件生成二进制文件。

这确保了可以为硬件创建优化过的代码,而不需要依赖任何一行手写的内核。驱动程序堆栈负责多主机多分区,这在前几张幻灯片中介绍过。

此外,在软件堆栈中还有分析器和调试器。
image.png
通过全栈的垂直整合优化的方式,打破了传统的扩展障碍。这就是ojo如何在堆栈上实现模块化,以提高性能的新的水准。

image.png

这将是一台速度最快的AI训练计算机,在相同的成本下,性能每瓦特提高1.3倍。这不仅节省了能源,而且节省了5倍的占地面积。

这将是Dojo超级电脑。

image.png

还没有结束。

特斯拉已经有了一个完整的下一代计划,正在考虑可以做的10倍以上的改进,从各个不同的层面:芯片、系统、集群、软件。

(正文完)

来源:软硬件融合
微信公众号:
 title=

相关文章推荐

更多软硬件技术干货请关注软硬件融合专栏。
2 阅读 1.1k
推荐阅读
关注数
2752
内容数
82
软硬件融合
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息