數據中心架構克服AI內存瓶頸的方法

人工智能的懷疑論者批評了當前技術中存在的內存瓶頸，認為無法加速處理器和內存之間的數據移動阻礙了有用的實際應用程序。

用于在數據中心訓練 AI 模型的 AI 加速器需要可用的最高內存帶寬。雖然將整個模型存儲在處理器中會消除等式中的片外存儲器，但這不是一個可行的解決方案，因為最大的模型測量數十億或數萬億個參數。

過去的系統內存受限，而今天的數據中心架構使用各種技術來克服內存瓶頸。

高帶寬內存

一種流行的解決方案是使用高帶寬內存 (HBM)，其中包括通過硅中介層將 4、8 或 12 個 DRAM 裸片的 3D 堆棧連接到處理器。該技術的最新版本 HBM2E 具有比其前身更快的每針信號速率，每針高達 3.6 Gb/s，從而提高了帶寬。三星和 SK Hynix 各自提供 8 個芯片 HBM2E 堆棧，總容量為 16 GB，提供 460 GB/s 帶寬（據 SK Hynix 稱，DDR5 為 2.4 GB/s，GDDR6 為 64 GB/s）。HBM3 旨在將速度和容量推得更高。

最新版本的 Nvidia 旗艦數據中心 GPU A100 提供 80 GB 的 HBM2E 性能和 2 TB/s 的內存帶寬。A100 包含五個 16-GB DRAM 堆棧，加入一個使用 HBM2 的 40-GB 版本，總帶寬為 1.6 TB/s。兩者之間的差異使深度學習推薦模型的 AI 模型訓練速度提高了三倍，這是一種已知的內存消耗。

同時，數據中心 CPU 正在利用 HBM 帶寬。Intel 的下一代 Xeon 數據中心 CPU Sapphire Rapids 將 HBM 引入 Xeon 系列。它們是英特爾首款使用專為 AI 等矩陣乘法工作負載設計的新 AMX 指令擴展的數據中心 CPU。他們還將能夠使用片外 DDR5 DRAM 或 DRAM 加 HBM?！巴ǔＧ闆r下，CPU 針對容量進行了優化，而加速器和 GPU 針對帶寬進行了優化，”英特爾高級首席工程師 Arijit Biswas 在最近的一次熱芯片介紹?！叭欢?，隨著模型大小呈指數級增長，我們看到對容量和帶寬的持續需求沒有權衡取舍。Sapphire Rapids 通過原生支持這兩者來做到這一點?！?/p>

Biswas 補充說，該方法通過內存分層得到了增強，“其中包括對軟件可見的 HBM 和 DDR 的支持，以及使用 HBM 作為 DDR 支持的緩存的軟件透明緩存，”Biswas 補充道。

然而，Sapphire Rapids 的總工程師 Nevine Nassif 告訴EE Times ，HBM 版本是以犧牲芯片面積為代價的。

“[與 HBM 兼容的]模具略有不同，”她說?！斑€有一個不同于 DDR5 控制器的 HBM 控制器。在沒有 HBM 的 Sapphire Rapids 版本中，我們在芯片的一個區域添加了用于加密、壓縮等的加速器。所有這些都消失了——除了數據流加速器——而是 HBM 控制器進入了。

“最重要的是，我們必須對網格進行一些更改，以支持 HBM 的帶寬要求，”Nassif 補充道。

除了 CPU 和 GPU，HBM 在數據中心 FPGA 中也很受歡迎。例如，英特爾的 Stratix 和賽靈思 Versal FPGA 都有 HBM 版本，一些 AI ASIC 也使用它。騰訊支持的數據中心 AI ASIC 開發商 Enflame Technology 將 HBM 用于其 DTU 1.0 設備，該設備針對云 AI 訓練進行了優化。80-TFLOPS (FP16/BF16) 芯片使用兩個 HBM2 堆棧，提供通過片上網絡連接的 512-GB/s 帶寬。

每美元的性能

