智东西
作者 ZeR0
编辑 漠影

智东西4月11日报道，美国AI三巨头不仅在大模型赛道争奇斗艳，还纷纷卷起自研AI芯片。昨天“AI界汪峰”谷歌刚推出新款云端定制芯片，今天Meta就将硅谷注意力吸引走，重磅宣布推出第二代自研AI训练和推理芯片MTIA v2。

//oss.zhidx.com/uploads/2024/04/661727db37b48_661727db33fde_661727db33fb8_1.mp4

Meta去年5月推出的第一代MTIA v1采用台积电7nm芯片。MTIA v2则换上了台积电5nm，每PE缓存从上一代的128KB增加到384KB，频率从800MHz提升到1.35GHz，INT8精度下的稠密算力达到上一代的近3.5倍，稀疏算力达到上一代的近7倍，达到708TFLOPS。

有意思的是，MTIA v2的面积仅变大13%，不过功耗增加到上一代的3.6倍，达到90W，上一代只有25W。

相比之下，英伟达H100 PCIe的最大功耗约在350~500W区间（SXM版高达700W），INT8精度下稀疏算力为3026TFLOPS。

在平台层，通过2倍的设备数量和2插槽CPU，MTIA v2能够实现6倍的模型服务吞吐量、比第一代MTIA系统提高多达1.5倍的每瓦性能。

与MTIA v1相比，MTIA v2性能显著改进，将Meta以前的解决方案的计算和内存带宽提高了1倍以上，同时保持了与工作负载的紧密联系，旨在增强Meta的推荐广告模型。

新MTIA芯片可处理低复杂性（LC）和高复杂性（HC）的排名和推荐模型，这些模型是Meta产品的组成部分。这些模型中，模型大小和每个输入样本的计算量可能相差约10-100倍。

Meta控制全栈，因此相比商用GPU，可实现更高的效率。Meta在优化内核、编译器、运行时和主机服务堆栈方面取得了重大进展。

MTIA v2与为MTIA v1开发的代码兼容，因此能快速开发新一代MTIA芯片，从第一块芯片到在16个地区运行的生产模型仅用时不到9个月。

MTIA已部署在数据中心中，正在为生产中的模型提供服务。

“事实证明，它在提供特定于Meta工作负载的性能和效率的最佳组合方面，是对商用GPU的高度补充。”Meta博客文章中写道。

Meta的下一代大规模基础设施将支持生成式AI产品和服务、推荐系统和先进的AI研究。Meta预计这项投资将在未来几年增长，并将致力于继续提高每瓦性能。

下周Meta将推出Llama 3大语言模型的两个小版本，为今年夏季推出被预期为最强开源大语言模型的Llama 3旗舰版本蓄力。

据外媒报道，Llama 3最大版本可能拥有超过1400亿个参数，比上一代的2倍还多。

一、翻开MTIA v2内部架构：超大SRAM容量，计算、带宽翻倍提升

该芯片的架构基本上专注于为服务排名和推荐模型提供计算、内存带宽和内存容量的适当平衡。

因此，即便其batch size大小相对较低，也需能提供相对较高的利用率。

相对于典型GPU，MTIA v2提供了超大的SRAM容量，可在batch size有限的情况下提供高利用率，并在遇到大量潜在并发工作时提供足够的计算。

该加速器由一个8×8的处理单元网格（PE）组成。这些PE显著提高了稠密计算性能（比上一代提高256%）和稀疏计算性能（比上一代提高591%）。

这部分来自于与稀疏计算的流水线相关的体系结构的改进。它还来自于如何提供PE网格：Meta芯片团队将本地PE存储的大小增加到2倍，将片上SRAM增加到1倍，并将其带宽增加到3.5倍，并将LPDDR5的容量翻倍。

新MTIA设计还拥有改进的片上网络（NoC）架构，可将带宽翻倍，并支持以低延迟在不同的PE之间进行协调，增强可扩展性。

二、打造大型机架式系统，最多容纳72个MTIA v2芯片

硬件系统和软件栈与芯片的协同设计对于整个推理解决方案的成功至关重要。

