运算性能突破360 TFLOPS,特斯拉展示全自动驾驶汽车AI训练芯片Dojo D1

不让Nvidia、Google专美于前,在今年Tesla AI Day活动上,美国电动汽车大厂特斯拉(Tesla)展现自家最新自动驾驶汽车应用同时,首度公开披露一款自制AI训练芯片Dojo D1,强调拥有GPU的强大运算力,更兼具CPU的运用弹性,以及超高速传输带宽,甚至在AI算力表现也优于Google的TPU v3,可提供362 TFLOPS运算性能,特斯拉不只用它加速AI训练创建全自动驾驶汽车,未来也将作为首款人型机器人Tesla Bot的AI训练使用。

这颗自动驾驶汽车AI芯片D1,是由Tesla工程团队一手开发完成,从芯片架构、制作到封装,采用先进7纳米制程,其内置高达500亿个晶体管。根据Tesla的介绍,D1处理器芯片是由354个高性能训练节点所组成的,每个训练节点都可视为是一个个运算芯片,每个训练芯片皆采用定制ISA指令集架构设计完成,特别对于ML工作负载执行加以优化,并内置1.25MB高速的SRAM以及低延迟和高带宽的网络fabric,使得单一芯片,在BF16/CFP8测试基准下,其运算性能可达1.024 TFLOPS。

由354个训练节点所组成的D1模块芯片,其运算力更一举达到362 TFLOPS(每秒1万亿次浮点运算),若以目前市面已知的ML芯片 (TPU v3、GPU(HBM-Links互联)或其他创业公司ML芯片)性能来做比较,Tesla指出,D1算力表现还优于其他市面ML芯片,甚至比Google的TPU v3表现都还好。另外这颗处理器本身的热设计功耗(TDP)仅有400瓦,相较之下,TPUv3有达到450瓦。

不只运算性能,Tesla也特别强调这颗AI处理器,在设计上采用新的芯片互联架构,可提供内部高速互联,总带宽可达到每秒40TB的传输能力,而且每个D1芯片,能支持最多576信道,可用于高速I/O应用,跟当前最先进网络交换机相比,Tesla表示,D1芯片可提供高达两倍的传输带宽。

Tesla Dojo项目主要负责人Genesh Venugopa则说:“Dojo D1处理器芯片是一款纯机器学习训练机器,完全专为ML训练和高速带宽而设计。”

不单是自动驾驶汽车AI训练上会用到,Tesla首席执行官马斯克在会中也透露,明年将展示一款人型机器人原型Tesla Bot,也将运用到这颗AI训练芯片,作为机器人AI训练来使用。他表示,Tesla是世上最大的机器人公司,“因为我们的汽车就像装了轮子的半感知机器人(semi-sentient robot),所以运用到人型机器人身上也说得通。”

Genesh Venugopa表示,该研究团队最新还以约1,500个D1芯片互联打造一套ML训练引擎系统,称为训练瓦片(Training Tile),每个Training Tile系统,可提供运算性能更高,达到9 petaFLOPS,等于是一台AI超级计算机。目前他们已完成一台Training Tile原型。

为了打造这套AI训练系统,研究团队费尽心思,不仅设计Dojo接口处理器作为host bridge,用于与PCIe Gen4的连接,还加入高速带宽的DRAM共享内存,并以Radix网络连接创建低延迟Mesh架构。另在电源方面,则重新定制一个电力调整器模块,搭配52伏特直流电输入,可以做到更省功耗,更省电。不只硬件,在软件面向,他们也结合了DPU (Dojo Processing Unit)、Dojo Compiler Engine、Neural Net models等软件组件,创建完整软件堆栈,借此来实现AI运算规模化。

接下来,他们还要以120个Training Tile系统创建一个超大型ExaPOD运算集群,其运算性能更可达1.1 ExaFLOPS(每秒一百万万亿次运算),还具有多达每秒36TB对外高速带宽。待完成以后,Genesh Venugopa表示,它将是世上最快的AI训练超级计算机,预期在相同成本下,可提供多4倍性能,而且更省电,每瓦性能将能提升1.3倍之多。未来还会开发下一代Dojo硬件,性能更提升10倍。