深度解讀亞馬遜AI芯片核心技術(shù)

亞馬遜網(wǎng)絡(luò)服務(wù)可能不是第一家創(chuàng)建自己的定制計(jì)算引擎的超大規(guī)模提供商和云構(gòu)建商，但它緊隨谷歌之后發(fā)布了自研的AI芯片——谷歌于 2015 年開始使用其自主研發(fā)的 TPU 加速器來處理人工智能工作負(fù)載。

AWS 于 2017 年開始使用“Nitro”DPU，并很快決定如果要在未來繼續(xù)在服務(wù)器基礎(chǔ)設(shè)施方面進(jìn)行創(chuàng)新，就需要在所有方面對(duì)計(jì)算引擎進(jìn)行創(chuàng)新?，F(xiàn)在我們已經(jīng)看到了多代 Nitro DPU、四代 Graviton Arm 服務(wù)器 CPU、兩代 Inferentia AI 推理加速器，以及現(xiàn)在的第二代 Trainium AI 訓(xùn)練加速器。Trainium2 芯片與 Graviton4 服務(wù)器 CPU 一起在 AWS 最近于拉斯維加斯舉辦的 re:Invent 2023 大會(huì)上亮相，我們花了一些時(shí)間嘗試了解這個(gè)新的 AI 訓(xùn)練引擎以及它與 Inferentia 系列的關(guān)系。

AWS 尚未公布有關(guān)這些 AI 計(jì)算引擎的大量詳細(xì)信息，但我們（指代nextplatform）設(shè)法與 Gadi Hutt 進(jìn)行了一些交流，他是 AWS Annapurna Labs 部門的業(yè)務(wù)開發(fā)高級(jí)總監(jiān)，負(fù)責(zé)設(shè)計(jì)其計(jì)算引擎并通過代工廠和指導(dǎo)它們。深入了解 AWS 系統(tǒng)，更深入地了解 Inferentia 和 Trainium 之間的關(guān)系以及對(duì) Trainium2 的期望；我們還對(duì)技術(shù)文檔中的規(guī)格進(jìn)行了一些挖掘，并嘗試填補(bǔ)空白，就像我們發(fā)現(xiàn)信息缺乏時(shí)所做的那樣。

不過，為了做好準(zhǔn)備，我們先做一點(diǎn)數(shù)學(xué)計(jì)算，然后再了解 AWS AI 計(jì)算引擎的數(shù)據(jù)源和速度。

在 AWS 首席執(zhí)行官 Adam Selipsky 的 re:Invent 主題演講中，Nvidia 聯(lián)合創(chuàng)始人兼首席執(zhí)行官黃仁勛是一位驚喜嘉賓，他在講話中表示，在“Ampere”A100 和“Hopper”H100 期間，AWS 購買了兩百萬個(gè)這樣的設(shè)備。

有傳言稱，AWS 將在 2023 年完成大約 50,000 個(gè) H100 訂單，我們假設(shè)去年可能有 20,000 個(gè)訂單。以每臺(tái) 30,000 美元的價(jià)格——考慮到需求，Nvidia 幾乎沒有動(dòng)力打折，而且最近幾個(gè)季度的凈利潤率遠(yuǎn)高于其數(shù)據(jù)中心收入的 50% ——即 21 億美元。這還剩下 193 萬美元的 A100，按照 2020 年至今的平均價(jià)格約為 12,500 美元計(jì)算，總計(jì) 241.3 億美元。

在如此巨大的投資流中，性價(jià)比曲線顯然還有彎曲的空間，而且 AWS 創(chuàng)建了自己的 Titan 模型，以供母公司 Amazon 和數(shù)以萬計(jì)的企業(yè)客戶使用，并提供其他模型，重要的是來自 Anthropic 的Claude 2，在其本土開發(fā)的 Inferentia 和 Trainium 上運(yùn)行。

我們認(rèn)為這條曲線看起來與 AWS 使用 Graviton 服務(wù)器 CPU 所做的沒有太大不同。AWS 非常樂意銷售 Intel 和 AMD 的 CPU，但它的價(jià)格/性能比“傳統(tǒng)”高出 30% 到 40%。由于它通過 Graviton 省去了中間商，因此可以以更低的價(jià)格提供 Arm CPU 實(shí)例，這對(duì)越來越多的客戶有吸引力。我們預(yù)計(jì) Nvidia 和 AMD GPU 以及 AWS 制造的 Inferentia 和 Trainium 設(shè)備之間也會(huì)存在同樣的傳統(tǒng)定價(jià)差距。

