近日，Meta 透露了其在人工智能方面取得的最新進展。

人們提起 Meta 時，通常會想到其應(yīng)用程序，包括 Facebook、Instagram、WhatsApp 或即將推出的元宇宙。但許多人不知道的是這家公司設(shè)計和構(gòu)建了非常復(fù)雜的數(shù)據(jù)中心來運營這些服務(wù)。

與 AWS、GCP 或 Azure 等云服務(wù)提供商不同，Meta 不需要披露有關(guān)其硅芯選擇、基礎(chǔ)設(shè)施或數(shù)據(jù)中心設(shè)計的細節(jié)，除了其 OCP 設(shè)計用來給買家留下深刻印象。Meta 的用戶希望獲得更好、更一致的體驗，而不關(guān)心它是如何實現(xiàn)的。

在 Meta，AI 工作負載無處不在，它們構(gòu)成了廣泛用例的基礎(chǔ)，包括內(nèi)容理解、信息流、生成式 AI 和廣告排名。這些工作負載在 PyTorch 上運行，具有一流的 Python 集成、即時模式（eager-mode）開發(fā)和 API 簡潔性。特別是深度學(xué)習(xí)推薦模型（DLRMs），對于改善 Meta 的服務(wù)和應(yīng)用體驗非常重要。但隨著這些模型的大小和復(fù)雜性的增加，底層的硬件系統(tǒng)需要在保持高效的同時提供指數(shù)級增長的內(nèi)存和計算能力。

Meta 發(fā)現(xiàn)，對于目前規(guī)模的 AI 運算和特定的工作負載，GPU 的效率不高，并不是最佳選擇。因此，該公司提出了推理加速器 MTIA，幫助更快地訓(xùn)練 AI 系統(tǒng)。

MTIA V1

MTIA v1（推理）芯片（die）

2020 年，Meta 為其內(nèi)部工作負載設(shè)計了第一代 MTIA ASIC 推理加速器。該推理加速器是其全棧解決方案的一部分，整個解決方案包括芯片、PyTorch 和推薦模型。

MTIA 加速器采用 TSMC 7nm 工藝制造，運行頻率為 800 MHz，在 INT8 精度下提供 102.4 TOPS，在 FP16 精度下提供 51.2 TFLOPS。它的熱設(shè)計功耗 (TDP) 為 25 W。

MTIA 加速器由處理元件 (PE)、片上和片外存儲器資源以及互連組成。該加速器配備了運行系統(tǒng)固件的專用控制子系統(tǒng)。固件管理可用的計算和內(nèi)存資源，通過專用主機接口與主機通信，協(xié)調(diào)加速器上的 job 執(zhí)行。

內(nèi)存子系統(tǒng)使用 LPDDR5 作為片外 DRAM 資源，可擴展至 128 GB。該芯片還有 128 MB 的片上 SRAM，由所有 PE 共享，為頻繁訪問的數(shù)據(jù)和指令提供更高的帶寬和更低的延遲。

MTIA 加速器網(wǎng)格包含以 8x8 配置組織的 64 個 PE，這些 PE 相互連接，并通過網(wǎng)狀網(wǎng)絡(luò)連接到內(nèi)存塊。整個網(wǎng)格可以作為一個整體來運行一個 job，也可以分成多個可以運行獨立 job 的子網(wǎng)格。

每個 PE 配備兩個處理器內(nèi)核（其中一個配備矢量擴展）和一些固定功能單元，這些單元經(jīng)過優(yōu)化以執(zhí)行關(guān)鍵操作，例如矩陣乘法、累加、數(shù)據(jù)移動和非線性函數(shù)計算。處理器內(nèi)核基于 RISC-V 開放指令集架構(gòu) (ISA)，并經(jīng)過大量定制以執(zhí)行必要的計算和控制任務(wù)。

每個 PE 還具有 128 KB 的本地 SRAM 內(nèi)存，用于快速存儲和操作數(shù)據(jù)。該架構(gòu)最大限度地提高了并行性和數(shù)據(jù)重用性，這是高效運行工作負載的基礎(chǔ)。

該芯片同時提供線程和數(shù)據(jù)級并行性（TLP 和 DLP），利用指令級并行性 (ILP)，并通過允許同時處理大量內(nèi)存請求來實現(xiàn)大量的內(nèi)存級并行性 (MLP)。

MTIA v1 系統(tǒng)設(shè)計

MTIA 加速器安裝在小型雙 M.2 板上，可以更輕松地集成到服務(wù)器中。這些板使用 PCIe Gen4 x8 鏈接連接到服務(wù)器上的主機 CPU，功耗低至 35 W。

帶有 MTIA 的樣品測試板

托管這些加速器的服務(wù)器使用來自開放計算項目的 Yosemite V3 服務(wù)器規(guī)范。每臺服務(wù)器包含 12 個加速器，這些加速器連接到主機 CPU，并使用 PCIe 交換機層級相互連接。因此，不同加速器之間的通信不需要涉及主機 CPU。此拓撲允許將工作負載分布在多個加速器上并并行運行。加速器的數(shù)量和服務(wù)器配置參數(shù)經(jīng)過精心選擇，以最適合執(zhí)行當前和未來的工作負載。

MTIA 軟件棧

MTIA 軟件（SW）棧旨在提供給開發(fā)者更好的開發(fā)效率和高性能體驗。它與 PyTorch 完全集成，給用戶提供了一種熟悉的開發(fā)體驗。使用基于 MTIA 的 PyTorch 與使用 CPU 或 GPU 的 PyTorch 一樣簡單。并且，得益于蓬勃發(fā)展的 PyTorch 開發(fā)者生態(tài)系統(tǒng)和工具，現(xiàn)在 MTIA SW 棧可以使用 PyTorch FX IR 執(zhí)行模型級轉(zhuǎn)換和優(yōu)化，并使用 LLVM IR 進行低級優(yōu)化，同時還支持 MTIA 加速器自定義架構(gòu)和 ISA。

下圖為 MTIA 軟件棧框架圖：

作為 SW 棧的一部分，Meta 還為性能關(guān)鍵型 ML 內(nèi)核開發(fā)了一個手動調(diào)整和高度優(yōu)化的內(nèi)核庫，例如完全連接和嵌入包運算符。在 SW 棧的更高層級可以選擇在編譯和代碼生成過程中實例化和使用這些高度優(yōu)化的內(nèi)核。

此外，MTIA SW 棧隨著與 PyTorch 2.0 的集成而不斷發(fā)展，PyTorch 2.0 更快、更 Python 化，但一如既往地動態(tài)。這將啟用新功能，例如 TorchDynamo 和 TorchInductor。Meta 還在擴展 Triton DSL 以支持 MTIA 加速器，并使用 MLIR 進行內(nèi)部表示和高級優(yōu)化。

MTIA 性能

Meta 比較了 MTIA 與其他加速器的性能，結(jié)果如下：

Meta 使用五種不同的 DLRMs（復(fù)雜度從低到高）來評估 MTIA

此外，Meta 還將 MTIA 與 NNPI 以及 GPU 進行了比較，結(jié)果如下：

評估發(fā)現(xiàn)，與 NNPI 和 GPU 相比，MTIA 能夠更高效地處理低復(fù)雜度（LC1 和 LC2）和中等復(fù)雜度（MC1 和 MC2）的模型。此外，Meta 尚未針對高復(fù)雜度（HC）模型進行 MTIA 的優(yōu)化。

審核編輯 ：李倩