Cerebras 推出萬億引數 AI 模型部署方式

我們最常收到的問題之一是如何在 Cerebras 晶圓級晶片上部署大型模型。單一 Cerebras WSE-3T 晶圓擁有 44 GB 的片上記憶體，而兆引數前沿模型僅權重就需要 TB 量級的記憶體。

在這篇文章中，我們將說明如何將大型前沿模型拆分到多個WSE-3T晶圓上，以每秒 1000-2000+ 令牌執行多兆引數模型。

我們將使用開源模型 Kimi K3（約 3 兆引數）作為說明範例。

最後，您將瞭解：

• 如何將完整模型放入片上記憶體
• 如何將模型對映到多個晶片上
• 如何處理大量併發請求

以 Kimi K3 為例拆解說明

Kimi K3 是一個開源的 3T 引數模型，擁有 69 個線性注意力層和 24 個全連線層。

對於此說明範例，我們使用：

• 假設平均請求大小 256k
• 4 位元權重量化（類似 Hugging Face 實作）
• 8 位元 KV 快取

#1：每層大小

為計算支援 Kimi K3 所需的 Cerebras WSE-3T 晶圓總數，先計算每層模型的大小。

Kimi K3 擁有 3 兆引數。若所有 3T 引數以每個 4 位元儲存，原始引數載荷為：

將該載荷除以 93 層後得到：

#2：晶圓數量

現在知道每層大約需要 16 GB，我們可以算出每個 44 GB 晶圓能容納多少層。答案取決於層型別：

• 全連線層需要額外空間用於 KV 快取，因為每個全連線（全注意力）層為每個活躍請求的每個令牌儲存鍵和值。因此，我們每個晶圓放置一層全連線層，並將剩餘 SRAM 留給其 KV 快取。
• 線性注意力層不需要額外的 KV 快取空間。它不為每個令牌儲存鍵和值，而是為每個請求保留一個小型、固定大小的狀態，與上下文長度無關。這樣每個晶圓可容納兩層線性層（2 × 16 GB = 32 GB）。

因此，使用 Kimi K3 的 69 個線性層和 24 個全連線層，我們得到：

總晶圓數為 35 + 24 = 59 晶圓。

#3：將晶圓連線在一起

Cerebras 以流水線方式執行模型。Cerebras 編譯器將模型層對映到晶圓區域後，執行時期會為令牌在晶圓之間流動對映路徑。為了極快的延遲，只有啟用值需要在晶圓之間傳遞，並可透過乙太網路 RDMA 完成。每跳大約 2 微秒。

重要的是要強調，因為每個 Cerebras 晶圓擁有 44 GB 的片上 SRAM，而每層僅 16 GB，我們通常能將整個層直接放在晶片上。

這意味著在大多數情況下，模型權重不需要移動，層也不需要跨晶圓拆分。只有啟用值從一個晶圓傳遞到下一個。

#4：如何處理 1 個請求

現在我們已將晶圓連線並佈局模型，我們可以把每個晶圓視為組裝線上的一個階段：每個階段的模型權重保持不變，請求的中間結果在各層之間傳遞。

例如，晶圓 1 執行模型層 1–2，並將其啟用值傳遞給晶圓 2；晶圓 2 執行層 3–4，並將結果傳遞給晶圓 3，依此類推，直到最後一個晶圓產生下一個令牌。

#5：如何處理多個請求

這引起了一個常見的誤解。若一個 token 必須按順序經過每片晶圓，是否每片晶圓在等待時都會閒置？實際上，系統是否會以單筆批次運作？

不！如果我們一次只傳送一個請求並等它完成後再開始下一個，這將是對管線的低效利用。相反，我們保持整條管線都充滿不同的請求。由於模型已經被拆分成管線階段，每個階段可以同時處理一個 不同的請求同時處理。

當第 3 階段為請求 A 處理一個 token 時，第 2 階段為請求 B 處理一個，第 1 階段為請求 C 處理一個。在我們的例子中，每個階段都是一層。實際上，Cerebras 使用更細粒度的階段，且數量更多。

Loading player...

每個階段也可以處理一個 微批次，即一小組請求，而不是單個請求。因此，正在進行中的請求數量大約等於階段數乘以微批次大小。

Loading player...

為了了解權衡，較小的微批次會使每個請求更快，而較大的微批次則提升總吞吐量。

#6: 我們能提供 多少 個請求？

分配給 KV 快取的晶片記憶體量決定了同時請求的最大數量。每個正在進行的請求在每片密集層晶圓上都需要自己的 KV 快取，因此晶片記憶體大小設定了同時請求的上限。如上所述，我們只需要計算密集層的 KV 快取需求。

對於每個已經擁有一層密集層的 44 GB 晶圓，我們可以用以下公式計算 KV 快取剩餘空間：

總晶圓儲存空間 - 一層密集層的大小。以此為例：

計算同時請求數量

要找出任何特定模型對映所支援的同時請求數，我們首先計算單個 256k 請求佔用的快取量：

接著計算此配置可支援的同時請求總數，我們將其除以：

同時請求總數代表任何時刻正在管線中移動的同時請求數量。與 GPU 批次不同，這些請求不會相互等待完成，而是連續地流經管線。

當然，並非同時性或批次大小能準確告訴你推論系統的效率。評估推論系統吞吐量時，更有意義的指標是每單位功率的總 token 吞吐量（例如每秒 token 數 / 每兆瓦），在相同的輸出生成速度下測量。

此規模如何擴充套件至超過 3T 的引數

同樣的方法可擴充套件到更大的模型：計算權重和請求狀態所需的記憶體，將模型分佈到多個晶圓上，並將它們連線成管線。每增加一片晶圓都提供計算與晶片記憶體，讓我們能容納超過單片晶圓容量的模型。

感謝你與我們一起學習！有興趣在 Cerebras 上部署模型或使用我們的推論雲服務嗎？

• 從我們的免費 API 開始：cloud.cerebras.ai
• 聯絡我們的銷售團隊：sales@cerebras.ai

設計由 Halley Chang 製作。感謝 Atharva Talpade、Zhenwei Gao、Angela Yeung、Jessica Liu、Joyce Er、Emma Call 與 Rohan Gupta 的審查與回饋。