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

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

AWS 于 2017 年開始使用“Nitro”DPU，并很快決定如果要在未來繼續(xù)在服務(wù)器基礎(chǔ)設(shè)施方面進行創(chuàng)新，就需要在所有方面對計算引擎進行創(chuàng)新。現(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 大會上亮相，我們花了一些時間嘗試了解這個新的 AI 訓(xùn)練引擎以及它與 Inferentia 系列的關(guān)系。

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

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

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

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

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

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

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

推斷 TRAINIUM 的架構(gòu)

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

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

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

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

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

Trainium1 芯片于 2020 年 12 月發(fā)布，并以兩個不同的實例（Trn1 和 Trn1n）發(fā)貨。我們當(dāng)時對 Trainium1 和2021 年 12 月的這些實例進行了盡可能多的分析，坦率地說，AWS 沒有提供大量有關(guān)這些本土 AI 計算引擎的數(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 倍，同時帶寬也提高了 2 倍。

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

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

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

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

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

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

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

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

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

我們還認為，鑒于 AWS 希望使用 Trainium2 將 UltraCluster 擴展到 100,000 個設(shè)備，它將減少 Trainium2 上的實際核心數(shù)量，使其比我們在上表中顯示的 64 個核心少很多。

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

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

以下是使用 Inferentia 和 Trainium 芯片可用的實例：

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

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

審核編輯：黃飛

?