Build on Trainium

“Build on Trainium” 是一项耗资 1.1 亿美元的投资计划，专注于人工智能研究和大学教育，旨在支持 AWS Trainium 的下一代创新和发展。AWS Trainium 是一款专为推进最先进的人工智能理念和应用而设计的人工智能脉动阵列芯片。“Build on Trainium” 资助在 Trainium 上进行的创新人工智能研究，投资领先的学术团队，在关键领域推动创新，包括新模型架构、机器学习库、优化、大规模分布式系统等。这项为期多年的计划激励学术界使用、投资和参与围绕 Trainium 的开源社区，为人工智能的未来奠定了基础。将这些优势与 Neuron 软件开发工具包 (SDK) 和最近推出的 Neuron Kernel 接口 (NKI) 相结合，Trainium 客户现在可以在云端进行大规模创新。

我们已经创建了一个包含多达 40,000 个 Trainium 芯片的专用 Trainium 研究集群，用户可以通过 Amazon EC2 Trn1 实例访问该集群，而这些实例通过 Amazon EC2 UltraClusters 连接在一个单一的非阻塞 PB 级网络上。研究团队和学生可以使用适用于 ML 的 Amazon EC2 容量块进行自管理容量块预留来访问这些芯片。

我们正在向广泛的研究界进行多轮亚马逊研究奖 (ARA) 提案征集 (CFP)，入选的提案将获得 AWS Trainium 积分和可访问 Trainium 研究集群。“Build on Trainium” 欢迎那些利用流行开源机器学习库和框架的研究提案，并将研究成果回馈给开源社区，以增加机器学习开发者社区的资源。

Neuron Kernel 接口 (NKI) 是适用于 AWS 人工智能芯片 Trainium 和 Inferentia 的全新编程接口。NKI 提供对 AWS Trainium 和 Inferentia 上可用硬件原语和指令的直接访问，使研究人员能够构建和调整计算内核以实现最佳性能。它是一个基于 Python 的编程环境，采用常用的类似 Triton 的语法和图块级语义。研究人员可以使用 NKI 为深度学习模型增加新功能、优化性能，并推动科学创新。要了解更多信息，请访问 NKI 文档页面。

在麻省理工学院的设备实现实验室，我们正在使用 AWS Trainium 来突破医疗人工智能研究的极限。我们的三维超声分割和声速估算模型比以往任何时候都更快、更高效地训练，将实验时间缩短了一半以上，同时实现了最先进的精度。AWS Trainium 使我们能够以传统 GPU 系统无法实现的方式扩大研究规模。通过在 AWS Trainium (trn.32xlarge) 上训练我们的三维全卷积神经网络，与 NVIDIA A100 实例相比，我们实现了最先进的性能，吞吐量提高了 50%，成本更低。使用 32 节点的 Trainium 集群，我们高效地进行了 180 多项消融实验，将总训练时间从几个月缩短到几周，并加速了麻省理工学院的医疗人工智能创新。将来，我们计划使用 Trainium 来训练人工智能代理模型，这些模型可以操作和自动化数字超声工作流程，从而为临床医生节省大量时间并为患者提供更好的护理。”

“CMU Catalyst 研究小组致力于优化机器学习系统。我们的项目旨在简化不同机器学习系统的优化。Trainium 的独特之处在于通过神经元内核接口 (NKI) 提供低级控制和可访问的编程接口。

在 AWS 通过 Build on Trainium 计划的支持下，我们的研究人员得以在关键内核——FlashAttention 上探索高级优化。最令我们惊讶的是迭代的速度：使用公开的 NKI、Neuron profiler 和架构文档，我们在短短一周内在先前技术的基础上取得了有意义的改进。强大的工具和清晰的硬件洞察相结合，使我们的团队可以进行复杂的低级优化。

AWS Trainium 和 Neuron Kernel 接口 (NKI) 使像我们这样的研究人员能够更快地创新，消除了通常会减慢特定硬件优化工作的障碍。”

 

“通过 Build on Trainium 计划，他的团队获得了对 AWS Neuron 新的 NKI 开源编译器堆栈的完全访问权限，包括直接查看 Trainium ISA 和 API，以进行精确的调度和内存分配。这种可见性和控制水平使他的学生能够更轻松地分析优化机会，更有效地发现性能实现。”

Christopher Fletcher，加州大学伯克利分校计算机科学副教授
 

“访问 AWS Trainium 和 Inferentia 在推进我们对大规模、高效人工智能系统的研究和教育方面发挥了重要作用。我们使用这些平台进行专家混合训练和推理优化，对新的运行时和调度技术进行原型设计，以提高新兴加速器架构的可扩展性、效率和可移植性。通过利用 Neuron Developer 堆栈，UIUC 研究人员正在开发新的运行时和调度技术，以提高人工智能工作负载的效率和可移植性。Neuron Developer 堆栈的开放性给团队留下了特别深刻的印象，这使得这些平台对运行时研究具有价值，并实现了超越传统 GPU 架构的稀疏性、内存层次结构和通信效率创新。”

“通过利用 AWS Trainium 和 Build on Trainium 计划，我和我的学生得以显著加快我们的量子电路模拟速度。该项目汇集了一群强大的学生，他们共同构建了高性能模拟器，从而能够以前所未有的规模进行更深入的实验和动手学习。”

“我们在悉尼科技大学的研究团队正在通过开发自定义的 Neuron NKI 内核来探索树环水印算法的集成。通过 Build on Trainium 计划访问开源 Neuron 堆栈具有变革性。它使我们能够前所未有地了解 Trainium 架构，并能够直接在硬件层面工作。访问 Trainium 使我们的团队能够显著加快我们的水印工作负载，缩短了迭代周期，使我们能够探索更复杂的模型和技术。这种访问深度使我们的研究人员能够对新想法进行原型设计，尝试低级优化，并突破水印系统在现代 AI 加速器上可以实现的极限。”