什么是 AWS Neuron？
AWS Neuron 是开发人员堆栈，用于在 AWS Trainium 和 AWS Inferentia 上运行深度学习和生成式人工智能工作负载。Neuron 建立在开源基础上，可让开发人员使用 PyTorch 和 JAX 框架以及 HuggingFace、vLLM、PyTorch Lightning 等机器学习库进行本地构建、部署和探索，而无需修改您的代码。 它包括编译器、运行时、训练和推理库，以及用于监控、分析和调试的开发人员工具。Neuron 支持端到端机器学习（ML）开发生命周期，从构建和部署深度学习和人工智能模型、进行优化以实现最高性能和最低成本，以及更深入地了解模型行为。
Neuron 支持快速实验、前沿模型的生产规模训练、通过自定义内核的 Neuron Kernel 接口（NKI）进行低级别性能优化、代理式人工智能和强化学习工作负载的成本优化推理部署，以及使用 Neuron Explorer 进行全面的分析和调试。
为研究人员而打造
Neuron 通过在 Trainium 上原样运行 PyTorch 代码，实现快速的人工智能研究。借助 PyTorch Eager 模式支持，研究人员可以尝试新想法并快速迭代。使用 PyTorch 分布式库（如 FSDP、DDP 和 DTensor）可以轻松实现跨芯片分片模型或扩展到多个节点的扩展。Neuron 支持 torch.compile，诸如 TorchTitan 和 HuggingFace Transformers 这样的库现在可以直接在 Trainium 上运行，而无需进行修改。此外，JAX 开发人员可以利用 Neuron 轻松开发、优化他们的模型并将其部署到 Inferentia 和 Trainium。
为提高工作效率而打造
Neuron 优化代理式人工智能和强化学习工作负载的推理经济学。标准 vLLM V1 API 在 Trainium 和 Inferentia 上运行，具有开箱即用的高性能，并且专家并行性、分解推理和推测性解码等功能，以及来自 Neuron Kernel Library 的优化内核，从而最大限度地提高令牌经济性。 机器学习开发人员可以使用 HuggingFace Optimum Neuron、PyTorch Lightning 和 TorchTitan 进行训练，然后使用标准 vLLM API 部署推理。
为创新而打造
构建 AI 模型需要快速创新和性能优化。虽然像 PyTorch 这样的标准框架可以轻松扩展实验，但突破性能界限需要优化完整的堆栈（芯片、服务器和 UltraServer）。Neuron 通过 Neuron Kernel 接口（NKI）、Neuron Explorer 提供的更深入见解以及称为 Neuron Kernel Library（NKILib）的优化内核库，为机器学习性能工程师提供对 AWS AI 芯片的空前访问权限。NKI 提供用于内存分配、执行计划和直接访问 Trainium ISA 的 API，从而实现对指令级编程的控制。NKI Compiler 是开源的工具，基于 MLIR 构建，可让开发人员深入了解完整的编译器管道。开源的 Neuron Kernel Library 通过源代码、文档和基准测试提供优化的实施。Neuron Explorer 提供一套统一的工具，用于指导开发人员完成性能优化和调试之旅。性能工程师可以跟踪从源代码到硬件操作的执行情况，分析单节点和分布式应用程序，并获得基于人工智能支持的见解以及切实可行的内核优化和性能改进建议。
为开源而打造
人工智能创新在开放社区中蓬勃发展，开发人员可以在其中检查、修改和做出贡献。Neuron 致力于打造开源社区并促进创新。虽然我们正在将堆栈的更多部分移至开源环境，但 NKI Compiler、Neuron Kernel Driver、Neuron Kernel Library、NxD Inference、Neuron Explorer 以及 PyTorch、JAX 和 vLLM 集成现已完全开源。开源库和工具可让开发人员不受任何阻碍地检查编译器实施，为优化做出贡献，以及调整内核代码。与我们共同构建。
认识 Neuron
Neuron 提供与 PyTorch 的原生集成，使研究人员和机器学习开发人员能够在 Trainium 上原样运行现有代码。包括 FSDP、DDP 和 DTensor 在内的标准 API 可无缝地用于跨多节点设置的分布式训练。诸如 TorchTitan、HuggingFace Optimum Neuron、PyTorch Lightning 等流行的机器学习库，只需进行少量修改即可直接运行。使用熟悉的工作流程和工具训练模型（从预训练到借助强化学习的训练后），同时利用 Trainium 的性能和成本优势进行实验和生产规模的训练。
Neuron 支持在 Trainium 和 Inferentia 上使用标准框架和 API 进行生产环境推理部署。vLLM 与标准 API 的集成利用 Neuron Kernel Library 中优化的内核提供高性能服务。包括专家并行度、分解推理和推测性解码的高级功能可最大限度地提高每秒令牌数，同时最大限度地降低每个令牌的成本。通过即用的性能优化，大规模部署代理式人工智能和强化学习工作负载。
对于追求最高硬件效率的性能工程师而言，Neuron 通过 Neuron Kernel 接口（NKI）提供完全控制，可直接访问 NeuronISA 指令集、内存分配和执行计划。开发人员可以创建标准框架中未提供的新操作，并使用自定义内核优化性能关键代码。基于 MLIR 构建的开源 NKI Compiler 为编译过程提供透明度。Neuron Kernel Library 提供可用于生产、经过优化的内核，包括完整的源代码、文档和基准测试。
Neuron Explorer 提供一套统一的工具，用于指导开发人员完成性能优化和调试之旅。通过将分析、调试、优化实施和改进验证整合到一个环境中，Neuron Explorer 消除因分散的工具而浪费的时间。通过代码链接对 PyTorch、JAX 和 NKI 进行分层分析，可以跟踪从源代码到硬件操作的执行过程。人工智能支持的建议可分析配置文件以确定瓶颈，并为分片策略和内核优化提供切实可行的见解。用户界面在 GitHub 上是开源的。
Neuron 提供全面的监控和可观测性功能，助力机器学习开发人员和 MLOps 团队保持生产部署的卓越运营。Amazon CloudWatch 原生集成支持对机器学习基础设施进行集中监控，并支持 Kubernetes 和 Amazon EKS 上的容器化应用程序。合作伙伴平台与 Datadog 等工具的集成通过统一监控、日志记录和提醒扩展可观测性。Neuron 提供一些实用程序，包括用于实时监控的 neuron-top、用于指标收集的 Neuron Monitor、用于获取设备列表的 neuron-ls 以及用于获取详细系统信息的 Neuron Sysfs。
Neuron 使用预配置的环境和基础设施工具简化机器学习开发人员和 MLOps 团队的部署。Neuron 深度学习 AMI（DLAMI）和深度学习容器（DLC）预装 Neuron 软件堆栈、流行框架和基本库。对于 Kubernetes 部署，Neuron Device Plugin 管理资源分配，Neuron Scheduler Extension 提供智能的工作负载布局，Dynamic Resource Allocation（DRA）驱动程序则通过直观的基于大小的资源选择来抽象化硬件拓扑的复杂性。Helm 图表简化容器化部署的编排。