亚马逊AWS的Trainium2 AI架构

AWS最新推出的Trainium2 AI训练引擎在re:Invent 2023主机上首次亮相，引起广泛关注，通过与AWS实验室的Gadi Hutt的交流和对技术文档的挖掘，可以试图深入了解Trainium2与之前Inferentia系列的关系以及对Trainium2的期望。
AWS自2017年推出Nitro DPU以来，不断创新，推出了多代Nitro DPU、Graviton Arm服务器CPU、Inferentia AI推理加速器等产品。Trainium2与Graviton4服务器CPU一同亮相，AWS似乎在全方位创新计算引擎。在与Gadi Hutt的交流中，我们努力填补由于缺乏详细信息而存在的空白。
在技术方面，Trainium2与Inferentia系列相比，有望正面竞争Nvidia的Hopper H100 GPU加速器。AWS首席执行官Adam Selipsky在re:Invent主题演讲中透露，AWS已经购买了数百万个Nvidia的A100和H100产品，表明其巨大的投资。AWS似乎通过创建自家的Titan模型，并在本土开发的Inferentia和Trainium上运行模型，实现了更高的性价比。我们预计Trainium2将与H100竞争，考虑到H100的高昂价格和供应难题。AWS在推动AI计算引擎方面取得了显著进展，但在定价上可能存在与Graviton服务器CPU类似的差距。AWS通过创造自有模型和设备，可能在价格/性能比上具有竞争优势，这类似于其在Arm CPU实例方面的策略。尽管AWS乐意销售Intel和AMD的CPU，但通过省去中间商，它能够以更具吸引力的价格提供Arm CPU实例。Trainium2的推出使得AWS在AI计算引擎领域更具竞争力，而其对自主研发和创新的执着，似乎在性价比上取得了一定优势。这将在未来的云计算市场中引发更多关注和探讨。
Trainium架构是由AWS持续创新的一部分，结合了Inferentia1、Trainium1的元素，并在Inferentia2的基础上进行了改进。架构的层次结构包括计算元素、存储元素和连接它们的网络，其中元素的抽象级别可能在不同架构中发生变化。

Inferentia1芯片由Annapurna Labs团队于2018年发布，具有四个NeuronCore内核，包括ScalarEngine和VectorEngine，类似于Nvidia GPU中的CUDA核心。NeuronCore还包括TensorEngine，用于加速矩阵数学，类似于Nvidia GPU中的TensorCore。Inferentia1的NeuronCore-v1内核在FP16/BF16粒度下提供16 teraflops的性能。

Trainium1芯片于2020年发布，使用NeuronCore-v2核心，添加了32 GB HBM堆叠DRAM内存以提高带宽。Trainium1相较于Inferentia1可能减少了NeuronCore数量，但在每个核心内的标量、矢量和张量引擎数量增加了一倍。此外，Trainium1引入了称为GPSIMD引擎的通用处理器，可以直接使用C和C++进行寻址。

Inferentia2芯片基本上是Trainium1芯片的变体，通过保留HBM带宽并可能未激活某些元件来适应推理工作负载。其架构与Trainium1相似，但具有更少的NeuronLink-v2互连端口。

Trainium2架构推断为两个Trainium1芯片相互连接，可能以单片或两个小芯片插槽的形式出现，通过高速互连进行连接。Trainium2在计算元素、存储元素和网络连接方面与之前的架构相似，但可能在性能和带宽方面进行了改进。Trainium架构代表了AWS在AI计算引擎领域的不断演进，利用先前架构的优势并通过新的创新不断提高性能和效率。