AI/ML トレーニングと推論のために Amazon EC2 Trn2 インスタンスと Trn2 UltraServer が使用可能に

新しい Amazon Elastic Compute Cloud (Amazon EC2) Trn2 インスタンスと Trn2 UltraServers は、ML トレーニングと推論のための最も強力な EC2 コンピューティングオプションです。第 2 世代の AWS Trainium チップ (AWS Trainium2) を搭載した Trn2 インスタンスは、第 1 世代の Trn1 インスタンスと比較して、速度が 4 倍、メモリ帯域幅が 4 倍、メモリキャパシティが 3 倍になっています。Trn2 インスタンスは、現行世代の GPU ベースの EC2 P5e および P5en インスタンスよりも 30～40% 優れた料金パフォーマンスを提供します。

16 個の Trainium2 チップに加えて、各 Trn2 インスタンスは 192 vCPU、2 TiB のメモリ、3.2 Tbps の Elastic Fabric Adapter (EFA) v3 ネットワーク帯域幅を備えており、前世代よりも最大 35% 低いレイテンシーを提供します。

まったく新しいコンピューティングオファリングである Trn2 UltraServer は、高帯域幅で低レイテンシーの NeuronLink インターコネクトに接続された 64 個の Trainium2 チップを搭載しており、最先端の基盤モデルで極めて優れた推論およびトレーニングパフォーマンスを実現します。

数万の Trainium チップが既に Amazon および AWS サービスで使用されています。例えば、最近のプライムデーでは、80,000 個を超える AWS Inferentia および Trainium1 チップが Rufus ショッピングアシスタントをサポートしました。Trainium2 チップは、既に Amazon Bedrock での Llama 3.1 405B および Claude 3.5 Haiku モデルのレイテンシー最適化バージョンに採用されています。

スケールアップ、スケールアウト、スケールアップ
フロンティアモデルのサイズと複雑さの持続的な成長は、革新的なコンピューティング性能の形態によって実現され、同様に革新的なアーキテクチャの形態にまとめられています。物事がよりシンプルだった時代には、高いスケーラビリティを実現するための設計について、スケールアップ (より大きなコンピュータを使用) とスケールアウト (より多くのコンピュータを使用) の 2 つの方法で説明できました。今日、Trainium2 チップ、Trn2 インスタンス、および後ほど説明するさらに大規模なコンピューティングオファリングを見ると、両方のモデルが適用されるように見えますが、これらは全体的な階層において異なるレベルに当てはまります。NeuronCore から UltraCluster まで拡張した Trn2 の構成要素を見てみましょう。

NeuronCore は、Trainium2 チップの中核です。第 3 世代の各 NeuronCore には、スカラーエンジン (1 個の入力から 1 個の出力)、ベクトルエンジン (複数の入力から複数の出力)、テンソルエンジン (シストリックアレイの乗算、畳み込み、転置)、および GPSIMD (汎用単一命令複数データ) コアが含まれています。

各 Trainium2 チップには、8 個の NeuronCore と 96 GiB の高帯域幅メモリ (HBM) が搭載されており、2.9 TB/秒の HBM 帯域幅をサポートしています。コアは個別にアドレス指定して使用することも、物理コアのペアを単一の論理コアにグループ化することもできます。単一の Trainium2 チップは、最大 1.3 PFLOPS の高密度 FP8 コンピューティングと最大 5.2 PFLOPS のスパース FP8 コンピューティングを提供し、HBM キューの自動並べ替えによりメモリ帯域幅の使用率を 95% まで高めることができます。

一方、各 Trn2 インスタンスには、16 個の Trainum2 チップが搭載されています。合計で、128 個の NeuronCore、1.5 TiB の HBM、および 46 TB/秒の HBM 帯域幅となります。これらを掛け合わせると、最大 20.8 PFLOPS の高密度 FP8 コンピューティングと最大 83.2 PFLOPS のスパース FP8 コンピューティングとなります。Trainium2 チップは、2D トーラスの NeuronLink を介して接続され、1 GB/秒の高帯域幅かつ低レイテンシーのチップ間通信を実現します。

UltraServer には、低レイテンシーかつ高帯域幅の NeuronLink に接続された 4 個の Trn2 インスタンスがあります。512 個の NeuronCore、64 個の Trainium2 チップ、6 TiB の HBM、および 185 TB/秒の HBM 帯域幅となります。計算すると、最大 83 PFLOPS の高密度 FP コンピューティングと最大 332 PFLOPS のスパース FP8 コンピューティングが実現されます。インスタンス内の NeuronCore を接続する 2D トーラスに加えて、4 個のインスタンスのそれぞれで対応する XY 位置にあるコアがリング状に接続されます。推論では、UltraServer は業界をリードする応答時間を実現し、極めて優れたリアルタイムエクスペリエンスを生み出すのに役立ちます。トレーニングでは、UltraServer はスタンドアロンインスタンスと比較して、モデルの並列処理のための集合通信を高速化することで、モデルトレーニングの速度と効率を高めます。UltraServer は、1 兆パラメータレベル以上でのトレーニングと推論をサポートするように設計されています。プレビュー形式で提供されています。プレビューに参加するには、当社までお問い合わせください。

Trn2 インスタンスと UltraServer は、EC2 UltraClusters にデプロイされ、単一の Pb 規模の非ブロッキングネットワークで数万の Trainium チップにわたるスケールアウト分散トレーニングを可能にし、Amazon FSx for Lustre の高性能ストレージにアクセスできます。

Trn2 インスタンスの使用
Trn2 インスタンスは、米国東部 (オハイオ) AWS リージョンで本番での使用のために現在使用可能で、Amazon EC2 Capacity Blocks for ML を使用して予約できます。最大 64 個のインスタンスを最大 6 か月間予約できます。予約は最大 8 週間前まで受け付けています。すぐに開始することもできるほか、必要に応じて予約を延長できます。詳細については、「機械学習ワークロードの GPU 容量を予約するための Amazon EC2 Capacity Blocks for ML の発表」をお読みください。

ソフトウェア側では、AWS Deep Learning AMI を使用して開始できます。これらのイメージは、おそらく既にご存知であり、使用しているフレームワークとツール (PyTorch、JAX など) で事前設定されています。

AWS Neuron SDK を使用してアプリケーションを構築した場合は、Trn2 インスタンスで使用するために、それらを移行して再コンパイルできます。この SDK は、JAX、PyTorch、および Hugging Face、PyTorch Lightning、NeMo などの重要なライブラリとネイティブに統合します。Neuron には、オープンソースの PyTorch ライブラリ NxD Training および NxD Inference を使用した分散トレーニングと推論のためのすぐに使用できる最適化が含まれており、プロファイリングとデバッグのための詳細なインサイトが提供されます。また、Neuron は、安定した HLO と GSPMD を含む OpenXLA もサポートしているため、PyTorch/XLA および JAX デベロッパーは Trainium2 のために Neuron のコンパイラ最適化を利用できます。

– Jeff;

原文はこちらです。