AWS 通过推出自主研发的 Trainium2 处理器和基于其的 ExaFLOPS 超级计算机,开辟了一条与众不同的 AI 路径。
Trainium2 提供高达 1.3 FP8 PetaFLOPS 的性能,支持大规模生成式 AI 模型训练和推理。基于 Trainium2 的超大规模集群,AWS 迈向 130 FP8 ExaFLOPS 的计算能力,直接挑战 Nvidia GPU 的主导地位。与此同时,AWS 宣布即将推出性能提升四倍的 Trainium3 芯片。
这一趋势表明,自研 AI 加速器正逐渐成为云计算巨头优化成本和性能的核心竞争力。
Part 1
● Trainium2 作为亚马逊的第二代 AI 加速器,展现出了独特的架构设计与强大的性能表现。
◎ 其采用多块系统级封装,包含两个计算块、使用四个堆栈的 96GB HBM3 内存以及两个静态芯片组以确保封装一致性。
◎ 在计算核心方面,由少量大型 NeuronCore 组成,这种设计与传统 GPU 使用大量较小张量核心形成鲜明对比,大型核心在处理生成式 AI 工作负载时能够有效减少控制开销。
● 每个 NeuronCore 内部集成了张量引擎、矢量引擎、标量引擎和 GPSIMD 引擎,各司其职协同工作。
◎ 张量引擎作为核心计算单元,以 128×128 脉动阵列的形式,从 SRAM 缓冲区收集输入并输出结果到部分和 SRAM 缓冲区,承担了现代 LLM 工作负载中 80%以上的功率和 FLOPS 消耗。
◎ 矢量引擎专注于加速矢量运算,如在注意层中计算 softmax 或在层/批处理规范化层中计算移动平均值和方差时发挥重要作用;
◎ 标量引擎则用于执行元素级操作等简单映射操作;
◎ 而 GPSIMD 引擎具备图灵完备性,可运行任意 C++代码,方便开发人员快速实现自定义操作。
Trainium2 还创新性地配备了专用的集体通信核心,专门用于与其他芯片通信,实现了计算通信重叠且避免了资源争用。
这一设计相较于 Nvidia 和 AMD GPU 在通信与计算资源共享同一核心的模式具有显著优势,大大降低了用户在平衡通信与计算操作时的复杂性。
Trainium2 由 Annapurna Labs 设计,专为支持大规模生成式 AI 模型(如大型语言模型 LLM 和基础模型 FM)设计。
● 具备以下特点:
◎ 计算能力:每颗芯片提供 1.3 FP8 PetaFLOPS 性能,支持 BF16 和 FP8 精度。
◎ 存储性能:搭载 96GB HBM3,提供高达 46 TB/s 的带宽。
◎ 扩展性:通过 NeuronLink 网络实现多芯片互连,支持多达 64 芯片的服务器配置,并扩展到超大规模集群。
● ExaFLOPS 级超级计算机的构建与意义
ExaFLOPS(百亿亿次浮点运算每秒)代表了当前 AI 超级计算的顶级性能。AWS 基于数十万个 Trainium2 构建的 Rainier 项目,将实现 130 FP8 ExaFLOPS,这一性能级别相当于约 32,768 块 Nvidia H100 GPU 的总和。
通过将大量 Trainium2 处理器进行组合,如在 EC2 Trn2 UltraServer 中采用 64 个互连的 Trainium2 芯片,实现了高达 83.2 FP8 PetaFLOPS 的性能以及 6 TB 的 HBM3 内存,峰值带宽达到 185 TB/s,并借助 12.8 Tb/s Elastic Fabric Adapter (EFA) 网络进行互连,代号为“Rainier 项目”的 EC2 UltraCluster 更是由数十万个 Trainium2 处理器驱动,预计可提供约 130 FP8 ExaFLOPS 的超强性能。
如此强大的计算能力为大规模 AI 模型的训练和推理提供了坚实的基础。在当今 AI 领域,随着模型参数数量呈指数级增长,如 GPT-4 等模型已经拥有数万亿参数,对计算资源的需求也随之急剧攀升。
ExaFLOPS 级超级计算机的出现,使得研究人员能够在更短的时间内训练出更加复杂、精准的 AI 模型,加速了 AI 技术在自然语言处理、计算机视觉、推荐系统等众多领域的创新与应用进程,推动了整个 AI 行业从理论研究向实际大规模应用的跨越。
● 行业趋势驱动因素
基础模型和大型语言模型的快速发展使得模型参数数量不断逼近数万亿级别。这种复杂度的提升对计算资源提出了前所未有的挑战,传统的计算架构难以满足大规模模型训练和推理所需的海量计算需求。
例如,在训练 GPT-4 这样的超大型模型时,需要耗费巨大的计算资源和漫长的时间,促使企业寻求更强大、更高效的计算解决方案,这成为了 AWS 研发 Trainium2 处理器并构建 ExaFLOPS 级超级计算机的主要驱动力之一。
在大规模 AI 计算场景中,成本和能源效率成为了关键因素。AWS 通过自研 Trainium2 处理器,旨在实现更高的性价比和能源效率。
Trainium2 在设计上针对 AI 工作负载进行了优化,例如采用相对较低的算术强度设计,以适应由于 ML 研究进展导致模型算术强度增长较慢的趋势,如混合专家模型(MoE)中分组 GEMM 对内存带宽需求较大的情况,通过优化芯片架构和系统集成,AWS 期望在提供强大计算能力的同时,降低训练和部署 AI 模型的总体成本,提高能源利用效率,以满足企业在大规模 AI 应用中的经济和环保需求。
在全球科技竞争日益激烈的背景下,企业对于关键技术的自主可控性愈发重视。
AWS 作为全球领先的云计算服务提供商,通过自主研发 Trainium2 处理器,减少了对外部供应商(如 Nvidia)的依赖,实现了在 AI 芯片领域的自主可控,有助于保障其云计算服务的稳定性和安全性。
还能够根据自身对市场需求的理解和技术发展趋势的判断,灵活地进行芯片技术的迭代和创新,形成独特的技术差异化优势。
Trainium2 的扩展网络 NeuronLink 采用与 Nvidia NVLink 不同的拓扑结构,更接近于 TPU 类拓扑,通过点对点连接实现了特定的扩展功能,为用户提供了另一种选择,满足了不同用户在 AI 计算架构上的多样化需求。
Part 2
AWS 的这一举措将对 AI 基础设施市场格局产生深远的影响,在芯片层面,Trainium2 处理器的出现打破了 Nvidia 在高性能 AI 芯片领域的长期主导地位,为市场引入了新的竞争力量。
随着 AWS 不断优化和推广 Trainium2 及其后续产品,其他云计算服务提供商和企业可能会重新审视其芯片采购策略,考虑采用更多样化的芯片解决方案,以避免过度依赖单一供应商。
这将促使芯片市场竞争更加激烈,推动各芯片制造商加大研发投入,加速技术创新,从而降低芯片成本,提高性能和能效。
在超级计算机领域,AWS 基于 Trainium2 构建的 ExaFLOPS 级超级计算机为企业和研究机构提供了一种全新的、可替代传统 GPU 集群的高性能计算选择。
这种大规模、低成本、高效率的超级计算机将吸引更多用户将其 AI 工作负载迁移到 AWS 平台,改变当前 AI 超级计算机市场的份额分布。
对于那些正在计划构建或扩展其 AI 基础设施的企业来说,AWS 的方案可能会成为一个极具吸引力的选项,从而影响整个 AI 基础设施建设的市场走向,推动行业朝着多元化、高性能、低成本的方向发展。
为了巩固其在 AI 领域的竞争优势,AWS 将继续沿着自主创新的道路前行,在芯片技术研发上,AWS 已经推出了 Trainium3 处理器,该处理器采用台积电 3nm 级制程工艺制造,预计 2025 年面向客户上市。Trainium3 的性能将是 Trainium2 的四倍,这将进一步提升 AWS 在 AI 芯片性能方面的竞争力。
AWS 可能会继续加大在芯片研发上的投入,探索新的架构设计、制程工艺优化以及与其他技术(如量子计算、光子计算等新兴技术)的融合可能性,以保持其在 AI 芯片技术前沿的地位。
AWS 将充分利用其全球领先的云计算服务平台优势,将 Trainium 系列处理器与其他 AWS 服务(如 Amazon S3 存储服务、Amazon EC2 弹性计算服务等)进行深度整合,为用户提供一站式的 AI 解决方案。
例如,通过优化 Trainium 与 Amazon S3 之间的数据传输通道,实现更快的数据加载和存储,提高整个 AI 工作流程的效率。
此外,AWS 还将加强与合作伙伴的合作关系,包括 AI 软件开发商、系统集成商、企业客户等,共同探索 Trainium 系列处理器在不同行业和应用场景中的最佳实践,拓展其市场应用范围。
例如,与 AI 软件开发商合作优化深度学习框架对 Trainium 处理器的支持,与企业客户合作开发特定行业的 AI 应用解决方案,如医疗保健领域的疾病诊断辅助系统、金融领域的风险预测模型等,从而进一步扩大 Trainium 系列处理器在全球市场的影响力和占有率。
AWS 利用 Trainium2 处理器构建 ExaFLOPS 级超级计算机无疑是 AI 领域的一项重大战略布局,其对技术创新、行业趋势以及市场格局的影响深远而持久。
AWS 在满足 AI 计算需求日益增长的背景下,对成本、能效和自主可控性的深度考量与战略追求,推动了 AI 基础设施朝着多元化、高性能、低成本的方向发展。