AWS Neuron

用于在 AWS Trainium 和 AWS Inferentia 上优化人工智能和深度学习的软件开发工具包(SDK)

什么是 AWS Neuron?

AWS Neuron 是一款软件开发工具包(SDK),用于在 AWS Inferentia 和 AWS Trainium 驱动的 Amazon Elastic Compute Cloud(Amazon EC2)实例上运行深度学习和生成式人工智能工作负载。它包括编译器、运行时、训练和推理库,以及用于监控、分析和调试的开发人员工具。Neuron 支持端到端机器学习(ML)开发生命周期,包括构建和部署深度学习和人工智能模型、进行优化以实现最高性能和最低成本,以及更深入地了解模型行为。

背景图案

与常用的机器学习框架和库进行原生集成

Neuron 与 PyTorch 和 JAX,以及 Hugging Face Optimum Neuron、PyTorch Lightning 和 AXLearn 等基本 ML 库原生集成。Neuron 还支持 OpenXLA,包括 StableHLO 和 GSPMD,让 PyTorch、XLA 和 JAX 开发人员能够使用 Neuron 针对 Inferentia 和 Trainium 的编译器优化。通过 Neuron,您可以在 Amazon SageMaker、Amazon EKS、Amazon ECS、AWS ParallelCluster 和 AWS Batch 等服务以及 Ray(Anyscale)、Domino Data Lab、Datadog 和 Weights & Biases 等第三方服务中使用基于 Trainium 和 Inferentia 的实例。

背景图案

分布式训练和推理库

Neuron 使用开源 PyTorch 库 NxD Training 和 NxD Inference,对分布式训练和推理进行了开箱即用的优化。NxD Training 可简化和优化大规模分布式训练,并支持各种模型架构、并行策略和训练工作流程。NxD Inference 为优化模型推理提供了全面的解决方案,其主要功能包括设备端采样、QKV 权重融合、连续批处理、推测解码、动态分桶和分布式推理。NxD Inference 还集成了 vLLM 和 Hugging Face TGI 等服务解决方案。它们都包括用于不同模型架构的模型中心。

背景图案

先进的应用科学能力

Neuron 拥有多项应用科学能力,让科学家和研究人员能够在 Trainium 和 Inferentia 上推动开源人工智能研究和创新的发展。Neuron Kernel 接口(NKI)提供对 Trainium 和 Inferentia 上可用硬件原语和指令的直接访问,使研究人员能够构建和调整计算内核以实现最佳性能。它是一个基于 Python 的编程环境,采用常用的类似 Triton 的语法和图块级语义。研究人员可以使用 NKI 为深度学习模型增加新功能、优化性能,并推动科学创新。通过 Neuron 的自定义 C++ 运算符,开发人员可以创建自己的运算符,并针对 Inferentia 和 Trainium 进行优化,从而扩展 SDK 的功能。

背景图案

功能强大的开发人员工具

AWS Neuron SDK 提供全面的工具集,可深入了解在 AWS Inferentia 和 Trainium 驱动的 EC2 实例上监控、管理和优化深度学习模型的情况。它提供 neuron-top、neuron-monitor 和 Neuron Sysfs 等实用程序,以监控硬件资源、模型执行和系统细节。对于 Kubernetes 和 EKS 上的容器化应用程序,Neuron 可通过 Amazon CloudWatch 集成以及 Data Dog 和 Weights & Biases 等其他常用的可观测性工具简化监控。此外,neuron-profile 工具还有助于识别和解决单节点和分布式应用程序的性能瓶颈,并为常见的 ML 框架提供原生分析功能。

背景图案

入门

Neuron Deep Learning Amazon 系统映像(Neuron DLAMI)预配置了 Neuron SDK、常用框架和有用的库,让您能够在 AWS Inferentia 上快速开始训练和运行推理。Neuron DLAMI 可简化工作流程并优化性能,免除复杂的设置,让您可以专注于构建和部署人工智能模型。开始使用 Neuron DLAMI

使用预配置的 AWS Neuron Deep Learning Containers(Neuron DLC)以及 Trainium 和 Inferentia 的优化框架,以快速部署模型。对于自定义解决方案,您可以构建自己的容器,并利用 Kubernetes 功能,如 Neuron Device Plugin、Neuron Scheduler Extension 和 Helm 图表。与 Amazon EKS、AWS Batch 和 Amazon ECS 等 AWS 服务无缝集成,实现可扩展部署。开始使用 Neuron DLC

Optimum Neuron 对 Hugging Face Transformers 和 AWS Neuron SDK 进行桥接,为 Trainium 和 Inferentia 提供标准的 Hugging Face API。它为训练和推理提供解决方案,包括对大规模的模型训练和人工智能工作流程的部署支持。Optimum Neuron 支持 Amazon SageMaker 和预构件的 Deep Learning Containers,以简化适用于机器学习的 Trainium 和 Inferentia 的使用。借助这种集成,开发人员能够使用熟悉的 Hugging Face 界面,同时利用 Trainium 和 Inferentia 来开发基于转换器的项目。开始使用 Hugging Face Optimum Neuron

利用 Neuron,您可以使用 Amazon SageMaker JumpStart 训练和部署模型。JumpStart 支持对 Meta 的 Llama 系列模型等常用模型进行微调和部署。开始使用 SageMaker JumpStart