AWS Neuron がトレーニング向け Neuron Kernel Interface (NKI)、NxD Training、JAX のサポートを提供
本日、AWS は Neuron 2.20 のリリースを発表しました。Neuron 2.20 は、AWS Trainium と Inferentia 用のプログラミングインターフェイスである Neuron Kernel Interface (NKI) (ベータ版) を搭載しています。これにより、開発者は新しい機能、最適化、科学革新に適したコンピューティングカーネルを構築できます。さらに、このリリースでは、NeMo と互換性のあるユーザーフレンドリーなインターフェイスを備えた、効率的な分散トレーニングを可能にする PyTorch ベースのライブラリである NxD Training (ベータ版) が搭載されています。また、JAX フレームワーク (ベータ版) もサポートされています。
AWS Neuron は、生成 AI 向けに構築された AWS Inferentia および Trainium ベースのインスタンス用の SDK です。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) および Pixart-alpha や Pixart-sigma などのディフュージョントランスフォーマー (DiT) モデルのサポートが追加されました。さらに、デバイスでの top-p sampling、フラッシュデコードでの 128K コンテキスト長による推論サポートが追加されています。また、Rocky 9.0 オペレーティングシステムと、Neuron コンパイラでの RMSNorm 演算子と RMSNormDx 演算子のサポートも追加されました。
詳細については、Neuron リリースノートを参照してください。