AWS上での基盤モデルのトレーニングと推論のための構成要素

インフラストラクチャ：コンピュート、ネットワーク、ストレージ

加速コンピュート

  • P5 Instances: NVIDIA H100 GPU (p5.48xlarge) と H200 GPU (p5e.48xlarge/p5en.48xlarge) を搭載。
  • P6 Instances: NVIDIA Blackwell B200 (p6-b200.48xlarge) と B300 (p6-b300.48xlarge) アーキテクチャを導入。

ネットワーキングと相互接続

  • Internal Scale-up (NVLink/NVSwitch): 単一ノード内の高帯域・低遅延接続。
  • External Scale-out (EFA): Elastic Fabric Adapter (EFA) は、Scalable Reliable Datagram (SRD) プロトコルを使用した OS バイパス RDMA を提供し、ノード間通信を実現。EFAv3 は EFAv2 に比べパケット遅延を約 35% 削減し、EFAv4 は EFAv3 に対して集合通信性能をさらに 18% 向上させます。
  • UltraServers: Amazon EC2 UltraServers（例：P6e-GB200）は NVLink ドメインを単一インスタンスを超えて拡張し、最大 72 台の Blackwell GPU を同一 NVLink ドメイン内に配置でき、NVLink ファブリックから外部への通信頻度を低減します。

階層型ストレージ

  1. Local NVMe SSD: ホットデータ用のエフェメラルインスタンスストア（最大 30.72 TB の生容量）。
  2. Amazon FSx for Lustre: 高集約スループットとサブミリ秒レイテンシを提供するマネージド並列ファイルシステム。
  3. Amazon S3: 永続的な保存と Data Repository Associations を介したデータセットの遅延ロードに使用。

リソースオーケストレーション：Slurm と Kubernetes

Slurm（HPC 中心）

Slurm は HPC ワークロードの主要なマネージャで、ジョブレベルでスケジューリングし、タスク開始前に必要なすべてのノードが割り当てられることを保証します。AWS は以下を通じて Slurm をサポートします。

  • AWS ParallelCluster: Slurm クラスタのデプロイを自動化するオープンソースツール。
  • AWS Parallel Computing Service (PCS): Slurm 用のマネージドコントロールプレーン。
  • Amazon SageMaker HyperPod: Slurm モードに継続的なノードヘルスモニタリングとジョブ自動再開機能を追加。

Kubernetes（クラウドネイティブ）

Kubernetes はデプロイに優れていますが、ジョブレベルの原子性やトポロジー認識がネイティブに欠けています。これらのギャップは以下で埋められます。

  • Kueue: ジョブレベルのギャングアドミッションとマルチテナントクォータを管理。
  • Volcano and NVIDIA KAI Scheduler: NVLink と EFA の使用を最適化するトポロジー認識ポッド配置を提供。
  • SageMaker HyperPod (EKS mode): マネージド Kueue と Karpenter を統合し、ジャストインタイムプロビジョニングを実現。また、checkpointless training を導入し、EFA を介したピアツーピア状態レプリケーションにより、ストレージからマルチテラバイトのチェックポイントを読み込まずに障害から復旧します。

ML ソフトウェアスタック

低レベル有効化とランタイム

  • Kernel Drivers: NVIDIA GPU ドライバは GPUDirect RDMA をサポートし、EFA ドライバは libfabric を介した OS バイパスネットワーキングを提供。
  • CUDA and Kernels: CUDA Toolkit 13.x は Blackwell アーキテクチャをサポート。FlashAttention のような融合カーネルや Triton、NVIDIA CuTe といったプログラマブルツールチェーンが性能をさらに向上させます。

通信基盤

  • NCCL: NVIDIA Collective Communications Library は all-reduce と all-gather のトポロジー認識アルゴリズムを実装。Mixture-of-Experts (MoE) モデルでは、トークンをエキスパート間でルーティングするために all-to-all 集合が重要です。
  • aws-ofi-nccl: NCCL のトランスポート API を libfabric にマッピングし、NCCL が EFA の SRD プロトコルを使用できるようにするプラグイン。
  • NIXL: NVIDIA Inference Xfer Library は、プレフィルとデコードフェーズを分離した分散推論アーキテクチャ向けにポイントツーポイント転送を提供。

フレームワーク

  • PyTorch: 分散ワークロードの主要フレームワークで、torch.distributed と FSDP2 を利用してパラメータとオプティマイザ状態をシャーディング。
  • Distributed Frameworks:
    • Hugging Face Transformers/Accelerate: 使いやすさと互換性を重視。
    • NVIDIA Megatron Core/NeMo: 3D パラレル（テンソル、パイプライン、エキスパート）で最大効率を追求。
    • veRL: ハイブリッドフローを使用してトレーニングバックエンド（FSDP2、Megatron）と推論エンジン（vLLM、SGLang）を単一ジョブで混在させる強化学習フレームワーク。
    • vLLM and SGLang: PagedAttention と RadixAttention を活用し、KV キャッシュ管理とリクエストスケジューリングを最適化する推論エンジン。

可観測性と障害検出

テレメトリースタック

  • Prometheus and Grafana: メトリクス収集と可視化の標準。AWS は Amazon Managed Service for Prometheus (AMP) と Amazon Managed Grafana (AMG) を提供し、運用負荷を削減。
  • DCGM-Exporter: NVIDIA GPU メトリクスを公開。SM アクティビティ（DCGM_FI_PROF_SM_ACTIVE）は、基本的な利用率よりも計算効率の正確な指標とされています。

ヘルスモニタリング

ハードウェア障害の事前検出はトレーニング中断を防ぎます。重要な指標は次のとおりです。

  • ECC Errors: 単一ビットエラー（SBE）率の上昇は、しばしば二重ビットエラー（DBE）に先行します。
  • XID Events: XID 63（行リマップ失敗）、XID 64（GPU がバスから外れた）、XID 94/95（制御/非制御エラー）などの特定エラーは即時のノード交換を引き起こします。

Sources