拋開這些數(shù)學(xué)問題，我們來談?wù)?Inferentia1，并概覽一下 Inferentia2 和 Trainium1，這樣我們就可以了解 Trainium2，它將與當(dāng)前在 Hopper H100 GPU 加速器上運(yùn)行的工作負(fù)載進(jìn)行正面競(jìng)爭(zhēng)。如您所知，H100 的價(jià)格幾乎與黃金一樣昂貴（目前 SXM5 版本每盎司的價(jià)格約為黃金的一半），并且像稀土礦物一樣難以獲得，并且是支撐人工智能經(jīng)濟(jì)的不可或缺的一部分。

推斷 TRAINIUM 的架構(gòu)

所有計(jì)算引擎都是計(jì)算元素、存儲(chǔ)元素和將它們連接在一起的網(wǎng)絡(luò)的層次結(jié)構(gòu)，并且圍繞這些元素的抽象級(jí)別可能會(huì)發(fā)生變化，特別是當(dāng)架構(gòu)是新的并且工作負(fù)載快速變化時(shí)。

Inferentia1 芯片由 Annapurna Labs 的人員創(chuàng)建，于 2018 年 11 月首次發(fā)布，一年后全面上市，是 AWS AI 引擎工作的基礎(chǔ)，以下是該設(shè)備的架構(gòu)：

該器件有四個(gè)內(nèi)核，具有四個(gè)不同的計(jì)算元件，以及用于片上存儲(chǔ)的片上 SRAM 存儲(chǔ)器和用于片外存儲(chǔ)的 DDR4 主存儲(chǔ)器，這與當(dāng)今許多 AI 芯片制造商一樣。AWS 沒有提供 Inferentia1 設(shè)備的 SRAM 內(nèi)存和緩存大小或時(shí)鐘速度的詳細(xì)規(guī)格，甚至沒有提供設(shè)備中使用的每個(gè) NeuronCore 內(nèi)的元件數(shù)量。不過在 Inferentia 和 Trainium 芯片的 Neuron SDK 中，AWS 確實(shí)討論了 Inferentia2 和 Trainium1 中使用的 NeuroCore-V2 核心的架構(gòu)，我們可以以此為基礎(chǔ)來弄清楚 Inferentia1 是什么，并推斷出 Trainium2 可能是什么是。

無論是哪代，NeuronCore 都有一個(gè)處理標(biāo)量計(jì)算的 ScalarEngine 和一個(gè)處理各種精度的整數(shù)和浮點(diǎn)數(shù)據(jù)向量計(jì)算的 VectorEngine。這些大致相當(dāng)于 Nvidia GPU 中的 CUDA 核心。根據(jù) Neuron SDK，NeuronCore-v1 ScalarEngine 每個(gè)周期處理 512 個(gè)浮點(diǎn)運(yùn)算，VectorEngine 每個(gè)周期處理 256 個(gè)浮點(diǎn)運(yùn)算。（我們認(rèn)為AWS的意思是ScalarEngine上每個(gè)周期有512位處理，VectorEngine上每個(gè)周期有256位處理，然后您通過這些以選擇的格式泵送數(shù)據(jù)以進(jìn)行特定類型的計(jì)算。

NeuronCore 架構(gòu)還包括一個(gè) TensorEngine，用于加速矩陣數(shù)學(xué)，超出了通過 VectorEngine 推動(dòng)代數(shù)矩陣數(shù)學(xué)的能力，而矩陣數(shù)學(xué)對(duì)于 HPC 和 AI 工作負(fù)載至關(guān)重要，它們通常會(huì)完成大量工作并提供最大的吞吐量。TensorEngine 大致類似于 Nvidia GPU 中的 TensorCore，在 NeuronCore-v1 內(nèi)核中，它們可以在 FP16/BF16 粒度下提供 16 teraflops。

Trainium1 芯片于 2020 年 12 月發(fā)布，并以兩個(gè)不同的實(shí)例（Trn1 和 Trn1n）發(fā)貨。我們當(dāng)時(shí)對(duì) Trainium1 和2021 年 12 月的這些實(shí)例進(jìn)行了盡可能多的分析，坦率地說，AWS 沒有提供大量有關(guān)這些本土 AI 計(jì)算引擎的數(shù)據(jù)。Trainium1 使用了新的 NeuroCore-v2 核心，它具有相同的元素，但核心數(shù)量更少，如下所示：

