在AI的加速操作
芝能科技出品
AWS最新推出的Trainium2 AI訓(xùn)練引擎在re:Invent 2023主機(jī)上首次亮相，引起廣泛關(guān)注，通過與AWS實(shí)驗(yàn)室的Gadi Hutt的交流和對(duì)技術(shù)文檔的挖掘，可以試圖深入了解Trainium2與之前Inferentia系列的關(guān)系以及對(duì)Trainium2的期望。
AWS自2017年推出Nitro DPU以來，不斷創(chuàng)新，推出了多代Nitro DPU、Graviton Arm服務(wù)器CPU、Inferentia AI推理加速器等產(chǎn)品。Trainium2與Graviton4服務(wù)器CPU一同亮相，AWS似乎在全方位創(chuàng)新計(jì)算引擎。在與Gadi Hutt的交流中，我們努力填補(bǔ)由于缺乏詳細(xì)信息而存在的空白。
在技術(shù)方面，Trainium2與Inferentia系列相比，有望正面競(jìng)爭(zhēng)Nvidia的Hopper H100 GPU加速器。AWS首席執(zhí)行官Adam Selipsky在re:Invent主題演講中透露，AWS已經(jīng)購買了數(shù)百萬個(gè)Nvidia的A100和H100產(chǎn)品，表明其巨大的投資。AWS似乎通過創(chuàng)建自家的Titan模型，并在本土開發(fā)的Inferentia和Trainium上運(yùn)行模型，實(shí)現(xiàn)了更高的性價(jià)比。我們預(yù)計(jì)Trainium2將與H100競(jìng)爭(zhēng)，考慮到H100的高昂價(jià)格和供應(yīng)難題。AWS在推動(dòng)AI計(jì)算引擎方面取得了顯著進(jìn)展，但在定價(jià)上可能存在與Graviton服務(wù)器CPU類似的差距。AWS通過創(chuàng)造自有模型和設(shè)備，可能在價(jià)格/性能比上具有競(jìng)爭(zhēng)優(yōu)勢(shì)，這類似于其在Arm CPU實(shí)例方面的策略。盡管AWS樂意銷售Intel和AMD的CPU，但通過省去中間商，它能夠以更具吸引力的價(jià)格提供Arm CPU實(shí)例。Trainium2的推出使得AWS在AI計(jì)算引擎領(lǐng)域更具競(jìng)爭(zhēng)力，而其對(duì)自主研發(fā)和創(chuàng)新的執(zhí)著，似乎在性價(jià)比上取得了一定優(yōu)勢(shì)。這將在未來的云計(jì)算市場(chǎng)中引發(fā)更多關(guān)注和探討。
Trainium架構(gòu)是由AWS持續(xù)創(chuàng)新的一部分，結(jié)合了Inferentia1、Trainium1的元素，并在Inferentia2的基礎(chǔ)上進(jìn)行了改進(jìn)。架構(gòu)的層次結(jié)構(gòu)包括計(jì)算元素、存儲(chǔ)元素和連接它們的網(wǎng)絡(luò)，其中元素的抽象級(jí)別可能在不同架構(gòu)中發(fā)生變化。
Inferentia1芯片由Annapurna Labs團(tuán)隊(duì)于2018年發(fā)布，具有四個(gè)NeuronCore內(nèi)核，包括ScalarEngine和VectorEngine，類似于Nvidia GPU中的CUDA核心。NeuronCore還包括TensorEngine，用于加速矩陣數(shù)學(xué)，類似于Nvidia GPU中的TensorCore。Inferentia1的NeuronCore-v1內(nèi)核在FP16/BF16粒度下提供16 teraflops的性能。
Trainium1芯片于2020年發(fā)布，使用NeuronCore-v2核心，添加了32 GB HBM堆疊DRAM內(nèi)存以提高帶寬。Trainium1相較于Inferentia1可能減少了NeuronCore數(shù)量，但在每個(gè)核心內(nèi)的標(biāo)量、矢量和張量引擎數(shù)量增加了一倍。此外，Trainium1引入了稱為GPSIMD引擎的通用處理器，可以直接使用C和C++進(jìn)行尋址。
Inferentia2芯片基本上是Trainium1芯片的變體，通過保留HBM帶寬并可能未激活某些元件來適應(yīng)推理工作負(fù)載。其架構(gòu)與Trainium1相似，但具有更少的NeuronLink-v2互連端口。
Trainium2架構(gòu)推斷為兩個(gè)Trainium1芯片相互連接，可能以單片或兩個(gè)小芯片插槽的形式出現(xiàn)，通過高速互連進(jìn)行連接。Trainium2在計(jì)算元素、存儲(chǔ)元素和網(wǎng)絡(luò)連接方面與之前的架構(gòu)相似，但可能在性能和帶寬方面進(jìn)行了改進(jìn)。Trainium架構(gòu)代表了AWS在AI計(jì)算引擎領(lǐng)域的不斷演進(jìn)，利用先前架構(gòu)的優(yōu)勢(shì)并通過新的創(chuàng)新不斷提高性能和效率。
