AWS Neuron 引入了 Neuron Kernel Interface (NKI)、NxD Training 和 JAX 对训练的支持

发布于: 2024年9月17日

今天，AWS 宣布发布 Neuron 2.20，该版本引入了 Neuron Kernel Interface (NKI)（测试版）。NKI 是一种适用于 AWS Trainium 和 Inferentia 的编程接口，让开发人员能够构建经过优化的计算内核，以便实现新的功能、优化和科学创新。此外，该版本还引入了 NxD Training（测试版），这是一个基于 PyTorch 的库，可实现高效的分布式训练，具有与 NeMo 兼容的用户友好界面。该版本还引入了对 JAX 框架（测试版）的支持。

AWS Neuron 是基于 AWS Inferentia 和 Trainium 的实例的 SDK，专为生成式人工智能而构建。Neuron 与 PyTorch 等热门 ML 框架集成。Neuron 包括编译器、运行时、工具和库，用于支持在 Trn1 实例和 Inf2 实例上对 AI 模型进行高性能训练和推理。

该版本增加了针对模型训练和推理的支持功能和性能改进。在训练方面，该版本增加了对 Llama 3.1 8B 和 70B 模型的支持，序列长度可达 32K。该版本还引入了 torch.autocast()，以支持原生 PyTorch 混合精度和 PEFT LoRA 技术。在推理方面，Neuron 2.20 增加了对 Llama 3.1（405b、70b、8b）以及 Pixart-alpha 和 Pixart-sigma 等扩散转换器 (DiT) 模型的支持。此外，该版本增加了推理支持，包括在设备上进行 top-p 采样和使用 Flash Decoding 实现 128K 的上下文长度。该版本还增加了对 Rocky 9.0 操作系统以及 Neuron 编译器中 RMSNorm 和 RMSNormDx 运算符的支持。

有关更多信息，请参阅 Neuron 发行说明。