Meta宣佈全新訓推一體加速器：完全整合PyTorch 2，效能3倍提升 机器之心 發表於2024-04-11
Meta 瘋狂砸入數十億美元，一部分招攬人才，一部分造晶片。 Meta 正在不遺餘力地想要在生成式 AI 領域趕上競爭對手，目標是投入數十億美元用於 AI 研究。這些巨資一部分用於招募 AI 研究員。但更大的一部分用於開發硬體，特別是用於執行和訓練 Meta AI 模型的晶片。 在英特爾 宣佈其最新人工智慧 加速器硬體的第二天，Meta 便迅速公佈了關於晶片研發的最新成果：下一代 MTIA（Meta Training and Inference Accelerator），其中 MTIA 是專為 Meta AI 工作負載而設計的定製晶片系列。分析認為，Meta 此舉意在減少對英偉達及其他外部公司晶片的依賴。
其實，Meta 在去年就推出了第一代 MTIA，與前代相比，最新版本顯著改進了效能，並有助於強化內容排名和推薦廣告模型。從 Meta 官方介紹中我們可以看出，MTIA v2 目前已在其 16 個資料中心投入使用，與 MTIA v1 相比，整體效能提高了 3 倍。不過最早的 MTIA v1 預計要到 2025 年才會釋出，MTIA v2 只能再等等了，但 Meta 表示這兩款 MTIA 晶片現已投入生產。 從下圖可以看到，下一代 MTIA 採用了臺積電 5nm 製程工藝。 雖然 MTIA 主要用於訓練排名和推薦演算法，但 Meta 表示，他們的最終目標是擴充套件晶片的功能，用來訓練如 Llama 等語言模型 的生成式 AI 。
新的 MTIA 晶片從根本上專注於提供計算、記憶體頻寬和記憶體容量的適當平衡。該晶片將擁有 256MB 片上記憶體，頻率為 1.3GHz，而 v1 的片上記憶體為 128MB 和 800GHz。有訊息稱，Meta 致力於 MTIA v2 研究已有一段時間了，該專案內部代號為 Artemis。 新一代晶片的最新加速器由 8x8 個處理元件 (PE) 組成。這些 PE 顯著提高了密集計算效能（比 MTIA v1 提高了 3.5 倍）和稀疏計算效能（提高了 7 倍）。這些提高一部分歸功於架構的改進，另一部分歸功於 PE 網格供電方式：Meta 將本地 PE 儲存的大小增加了兩倍，將片上 SRAM 增加了一倍，將其頻寬增加了 3.5 倍，並將 LPDDR5 的容量增加了一倍。 最新版本的 MTIA 設計方式還採用改進的片上網路 (NoC) 架構，使頻寬加倍，並允許以低延遲的方式在不同 PE 之間進行協調。 高效地為工作負載提供服務不單單是晶片的挑戰，共同設計硬體系統和軟體堆疊對於整體推理解決方案的成功同樣至關重要。 為了支援下一代 MTIA 晶片，Meta 開發了一個大型機架式系統，最多可容納 72 個加速器。該系統由三個機箱組成，每個機箱包含 12 個板，每個板上有兩個加速器。 為此，Meta 專門設計了該系統，以便將晶片的時脈頻率設定為 1.35G Hz（從 800MHz 開始提高），並以 90 瓦的功率執行，而第一代的設計功耗為 25 瓦。Meta 的設計確保提供更密集的功能以及更高的計算、記憶體頻寬和記憶體容量。這種密集性使得可以更輕鬆地適應各種複雜性和尺寸的模型。 此外，Meta 還將加速器之間、主機與加速器之間的結構升級到 PCIe Gen5，以提高系統的頻寬和可擴充套件性。如果選擇橫向擴充套件至機架之外，還能新增一個 RDMA NIC。 從投資 MTIA 之初起，軟體就一直是 Meta 重點關注的領域之一。作為 PyTorch 的最初開發者，Meta 非常重視可程式設計性和開發效率。 在設計時，MTIA 堆疊旨在與 PyTorch 2.0 以及 TorchDynamo 、TorchInductor 等功能完全整合。同時，前端圖形級捕獲、分析、轉換和提取機制（比如 TorchDynamo、torch.export 等）與 MTIA 無關，並正在被重新使用。 此外， MTIA 的較低階別編譯器從前端獲取輸出並生成高效且裝置特定的程式碼。該編譯器本身由幾個元件組成，分別負責為模型和核心生成可執行程式碼。 下圖為負責與驅動程式 / 韌體連線的執行時堆疊。MTIA Streaming 介面抽象提供了推理和（未來）訓練軟體管理裝置記憶體以及在裝置上執行運算元、執行編譯圖所需要的基本和必要操作。 最後，執行時與使用者空間中的驅動程式進行互動，Meta 做出這一決定是為了能夠在生產堆疊中更快地迭代驅動程式和韌體。 在許多方面，這一全新晶片系統執行軟體堆疊的方式與 MTIA v1 類似，使得團隊的部署速度更快，其中 Meta 已經完成了在該架構上執行應用程式所需的大部分必要的整合和開發工作 。 下一代 MTIA 旨在相容為 MTIA v1 開發的程式碼。由於 Meta 已經將完整的軟體堆疊整合到晶片中，因此在幾天內就可以使用這款新晶片啟動並執行流量。這使得 Meta 能夠快速落地下一代 MTIA 晶片，在不到 9 個月的時間內從推出首個晶片到在 16 個資料中心執行生產模型。 Meta 透過建立 Triton-MTIA 編譯器後端來為 MTIA 硬體生成高效能程式碼，從而進一步最佳化了軟體堆疊。作為一種開源語言和編譯器，Triton 用於編寫高效的機器學習 計算核心，可以提高開發人員編寫 GPU 程式碼的效率。Meta 發現 Triton 語言與硬體無關，足以適用於 MTIA 等非 GPU 硬體架構。 此外，Triton-MTIA 後端執行最佳化後，可以最大限度地提高硬體利用率並支援高效能核心。 Meta 實現了對 Triton 語言功能的支援並整合到 PyTorch 2 中，為 PyTorch 運算元提供了廣泛的覆蓋。例如，藉助 TorchInductor，開發人員可以在提前 (AOT) 和即時 (JIT) 工作流程中利用 Triton-MTIA。 Meta 還觀察到，Triton-MTIA 顯著提高了開發人員的效率，能夠擴充套件計算核心的編寫範圍，並大大擴充套件對 PyTorch 運算元的支援。 目前的結果表明，下一代 MTIA 晶片可以處理 Meta 產品元件中的低複雜性 (LC) 和高複雜性 (HC) 排序和推薦模型。在這些模型中，模型大小和每個輸入樣本的計算量可能存在 10 倍到 100 倍的差異。 由於 Meta 控制整個堆疊，因此相較於商用 GPU 可以實現更高的效率。實現這些收益需要持續的努力，隨著在系統中構建和部署 MTIA 晶片，Meta 將繼續提高每瓦效能。 早期結果表明，在 Meta 評估的四個關鍵模型中，下一代 MTIA 的效能比第一代晶片提高了 3 倍。在平臺層面，與第一代 MTIA 系統相比，憑藉 2 倍的裝置數量和強大的 2-socket CPU，Meta 能夠實現 6 倍的模型服務吞吐量和 1.5 倍的每瓦效能提升。 為了實現這些目標，Meta 在最佳化核心、編譯器、執行時和主機服務堆疊方面取得了重大進展。隨著開發者生態系統的成熟，最佳化模型的時間正在縮短，但未來提高效率的空間更大。 https://ai.meta.com/blog/next-generation-meta-training-inference-accelerator-AI-MTIA/