雖然 HBM 為數據中心 AI 加速器所需的片外內存提供了極高的帶寬，但仍然存在一些值得注意的問題。

Graphcore 對不同內存技術的容量和帶寬的比較。當其他人嘗試使用 HBM2E 解決這兩個問題時，Graphcore 在其 Colossus Mk2 AI 加速器芯片上使用了主機 DDR 內存和片上 SRAM 的組合。（來源：Graphcore）

Graphcore 就是其中之一。在他的 Hot Chips 演示中，Graphcore 首席技術官 Simon Knowles 指出，在大型 AI 模型中更快的計算需要內存容量和內存帶寬。雖然其他人使用 HBM 來提高容量和帶寬，但權衡包括 HBM 的成本、功耗和熱限制。

Graphcore 的第二代智能處理單元 (IPU) 改為使用其 896 MiB 的大型片上 SRAM 來支持其 1,472 個處理器內核所需的內存帶寬。Knowles 說，這足以避免卸載 DRAM 所需的更高帶寬。為了支持內存容量，太大而無法在芯片上安裝的 AI 模型使用服務器級 DDR 形式的低帶寬遠程 DRAM。該配置連接到主機處理器，允許中型模型分布在 IPU 集群中的 SRAM 上。

鑒于該公司以每美元性能為基礎推廣其 IPU，Graphcore 拒絕 HBM 的主要原因似乎是成本。

“與 AI 處理器集成的 HBM 的凈成本是服務器級 DDR 每字節成本的 10 倍以上，”Knowles 說?！凹词谷萘窟m中，HBM 也主導著處理器模塊的成本。如果 AI 計算機可以使用 DDR，它可以部署更多 AI 處理器，但總擁有成本相同?！?/p>

據 Knowles 稱，40 GB 的 HBM 有效地將封裝的標線大小處理器的成本提高了三倍。Graphcore 的 8 GB HBM2 與 8 GB DDR4 的成本細分估計 HBM 芯片的尺寸是 DDR4 芯片的兩倍（將 20-nm HBM 與 Knowles 認為是同時代的 18-nm DDR4 進行比較），從而增加了制造費用。然后是 TSV 蝕刻、堆疊、組裝和封裝的成本，以及內存和處理器制造商的利潤率。

Graphcore 對 HBM2 與 DDR4 內存的成本分析顯示，前者的成本是后者的 10 倍。（來源：Graphcore）

“DDR DIMM 不會發生這種邊距堆疊，因為用戶可以直接從內存制造商處采購，”Knowles 說?！笆聦嵣?，可插拔的計算機組件生態系統出現的一個主要原因是為了避免保證金堆積?！?/p>

走得更遠

從 Hot Chips 的隱形模式中脫穎而出，Esperanto Technologies 提供了另一種解決內存瓶頸問題的方法。該公司的 1000 核 RISC-V AI 加速器針對的是超大規模推薦模型推理，而不是上面提到的 AI 訓練工作負載。

世界語的創始人兼執行主席戴夫·迪策爾指出，數據中心推理不需要巨大的片上內存?！拔覀兊目蛻舨幌胍?250 MB 的芯片，”Ditzel 說?！八麄兿胍?100 MB——他們想用推理做的所有事情都適合 100 MB。任何比這更大的東西都需要更多?！?/p>

Ditzel 補充說，客戶更喜歡將大量 DRAM 與處理器放在同一張卡上，而不是在芯片上?！八麄兘ㄗh我們，‘只需將所有內容都放到卡上，然后使用您的快速接口。然后，只要你能比通過 PCIe 總線更快地獲得 100 GB 的內存，這就是勝利。'”

將 Esperanto 的方法與其他數據中心推理加速器進行比較，Ditzel 說其他人專注于消耗整個功率預算的單個巨型處理器。這家初創公司堅稱，Esperanto 的方法——安裝在雙 M.2 加速卡上的多個低功耗處理器——可以更好地使用片外內存。單芯片競爭對手“引腳數量非常有限，因此他們必須使用 HBM 之類的產品才能在少量引腳上獲得非常高的帶寬——但 HBM 確實很昂貴、難以獲得且功率很高，”迪策爾說。

