AWS Neuron とは何ですか?
AWS Neuron は、AWS Trainium および AWS Inferentia で深層学習と生成 AI ワークロードを実行するための開発者スタックです。Neuron はオープンソースの基盤上に構築されており、開発者はコードを変更することなく、PyTorch および JAX フレームワークや HuggingFace、vLLM、PyTorch Lightning などの ML ライブラリを使用してネイティブに構築、デプロイ、探索できます。 コンパイラー、ランタイム、トレーニングライブラリ、推論ライブラリに加えて、モニタリング、プロファイリング、デバッグ用のデベロッパーツールが含まれています。Neuron は、深層学習および AI モデルの構築とデプロイ、最高のパフォーマンスと最低のコストを実現するための最適化、モデルの動作に関するより深いインサイトの取得まで、エンドツーエンドの機械学習 (ML) 開発ライフサイクルをサポートします。
Neuron は、迅速な実験、最先端のモデルの生産規模のトレーニング、カスタム カーネルの Neuron Kernel Interface (NKI) による低レベルのパフォーマンス最適化、エージェンティック AI および強化学習ワークロード向けのコスト最適化された推論のデプロイメント、および Neuron Explorer による包括的なプロファイリングとデバッグを可能にします。
研究者向けに構築
Neuron は、ネイティブ PyTorch コードを変更せずに Trainium 上で実行することで、迅速な AI 研究を可能にします。研究者は、PyTorch Eager モードのサポートにより、新しいアイデアを試し、迅速に反復することができます。FSDP、DDP、DTensor などの PyTorch 分散ライブラリを使用すると、チップ間でモデルをシャードしたり、複数のノードにスケーリングしたりして、スケールするのが容易になります。Neuron は torch.compile をサポートしていて、TorchTitan や HuggingFace Transformers などのライブラリは変更なしで Trainium 上で直接動作するようになりました。また、JAX 開発者は Neuron を利用してモデルを簡単に開発、最適化し、Inferentia や Trainium にデプロイすることができます。
生産性を重視した設計
Neuron は、エージェンティック AI と強化学習ワークロードの推論経済性を最適化します。標準の vLLM V1 API は、エキスパート並列処理、分散推論、投機的デコードなどの機能と、大規模なトークン経済を最大化するために Neuron カーネルライブラリから最適化されたカーネルを備え、すぐに使用できる高性能で Trainium および Inferentia で動作します。 機械学習開発者は、HuggingFace Optimum Neuron、PyTorch Lightning、TorchTitan でトレーニングを行い、その後、標準 vLLM API を使用して推論をデプロイできます。
イノベーションを重視した設計
AI モデルの構築には、迅速なイノベーションとパフォーマンスの最適化の両方が必要です。PyTorch のような標準フレームワークでは実験の規模を簡単に拡張できますが、パフォーマンスの限界を押し広げるには、スタック全体 (チップ、サーバー、UltraServer) を最適化する必要があります。Neuron は、Neuron Kernel Interface (NKI) を通じて ML パフォーマンス エンジニアに AWS AI チップへの比類のないアクセス、Neuron Explorer によるより深いインサイト、および Neuron Kernel Library (NKILib) と呼ばれる最適化されたカーネルライブラリを提供します。NKI は、メモリ割り当て、実行スケジューリング、Trainium ISA への直接アクセス用の API を提供し、命令レベルのプログラミングを制御できるようにします。NKI Compiler はオープンソースで、MLIR に基づいて構築されており、開発者はコンパイラパイプライン全体を把握できます。オープンソースの Neuron Kernel Library は、ソースコード、ドキュメント、ベンチマークを含む最適化された実装を提供します。Neuron Explorerは、開発者がパフォーマンスの最適化とデバッグを行う際に役立つ統合ツールスイートを提供します。パフォーマンスエンジニアは、ソースコードからハードウェア操作までの実行状況を追跡し、単一ノードアプリケーションと分散アプリケーションをプロファイリングし、カーネルの最適化とパフォーマンス改善に関する AI を活用したインサイトと実用的な推奨事項を受け取ることができます。
オープンソース向けに構築
AI イノベーションは、開発者が調査し、修正し、貢献できるオープンコミュニティで繁栄します。Neuron はオープンソースコミュニティとイノベーションの促進に取り組んでいます。スタックの多くの部分をオープンソースに移行していて、NKI Compiler、Neuron Kernel Driver、Neuron Kernel Library、NxD Inference、Neuron Explorer、PyTorch、JAX、vLLM の統合は現在完全にオープンソースとして利用可能です。オープンソースのライブラリとツールにより、開発者はコンパイラの実装を検査し、最適化を行い、カーネルコードを適応させることができます。一緒に構築しましょう。
Neuron に会う
Neuron は PyTorch とネイティブに統合されているため、研究者や ML 開発者は既存のコードを変更せずに Trainium で実行できます。FSDP、DDP、DTensor などの標準 API は、マルチノード設定にわたる分散型トレーニングでシームレスに機能します。TorchTitan、HuggingFace Optimum Neuron、PyTorch Lightning などの人気の ML ライブラリは、最小限の変更で直接実行されます。事前トレーニングから強化学習による事後トレーニングまで、使い慣れたワークフローとツールを使用してモデルをトレーニングし、Trainium のパフォーマンスとコストの利点を実験と生産規模のトレーニングの両方に活用します。
Neuron は、Trainium と Inferentia 上の標準フレームワークと API を使用してプロダクション推論のデプロイメントを可能にします。vLLM を標準 API と統合することで、Neuron Kernel Library の最適化されたカーネルを使用して高性能なサービスを提供できます。エキスパート並列処理、非集約推論、投機的デコードなどの高度な機能により、トークンあたりのコストを最小限に抑えながら、1 秒あたりのトークンを最大化します。すぐに使えるパフォーマンス最適化機能により、エージェンティック AI と強化学習のワークロードを大規模にデプロイできます。
最大のハードウェア効率を求めるパフォーマンスエンジニア向けに、Neuron は NeuronISA 命令セット、メモリ割り当て、実行スケジューリングへの直接アクセスにより、NeuronKernel Interface (NKI) を介して完全な制御を提供します。開発者は、標準フレームワークでは利用できない新しい操作を作成したり、カスタムカーネルを使用してパフォーマンスが重要なコードを最適化したりできます。MLIR 上に構築されたオープンソースの NKI Compiler は、コンパイルプロセスの透明性を提供します。Neuron Kernel Library は、完全なソースコード、ドキュメント、ベンチマークを備えた、本番環境に対応した最適化されたカーネルを提供します。
Neuron Explorerは、開発者がパフォーマンスの最適化とデバッグを行う際に役立つ統合ツールスイートを提供します。Neuron Explorer は、プロファイリング、デバッグ、最適化の実装、改善の検証を 1 つの環境に統合することで、断片化されたツール間でタイムロストを排除します。PyTorch、JAX、NKI のコードリンクによる階層プロファイリングは、ソースコードからハードウェア操作までの実行をトレースします。AI を活用したレコメンデーションは、プロファイルを分析してボトルネックを特定し、シャーディング戦略とカーネル最適化のための実用的なインサイトを提供します。UI は GitHub で公開されているオープンソースです。
Neuron は包括的なモニタリングとオブザーバビリティ機能を提供することで、ML 開発者と MLOPs チームが本番環境へのデプロイにおいて優れた運用を維持できるようにします。ネイティブの Amazon CloudWatch 統合により、Kubernetes と Amazon EKS 上のコンテナ化されたアプリケーションをサポートしながら、ML インフラストラクチャ全体の一元的なモニタリングが可能になります。パートナープラットフォームと Datadog などのツールを統合することで、監視、ログ、アラートが一元化され、オブザーバビリティが拡張されます。Neuron には、リアルタイムモニタリング用の neuron-top、メトリクス収集用の Neuron Monitor、デバイスリスト用の neuron-ls、詳細なシステム情報用の Neuron Sysfs などのユーティリティが用意されています。
Neuron は、事前設定された環境とインフラストラクチャツールを使用して、ML 開発者と MLOps チームのデプロイメントを簡素化します。Neuron 深層学習 AMI (DLAMI) と深層学習コンテナ (DLC) には、Neuron ソフトウェアスタック、一般的なフレームワーク、および必須ライブラリが付属しています。Kubernetes のデプロイメントでは、Neuron Device Plugin がリソース割り当てを管理し、Neuron Scheduler Extension がインテリジェントなワークロード配置を提供し、動的リソース割り当て (DRA) ドライバーが直感的なサイズベースのリソース選択によりハードウェアトポロジーの複雑さを抽象化します。Helm チャートは、コンテナ化されたデプロイメントのオーケストレーションを効率化します。