通過 Trainium1 芯片，AWS 添加了 32 GB HBM 堆疊 DRAM 內(nèi)存以提高設(shè)備的帶寬，轉(zhuǎn)向 PCI-Express 5.0 外形尺寸和 I/O 插槽，并增加了 NeuronLink 芯片間互連鏈路的數(shù)量其帶寬提高了 2 到 4 倍，同時(shí)帶寬也提高了 2 倍。

我們沒有證據(jù)證明這一點(diǎn)，但我們認(rèn)為，通過 Trainium1，AWS 將每個(gè)芯片的 NeuronCore 數(shù)量比 Inferentia1 減少了一半（兩個(gè)而不是四個(gè)），然后每個(gè)核心內(nèi)的標(biāo)量、矢量和張量引擎的數(shù)量增加了一倍。當(dāng)時(shí)的變化實(shí)際上是緩存和內(nèi)存層次結(jié)構(gòu)抽象級(jí)別的變化，基本上使 NeuronCore 在每個(gè)計(jì)算元素類型中實(shí)現(xiàn)多線程。

有趣的是，Trainium 芯片中首次使用的 NeuronCore-v2 還包括稱為 GPSIMD 引擎的東西，它是一組 8 個(gè) 512 位寬的通用處理器。（確實(shí)非常有趣。）這些設(shè)備可以直接使用 C 和 C++ 進(jìn)行尋址，并且可以訪問片上 SRAM 和內(nèi)核上的其他三種類型的引擎，并用于實(shí)現(xiàn)需要加速且不需要加速的自定義操作。直接受其他引擎中的數(shù)據(jù)和計(jì)算格式支持。（我們必須查閱Flynn 的分類法，試圖弄清楚這個(gè) GPSIMD 引擎是如何適應(yīng)的，并且從文檔中并不清楚我們看到的是這是一個(gè)陣列處理器、管道處理器還是關(guān)聯(lián)處理器.)

用于推理的 Inferentia2 芯片基本上是一個(gè) Trainium1 芯片，其數(shù)量只有一半的 NeuronLink-v2 互連端口。芯片的某些元件也可能未激活。時(shí)鐘速度和性能似乎大致相同，HBM 內(nèi)存容量和帶寬也是如此。

“Inferentia2 和 Trainium1 之間的芯片架構(gòu)幾乎相同，”Hutt 告訴The Next Platform?！拔覀?yōu)?Inferentia2 保留了 HBM 帶寬，因?yàn)檫@對(duì)于推理非常重要，而不僅僅是訓(xùn)練。LLM 推理實(shí)際上受內(nèi)存限制，而不是計(jì)算限制。因此，我們實(shí)際上可以采用類似的硅架構(gòu)并盡可能降低成本 - 例如，我們不需要那么多的 NeuronLink。通過推理，當(dāng)我們從一個(gè)加速器移動(dòng)到另一個(gè)加速器時(shí)，我們只需要環(huán)形的鏈接來生成token。當(dāng)然，通過訓(xùn)練，您需要完整的網(wǎng)狀連接，最大限度地減少服務(wù)器內(nèi)部每個(gè)加速器之間的跳數(shù)。當(dāng)然，當(dāng)你訪問訓(xùn)練服務(wù)器時(shí)，你需要大量的網(wǎng)絡(luò)帶寬?！?/p>

使用 16 個(gè) Trainium1 芯片的服務(wù)器上的互連是一個(gè) 2D 環(huán)面或 3D 超立方體，根據(jù) Hutt 的說法，它們是同一件事，具體取決于你想如何談?wù)撍?，它看起來像這樣：

下面的表格匯集了我們所知道的 Inferentia1、Inferentia2 和 Trainium1 的饋送和速度，以及我們對(duì) Trainium2 的預(yù)測(cè)（以粗體紅色斜體顯示）。Neuron SDK 尚未使用有關(guān) Trainium2 的任何信息進(jìn)行更新。

根據(jù)如下所示的 Trainium2 芯片封裝圖片，我們認(rèn)為它本質(zhì)上是兩個(gè) Trainium1 芯片互連在一起，要么作為單片芯片，要么作為兩個(gè)小芯片插槽，通過某種高速互連將它們連接在一起：

Hutt 沒有透露 Trainium 芯片的任何數(shù)據(jù)和速度，但他確實(shí)確認(rèn) Trainium2 擁有更多的內(nèi)核和更多的內(nèi)存帶寬，并進(jìn)一步補(bǔ)充說，芯片的有效性能將擴(kuò)展到 4 倍——他稱這是一個(gè)保守的數(shù)字，也許超出了現(xiàn)實(shí)世界人工智能訓(xùn)練工作負(fù)載的 Trainium1 的數(shù)字，因?yàn)檫@些工作負(fù)載更多地受到內(nèi)存容量和內(nèi)存帶寬的限制，而不是受計(jì)算的限制。