Esperanto 聲稱通過使用六個較小的芯片而不是單個大芯片解決了內存瓶頸，留下了可用于連接 LPDDR4x 芯片的引腳。（來源：世界語技術）

Esperanto 的多芯片方法使更多引腳可用于與片外 DRAM 通信。除了六個處理器芯片外，該公司還使用了 24 個專為手機設計的廉價 LPDDR4x DRAM 芯片，在低電壓下運行，“每比特能量與 HBM 大致相同”，Ditzel 說。

“因為 [LPDDR4x] 的帶寬 [比 HBM] 低，我們可以通過更寬來獲得更多帶寬，”他補充道?！拔覀冊诩铀倏ㄉ系膬却嫦到y上使用 1,500 位寬，[而單芯片競爭對手] 買不起 1,500 位寬的內存系統，因為對于每個數據引腳，你必須有幾個電源還有幾個接地針，針太多了。

“之前處理過這個問題，我們說，'讓我們把它分開，'”Ditzel 說。

通過 822-GB/s 的內存帶寬訪問 192 GB 的總內存容量。所有 64 位 DRAM 芯片的總和計算出 1,536 位寬的內存系統，分成 96 個 16 位通道以更好地處理內存延遲。這一切都符合 120 W 的功率預算。

流水線權重

晶圓級 AI 加速器公司 Cerebras Systems 設計了一個規模遠端的內存瓶頸解決方案。在 Hot Chips 上，該公司發布了 MemoryX，這是一款用于其 CS-2 AI 加速器系統的內存擴展系統，旨在實現高性能計算和科學工作負載。MemoryX 旨在訓練具有一萬億或更多參數的大型 AI 模型。

Cerebras Systems 的 MemoryX 是其 CS-2 晶圓級引擎系統的片外內存擴展，其行為就像在片上一樣。（來源：Cerebras 系統）

MemoryX 是 DRAM 和閃存的組合，其行為就像片上一樣。該架構被宣傳為具有彈性，旨在容納 4 TB 到 2.4 PB（2000 億到 120 萬億個參數）——足以容納世界上最大的 AI 模型。

該公司的聯合創始人兼首席硬件架構師 Sean Lie 表示，為了使其片外內存表現得像片上一樣，Cerebras 優化了 MemoryX，以消除延遲影響的方式將參數和權重數據流式傳輸到處理器。

“我們將內存與計算分開，從根本上分解它們，”他說?！巴ㄟ^這樣做，我們使溝通變得優雅而直接。我們可以這樣做的原因是神經網絡對模型的不同組件使用不同的內存。因此，我們可以為每種類型的內存和每種類型的計算設計一個專門構建的解決方案?！?/p>

結果，這些組件被解開，從而“簡化了縮放問題，”Lie 說。

在訓練期間，必須立即訪問對延遲敏感的激活記憶。因此，Cerebras 將激活保留在芯片上。

點擊查看完整大小的圖片

Cerebras 使用流水線在 AI 訓練期間消除對延遲敏感的通信。（來源：Cerebras 系統）

Cerebras 將權重存儲在 MemoryX 上，然后根據需要將它們流式傳輸到芯片上。Lie 說，在沒有背靠背依賴性的情況下，權重記憶的使用相對較少。這可以用來避免延遲和性能瓶頸。粗粒度流水線還避免了層之間的依賴關系；層的權重在前一層完成之前開始流式傳輸。

同時，細粒度流水線避免了訓練迭代之間的依賴關系；后向傳播中的權重更新與同一層的后續前向傳播重疊。

“通過使用這些流水線技術，權重流執行模型可以隱藏外部權重的額外延遲，并且我們可以達到與權重在晶圓本地 [訪問] 相同的性能，”Lie 說。

本文最初發表于EE Times。

審核編輯 黃昊宇