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

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

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

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

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

MTIA V1

MTIA v1（推理）芯片（die）

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

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

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

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

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

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

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

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

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

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

帶有 MTIA 的樣品測(cè)試板

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

MTIA 軟件棧

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

下圖為 MTIA 軟件?？蚣軋D：

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

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

MTIA 性能

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

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

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

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

審核編輯 ：李倩