AI推理計算的可能終局：存算一體芯片！

生成式AI為提高效率帶來了令人難以置信的商機，但我們仍然面臨著一個阻礙廣泛采用的巨大挑戰：運行推理的高昂成本。

我們已經聽說過大型語言模型（LLM）訓練期間產生的驚人費用：大量的GPU，高昂的電費。分析師估計，Meta?2024年可能會在GPU上花費150億美元。生成式AI需要大量的內存和帶寬來進行權重計算和數據處理，這是大規模運行模型的主要障礙。就連OpenAI的Sam Altman也表示：“沒有突破，就不可能實現這一目標。這促使我們加大對核聚變的投資。”

迄今為止，CPU、GPU和許多定制設計的加速器一直是AI的主要選擇，但即使是最先進的解決方案也會被傳統的馮·諾伊曼架構所拖累。定制設計的ASIC、FPGA和APU提供具有更高內存帶寬的專用控制器，并且需要大量RAM來將模型保存在內存中，這使得生成式AI幾乎不可能經濟地交付。

盡管存在這些成本和計算限制，但很快企業將廣泛部署多種AI模型，擴大對推理的需求。這將需要計算能力的另一個顯著提高。生成式AI推理需要更多的計算和內存容量，因為今天的模型比過去的非生成式ML模型大得多。

推理的計算不僅取決于模型的大小，還取決于用戶輸入提示的大小。例如，當提示符長度從8K增加到32K時，OpenAI的GPT-4運行成本增加了一倍。另一方面，模型生成每個輸出token的速度是一個重要的用戶體驗度量。這里內存帶寬很重要，但是遺留架構的內存帶寬限制加劇了推斷成本和功耗的痛點。

AI社區正在積極尋找解決方案來應對這些挑戰，包括氣候友好型能源（風能、太陽能等）、新的硬件設計方法和算法優化。今天我想談談解決這個問題的方法。

存儲和能量墻

加州大學伯克利分校BAIR/SkyLab的研究科學家Amir Gholami及其同事在他們的文章“AI和內存墻”中描述了訓練和服務transformer模型的挑戰。“內存墻”是William Wulf和Sally McKee在1995年創造的一個術語，它既包括有限的容量，也包括存儲傳輸的帶寬。分布式推理可以避免單個加速器有限的內存容量和帶寬，但這種方法也面臨內存墻問題：在加速器之間移動數據的通信瓶頸，這比芯片上的數據移動更慢，效率更低。

從訓練過的權重中創建內容需要大量的微小計算。為了進行這些計算，使用了GEMM（General Matrix Multiply，通用矩陣乘法）操作。GEMM不需要大型處理器，但它需要快速高效的小型計算。

這就是記憶墻阻礙的地方。對于每個操作，數據必須在DDR RAM和處理器之間以及處理器之間傳輸。即使這段旅程只有毫米的距離，由于馮·諾依曼架構中存儲和計算處理器的分離，它也需要時間和能量。而且有這么多的計算！

每次數據在存儲器總線上移動時，DRAM訪問每字節消耗大約60皮焦耳，而每次操作的計算只需要50-60飛焦耳。這意味著來回移動數據所消耗的能量是實際使用數據所消耗能量的一千倍。隨著這些能源成本累積到數百萬用戶提示，每個GPU數百瓦，以及全球數千臺服務器和數據中心，你可以看到為什么Altman認為我們需要核能突破。

內存計算

如果在存儲和處理之間來回移動數據的效率如此之低，那么我們是否可以將計算移到內存中呢？近年來，內存計算（in-memory computing，簡稱IMC）通過直接在內存單元附近/內部執行MAC（乘-累加）操作而成為一種很有前途的替代方法。

普渡大學的研究表明，與機器學習推理的既定基線相比，內存計算架構的能耗降低至0.12倍。位于比利時KU Leuven研究型大學的MICAS中心支持IMC，強調其降低訪問開銷和實現大規模并行化機會的能力，可能導致能源效率和吞吐量的數量級提高。

該領域的早期探索將模擬IMC視為評估權重和在預訓練的LLM上運行推理的有效方法。然而，這種方法需要昂貴的數模轉換器和額外的錯誤檢查。

數字內存計算（DIMC）提供了一種替代方案，避開了模擬IMC的挑戰，提供無噪聲計算和更大的空間映射靈活性。與模擬相比，DIMC犧牲了一些面積效率，但為處理未來的AI需求提供了更多的靈活性和能力。KU Leuven的研究還支持SRAM作為IMC的首選解決方案，因為與基于NVM的解決方案相比，SRAM具有魯棒性和可靠性。

DIMC有望徹底改變AI推理，降低成本并提高性能。鑒于生成式AI的采用速度很快，只有通過將計算引入內存并提高性能來追求降低成本和功耗的新方法才有意義。通過減少不必要的數據移動，我們可以顯著提高AI的效率，并改善AI未來的經濟狀況。

Sree Ganesan是d-Matrix的產品副總裁，d-Matrix是一家為生成式AI推理開發AI芯片的初創公司。