什么是 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 框架提供原生分析功能。