[前沿科技]Meta首代AI推理加速器7nm制程

据天中威科技小编了解，近日，Meta透露了其在人工智能方面取得的最新进展。人们提起Meta时，通常会想到其应用程序，包括Facebook、Instagram、WhatsApp或即将推出的元宇宙。但许多人不知道的是这家公司设计和构建了非常复杂的数据中心来运营这些服务。

与AWS、GCP或Azure等云服务提供商不同，Meta不需要披露有关其硅芯选择、基础设施或数据中心设计的细节，除了其OCP设计用来给买家留下深刻印象。Meta的用户希望获得更好、更一致的体验，而不关心它是如何实现的。

在Meta，AI工作负载无处不在，它们构成了广泛用例的基础，包括内容理解、信息流、生成式AI和广告排名。这些工作负载在PyTorch上运行，具有一流的Python集成、即时模式（eager-mode）开发和API简洁性。特别是深度学习推荐模型（DLRMs），对于改善Meta的服务和应用体验非常重要。但随着这些模型的大小和复杂性的增加，底层的硬件系统需要在保持高效的同时提供指数级增长的内存和计算能力。

Meta发现，对于目前规模的AI运算和特定的工作负载，GPU的效率不高，并不是最佳选择。因此，该公司提出了推理加速器MTIA，帮助更快地训练AI系统。

MTIA V1

MTIA v1（推理）芯片（die）

2020年，Meta为其内部工作负载设计了第一代MTIA ASIC推理加速器。该推理加速器是其全栈解决方案的一部分，整个解决方案包括芯片、PyTorch和推荐模型。

MTIA加速器采用TSMC 7nm工艺制造，运行频率为800 MHz，在INT8精度下提供102.4 TOPS，在FP16精度下提供51.2 TFLOPS。它的热设计功耗 (TDP) 为25W。

MTIA加速器由处理元件 (PE)、片上和片外存储器资源以及互连组成。该加速器配备了运行系统固件的专用控制子系统。固件管理可用的计算和内存资源，通过专用主机接口与主机通信，协调加速器上的job执行。

内存子系统使用LPDDR5作为片外DRAM资源，可扩展至128 GB。该芯片还有128 MB的片上SRAM，由所有PE共享，为频繁访问的数据和指令提供更高的带宽和更低的延迟。

MTIA加速器网格包含以8x8配置组织的64个PE，这些PE相互连接，并通过网状网络连接到内存块。整个网格可以作为一个整体来运行一个job，也可以分成多个可以运行独立job的子网格。

每个PE配备两个处理器内核（其中一个配备矢量扩展）和一些固定功能单元，这些单元经过优化以执行关键操作，例如矩阵乘法、累加、数据移动和非线性函数计算。处理器内核基于RISC-V开放指令集架构 (ISA)，并经过大量定制以执行必要的计算和控制任务。

每个PE还具有128 KB的本地SRAM内存，用于快速存储和操作数据。该架构最大限度地提高了并行性和数据重用性，这是高效运行工作负载的基础。

该芯片同时提供线程和数据级并行性（TLP和DLP），利用指令级并行性 (ILP)，并通过允许同时处理大量内存请求来实现大量的内存级并行性 (MLP)。

MTIA v1系统设计

MTIA加速器安装在小型双M.2板上，可以更轻松地集成到服务器中。这些板使用PCIe Gen4 x8链接连接到服务器上的主机CPU，功耗低至35 W。

带有MTIA的样品测试板

托管这些加速器的服务器使用来自开放计算项目的Yosemite V3服务器规范。每台服务器包含12个加速器，这些加速器连接到主机CPU，并使用PCIe交换机层级相互连接。因此，不同加速器之间的通信不需要涉及主机CPU。此拓扑允许将工作负载分布在多个加速器上并并行运行。加速器的数量和服务器配置参数经过精心选择，以最适合执行当前和未来的工作负载。

MTIA软件栈

MTIA软件（SW）栈旨在提供给开发者更好的开发效率和高性能体验。它与PyTorch完全集成，给用户提供了一种熟悉的开发体验。使用基于MTIA的PyTorch与使用CPU或GPU的PyTorch一样简单。并且，得益于蓬勃发展的PyTorch开发者生态系统和工具，现在MTIA SW栈可以使用PyTorch FX IR执行模型级转换和优化，并使用LLVM IR进行低级优化，同时还支持MTIA加速器自定义架构和ISA。

下图为MTIA软件栈框架图：

作为SW栈的一部分，Meta还为性能关键型ML内核开发了一个手动调整和高度优化的内核库，例如完全连接和嵌入包运算符。在SW栈的更高层级可以选择在编译和代码生成过程中实例化和使用这些高度优化的内核。

此外，MTIA SW栈随着与PyTorch 2.0的集成而不断发展，PyTorch 2.0更快、更Python化，但一如既往地动态。这将启用新功能，例如TorchDynamo和TorchInductor。Meta还在扩展Triton DSL以支持MTIA加速器，并使用MLIR进行内部表示和高级优化。

MTIA性能

Meta比较了MTIA与其他加速器的性能，结果如下：

Meta使用五种不同的DLRMs（复杂度从低到高）来评估MTIA

此外，Meta还将MTIA与NNPI以及GPU进行了比较，结果如下：

评估发现，与NNPI和GPU相比，MTIA能够更高效地处理低复杂度（LC1和LC2）和中等复杂度（MC1和MC2）的模型。此外，Meta尚未针对高复杂度（HC）模型进行MTIA的优化。

参考链接：

https://ai.facebook.com/blog/meta-training-inference-accelerator-AI-MTIA/

文章来源：机器之心