为了支持下一代芯片，Meta开发了一个大型的机架式系统，最多可容纳72个加速器。

它由3个机箱组成，每个机箱包含12块电路板，每块电路板上有两个加速器。Meta特别设计了这个系统，这样就可以将芯片的时钟设置在1.35GHz（高于800mhz），功耗90W，而第一代设计的功耗为25W。

其设计确保提供更密集的功能，具有更高的计算、内存带宽和内存容量。这种密度使其更容易地适应广泛的模型复杂性和大小。

除此之外，Meta还将加速器之间以及主机和加速器之间的结构升级到PCIe Gen5，以增加系统的带宽和可扩展性。如果选择扩展到机架之外，还能选择添加RDMA网卡。

三、软件才是关键！看看完整的MTIA软件栈

从投资MTIA开始，软件一直是Meta关注的关键领域之一。

作为PyTorch的初始开发人员，Meta团队重视可编程性和开发人员效率。其MTIA软件栈旨在与PyTorch 2.0、TorchDynamo、TorchInductor等功能完全集成。

前端图形级捕获、分析、转换和提取机制（如TorchDynamo、torch）对MTIA是不可知的，并正在被重用。MTIA的低级编译器从前端获取输出，生成高效且特定于设备的代码。

下方是运行时堆栈，负责与驱动程序/固件接口。MTIA流接口抽象提供了基本和必要的操作，推理和（未来）训练软件都需要管理设备内存，以及在设备上运行操作符和执行编译图。

最后，运行时与驱动程序交互，驱动程序位于用户空间中——Meta团队做出这个决定是为了使其能够更快地迭代生产栈中的驱动程序和固件。

在许多方面，这个新的芯片系统运行的软件栈类似于MTIA v1，使团队部署的速度更快，因为Meta团队已经完成了许多必要的集成和开发工作，需要能够在这个架构上运行其应用程序。

新的MTIA被设计为与为MTIA v1开发的代码兼容。由于Meta已将完整的软件栈集成到芯片上，他们在几天内就启动并运行了这个新芯片。这使其能够快速开发下一代MTIA芯片，在不到9个月的时间里，从第一块芯片到在16个地区运行的生产模型。

四、Triton-MTIA编译器：为MTIA硬件生成高性能代码

通过创建Triton-MTIA编译器后端来为MTIA硬件生成高性能代码，Meta进一步优化了软件栈。

Triton是一种开源语言和编译器，用于编写高效的ML计算内核。它提高了开发人员编写GPU代码的效率，Meta发现Triton语言与硬件无关，足以适用于像MTIA这样的非GPU硬件架构。

Triton-MTIA后端执行优化以最大化硬件利用率并支持高性能内核。它还提供了利用Triton和MTIA自动调优基础设施来探索内核配置和优化空间的关键旋门。

Meta已实现对Triton语言特性的支持，并集成到PyTorch 2中，为PyTorch操作符提供了广泛的覆盖。例如，得益于TorchInductor，开发人员可以在提前（AOT）和即时（JIT）工作流程中利用Triton-MTIA。

Meta观察到使用Triton-MTIA极大地提高了开发人员的效率，它允许扩展计算内核创作并显着扩展PyTorch操作符的支持。

结语：Meta在定制芯片方面的持续投资

今年Meta正加大资本支出，其中大部分用于与AI相关的服务器和数据中心建设，包括向英伟达大量采购先进GPU芯片，也很可能包括新一代MTIA芯片的研发支出。

Meta去年资本支出为280亿美元，占其收入的 21%，今年其预计资本支出将高达370亿美元。

根据Meta博客文章，MTIA推理加速器是其更广泛的全栈开发计划的一部分，用于定制特定领域的芯片，以解决独特的工作负载和系统。

MTIA将是Meta长期路线图的重要组成部分，为Meta独特的AI工作负载构建和扩展最强大、最高效的基础设施。

为了实现其定制芯片的目标，Meta不仅要在计算芯片上投资，还要在内存带宽、网络和容量以及其他下一代硬件系统上投资。

目前Meta正在进行几个旨在扩大MTIA范围的项目，包括对生成式AI工作负载的支持。