特斯拉希望将机器学习芯片带到 Dojo

为了满足人们对更大的人工智能和机器学习模型的渴望，特斯拉在 Hot Chips 34 上透露了他们完全定制的超级计算架构 Dojo 的大量细节。

该系统本质上是一个巨大的可组合超级计算机，尽管与我们在 返回顶部，它由一个完全定制的架构构建而成，涵盖计算、网络和输入/输出 (I/O) 芯片到指令集架构 (ISA)、电力传输、封装和冷却。 所有这些都是为了大规模运行量身定制的特定机器学习训练算法的明确目的。

特斯拉硬件工程高级总监 Ganesh Venkataramanan 表示：“现实世界的数据处理只有通过机器学习技术才能实现，无论是自然语言处理、在为人类视觉设计的街道上驾驶，还是与日常环境交互的机器人技术。”在他的主题演讲中。

然而，他认为扩展分布式工作负载的传统方法未能以跟上机器学习需求所需的速度加速。 实际上，摩尔定律并没有削减它，也没有可用于大规模 AI/ML 训练的系统，即 CPU/GPU 的某种组合，或者在极少数情况下使用专用 AI 加速器。 

“传统上，我们制造芯片，将它们放在封装上，封装在 PCB 上，然后进入系统。 系统进入机架，”Venkataramanan 说。 问题是每次数据从芯片移动到封装并从封装移出时，都会导致延迟和带宽损失。

数据中心三明治

因此，为了克服这些限制，Venkataramanan 和他的团队从头开始。

“就在我对 Elon 的采访中，他问我你能做什么不同于用于 AI 的 CPU 和 GPU。 我觉得整个团队仍在回答这个问题。”

这导致了 Dojo 训练块的开发，这是一个独立的计算集群，占用一个半立方英尺，能够在 556kW 液冷封装中实现 32 TFLOPS 的 FP15 性能。

每个 tile 都配备了 11GB 的 SRAM，并在整个堆栈中使用自定义传输协议通过 9TB/s 结构连接。

“这个训练块代表了从计算机到内存到电力输送、再到通信的无与伦比的集成量，无需任何额外的开关，”Venkataramanan 说。

训练模块的核心是特斯拉的 D1，这是一个 50 亿个晶体管芯片，基于台积电的 7nm 工艺。 特斯拉表示，每个 D1 能够在 22W 的 TDP 下实现 32 TFLOPS 的 FP400 性能。 然而，特斯拉指出，该芯片能够运行广泛的浮点计算，包括一些自定义计算。

“如果你比较晶体管的平方毫米，这可能是现有技术的最前沿，”Venkataramanan 说。

特斯拉随后采用了 25 个 D1，将它们分类为已知良好的裸片，然后使用台积电的晶圆系统技术将它们封装，以“以极低的延迟和极高的带宽实现大量计算集成，”他说。

然而，晶圆上系统设计和垂直堆叠架构在供电方面带来了挑战。

根据 Venkataramanan 的说法，当今大多数加速器都将电源直接放置在硅附近。 他解释说，虽然经过验证，但这种方法意味着加速器的大部分区域必须专用于这些组件，这使得 Dojo 不切实际。 相反，特斯拉将他们的芯片设计为直接通过芯片底部提供电力。 

把它放在一起

“我们可以用这个训练块构建整个数据中心或整个建筑物，但训练块只是计算部分。 我们还需要喂养它，”Venkataramanan 说。

为此，Tesla 还开发了 Dojo 接口处理器 (DIP)，它充当主机 CPU 和训练处理器之间的桥梁。 DIP 还用作共享高带宽内存 (HBM) 的来源和高速 400Gbit/sec NIC。

每个 DIP 具有 32GB 的 HBM，最多可以将其中五张卡以 900GB/s 的速度连接到一个训练图块，以总计 4.5TB/s 的速度连接到主机，每个图块总共有 160GB 的 HBM。

Tesla 的 V1 配置对（或 150 个 D1 芯片）阵列支持四个主机 CPU，每个主机 CPU 配备五个 DIP 卡，以实现声称的 BF16 或 CFP8 性能 exaflop。

总而言之，Venkataramanan 说架构 - 详细 在这里深入 by 下一个平台 – 使特斯拉能够克服与 Nvidia 和 AMD 等传统加速器相关的限制。

“传统加速器是如何工作的，通常你会尝试将整个模型安装到每个加速器中。 复制它，然后将数据流过它们中的每一个，”他说。 “如果我们的模型越来越大，会发生什么？ 这些加速器可能会因为内存不足而失效。”

他指出，这不是一个新问题。 例如，Nvidia 的 NV-switch 使内存能够跨大量 GPU 进行池化。 然而，Venkataramanan 认为这不仅增加了复杂性，而且还引入了延迟和带宽妥协。

“我们从一开始就考虑到了这一点。 我们的计算模块和每个模具都是为安装大型模型而制造的，”Venkataramanan 说。

软件

这种专门的计算架构需要专门的软件堆栈。 然而，Venkataramanan 和他的团队认识到，可编程性将决定 Dojo 的成败。

“当我们设计这些系统时，软件对应物的易用性是最重要的，”他说。 “研究人员不会等待你的软件人员编写手写内核来适应我们想要运行的新算法。”

为此，Tesla 放弃了使用内核的想法，并围绕编译器设计了 Dojo 的架构。

“我们所做的是使用 PiTorch。 我们创建了一个中间层，它可以帮助我们并行化以扩展它下面的硬件。 一切都在下面是编译的代码，”他说。 “这是创建适用于所有未来工作负载的软件堆栈的唯一方法。”

尽管强调软件的灵活性，Venkataramanan 指出，目前在他们的实验室中运行的平台暂时仅限于特斯拉使用。

“我们首先关注的是我们的内部客户，”他说。 “埃隆已经公开表示，随着时间的推移，我们将把它提供给研究人员，但我们没有时间框架。 ®