AWS Neuron

AWS Neuron 是用于在基于 AWS Inferentia 和 AWS Trainium 的 Amazon Elastic Compute Cloud（Amazon EC2）实例上运行深度学习与生成式 AI 工作负载的软件开发工具包（SDK）。它包含编译器、运行时、训练与推理库以及用于监控、性能分析和调试的开发者工具。Neuron 支持端到端机器学习（ML）开发生命周期，涵盖模型构建与部署、优化以实现最高性能和最低成本，以及深入洞察模型行为。

与主流 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 等第三方服务结合使用。

分布式训练与推理库
Neuron 内置对开源 PyTorch 库 NxD Training 和 NxD Inference 的开箱即用优化，支持分布式训练与推理。NxD Training 简化并优化大规模分布式训练，支持多种模型架构、并行策略和训练流程。NxD Inference 提供全面的优化推理解决方案，具备关键特性如设备端采样、QKV 权重融合、连续批处理、推测解码、动态分桶和分布式推理。NxD Inference 还可与 vLLM 和 Hugging Face TGI 等推理服务集成，并提供针对不同模型架构的模型仓库。

先进的应用科学能力
Neuron 具备多项应用科学功能，助力科学家与研究人员在 Trainium 和 Inferentia 上推动开源 AI 研究与创新的边界。Neuron Kernel Interface（NKI）为 Trainium 和 Inferentia 上可用的硬件原语与指令提供直接访问，使研究人员能够构建并调优计算内核以实现最佳性能。NKI 是一种基于 Python 的编程环境，采用常见的 Triton 风格语法和层级块语义。研究人员可利用 NKI 为深度学习模型添加新功能、优化及科学创新。Neuron 的自定义 C++ 操作符使开发者可通过创建专为 Inferentia 和 Trainium 优化的操作符来扩展 SDK 功能。