AWS Neuron 導入神經元核心介面 (NKI)、NxD 訓練和 JAX 支援以進行訓練

張貼日期: 2024年9月17日

今天，AWS 宣布推出 Neuron 2.20，其導入神經元核心介面 (NKI) (Beta 版)，這是 AWS Trainium 和 Inferentia 的程式設計介面，讓開發人員能夠建構最佳化的運算核心，以實現新功能、最佳化和科學創新。此外，本版本還導入 NxD Training (Beta 版)，這是一個以 PyTorch 為基礎的程式庫，可實現高效率的分散式訓練，並具有與 NeMo 相容的簡單易用介面。此版本也導入對 JAX 框架 (Beta 版) 的支援。

AWS Neuron 是 AWS Inferentia 和 Trainium 執行個體的軟體開發套件，專為生成式 AI 打造。Neuron 與 PyTorch 等流行的 ML 架構整合。Neuron 包括編譯器、執行時期、工具和程式庫，以支援在 Trn1 和 Inf2 執行個體上進行的高效能 AI 模型訓練和推論。

此版本也新增模型訓練和推論的支援功能和效能改進。對於訓練，此版本新增 Llama 3.1 8B 和 70B 模型支援，可支援高達 32K 序列長度，以及用於原生 PyTorch 混合精度支援和 PEFT LoRa 技術的 torch.autocast()。對於推論，Neuron 2.20 新增 Llama 3.1 (405b、70b、8b) 和 Diffusion-Transformers (DiT) 模型 (如 Pixart-alpha 和 Pixart-sigma) 的支援。此外，此版本新增推論支援，其中具有對裝置的 top-p 採樣，以及使用 Flash 解碼的 128K 內容長度。此版本還新增對 Neuron 編譯器中 Rocky 9.0 作業系統以及 RMSNorm 和 RMSNormDx 運算子的支援。

如需詳細資訊，請參閱 Neuron 版本備註。