我們認(rèn)為Trainium2芯片將有32個(gè)核心，并將從Trainium 1使用的7納米工藝縮小到4納米工藝，這樣核心的加倍就可以在與Trainium1相同或稍高的功率范圍內(nèi)完成。我們還認(rèn)為，Trainium2 中的時(shí)鐘速度將從 Trainium1 中使用的 3 GHz（AWS 已透露的數(shù)字）適度提高到 3.4 GHz。我們還認(rèn)為（僅基于預(yù)感），Trainium2 上的總 NeuronLink 帶寬將增加 33%，達(dá)到每個(gè)端口 256 GB/秒，從 Trainium2 產(chǎn)生 2 TB/秒，并且仍然允許 2D環(huán)面互連。每個(gè)芯片上的 NeuronLink 端口數(shù)量可能會(huì)增加，以增加環(huán)面的維數(shù)，并減少設(shè)備共享數(shù)據(jù)時(shí)設(shè)備之間的一些跳數(shù)。2D 環(huán)面意味著集群中任意兩個(gè) NeuronCore 芯片之間有固定的兩跳。網(wǎng)絡(luò)似乎不太可能增加到全面配置，但這是可能的。（SGI 曾經(jīng)用其超級(jí)計(jì)算機(jī)芯片組來做到這一點(diǎn)。）

我們還認(rèn)為，鑒于 AWS 希望使用 Trainium2 將 UltraCluster 擴(kuò)展到 100,000 個(gè)設(shè)備，它將減少 Trainium2 上的實(shí)際核心數(shù)量，使其比我們?cè)谏媳碇酗@示的 64 個(gè)核心少很多。

很難說它會(huì)在哪里，假設(shè)大約 10% 的核心將是無用的，因此芯片的產(chǎn)量將會(huì)高得多。您可以打賭，AWS 將保留任何可以在一組單獨(dú)的機(jī)器中運(yùn)行所有核心的 Trainium2 設(shè)備，很可能供每個(gè)核心都很重要的內(nèi)部使用。這將使 Trainium2 擁有 56 個(gè)核心或 58 個(gè)核心，甚至可能高達(dá) 56 個(gè)核心，并且所有帶寬都可供它們使用。AWS 承諾的 96 GB 可能僅用于我們認(rèn)為在設(shè)備上看到的四個(gè)內(nèi)存堆棧中的三個(gè)，該設(shè)備可能具有 128 GB 的實(shí)際 HBM 內(nèi)存。我們強(qiáng)烈懷疑這將是 HBM3 內(nèi)存，但 Hutt 沒有證實(shí)任何事情。

但他多次說過，性能是由內(nèi)存驅(qū)動(dòng)的，而不是指望原始峰值理論計(jì)算的增長(zhǎng)速度快于內(nèi)存帶寬，如果我們是對(duì)的，內(nèi)存帶寬將增長(zhǎng) 5 倍位于 Trainium1 和 Trainium2 之間。

以下是使用 Inferentia 和 Trainium 芯片可用的實(shí)例：

任何人都在猜測(cè) Trn2 實(shí)例在價(jià)格或性能方面的比較，但根據(jù)暗示和預(yù)感，我們堅(jiān)信 Trainium2 將提供 Nvidia H100 大約 2 倍的性能，這意味著它將與之旗鼓相當(dāng)- Nvidia 剛剛發(fā)布的 H200 配備了更大、更快的 HBM3e 內(nèi)存，適用于許多型號(hào)。當(dāng)我們建議 AWS 可能會(huì)對(duì)基于 Trainium2 的 EC2 實(shí)例相對(duì)于使用 Nvidia 的 H100 和 H200 GPU 的實(shí)例進(jìn)行定價(jià)時(shí)，其比率與其自己的 Graviton CPU 與 AMD 和 Intle X86 處理器之間的比率相同 - 介于 30% 到 40% 之間物有所值——Hutt并沒有勸阻我們放棄這種想法。但Hutt也沒有做出任何承諾，只是說性價(jià)比肯定會(huì)更好。

這并非毫無意義，而是將 100,000 個(gè)設(shè)備以 FP16 精度以 65 exaflops 連接在一起，并且沒有任何稀疏技巧，而是真正的 FP16 分辨率，有機(jī)會(huì)成為世界上最大的人工智能集群。

審核編輯：黃飛

?