AWS Neuron

AWS Trainium と AWS Inferentia 上の AI と深層学習を最適化する SDK

AWS Neuron とは何ですか?

AWS Neuron は、AWS Inferentia および AWS Trainium 搭載の Amazon Elastic Compute Cloud (Amazon EC2) インスタンスで深層学習と生成 AI ワークロードを実行するために使用される Software Development Kit (SDK) です。コンパイラー、ランタイム、トレーニングライブラリ、推論ライブラリに加えて、モニタリング、プロファイリング、デバッグ用のデベロッパーツールが含まれています。Neuron は、エンドツーエンドの機械学習 (ML) 開発のライフサイクルをサポートします。これには、深層学習と AI モデルの構築とデプロイ、最高のパフォーマンスと最低コストを実現するための最適化、モデルの動作に関するより深いインサイトの取得が含まれます。

背景パターン

一般的な ML フレームワークやライブラリとのネイティブ統合

Neuron は PyTorch や JAX のほか、Hugging Face Optimum Neuron、PyTorch Lightning、AXLearn などの重要な ML ライブラリとネイティブに統合されています。Neuron は StableHLO や GSPMD を含む OpenXLA もサポートしているため、PyTorch、XLA、JAX のデベロッパーは Neuron のコンパイラー最適化を Inferentia と Trainium に使用できます。Neuron を使用すると、Trainium ベースおよび Inferentia ベースのインスタンスを Amazon SageMaker、Amazon EKS、Amazon ECS、AWS ParallelCluster、AWS Batch などのサービスや、Ray (Anyscale)、Domino Data Lab、Datadog、Weights & Biases などのサードパーティーサービスで使用できます。

背景パターン

分散型のトレーニングライブラリと推論ライブラリ

Neuron には、オープンソースの PyTorch ライブラリである NxD Training と NxD Inference による分散トレーニングと推論のためのすぐに使える最適化機能が含まれています。NxD Training は、大規模な分散トレーニングを簡素化および最適化し、さまざまなモデルアーキテクチャ、並列処理戦略、およびトレーニングワークフローをサポートします。NxD Inference は、オンデバイスサンプリング、QKV ウェイトフュージョン、連続バッチ処理、投機的デコーディング、動的バケット、分散推論などの主要機能を備えた、最適化されたモデル推論のための包括的なソリューションを提供します。NxD Inference は、vLLM や Hugging Face TGI などのサービングソリューションとも統合されています。どちらにも、さまざまなモデルアーキテクチャ用のモデルハブが含まれています。

背景パターン

高度な応用科学能力

Neuron には、科学者や研究者が Trainium と Inferentia に関するオープンソースの AI 研究とイノベーションの限界を押し広げるための応用科学機能がいくつかあります。Neuron Kernel Interface (NKI) では、Trainium や Inferentia で利用できるハードウェアプリケーションミティブや命令に直接アクセスできるため、研究者はコンピューティングカーネルを構築およびチューニングして最適なパフォーマンスを得ることができます。Python ベースのプログラミング環境で、一般的に使用されている Triton のような構文とタイルレベルのセマンティクスを採用しています。研究者は NKI を使用して、新しい機能、最適化、科学的な革新によって深層学習モデルを強化できます。Neuron のカスタム C++ 演算子を使用すると、デベロッパーは Inferentia と Trainium 用に最適化された独自の演算子を作成することで、SDK の機能を拡張できます。

背景パターン

強力なデベロッパーツール

AWS Neuron SDK は、AWS Inferentia および Trainium を利用した EC2 インスタンスでの深層学習モデルのモニタリング、管理、最適化に関する深いインサイトを得るための包括的なツールセットを提供します。ハードウェアリソース、モデル実行、システム詳細を監視するためのニューロントップ、ニューロンモニター、Neuron Sysfs などのユーティリティを提供します。Kubernetes と EKS 上のコンテナ化されたアプリケーションでは、Neuron は Amazon CloudWatch 統合および Data Dog や Weights & Biases などの他の一般的なオブザーバビリティツールを通じてモニタリングを簡素化します。さらに、ニューロンプロファイルツールは、シングルノードアプリケーションと分散アプリケーションの両方でパフォーマンスのボトルネックを特定して対処するのに役立ち、一般的な ML フレームワークにネイティブプロファイリング機能をもたらします。

背景パターン

開始方法

Neuron Deep Learning Amazon マシンイメージ (Neuron DLAMI) には、Neuron SDK、一般的なフレームワーク、および便利なライブラリがあらかじめ設定されているため、AWS Inferentia でのトレーニングと推論の実行をすばやく開始できます。Neuron DLAMI はワークフローを合理化し、パフォーマンスを最適化し、セットアップの複雑さをなくします。これにより、AI モデルの構築とデプロイに集中できます。Neuron DLAMI の使用を開始する。

Trainium と Inferentia 向けに最適化されたフレームワークを備えた事前設定済みの AWS Neuron Deep Learning Containers (Neuron DLC) を使用してモデルを迅速にデプロイします。カスタムソリューションの場合は、独自のコンテナを構築し、Neuron Device Plugin、Neuron Scheduler Extension、Helm チャートなどの Kubernetes 機能をご活用ください。Amazon EKS、AWS Batch、Amazon ECS などの AWS サービスとシームレスに統合することで、スケーラブルなデプロイを実現します。Neuron DLC の使用を開始する。

Optimum Neuron は Hugging Face Transformer と AWS Neuron SDK をつなぎ、Trainium と Inferentia 向けの標準 Hugging Face API を提供します。大規模なモデルトレーニングや AI ワークフローのデプロイのサポートなど、トレーニングと推論の両方のソリューションを提供します。Amazon SageMaker と事前構築みの Deep Learning Containers をサポートする Optimum Neuron は、ML 向けの Trainium と Inferentia の使用を簡素化します。この統合により、デベロッパーはトランスフォーマーベースのプロジェクトに Trainium と Inferentia を活用しながら、使い慣れた Hugging Face インターフェイスで作業することができます。Hugging Face Optimum Neuron の使用を開始する。

Amazon SageMaker JumpStart を使用すると、Neuron を使用してモデルのトレーニングとデプロイを行うことができます。JumpStart は、Meta の Llama ファミリーモデルなどの一般的なモデルのファインチューニングとデプロイをサポートします。SageMaker JumpStart の使用を開始する