AWS における基盤モデルのトレーニングと推論のための構築ブロック
Amazon と Hugging Face の共同執筆記事は、AWS 上で大規模言語モデルのトレーニングと推論を効率的に実行するための具体的なアーキテクチャパターンとベストプラクティスを提示している。
キーポイント
AWS 上での分散トレーニング最適化
FSDP(Fully Sharded Data Parallel)や DeepSpeed を活用し、EC2 インスタンス群を効率的に構成して大規模モデルの学習コストと時間を削減する手法が解説されている。
推論パフォーマンスの最大化
Amazon EC2 P5 インスタンス(H100 GPU)や Inferentia2 を使用した低遅延・高スループットな推論環境の構築方法と、vLLM などの高速推論フレームワークとの連携が示されている。
コスト効率とスケーラビリティ
スポットインスタンスや自動スケーリングを活用して、トレーニングジョブのコストを最適化しつつ、需要に応じた柔軟なリソース拡張を実現する戦略が提案されている。
大規模モデルのボトルネックは通信とメモリ転送
モデルがスケールするにつれ、ステップ時間は計算スループットよりも集合通信やメモリアクセスに支配されるため、帯域幅の明示的な見積もりが必要となる。
GPU 間通信には内部と外部の2つのレジームが存在
ノード内では NVLink/NVSwitch を用いた高速なスケールアップで集合演算を処理し、ノード間では AWS の EFA(Elastic Fabric Adapter)による OS バイパスネットワークでスケールアウトを実現する。
AWS は EFA を基盤とした EC2 UltraClusters を提供
通信集約型の集合演算を数千インスタンスにわたって実行するために、EFA を利用した Amazon EC2 UltraClusters が構築されている。
EFA の世代別性能向上
P5en インスタンスの EFAv3 は EFAv2 よりパケット遅延を約 35% 削減し、P6 インスタンスの EFAv4 は EFAv3 より集合通信パフォーマンスがさらに 18% 向上しています。
影響分析・編集コメントを表示
影響分析
この記事は、企業が独自の大規模言語モデルを開発・運用する際の具体的な技術的ロードマップを提供しており、AWS と Hugging Face の連携によるエコシステム強化を象徴しています。実装レベルのベストプラクティスを公開することで、開発者の参入障壁を下げ、クラウド上での AI 開発標準化に寄与すると考えられます。
編集コメント
AWS と Hugging Face の連携による具体的な実装ガイドであり、大規模モデル開発の現場で即座に活用できる実践的な知見が得られる貴重な記事です。
長い間、ファウンデーションモデルにおける「スケーリング」とは主に一つの意味を持っていました:事前学習により計算リソースをより多く投入すれば、能力が向上するというものです。この直感は、モデルパラメータ、データセットサイズ、およびトレーニング計算量をスケールさせた際に損失に予測可能なべき乗則の傾向が見られると報告した実証的研究である Kaplan et al. (2020) によって裏付けられていました。実際、これらの傾向は、大規模なアクセラレータ容量への持続的な投資と、それを効率的に活用するために必要な周辺分散インフラストラクチャを正当化するものでした。
しかし最前線は進化しており、スケーリングはもはや単一の曲線ではありません。NVIDIA の「1 から 3 へのスケーリング法則」という枠組みは有用にも強調しています。事前学習を超えて、パフォーマンスはポストトレーニング(例えば、教師あり微調整 (SFT) や強化学習 (RL) ベースの手法)および推論時の計算量(「長い思考」、検索/検証、マルチサンプル戦略など)を通じて、ますますスケーリングされるようになります。
図: NVIDIA ブログ "AI's Three Scaling Laws, Explained" から改変。
これらスケーリングレジームを総合すると、ファウンデーションモデルのライフサイクル(事前学習、事後学習、推論)は、密結合したアクセラレータ計算、高帯域幅・低遅延ネットワーク、分散ストレージバックエンドという収束するインフラ要件へと向かいます。また、リソース管理におけるオーケストレーションの重要性が高まり、クラスターヘルスの維持や大規模環境でのパフォーマンス病理の診断には、アプリケーションレベルおよびハードウェアレベルの観測性(observability)が不可欠となります。
もう一つの重要なトレンドは、モデル開発フレームワーク、クラスターリソース管理、運用ツールにわたるオープンソースソフトウェア(OSS)エコシステムへの、ファウンデーションモデルのライフサイクルにおける依存度の高まりです。クラスター層では、リソース管理は通常、Slurm や Kubernetes などのシステムによって提供されます。モデル開発と分散トレーニングは、一般的に PyTorch や JAX などのフレームワークで実装されています。監視および可視化、つまり観測可能性(observability)は、メトリクス収集には Prometheus を、可視化とアラート通知には Grafana を使用して達成されることが多く、これらはインフラストラクチャおよびリソース管理の上に位置する運用層として機能します。図 1 はこの階層型アーキテクチャを示しており、ハードウェアインフラストラクチャがリソースオーケストレーションを支え、それがさらに ML フレームワークを可能にし、観測可能性がすべての層にわたって展開されている様子を描いています。
*Figure 1: ファウンデーションモデルのトレーニングおよび推論におけるオープンソースソフトウェアスタックの階層型アーキテクチャ*
本記事は、基盤モデルのトレーニングおよび推論に関わる機械学習エンジニアおよび研究者を対象としており、特にオープンソース(OSS)フレームワーク上に構築されたワークフローに重点を置いています。ここでは、AWS インフラストラクチャ(マルチノードアクセラレータ計算、高帯域・低遅延ネットワーク、分散共有ストレージ、および関連するマネージドサービス)が、基盤モデルのライフサイクル全体を通じて一般的な OSS スタックとどのように相互作用するかを分析します。主な目的は、事前トレーニング、事後トレーニング、推論にわたるシステムのボトルネックとスケーリング特性を理解するための技術的基盤を提供することです。本導入記事では、大規模分散トレーニングおよび推論を支える AWS インフラストラクチャコンポーネントと OSS ツールの統合ポイントに重点を置きつつ、全体システムアーキテクチャを紹介しています。
The AWS Building Blocks
このシリーズの残りの部分では、この階層型アーキテクチャが AWS 上でどのように実現されるかを、インフラストラクチャ、リソースオーケストレーション、ML ソフトウェアスタック、および観測可能性(Observability)を通じて順を追って検証します。以下のセクションで各レイヤーをプレビューします。
Infrastructure: Compute, Network, and Storage
図 1 に示すように、インフラストラクチャは、大容量デバイスメモリを備えたアクセラレータ計算、集合通信のための広帯域インターコネクト、データおよびチェックポイント用のスケーラブルな分散ストレージという、3 つの結合されたビルディングブロックによって支えられています。
加速されたコンピューティングは、大規模なファウンデーションモデルの事前学習、事後学習、推論の基盤を形成しています。AWS は、Amazon EC2 P インスタンスファミリーを含む、Amazon EC2 加速コンピューティングインスタンスの一部として、複数の世代にわたる NVIDIA GPU を提供しています。P5 インスタンスファミリーには、8 つの NVIDIA H100 GPU を搭載した p5.48xlarge、小規模なワークロード向けの単一の H100 GPU を搭載した p5.4xlarge、および NVIDIA H200 GPU を搭載した p5e.48xlarge/p5en.48xlarge バリアントが含まれます。P6 インスタンスファミリーでは、p6-b200.48xlarge に NVIDIA Blackwell B200 アーキテクチャが導入され、Blackwell Ultra B300 を搭載した p6-b300.48xlarge が登場しています。
これらの世代全体を通じて、主要なスケーリング軸は、ピーク Tensor スループット、HBM(High Bandwidth Memory)の容量と帯域幅、およびノード内およびノード間の相互接続帯域幅です。
一次近似として、浮動小数点演算毎秒(FLOPS)で測定されるピーク Tensor Core スループットは、これらのアクセラレータを共通の軸上に位置づけるのに役立ちます。以下の表は、NVSwitch/NVLink ベースのマルチ GPU ノードと整合する SXM/HGX クラス仕様を用いた、密集型 BF16/FP16 および FP8 Tensor 演算ごとの GPU 別ピークスループット、HBM(High Bandwidth Memory)容量、および HBM バンド幅を要約したものです。
GPU(代表的なバリアント)
BF16/FP16 Tensor ピーク(密集型)
FP8 Tensor ピーク(密集型)
FP4 Tensor ピーク(密集型)
HBM 容量
HBM バンド幅
0.9895 PFLOPS
1.979 PFLOPS
—
80 GB HBM3
3.35 TB/s
0.9895 PFLOPS
1.979 PFLOPS
—
141 GB HBM3e
4.8 TB/s
2.25 PFLOPS
4.5 PFLOPS
9 PFLOPS
180 GB HBM3e
8 TB/s
2.25 PFLOPS
4.5 PFLOPS
13.5 PFLOPS
288 GB HBM3e
8 TB/s
*注記:NVIDIA の製品表では、スパース性を考慮したテンソルスループットを報告することが多いですが、本表は密な(dense)スループットを報告しています。該当する箇所では、NVIDIA が HGX クラスプラットフォームに対して示唆している通り、密なスループットはスパーススループットの半分として算出されています (NVIDIA)。DGX の数値はシステムレベルのものです。B200 の HBM 容量および帯域幅の数値は、DGX の総量を 8 で割ることで GPU 単位の値として表されています (NVIDIA)。
モデルがスケールするにつれ、ステップ時間は生計算スループットよりも集合通信やメモリアクセスによって支配されることが多く、明示的なスケールアップおよびスケールアウトの帯域幅の考慮が必要となります。マルチ GPU インスタンスにおける GPU 間の通信は、2 つの領域にまたがります。内部スケールアップ(NVLink/NVSwitch) は、ノード内の GPU と GPU の間を高帯域・低遅延で接続し、ホストネットワークスタックを迂回して all-reduce や all-gather などの集合演算を実行可能にします。外部スケールアウト(EFA) は、OS をバイパスしたネットワークをノード間で提供し、AWS ではこれを Amazon EC2 UltraClusters の構築ブロックとして利用しており、ここでは数千のインスタンスにわたって通信集約型の集合演算が行われます。以下の表は、これらのインスタンスタイプにおける主要な仕様を要約したものです:
Instance Type
GPU
GPUs
GPU Memory
NVLink
NVLink BW (aggregate)
EFA
EFA BW (aggregate)
H100
1
80 GB HBM3
—
—
v2
12.5 GB/s
H100
8
640 GB HBM3
第 4 世代
7.2 TB/s
v2
400 GB/s
H200
8
1,128 GB HBM3e
第 4 世代
7.2 TB/s
v2
400 GB/s
H200
8
1,128 GB HBM3e
第 4 世代
7.2 TB/s
v3
400 GB/s
B200
8
1,440 GB HBM3e
第 5 世代
14.4 TB/s
v4
400 GB/s
B300
8
2,100 GB HBM3e
第 5 世代
14.4 TB/s
v4
800 GB/s
*注:EFA(Elastic Fabric Adapter)の帯域幅は、他の帯域幅指標との整合性を保つため、Gbps から GB/s に換算されています(÷8)。詳細は EC2 高速コンピューティングネットワーク仕様 を参照してください。NVLink および EFA の帯域幅数値は、リンクごとの値ではなく、インスタンスあたりの集計値として表示されています。対応するノード間相互接続およびネットワーク特性については、P5 インスタンスファミリーページ および P6 インスタンスファミリーページ を参照してください。
Elastic Fabric Adapter (EFA) は、Amazon EC2 向けのネットワークインターフェースであり、Scalable Reliable Datagram (SRD) プロトコルを使用して、OS をバイパスしたリモート直接メモリアクセス(RDMA)機能を提供します。Libfabric API を介してアプリケーションがネットワークデバイスと直接通信し、オペレーティングシステムカーネルをバイパスできるようにすることで、EFA は分散学習における集合演算のレイテンシを低減し、スループットを向上させます。
異なるインスタンスファミリーでは複数の世代の EFA が利用可能です。Amazon EC2 P5 および P5e インスタンスには EFA バージョン 2 (EFAv2) が搭載されています。P5en インスタンスで提供される EFA バージョン 3 (EFAv3) は、EFAv2 と比較してパケットレイテンシを約 35% 削減します。P6 インスタンスで利用可能な EFA バージョン 4 (EFAv4) は、EFAv3 と比較して集合通信のパフォーマンスをさらに 18% 向上させます。
スケーラビリティの観点では、分散トレーニング(ストリーミングコーパスとマルチテラバイト単位のチェックポイント書き込み)および大規模推論(重みのステージングと KV キャッシュ成長の管理)の両方が、階層型ストレージ構造を必要とします。具体的には、ホットデータ用としてローカル NVMe SSD を、共有高スループットアクセス用として Lustre を、永続的保存用として Amazon S3 をそれぞれ利用します。
本シリーズの主要なマルチ GPU インスタンスでは、ローカル NVMe は インスタンスストア(エフェメラル) として提供され、30.72 TB の生容量(8 × 3.84 TB の NVMe SSD) を備えています。詳細は EC2 アクセラレーテッドコンピューティング インスタンスのストレージ仕様 をご覧ください。
Lustre は、オープンソースで POSIX 準拠の分散ファイルシステムであり、高性能計算(HPC)分野で広く利用されています。これは多数のクライアント間で高い集約スループットを備えた共有名前空間を提供するものです。Amazon FSx for Lustre は、Lustre をフルマネージドサービスとして提供し、秒間テラバイト単位のデータ転送能力、数百万の IOPS(入出力オペレーション数)、サブミリ秒のレイテンシを実現する並列ファイルシステムとして公開しています。Data Repository Associations(データリポジトリ関連付け)により Amazon S3 と統合が可能で、トレーニングデータの遅延読み込みや、永続性を確保するための自動チェックポイントエクスポートをサポートします。
クラスター規模では、これらのインスタンスは Amazon EC2 UltraClusters にデプロイされており、これは 1 つの可用性ゾーン内に数千台のアクセラレータ付きインスタンスを単一の密接に配置されたクラスターとしてプロビジョニングし、ペタビット規模のノンブロッキングネットワークを使用して相互接続するものです。
図: 第 2 世代 Amazon EC2 UltraClusters(例:P5 UltraCluster)。
ステップごとの通信強度が高いワークロード(例えば、すべての GPU にわたってトークンの分散を行う MoE モデルにおけるエキスパート並列処理など)の場合、NVLink ドメインのサイズが主要な制約要因となる可能性があります。内部スケールアップ軸の拡張として NVLink ドメインを拡大することで、パフォーマンスに重要な通信が NVLink ファブリックから外部に出る頻度を低減できます。
Amazon EC2 UltraServers は、専用アクセラレータ相互接続を介して複数のコンポーネントインスタンスを接続することで、NVLink ドメインを単一の EC2 インスタンスを超えて拡張します。AWS によると、P6e-GB200 UltraServers は NVIDIA GB200 NVL72 プラットフォームを基盤として構築されており、1 つの NVLink ドメイン内で最大 72 個の Blackwell GPU と合計 13.4 TB の HBM3e を提供します。より大規模なスケールでは、EFA がマルチ UltraServer ジョブにおけるノード間ファブリックとして引き続き機能しますが、ドメイン内の GPU 数を増やすことで、パフォーマンスに重要な通信が NVLink ファブリックから外部に出る頻度を低減できます。
これらのシステムは、NVIDIA Grace–Blackwell スーパーチップで構築されており、キャッシュ整合性のある NVLink-C2C を介して Grace CPU メモリと Blackwell GPU の HBM を結合しています。これにより、明示的なホスト - デバイス間のコピーを行わずに、CPU 接続メモリと GPU 接続メモリの間を直接アクセスすることが可能になります。実際には、これは GPU ワークロードに対して利用可能な有効メモリを拡張する効果があります(例えば、より冷たいモデル状態や KV キャッシュを CPU 接続メモリに配置するなど)。ただし、ローカル HBM に比べてレイテンシが高く、帯域幅が低いというトレードオフはあります。
P6e-GB200 UltraServers のコンポーネントインスタンスタイプは p6e-gb200.36xlarge であり、4 つの GPU と Elastic Fabric Adapter (EFA) v4 ネットワークを提供します。以下の表は、各インスタンスおよび構成された UltraServer の構成を要約したものです。
Instance Type | GPU | GPUs | GPU Memory | Memory BW | NVLink | NVLink BW | EFA | EFA BW
---|---|---|---|---|---|---|---|---
p6e-gb200.36xlarge | GB200 NVL72 | 4 | 740 GB HBM3e | — | — | — | v4 | 200 GB/s
*注:p6e-gb200.36xlarge の EFA バンド幅は、公開された集約 EFA ネットワーク (4 × 400 Gbps) から GB/s に変換されています (÷8)。詳細は EC2 アクセラレーテッドコンピューティングネットワーク仕様 を参照してください。
UltraServer | Component instance type | GPUs (NVLink domain) | HBM3e (aggregate) | EFA | EFA BW
---|---|---|---|---|---
u-p6e-gb200x36 | p6e-gb200.36xlarge | 36 | 6.7 TB | v4 | 1,800 GB/s
u-p6e-gb200x72 | p6e-gb200.36xlarge | 72 | 13.4 TB | v4 | 3,600 GB/s
*注記:UltraServer の EFA バンド幅は、AWS が報告したテラビット毎秒 (Tbps) からギガバイト毎秒 (GB/s) に換算されています(÷8);詳細は P6e-GB200 UltraServers 発表 および P6 インスタンスファミリーページ をご参照ください。*
リソースオーケストレーション:Slurm と Kubernetes
トレーニングが数百から数千のアクセラレーターにまたがる場合、手動によるリソース管理は実質的に不可能になります。例えば、512 個の GPU を必要とするトレーニングジョブでは、64 台の 8-GPU ノード(P インスタンス)を同時にスケジューリングし、完了または失敗時にリソースを原子的に解放する必要があります。Slurm と Kubernetes の両方は、この課題に対処するためにコントロールプレーンアーキテクチャを採用しています:集中型スケジューラーがクラスターの状態を維持し割り当て決定を行い、ワーカーノードは割り当てられたワークロードを実行します。
*図 2: AWS における Slurm ベースおよび Kubernetes ベースのリソースオーケストレーションのハイレベルアーキテクチャ*
Slurm(Simple Linux Utility for Resource Management)は、高パフォーマンスコンピューティングにおける支配的なワークロードマネージャーであり、モジュール型プラグインアーキテクチャに基づいて構築されています。これにより、スケジューリングアルゴリズム、トポロジーモデル、リソースタイプ、およびアカウンティングバックエンドを独立して構成することが可能です。そのスケジューリングモデルは、リソースをパーティション(ノードの論理的なグループ化)に整理し、sbatch を介してジョブの提出を受け付け、割り当てられたノード間で同期起動を行いながら srun を通じて並列タスクを開始します。分散トレーニングにおいて特に重要なのは、Slurm がジョブレベルでスケジューリングを行う点です。つまり、タスクが開始される前に、マルチノジョブ全体を原子操作として割り当てます。バックフィルスケジューラー は、優先度の低いジョブをアイドル状態のスロットに開始し、より高い優先度のジョブの遅延を引き起こすことなく稼働させます。一方、多要素優先度 システムは、フェアシェア利用状況、ジョブの経過時間、QOS(Quality of Service)ティアを重み付けして、テナント間でのキュー順序を決定します。Slurm はまた、ネットワークスイッチ階層をモデル化するプラグインを通じて トポロジー認識配置 をサポートしており、AWS 環境では EFA ファブリックのトポロジーをエンコードして、最小限のスイッチホップ数でノード上にジョブを共置します。さらに、Generic Resource(GRES)インターフェースを通じてネイティブな GPU スケジューリング をサポートしており、これは GPU タイプを追跡し、デバイスアフィニティを強制する機能を持っています。
AWS は、Slurm ベースのオーケストレーションに対して複数のデプロイオプションを提供しています。AWS ParallelCluster は、EC2 上で Slurm クラスターを自動化して展開するオープンソースのクラスター管理ツールであり、ヘッドノードのプロビジョニング、計算リソース群のスケーリング、共有ストレージとの統合を処理します。AWS Parallel Computing Service (PCS) は、マネージドされたコントロールプレープを提供する代替手段です。特に分散型トレーニングワークロードについては、Amazon SageMaker HyperPod が Slurm モードをサポートしており、大規模なトレーニング向けに特別に設計された追加機能として、継続的なノードヘルスモニタリング や ジョブの自動再開機能 を提供しています。
Kubernetes は宣言的かつ API ドライブ型のアプローチを採用しており、ユーザーはリソースマニフェストを通じて望ましい状態を指定し、コントローラーが実際の状態をその状態に一致させるように調整します。Kubernetes はモデル展開において卓越した能力を発揮しますが、そのネイティブなスケジューリングモデルには、密結合型分散トレーニングにおいていくつかの課題が存在します。Kubernetes はポッドレベルでスケジューリングを行うため、ジョブレベルでの原子性が欠如しており、マルチノード構成のトレーニングジョブが部分的に開始されることがあります。具体的には、一部のランクは実行中である一方で他のランクは Pending 状態のままとなり、GPU が無駄になったりデッドロックが発生したりする可能性があります。また、標準的な Kubernetes には、優先度ベースのバックフィルを備えたバッチキューのセマンティクスや、通信集約型の集合操作(collectives)の配置のために必要なネットワークファブリックトポロジ(NVLink ドメイン、EFA インターコネクトなど)への組み込み認識機能も不足しています。
いくつかの Kubernetes ネイティブプロジェクトが、異なるレイヤーにおいてこれらのギャップを埋めています。Kueue はデフォルトスケジューラの上にアドミッションコントローラーとして動作し、ジョブレベルのガングアッドミッション、階層的公平共有によるマルチテナントクォータ管理、優先度ベースのプリエンプションを処理しますが、ポッド配置は基盤となるスケジューラに委譲します。Volcano と NVIDIA KAI Scheduler は異なるアプローチを採用し、デフォルトスケジューラを置き換えまたは拡張して、ガンスケジューリングをトポロジ認識型ポッド配置と直接統合します。Volcano は汎用バッチスケジューラとして機能し、KAI Scheduler は GPU 最適化配置のための深い NVLink/NVSwitch 認識機能を備えています。これらのレイヤーは補完関係にあります:Kueue がアドミッションとクォータポリシーを管理しつつ、承認されたジョブをトポロジ認識型スケジューラに渡して配置を行うことができます。
AWS における Kubernetes ベースのオーケストレーションには、Amazon Elastic Kubernetes Service (EKS) が、NVIDIA デバイスプラグイン を介して GPU スケジューリング機能を備えたマネージド Kubernetes を提供しています。Amazon SageMaker HyperPod も EKS モード をサポートしており、Kubernetes のオーケストレーション能力と HyperPod のトレーニング特化型機能を組み合わせています。HyperPod EKS は、大規模なファウンデーションモデルのトレーニング向けに設計された機能で EKS を拡張します。タスクガバナンス により、チーム間での計算リソースの割り当てとポリシーの強制が可能となり、管理された Kueue を導入して admission control(入場制御)を行い、Karpenter を用いてジョブ発生時に必要なノードを即時プロビジョニングします。チェックポイントレストレーニング は、従来のチェックポイントベースのフォールトトレランスに内在する回復遅延の問題に対処します。モデル状態を定期的に共有ストレージへシリアライズするのではなく、チェックポイントレストレーニングでは GPU 間で継続的なピアツーピアの状態複製を維持します。障害が発生した場合、生存しているノードは FSx for Lustre や S3 から数テラバイトのチェックポイントを読み込むのではなく、EFA ベースの通信を通じて失われた状態を再構築します。エラスティックトレーニング により、ジョブはリソースの利用状況に基づいて自動的にスケールします。追加のアクセラレーターが利用可能になった場合(例えば、完了したジョブから解放されたものや新たにプロビジョニングされた容量)、エラスティックジョブはそれらを活用して拡張できます。一方、より優先度の高いワークロードがリソースを必要とする場合は、トレーニングの進捗を維持しながらジョブを縮小できます。
ML ソフトウェアスタック
分散トレーニングと推論には、正しく構成され調整される複数のソフトウェア層が含まれます。有用なモデルは、ランタイムスタックを 5 つの層として捉えます。これらは、ハードウェアに隣接するコンポーネント(何らかの実行を行うために正しく機能する必要がある)から、プログラマの生産性とモデルのスループットを決定するフレームワークレベルの抽象化まで順に並べられています。具体的には、ハードウェア対応、アクセラレータランタイムおよび数値計算ライブラリ、通信基盤、ML フレームワーク、分散トレーニング/推論フレームワークです。
*図 3: EC2 インスタンス上での分散トレーニングと推論における ML ソフトウェアスタック*
ハードウェア対応:カーネルドライバ
基盤となる Linux カーネルドライバは、ハードウェアへの直接アクセスを提供します。NVIDIA GPU ドライバは計算機能の公開を行い、GPU とネットワークアダプタ間のデータ転送を可能にする GPUDirect RDMA をサポートしています。GDRCopy ドライバ(gdrdrv)は、GPU メモリへの CPU による低遅延コピーを可能にし、NCCL が小メッセージ転送に使用します。EFA driver は libfabric API を通じて OS バイパスネットワークを提供し、Lustre client ドライバは FSx for Lustre パラレルファイルシステムへの POSIX アクセスを可能にします。
アクセラレータランタイム、コンパイラ、およびカーネルライブラリ
CUDA プラットフォームは、GPU 計算のためのプログラミングモデルとランタイムを提供します。CUDA に基づいてコンパイルされたアプリケーションは、NVIDIA GPU でカーネルを実行し、デバイスメモリを管理し、複数のデバイスにわたる実行を調整できます。現在のリリースは CUDA Toolkit 13.x で、Blackwell アーキテクチャ(計算能力 10.x)をサポートしています。
現代のトレーニングおよび推論のパフォーマンスは、汎用ベンダープリミティブだけでなく、専門的な最適化ライブラリとカスタムカーネルによってますます駆動されています。FlashAttention などのカーネルは、アテンションを単一のメモリ効率的なパスに融合させ、HBM(High Bandwidth Memory)のトラフィックを削減し、スループットを向上させます。多くのチームは、形状や精度に特化した融合カーネル(例:レイヤーノーマライゼーション/リジューアル/活性化、量子化 GEMM、MoE 分散、KV キャッシュ操作など)を、自社のモデルに合わせてチューニングして記述しています。これは、Triton(Python GPU カーネルコンパイラ)や NVIDIA の CuTe(テンソルレイアウトおよびワープレベル DSL)といったプログラム可能なツールチェーンによって可能になり、CUTLASS などのライブラリが高度に最適化された GEMM および融合のビルディングブロックを提供しています。実際には、このカーネルとコンパイラ層が ML フレームワークと同程度にエンドツーエンドのパフォーマンスを決定することが多いです。
コミュニケーション基盤:NCCL とトランスポートプラグイン
マルチ GPU 学習は、効率的な集合通信に依存しています。NVIDIA Collective Communications Library (NCCL) は、トポロジ認識アルゴリズムを実装した集合演算(オールリダクション、オールギャザ、リダックスキャッター、オールツーオール、ブロードキャスト、およびポイントツーポイント送受信)を提供し、ノード内通信には NVLink を、ノード間トラフィックにはネットワーク転送を活用します。NCCL は動的に通信トポロジを検出し、メッセージサイズと利用可能な帯域幅に応じてリングまたはツリーアルゴリズムを選択します。データ並列およびテンソル並列戦略は主にオールリダクションとオールギャザに依存しますが、エキスパート並列性を備えた Mixture-of-Experts (MoE) モデルでは、トークンを GPU 間でルーティングするためにオールツーオールの集合通信が必要です:ディスパッチオールツーオールは各トークンを割り当てられたエキスパートをホストする GPU に送信し、コンバインオールツーオールはエキスパートの出力を発信元の GPU に返します (NVIDIA Developer Blog)。エキスパート並列グループ内のすべての GPU が他のすべての GPU とデータを交換するため、オールツーオールの通信量はエクスパート数に比例して増加し、高い並列度では主要なボトルネックとなる可能性があります。
AWS では、NCCL のノード間通信は aws-ofi-nccl プラグインを介して有効化されており、これは NCCL のトランスポート API を libfabric インターフェースにマッピングします。これにより、NCCL はアプリケーションの変更なしで EFA の OS バイパス機能とスケーラブル・リライアブル・データグラム(SRD)プロトコルを活用できます。
推論ワークロードにおいては、集合演算がすべての通信パターンを網羅するわけではありません。プリフェッチフェーズとデコードフェーズを別々の GPU プールに分離した非集約型推論アーキテクチャでは、特にインスタンス間で KV キャッシュ状態を転送する際に、効率的なポイント・ツー・ポイントデータ移動が必要です。NVIDIA Inference Xfer Library (NIXL) は、メモリ階層(HBM, DRAM, NVMe, 分散ストレージ)および相互接続(NVLink, InfiniBand, Ethernet)にわたるポイント・ツー・ポイント転送のための統一 API を提供することで、この要件に対応します。NIXL は NVIDIA Dynamo などの推論フレームワークと統合され、UCX や GPUDirect Storage などのバックエンドをサポートしています。
ML フレームワーク:PyTorch
ファウンデーションモデル開発における2 つの主要なフレームワークは、PyTorch と JAX です。JAX は XLA を通じて SPMD(Single Program Multiple Data)アプローチを採用しており、同じプログラムがデバイス間で実行され、データ分布と集合的な lowering が自動的に行われます。本シリーズでは、オープンソースエコシステムでより広く採用されており、以下で議論される分散トレーニングおよび推論フレームワークの基礎となっている PyTorch に焦点を当てます。
PyTorch は、GPU によるアクセラレーション、自動微分、柔軟なイーグゼキューションモデルを備えたテンソル計算を提供します。分散ワークロードにおいては、PyTorch の torch.distributed モジュールが中核となるプリミティブを提供します。これには、集合通信のためのプロセスグループと、Distributed Data Parallel (DDP) や Fully Sharded Data Parallel (FSDP2) を含む分散データ並列化の抽象化が含まれます。DDP はモデルを GPU 間で複製し、all-reduce を通じて勾配を同期します。一方、FSDP2 は ZeRO アルゴリズムの技術を用いて、パラメータ、勾配、オプティマイザの状態をワーカー間でシャード(分割)することで、単一の GPU メモリ容量を超えるモデルのトレーニングを可能にします。
分散トレーニングおよび推論フレームワーク
最上層には、PyTorch を基盤とし、大規模な分散トレーニングと推論のためのより高レベルの抽象化を提供するフレームワークが含まれています。トレーニングにおいては、3 つのカテゴリのフレームワークが、複雑性とパフォーマンスのトレードオフにおける異なるポイントに対応しています。以下にいくつかの例を示します。
Hugging Face Transformers は、Trainer クラスを提供しており、Accelerate を介して DDP(Data Parallelism)、FSDP(Fully Sharded Data Parallel)、DeepSpeed に対する抽象化を内蔵した形でサポートしています。このアプローチは使いやすさと広範なモデル互換性を優先しており、設定の簡素さが最大のスループットよりも重要となるファインチューニングや中規模のトレーニングに適しています。
NVIDIA Megatron Core は、Transformer Engine を介した FP8 混合精度計算などの最適化を備えた 3D パラレリズム(テンソル並列、パイプライン並列、エキスパート並列)を実装することで、大規模における最大効率性を目標としています。NeMo Framework は Megatron Core を基盤としており、事前トレーニングとファインチューニングのためのエンドツーエンドのワークフローを提供します。
人間フィードバックからの強化学習(RLHF)および関連するポストトレーニング手法においては、veRL(Volcano Engine Reinforcement Learning)が、PPO、GRPO、REINFORCE++ などのアルゴリズムを実装した柔軟なフレームワークを提供します。veRL の HybridFlow アーキテクチャでは、トレーニングバックエンド(FSDP2, Megatron)と推論エンジン(vLLM, SGLang)を同じジョブ内で組み合わせることが可能であり、アクターコンポーネントとロールアウトコンポーネント間でモデル重みをメモリ上で共有することで、重みの同期に伴うオーバーヘッドを回避します。
推論サービングにおいては、vLLM が PagedAttention を実装し、KV キャッシュ(Key-Value Cache)をページ化された仮想メモリとして管理することで断片化を削減し、より大きなバッチサイズを実現します。SGLang はこれに RadixAttention を追加してリクエスト間での自動プレフィックス再利用を実現し、CPU のスケジューリングと GPU 計算を重畳させるゼロオーバーヘッドのバッチスケジューラ、および予測されたキャッシュヒット率に基づいてリクエストをルーティングするキャッシュ認識型ロードバランサを提供します。両方のフレームワークは、単一 GPU メモリを超えるモデルをサービングするためのテンソル並列化をサポートしており、NVIDIA Dynamo と統合して、プリフェッチ(prefill)フェーズとデコード(decode)フェーズを分離する非集約型サービングアーキテクチャを実現します。
観測可能性
観測可能性は、大規模な分散型トレーニングシステムのデバッグと運用に不可欠な前提条件です。トレーニングジョブが停止したりスループットが低下したりした場合、実践者はその原因がハードウェアの故障、ネットワークの輻輳、ストレージのボトルネック、あるいはアプリケーションレベルの非効率性のいずれにあるかを可視化する必要があります。本シリーズで議論されているインフラスケール(数千個の GPU、ペタビット規模の相互接続帯域幅、テラバイト規模のチェックポイントデータ)において、課題は単純な監視から、体系的なテレメトリ収集、保存、分析へと移行します。観測可能性には 3 つのテレメトリカテゴリが含まれます:インフラストラクチャ指標(GPU、ネットワーク、ストレージ)、ワークロード指標(トレーニングスループット、キュー待ち時間)、および予防的な障害検出のためのアラートです。
コアスタック:Prometheus と Grafana
Kubernetes および HPC 環境における観測可能性の事実上の標準は、メトリクス収集に Prometheus を、可視化とアラートに Grafana を組み合わせたものです。Prometheus はプルベースモデルで動作し、メトリクスエクスポート者が公開する HTTP エンドポイントを定期的にスクレイピングします。収集されたメトリクスは時系列データベース(TSDB)に保存され、集約、フィルタリング、アラートルールの評価を行うための柔軟なクエリ言語である PromQL を介して照会されます。Grafana は Prometheus をデータソースとして取り込み、PromQL 式に基づいてダッシュボードを描画し、アラートをトリガーします。
本番環境でのデプロイメントにおいては、Amazon Managed Service for Prometheus (AMP) が、ストレージ、レプリケーション、高可用性の管理をオペレーターが行う必要なく、1 秒間に数百万サンプルの取り込みに対応できるスケーラブルなフルマネージド型の Prometheus 互換時系列データベースを提供します。Amazon Managed Grafana (AMG) は、AMP とのネイティブ統合および IAM Identity Center を介した AWS 認証機能を備えたマネージド型の Grafana ワークスペースを提供します。これら 2 つのサービスを組み合わせることで、既存の Prometheus エクスポートヤーや Grafana ダッシュボードとの互換性を維持しつつ、運用上のオーバーヘッドを解消できます。
GPU、ネットワーク、およびアプリケーションのテレメトリー
DCGM-Exporter は、NVIDIA の GPU メトリクスを Prometheus 形式で公開します。これには利用率、メモリ使用量、電力消費、温度、ならびに ECC エラーや XID イベントといったハードウェア健全性の指標が含まれます。トレーニングワークロードにおいては、SM アクティビティ(DCGM_FI_PROF_SM_ACTIVE)が、基本的な利用率メトリクスよりも計算効率をより正確に測定する手段となることがよくあります。
EFA は、分散トレーニングにおける集合操作のボトルネックを診断するのに役立つ、ドライバーレベルの統計情報(バイト数、パケット数、再送信回数、タイムアウト)を公開します。aws-ofi-nccl プラグインは NCCL を libfabric インターフェースに橋渡しし、オペレーターは EFA カウンターと NCCL 診断情報(NCCL_DEBUG=INFO)を組み合わせて、ネットワーク層の問題を特定できます。
Amazon FSx for Lustre は、スループットやメタデータレイテンシなどのクライアントサイドのメトリクスを公開します。一方、アプリケーションレベルのメトリクス(トレーニングにおけるステップ時間、1 秒あたりのトークン数、損失値;推論における TTFT、トークン間レイテンシ)は、Prometheus クライアントライブラリ経由でエクスポートできます。
GPU ヘルスモニタリングとアラート
予防的な障害検知により、ハードウェアの問題が拡張されたトレーニングの中断へと波及するのを防ぎます。一般的なワークフローでは、DCGM のヘルスメトリクスを監視し、エラーカウントが閾値を超えた場合にアラートを発令します。ECC シングルビットエラー(SBE)は少量であれば許容される場合もありますが、SBE 発生率の加速は、ダブルビットエラー(DBE)やその他の障害の前兆であることがよくあります。XID 63(行再マッピング失敗)、XID 64(GPU がバスから脱落)、および XID 94/95(封じ込め済み/封じ込め不可能なエラー)は、通常、即座にノードの交換を要します。
【GPU ヘルス - クラスターダッシュボード】(Grafana dashboard ID 21645) は、一般的な GPU エラーパターンの参照用可視化を提供します。このダッシュボードは、クラスター内のすべてのノードにわたる ECC エラー、XID イベント、熱違反、および行リマップステータスを集約し、オペレーターがトレーニングジョブに影響を与える前に故障するハードウェアを特定できるようにします。
*図 4: GPU エラーパターンとインスタンスレポートを示す GPU ヘルス - クラスターダッシュボード*
本シリーズのパート 5 では、大規模なクラスターの健全性を維持するためのメトリクス収集戦略、ダッシュボード設定、アラートパターンを含む、観測可能性アーキテクチャについて包括的に解説します。
結論
単一の事前トレーニングのスケーリング法則から、相互補完的な 3 つのレジーム(事前トレーニング、事後トレーニング、テスト時計算)への移行は、インフラストラクチャ要件を分断したのではなく、むしろ強化しました。これら 3 つのレジームすべてが、密結合されたアクセラレータ計算、高帯域幅・低遅延ネットワーク、スケーラブルな分散ストレージを必要としますが、主にワークロードプロファイルとリソーススケジューリングのパターンにおいて異なります。
本記事では、AWS 上でこれらの要件を満たすための 4 レイヤー構成を紹介しました。それは、インフラストラクチャのビルディングブロック(EC2 P インスタンス、EFA ネットワーク、階層型ストレージ)、リソースオーケストレーション(Slurm と SageMaker HyperPod を用いた Kubernetes)、ML ソフトウェアスタック(カーネルドライバや CUDA から NCCL、PyTorch まで)、そして観測性(Prometheus、Grafana、GPU ヘルスモニタリング)です。各レイヤーは上位のレイヤーを制約し、同時に可能にします。設定ミスをしたドライバや飽和したネットワークリンクが、最適化された並列処理戦略よりも効果的に、それ以外の点で適切に調整されたトレーニング実行のボトルネックとなる可能性があります。
これらの統合ポイントを理解することは、ファウンデーションモデルのライフサイクル全体を通じてパフォーマンスのボトルネックを診断し、根拠のあるスケーリング判断を下すための基盤となります。
著者
Aman Shanbhag は NVIDIA の MARS MLOps チームに所属する AI パフォーマンスおよびインフラストラクチャエンジニアです。研究チームがスケーラブルで高性能な ML トレーニングおよび推論システムを構築できるよう支援しています。以前は AWS でスペシャリストソリューションアーキテクトとして勤務し、世界中の顧客に対して AWS 上での ML トレーニングと推論の最適化をサポートしました。Aman はライス大学でコンピュータサイエンス、数学、そして起業家精神の学位を取得しており、AI インフラストラクチャやパフォーマンス最適化に注力しています。
原文を表示
For a long time, "scaling" in foundation models mostly meant one thing: spend more compute on pre-training and capabilities rise. That intuition was supported by empirical work such as Kaplan et al. (2020), which reported predictable power-law trends in loss as you scale model parameters, dataset size, and training compute. In practice, these trends justified sustained investment in large-scale accelerator capacity and the surrounding distributed infrastructure needed to keep it efficiently utilized.
But the frontier has evolved—and scaling is no longer a single curve. NVIDIA's "from one to three scaling laws" framing usefully emphasizes that, beyond pre-training, performance increasingly scales through post-training (e.g., supervised fine-tuning (SFT) and reinforcement learning (RL)-based methods) and through test-time compute ("long thinking," search/verification, multi-sample strategies).
Figure: Adapted from "AI's Three Scaling Laws, Explained" (NVIDIA Blog).
Taken together, these scaling regimes push the foundation-model lifecycle—pre-training, post-training, and inference—toward convergent infrastructure requirements: tightly coupled accelerator compute, a high-bandwidth low-latency network, and a distributed storage backend. They also raise the importance of orchestration for resource management, and of application- and hardware-level observability to maintain cluster health and diagnose performance pathologies at scale.
Another key trend is the increasing reliance of the foundation-model lifecycle on an open-source software (OSS) ecosystem that spans model development frameworks, cluster resource management, and operational tooling. At the cluster layer, resource management is typically provided by systems such as Slurm and Kubernetes. Model development and distributed training are commonly implemented in frameworks such as PyTorch and JAX. Monitoring and visualization—that is, observability—are often achieved using Prometheus for metrics collection and Grafana for visualization and alerting, positioned as an operational layer atop infrastructure and resource management. Figure 1 illustrates this layered architecture, showing how hardware infrastructure supports resource orchestration, which in turn enables ML frameworks, with observability spanning across all layers.
*Figure 1: The layered architecture of open-source software stacks for foundation model training and inference*
This post is intended for machine learning engineers and researchers involved in foundation model training and inference, with particular attention to workflows built atop OSS frameworks. It analyzes how AWS infrastructure—including multi-node accelerator compute, high-bandwidth low-latency networking, distributed shared storage, and associated managed services—interacts with common OSS stacks across the foundation model lifecycle. The primary goal is to provide a technical foundation for understanding systems bottlenecks and scaling characteristics spanning pre-training, post-training, and inference. This introductory post surfaces the overall system architecture, emphasizing the integration points between AWS infrastructure components and OSS tools that underpin large-scale distributed training and inference.
The AWS Building Blocks
The remainder of this series examines how this layered architecture is realized on AWS, progressing through infrastructure, resource orchestration, the ML software stack, and observability. The following sections preview each layer.
Infrastructure: Compute, Network, and Storage
As illustrated in Figure 1, infrastructure is anchored by three coupled building blocks—accelerated compute with large device memory, wide-bandwidth interconnect for collective communication, and scalable distributed storage for data and checkpoints.
Accelerated compute forms the foundation of large-scale foundation model pre-training, post-training, and inference. AWS offers several generations of NVIDIA GPUs as part of its Amazon EC2 accelerated computing instances, including the Amazon EC2 P instance family. The P5 instance family includes p5.48xlarge with eight NVIDIA H100 GPUs, p5.4xlarge with a single H100 GPU for smaller-scale workloads, and p5e.48xlarge/p5en.48xlarge variants with NVIDIA H200 GPUs. The P6 instance family introduces NVIDIA Blackwell B200 architecture with p6-b200.48xlarge and Blackwell Ultra B300 with p6-b300.48xlarge.
Across these generations, the dominant scaling axes are peak Tensor throughput, HBM capacity and bandwidth, and interconnect bandwidth (within and across nodes).
As a first-order approximation, peak Tensor Core throughput—measured in floating point operations per second (FLOPS)—helps situate these accelerators on a common axis. The table below summarizes per-GPU peak throughput for dense BF16/FP16 and FP8 Tensor operations, along with HBM capacity and HBM bandwidth, using SXM/HGX-class specifications that align with NVSwitch/NVLink-based multi-GPU nodes.
GPU (representative variant)
BF16/FP16 Tensor peak (dense)
FP8 Tensor peak (dense)
FP4 Tensor peak (dense)
HBM capacity
HBM bandwidth
0.9895 PFLOPS
1.979 PFLOPS
—
80 GB HBM3
3.35 TB/s
0.9895 PFLOPS
1.979 PFLOPS
—
141 GB HBM3e
4.8 TB/s
2.25 PFLOPS
4.5 PFLOPS
9 PFLOPS
180 GB HBM3e
8 TB/s
2.25 PFLOPS
4.5 PFLOPS
13.5 PFLOPS
288 GB HBM3e
8 TB/s
*Note: NVIDIA product tables often report Tensor throughput “with sparsity”; this table reports dense throughput. Where applicable, dense throughput is taken as half of sparse throughput, following NVIDIA’s guidance for HGX-class platforms (NVIDIA). DGX figures are system-level; the B200 HBM capacity and bandwidth values are expressed per GPU by dividing DGX totals by eight (NVIDIA).*
As models scale, step time is often dominated by collective communication and memory movement rather than raw compute throughput, motivating explicit scale-up and scale-out bandwidth accounting.
For the multi-GPU instances, GPU communication spans two regimes. Internal scale-up (NVLink/NVSwitch) provides high-bandwidth, low-latency GPU-to-GPU connectivity within a node, enabling collectives such as all-reduce and all-gather to execute without traversing the host networking stack. External scale-out (EFA) provides OS-bypass networking across nodes, which AWS uses as a building block for Amazon EC2 UltraClusters where communication-heavy collectives span thousands of instances. The following table summarizes key specifications across these instance types:
Instance Type
GPU
GPUs
GPU Memory
NVLink
NVLink BW (aggregate)
EFA
EFA BW (aggregate)
H100
1
80 GB HBM3
—
—
v2
12.5 GB/s
H100
8
640 GB HBM3
4th
7.2 TB/s
v2
400 GB/s
H200
8
1,128 GB HBM3e
4th
7.2 TB/s
v2
400 GB/s
H200
8
1,128 GB HBM3e
4th
7.2 TB/s
v3
400 GB/s
B200
8
1,440 GB HBM3e
5th
14.4 TB/s
v4
400 GB/s
B300
8
2,100 GB HBM3e
5th
14.4 TB/s
v4
800 GB/s
*Note: EFA bandwidth is converted from Gbps to GB/s (÷8) for consistency with other bandwidth metrics; see the EC2 accelerated computing networking specifications. NVLink and EFA bandwidth figures are shown as aggregate per-instance values rather than per-link values; see the P5 instance family page and the P6 instance family page for the corresponding intra-node interconnect and networking characteristics.*
Elastic Fabric Adapter (EFA) is a network interface for Amazon EC2 that provides OS-bypass remote direct memory access (RDMA) capability using the Scalable Reliable Datagram (SRD) protocol. By enabling applications to communicate directly with the network device through the Libfabric API—bypassing the operating system kernel—EFA reduces latency and improves throughput for collective operations in distributed training.
Multiple generations of EFA are available on different instance families. Amazon EC2 P5 and P5e instances are equipped with EFA version 2 (EFAv2). EFA version 3 (EFAv3), provided on P5en instances, reduces packet latency by approximately 35% compared to EFAv2. EFA version 4 (EFAv4), available on P6 instances, delivers an additional 18% improvement in collective communication performance relative to EFAv3.
At scale, both distributed training (streaming corpora and writing multi-terabyte checkpoints) and large-scale inference (staging weights and managing KV cache growth) motivate a tiered storage hierarchy—local NVMe SSD for hot data, Lustre for shared high-throughput access, and Amazon S3 for durable persistence.
In this series’ primary multi-GPU instances, local NVMe is provided as instance store (ephemeral) with 30.72 TB raw capacity (8 × 3.84 TB NVMe SSD); see the EC2 accelerated-computing instance store specifications.
Lustre is an open-source, POSIX compliant distributed file system widely used in high-performance computing (HPC) to provide a shared namespace with high aggregate throughput across many clients. Amazon FSx for Lustre provides Lustre as a fully managed service and exposes it as a parallel file system capable of terabytes per second of throughput, millions of IOPS, and sub-millisecond latencies. Data Repository Associations enable integration with Amazon S3, supporting lazy loading of training datasets and automatic checkpoint export for durability.
At cluster scale, these instances are deployed in Amazon EC2 UltraClusters, which provision thousands of accelerated instances as a single, tightly placed cluster within an Availability Zone and interconnect them using a petabit-scale nonblocking network.
Figure: 2nd-generation Amazon EC2 UltraClusters (example P5 UltraCluster).
For workloads with high per-step communication intensity (e.g., expert parallelism in MoE models where all-to-all token dispatch spans many GPUs), the size of the NVLink domain can become a first-order constraint. As an extension of the internal scale-up axis, increasing the NVLink domain reduces how often performance-critical communication must leave the NVLink fabric.
Amazon EC2 UltraServers extend the NVLink domain beyond a single EC2 instance by connecting multiple component instances through a dedicated accelerator interconnect. AWS reports that P6e-GB200 UltraServers are built on the NVIDIA GB200 NVL72 platform and expose up to 72 Blackwell GPUs and 13.4 TB of aggregate HBM3e within one NVLink domain. At larger scales, EFA remains the cross-node fabric for multi-UltraServer jobs, but increasing the intra-domain GPU count can reduce how often performance-critical communication must leave the NVLink fabric.
These systems are built from NVIDIA Grace–Blackwell superchips, which couple Grace CPU memory and Blackwell GPU HBM via cache-coherent NVLink-C2C, enabling direct access across CPU- and GPU-attached memory without explicit host–device copies. In practice, this can extend the effective memory available to GPU workloads (e.g., by placing colder model state or KV cache in CPU-attached memory) while avoiding PCIe-scale copy overheads, albeit with higher latency and lower bandwidth than local HBM.
The component instance type for P6e-GB200 UltraServers is p6e-gb200.36xlarge, which provides four GPUs and Elastic Fabric Adapter (EFA) v4 networking. The tables below summarize the per-instance and composed UltraServer configurations.
Instance Type
GPU
GPUs
GPU Memory
Memory BW
NVLink
NVLink BW
EFA
EFA BW
4
—
—
—
v4
*Note: The p6e-gb200.36xlarge EFA bandwidth is converted from the published aggregate EFA networking (4 × 400 Gbps) to GB/s (÷8); see the EC2 accelerated computing networking specifications.*
UltraServer
Component instance type
GPUs (NVLink domain)
HBM3e (aggregate)
EFA
EFA BW
u-p6e-gb200x36
36
6.7 TB
v4
1,800 GB/s
u-p6e-gb200x72
72
13.4 TB
v4
3,600 GB/s
*Note: UltraServer EFA bandwidth is converted from terabits per second (Tbps), as reported by AWS, to GB/s (÷8); see the P6e-GB200 UltraServers announcement and the P6 instance family page.*
Resource Orchestration: Slurm and Kubernetes
When training spans hundreds or thousands of accelerators, manual resource management becomes intractable. For example, a training job requiring 512 GPUs must co-schedule 64 eight-GPU nodes (P-instances) simultaneously, and release resources atomically upon completion or failure. Both Slurm and Kubernetes address this challenge through a control-plane architecture: a centralized scheduler maintains cluster state and makes allocation decisions, while worker nodes execute assigned workloads.
*Figure 2: High-level architecture of Slurm-based and Kubernetes-based resource orchestration on AWS*
Slurm (Simple Linux Utility for Resource Management) is the dominant workload manager in high-performance computing, built on a modular plugin architecture that allows the scheduling algorithm, topology model, resource types, and accounting backend to be configured independently. Its scheduling model organizes resources into partitions (logical groupings of nodes), accepts job submissions via sbatch, and launches parallel tasks via srun with synchronized startup across allocated nodes. Critically for distributed training, Slurm schedules at the job level—allocating entire multi-node jobs atomically before any task launches. A backfill scheduler starts lower-priority jobs in idle slots without delaying higher-priority ones, while a multi-factor priority system weighs fair-share usage, job age, and QOS tiers to order the queue across tenants. Slurm also supports topology-aware placement through plugins that model network switch hierarchies—on AWS, encoding the EFA fabric topology to co-locate jobs on nodes with minimal switch hops—and native GPU scheduling through its Generic Resource (GRES) interface, which tracks GPU types and enforces device affinity.
AWS provides multiple deployment options for Slurm-based orchestration. AWS ParallelCluster is an open-source cluster management tool that automates the deployment of Slurm clusters on EC2, handling head node provisioning, compute fleet scaling, and integration with shared storage. AWS Parallel Computing Service (PCS) offers an alternative that provides the managed control plane. For distributed training workloads specifically, Amazon SageMaker HyperPod supports Slurm mode with additional capabilities tailored to large-scale training, such as continuous node health monitoring and job auto-resume functionality.
Kubernetes takes a declarative, API-driven approach: users specify desired state through resource manifests, and controllers reconcile actual state to match. While Kubernetes excels at model deployment, its native scheduling model exposes several gaps for tightly coupled distributed training. Kubernetes schedules at the pod level; without job-level atomicity, a multi-node training job can partially start—some ranks running while others remain Pending—wasting GPUs or causing deadlocks. Vanilla Kubernetes also lacks batch queue semantics with priority-based backfill, built-in awareness of network fabric topology (NVLink domains, EFA interconnects) for placement of communication-heavy collectives.
Several Kubernetes-native projects address these gaps at different layers. Kueue operates as an admission controller atop the default scheduler, managing job-level gang admission, multi-tenant quotas with hierarchical fair sharing, and priority-based preemption—while delegating pod placement to the underlying scheduler. Volcano and NVIDIA KAI Scheduler take a different approach, replacing or augmenting the default scheduler to integrate gang scheduling directly with topology-aware pod placement—Volcano as a general-purpose batch scheduler, KAI Scheduler with deep NVLink/NVSwitch awareness for GPU-optimized placement. These layers are complementary: Kueue can manage admission and quota policy while passing admitted jobs to a topology-aware scheduler for placement.
For Kubernetes-based orchestration on AWS, Amazon Elastic Kubernetes Service (EKS) provides managed Kubernetes with GPU scheduling via the NVIDIA device plugin. Amazon SageMaker HyperPod also supports EKS mode, combining Kubernetes orchestration with HyperPod's training-specific capabilities. HyperPod EKS extends EKS with features designed for foundation model training at scale. Task governance provides compute allocation and policy enforcement across teams, integrating managed Kueue for admission control and Karpenter for just-in-time node provisioning. Checkpointless training addresses the recovery latency inherent in traditional checkpoint-based fault tolerance. Rather than periodically serializing model state to shared storage, checkpointless training maintains continuous peer-to-peer state replication across GPUs. When a failure occurs, surviving nodes reconstruct the lost state through EFA-based communication rather than reading multi-terabyte checkpoints from FSx for Lustre or S3. Elastic training enables jobs to automatically scale based on resource availability. When additional accelerators become available (e.g., from completed jobs or newly provisioned capacity), elastic jobs can expand to utilize them; when higher-priority workloads require resources, jobs can contract while maintaining training progress.
ML Software Stack
Distributed training and inference involve multiple software layers that must be correctly configured and coordinated. A useful model treats the runtime stack as five layers, ordered from hardware-adjacent components (which must function correctly for anything to run) to framework-level abstractions (which determine programmer productivity and model throughput): hardware enablement, accelerator runtime and math libraries, communication substrate, ML frameworks, and distributed training/inference frameworks.
*Figure 3: The ML software stack for distributed training and inference on EC2 instances*
Hardware enablement: kernel drivers
At the foundation, Linux kernel drivers provide direct hardware access. The NVIDIA GPU driver exposes compute capabilities and supports GPUDirect RDMA for direct data transfers between GPUs and network adapters. The GDRCopy driver (gdrdrv) enables low-latency CPU-initiated copies to and from GPU memory, used by NCCL for small-message transfers. The EFA driver provides OS-bypass networking through the libfabric API, and the Lustre client driver enables POSIX access to FSx for Lustre parallel file systems.
Accelerator runtime, compilers, and kernel libraries
The CUDA platform provides the programming model and runtime for GPU compute. Applications compiled against CUDA can launch kernels on NVIDIA GPUs, manage device memory, and coordinate execution across multiple devices. The current release is CUDA Toolkit 13.x, with support for Blackwell architecture (compute capability 10.x).
Modern training and inference performance is increasingly driven by specialized optimization libraries and custom kernels, not just general-purpose vendor primitives. Kernels like FlashAttention fuse attention into a single memory-efficient pass, cutting HBM traffic and improving throughput. Many teams also write shape- and precision-specialized fused kernels (e.g., layernorm/residual/activation, quantized GEMMs, MoE dispatch, KV-cache ops) tuned to their exact models. This is enabled by programmable toolchains such as Triton (Python GPU kernel compiler) and NVIDIA's CuTe (tensor layout and warp-level DSL), with libraries like CUTLASS providing highly optimized GEMM and fusion building blocks. In practice, this kernel and compiler layer often determines end-to-end performance as much as the ML framework.
Communication substrate: NCCL and transport plugins
Multi-GPU training depends on efficient collective communication. NVIDIA Collective Communications Library (NCCL) implements collective operations—all-reduce, all-gather, reduce-scatter, all-to-all, broadcast, and point-to-point send/receive—with topology-aware algorithms that exploit NVLink for intra-node communication and network transports for inter-node traffic. NCCL dynamically detects the communication topology and selects ring or tree algorithms depending on message size and available bandwidth. While data-parallel and tensor-parallel strategies rely primarily on all-reduce and all-gather, Mixture-of-Experts (MoE) models with expert parallelism depend on all-to-all collectives to route tokens between GPUs: a dispatch all-to-all sends each token to the GPU hosting its assigned expert, and a combine all-to-all returns expert outputs to the originating GPUs (NVIDIA Developer Blog). Because every GPU exchanges data with every other GPU in the expert-parallel group, all-to-all communication volume scales with the number of experts and can become a dominant bottleneck at high expert-parallelism degrees.
On AWS, NCCL's inter-node communication is enabled through the aws-ofi-nccl plugin, which maps NCCL's transport APIs to libfabric interfaces. This allows NCCL to leverage EFA's OS-bypass and Scalable Reliable Datagram (SRD) protocol without application changes.
For inference workloads, collective operations do not capture all communication patterns. Disaggregated inference architectures—which separate prefill and decode phases onto distinct GPU pools—require efficient point-to-point data movement, particularly for transferring KV cache state between instances. NVIDIA Inference Xfer Library (NIXL) addresses this requirement by providing a unified API for point-to-point transfers across memory tiers (HBM, DRAM, NVMe, distributed storage) and interconnects (NVLink, InfiniBand, Ethernet). NIXL integrates with inference frameworks such as NVIDIA Dynamo and supports backends including UCX and GPUDirect Storage.
ML frameworks: PyTorch
The two dominant frameworks for foundation model development are PyTorch and JAX. JAX takes an SPMD (Single Program Multiple Data) approach through XLA, where the same program executes across devices with automatic data distribution and collective lowering. This series focuses on PyTorch, which sees broader adoption in the open-source ecosystem and forms the basis for the distributed training and inference frameworks discussed below.
PyTorch provides tensor computation with GPU acceleration, automatic differentiation, and a flexible eager-execution model. For distributed workloads, PyTorch's torch.distributed module provides the core primitives: process groups for collective communication, and distributed data-parallel abstractions including Distributed Data Parallel (DDP) and Fully Sharded Data Parallel (FSDP2). DDP replicates models across GPUs and synchronizes gradients via all-reduce, while FSDP2 shards parameters, gradients, and optimizer states across workers using techniques from the ZeRO algorithm, enabling training of models that exceed single-GPU memory capacity.
Distributed training and inference frameworks
The top layer comprises frameworks that build on PyTorch to provide higher-level abstractions for distributed training and inference at scale. For training, three categories of frameworks address different points in the complexity-performance tradeoff. Below are few examples
Hugging Face Transformers provides the Trainer class with built-in support for distributed training via Accelerate, which abstracts over DDP, FSDP, and DeepSpeed. This path prioritizes ease of use and broad model compatibility, making it suitable for fine-tuning and moderate-scale training where configuration simplicity matters more than maximum throughput.
NVIDIA Megatron Core targets maximum efficiency at scale, implementing 3D parallelism (tensor, pipeline, and expert parallelism) with optimizations including FP8 mixed precision via Transformer Engine. The NeMo Framework builds on Megatron Core to provide end-to-end workflows for pre-training and fine-tuning.
For reinforcement learning from human feedback (RLHF) and related post-training methods, veRL (Volcano Engine Reinforcement Learning) provides a flexible framework that implements algorithms including PPO, GRPO, and REINFORCE++. veRL's HybridFlow architecture allows mixing training backends (FSDP2, Megatron) with inference engines (vLLM, SGLang) in the same job, avoiding weight synchronization overhead by sharing model weights in memory between actor and rollout components.
For inference serving, vLLM implements PagedAttention, managing the KV cache as paged virtual memory to reduce fragmentation and enable higher batch sizes. SGLang extends this with RadixAttention for automatic prefix reuse across requests, a zero-overhead batch scheduler that overlaps CPU scheduling with GPU computation, and a cache-aware load balancer that routes requests based on predicted cache hit rates. Both frameworks support tensor parallelism for serving models that exceed single-GPU memory, and both integrate with NVIDIA Dynamo for disaggregated serving architectures that separate prefill and decode phases.
Observability
Observability is a prerequisite for debugging and operating distributed training systems at scale. When a training job stalls or throughput degrades, practitioners need visibility into whether the cause is hardware failure, network congestion, storage bottlenecks, or application-level inefficiency. At the infrastructure scale discussed in this series—thousands of GPUs, petabits of interconnect bandwidth, and terabytes of checkpoint data—the challenge shifts from simple monitoring to systematic telemetry collection, storage, and analysis. Observability spans three telemetry categories: infrastructure metrics (GPU, network, storage), workload metrics (training throughput, queue latency), and alerting for proactive fault detection.
Core Stack: Prometheus and Grafana
The de facto standard for observability in Kubernetes and HPC environments combines Prometheus for metrics collection with Grafana for visualization and alerting. Prometheus operates on a pull-based model, periodically scraping HTTP endpoints exposed by metric exporters. Collected metrics are stored in a time-series database (TSDB) and queried via PromQL, a flexible query language for aggregation, filtering, and alerting rule evaluation. Grafana consumes Prometheus as a data source, rendering dashboards and triggering alerts based on PromQL expressions.
For production deployments, Amazon Managed Service for Prometheus (AMP) provides a fully managed, Prometheus-compatible time-series database that scales to ingest millions of samples per second without requiring operators to manage storage, replication, or high availability. Amazon Managed Grafana (AMG) offers a managed Grafana workspace with native integration to AMP and AWS authentication via IAM Identity Center. Together, these services eliminate operational overhead while preserving compatibility with existing Prometheus exporters and Grafana dashboards.
GPU, Network, and Application Telemetry
DCGM-Exporter exposes NVIDIA GPU metrics in Prometheus format, including utilization, memory usage, power, temperature, and hardware health indicators such as ECC errors and XID events. For training workloads, SM activity (DCGM_FI_PROF_SM_ACTIVE) often provides a more accurate measure of compute efficiency than basic utilization metrics.
EFA exposes driver-level statistics (bytes, packets, retransmits, timeouts) that help diagnose collective operation bottlenecks in distributed training. The aws-ofi-nccl plugin bridges NCCL to the libfabric interface, and operators can combine EFA counters with NCCL diagnostics (NCCL_DEBUG=INFO) to isolate network-layer issues.
Amazon FSx for Lustre exposes client-side metrics including throughput and metadata latency, while application-level metrics (step time, tokens per second, loss values for training; TTFT, inter-token latency for inference) can be exported via Prometheus client libraries.
GPU Health Monitoring and Alerting
Proactive fault detection prevents hardware issues from propagating into extended training interruptions. A typical workflow monitors DCGM health metrics and triggers alerts when error counts exceed thresholds. ECC single-bit errors (SBE) may be tolerable in small numbers, but accelerating SBE rates often precede double-bit errors (DBE) or other failures. XID 63 (row remap failure), XID 64 (GPU fallen off bus), and XID 94/95 (contained/uncontained errors) typically warrant immediate node replacement.
The GPU Health - Cluster dashboard (Grafana dashboard ID 21645) provides a reference visualization for common GPU error patterns. The dashboard aggregates ECC errors, XID events, thermal violations, and row remapping status across all cluster nodes, enabling operators to identify failing hardware before it impacts training jobs.
*Figure 4: GPU Health - Cluster dashboard showing GPU error patterns and instance reporting*
Part 5 of this series provides a comprehensive treatment of observability architecture, including metric collection strategies, dashboard configurations, and alerting patterns for maintaining cluster health at scale.
Conclusion
The shift from a single pre-training scaling law to three complementary regimes—pre-training, post-training, and test-time compute—has not fragmented infrastructure requirements; it has reinforced them. All three regimes demand tightly coupled accelerator compute, high-bandwidth low-latency networking, and scalable distributed storage, differing mainly in workload profile and resource scheduling patterns.
This post surfaced the four-layer architecture that addresses those requirements on AWS: infrastructure building blocks (EC2 P-instances, EFA networking, and tiered storage), resource orchestration (Slurm and Kubernetes with SageMaker HyperPod), the ML software stack (from kernel drivers and CUDA through NCCL to PyTorch), and observability (Prometheus, Grafana, and GPU health monitoring). Each layer constrains and enables the layers above it—a misconfigured driver or saturated network link can bottleneck an otherwise well-tuned training run just as effectively as a suboptimal parallelism strategy.
Understanding these integration points is the foundation for diagnosing performance bottlenecks and making informed scaling decisions across the foundation model lifecycle.
Authors
Aman Shanbhag is an AI Performance and Infrastructure Engineer on the MARS MLOps team at NVIDIA, where he helps research teams build scalable, high-performance ML training and inference systems. He previously worked as a Specialist Solutions Architect at AWS, supporting customers worldwide with ML training and inference optimization on AWS. Aman holds degrees in computer science, mathematics, and entrepreneurship from Rice University and focuses on AI infrastructure, performance optimization,
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み