2026-06-26 01:41 UTC+9サイト内リライト8 分で読了更新: 2026-06-26 02:09 UTC+9

NVIDIA Blackwell で Amazon SageMaker AI のモデルトレーニングを最適化

この記事では、Amazon SageMaker AI でトレーニングジョブを構成して、AWS 上の Blackwell アーキテクチャを最大限に活用する方法を説明します。Blackwell の拡張メモリを活用するバッチサイズとシーケンス長の選択、モデルサイズ（1B～64B パラメータ）に適した精度形式の選択、アクティベーションチェックポイントの戦略的適用について学びます。最後に、トレーニング構成を調整し、P6-B200 インスタンスで分散トレーニングジョブを起動するための実用的なフレームワークを提供します。

ソースAWS Machine Learning Blog著者: Andrea Gallo

記事インテリジェンス

エンジニア上級

要点

Blackwell の拡張メモリにより、より大きなバッチサイズ、より長いシーケンス長、簡素化されたモデルシャーディングが可能になります。
アクティベーションチェックポイントは、大規模モデル（〜14B+ パラメータ）の安定したトレーニングの前提条件です。
大規模モデルには FP8 がスループットとメモリ効率のバランスに優れ、MXFP8 は数値安定性に優れています。
P6-B200 インスタンスは、Flexible Training Plan を通じて SageMaker AI で利用できます。

重要な理由

このニュースが重要なのは、Blackwell の拡張メモリにより、より大きなバッチサイズ、より長いシーケンス長、簡素化されたモデルシャーディングが可能になりますためです。

技術的影響

モデル選定、推論コスト、プロダクト能力、評価基準に影響する可能性があります。

このパネルは AI が生成し、正確性を確認済みです。

NVIDIA Blackwell GPU を使用した Amazon SageMaker AI でのモデルトレーニングの最適化により、大規模 AI モデルの実用性が変わります。現在大規模モデルをトレーニングしている場合、GPU メモリに制限されたバッチサイズ、メモリ不足エラーを回避するための短いシーケンス長、スケールに伴う通信オーバーヘッドを増加させるモデルシャーディングなど、おなじみの制約に対処している可能性があります。Blackwell の拡張メモリと新しい精度形式は、これらの制約を直接軽減します。8 基の Blackwell GPU を搭載した P6-B200 インスタンスは、Amazon SageMaker AI トレーニングジョブで利用可能であり、Flexible Training Plan を使用してキャパシティを予約し、予測可能なアクセス、コスト管理、自動化されたリソース管理を実現できます。Amazon SageMaker AI トレーニングジョブは、基盤となるコンピューティングインフラストラクチャとリソースを自動的にプロビジョニングおよび管理することで、大規模な ML モデルトレーニングを可能にし、インフラストラクチャ運用ではなくデータとアルゴリズムに集中できます。

この記事では、Amazon SageMaker AI でトレーニングジョブを構成して Blackwell アーキテクチャを AWS 上で最大限に活用する方法を説明します。Blackwell の拡張メモリを活用するバッチサイズとシーケンス長の選択、モデルサイズ（1B～64B パラメータ）に適した精度形式の選択、アクティベーションチェックポイントの戦略的適用について学びます。最後に、トレーニング構成を調整し、P6-B200 インスタンスで分散トレーニングジョブを起動するための実用的なフレームワークを提供します。

適切に構成された Blackwell トレーニングジョブは、積極的なシャーディングなしでより大きなバッチサイズを処理できるため、通信オーバーヘッドが削減され、スループットが向上します。より長いシーケンス長が、長距離依存タスクで実現可能になります。適切な精度形式を使用すると、以前はマルチノードセットアップが必要だったモデルを 1 台の 8-GPU ノードで実行でき、迅速な反復サイクル、ネットワークオーバーヘッドの削減、インフラストラクチャコストの削減が可能になります。

NVIDIA Blackwell を理解する

トレーニングジョブを構成する前に、Blackwell が以前の GPU 世代と何が異なるかを理解することが役立ちます。Blackwell のデュアルチップアーキテクチャと第 5 世代 Tensor Core は、マルチ GPU トレーニングにすぐに使用できる顕著なメリットをもたらします。NVLink 5 相互接続は最大 1.8 TB/s の双方向 GPU 間帯域幅を提供し、B200 のより大きな HBM 容量と高いメモリ帯域幅は、大規模バッチ、長いシーケンス、分散トレーニングワークロードのメモリプレッシャーを軽減するのに役立ちます。

この記事の例では、1B から 64B パラメータのトランスフォーマーモデルを使用したシングルノード 8-GPU トレーニングを使用しています。トレーニング構成には、モデルパラメータ、勾配、オプティマイザ状態を複数の GPU にシャーディングして、単一 GPU メモリよりも大きなモデルをトレーニングする分散トレーニング手法である PyTorch Fully Sharded Data Parallel (FSDP) を使用しています。結果には、さまざまなバッチサイズ、シーケンス長、精度形式の複数の構成が含まれており、異なるアプローチがいつ最適な結果をもたらすかを示しています。

メモリ管理

Blackwell の拡張メモリ（B200 で 180 GB、B300 で 268 GB）により、次の 3 つの分野で最適化の余地が生まれます。

より大きなバッチサイズ：GPU 間の勾配同期ステップ数を減らし、全体的なスループットを向上させます。
簡素化されたモデルシャーディング：GPU あたりのメモリが増えるため、モデル並列度を下げたり、一部のモデルでは完全に排除したりできる可能性があります。シャーディングが少ないほど、GPU 間通信オーバーヘッドが減少します。
より長いシーケンス長：モデルが 1 回の処理でより多くのコンテキストを処理できるようになり、長距離依存タスクに重要です。

スループットが主な目標の場合は、バッチサイズの調整から始めてください。通信オーバーヘッドがボトルネックの場合は、最初にシャーディングを簡素化してください。タスクに長距離コンテキストが必要な場合は、シーケンス長を優先してください。バッチサイズとシーケンス長は両方ともメモリ消費を増加させるため、効果的なバランスを見つけることが重要です。

アクティベーションチェックポイントは、メモリ使用量と計算のバランスをとるのに役立ちます。これは、中間活性化を保存する代わりにバックワードパス中に再計算することで、計算時間の増加（モデルアーキテクチャに応じて通常 10～30% のオーバーヘッド）と引き換えに GPU メモリ使用量を削減します。解放されたメモリは、より大きなバッチサイズやより長いシーケンスに再投資できます。計算オーバーヘッドはワークロードによって異なるため、特定の構成をベンチマークしてトレードオフを理解してからチェックポイントを採用することをお勧めします。

たとえば、図 1 では、MXFP8 精度、8K シーケンス長の 1B パラメータ LLM について、3 つのトレーニング構成を比較しました。アクティベーションチェックポイントなし（BS=1）の場合、スループットは約 6K トークン/秒ですが、ピークメモリは 15.5 GB と高くなっています。同じバッチサイズでアクティベーションチェックポイントを有効にすると、メモリは 2.3 GB に劇的に低下します（中間活性化が再計算されるため）が、再計算オーバーヘッドのためにスループットもわずかに低下します。重要な見返りは 3 番目のバーにあります。アクティベーションチェックポイントを有効にしてバッチサイズを 16 に引き上げると、解放されたメモリによりはるかに大きなバッチが可能になり、スループットは約 51K トークン/秒（ベースラインの約 8 倍）に押し上げられ、ピークメモリは 22.8 GB に上昇しますが、それでも GPU 制限内です。

アクティベーションチェックポイントがワークロードに適しているかどうかを判断するには、モデルサイズとメモリ使用量を考慮してください。

小規模モデル（最大約 14B パラメータ）：通常、アクティベーションチェックポイントは必要ありません。Blackwell の拡張メモリにより、ほとんどの小規模モデルはチェックポイントなしで快適に収まります。この範囲の上限で実行していてメモリプレッシャーが発生している場合、アクティベーションチェックポイントは計算オーバーヘッドと引き換えに意味のあるメモリ節約をもたらし、それをより大きなバッチサイズに再投資できます。
大規模モデル（約 14B+ パラメータ）：このモデルサイズでは、メモリ消費はバッチサイズとシーケンス長に応じて 87～171 GB の範囲になります。アクティベーションチェックポイントなしでは、ほとんどの構成が CUDA メモリ不足 (OOM) エラーで失敗します。チェックポイントを追加すると、解放されたメモリによりバッチサイズを十分に増やせるため、追加の計算オーバーヘッドにもかかわらずスループットが向上します。大規模モデルでは、チェックポイントはオプションではなく、安定したトレーニングの前提条件です。

精度形式

Blackwell の第 5 世代 Tensor Core は、低精度形式（FP8、MXFP8、NVFP4）のハードウェアアクセラレーションを提供し、主にスループットの最適化であり、メモリ節約技術ではありません。低精度を使用するとメモリ帯域幅の要件が減り、GPU がサイクルごとに実行できる演算数が増加します。ただし、低精度トレーニングはデフォルトではほぼメモリニュートラルであり、Transformer Engine は高精度のプライマリウェイト（オプティマイザ更新用）と量子化コピーの両方を維持するため、低精度形式が直接メモリ使用量の削減につながるわけではありません。量子化自体にはオーバーヘッド（精度形式間の変換とメモリ内の複数のウェイトコピーの維持）が伴うため、正味のメリットはモデルサイズと、トレーニングが計算バウンドかメモリバウンドかによって異なります。NVFP4 は最高のスループットを提供しますが、そのパフォーマンス上の利点は主に大規模モデルと推論ワークロードに拡大します。このワークロードではプライマリウェイトは必要ありません。

計算バウンドのワークロード（通常は小規模モデル）では、計算速度が制限要因であり、量子化オーバーヘッドが低精度によるスループット向上を部分的に相殺します。メモリバウンドのワークロード（通常は大規模モデル）では、データ移動がボトルネックであり、低精度形式によるメモリフットプリントの削減が制約に直接対処し、より大きなメリットをもたらします。

小規模モデル（最大約 14B パラメータ）：このモデルサイズでは、低精度形式（FP8、MXFP8、NVFP4）はすべて、FP16 に対して同様の控えめなスループット向上をもたらします。これは、量子化オーバーヘッドが速度の利点を食いつぶすためです。バッチサイズの調整は、精度形式の選択よりも有意義なゲインをもたらす傾向があります。まずは FP8 から始めてスループットを向上させてください。これは MXFP8 や NVFP4 よりもオーバーヘッドが低く、ほとんどの小規模モデルワークロードの適切なデフォルトとなります。デフォルトの TransformerEngine 設定では、低精度形式は FP16 よりも多くのメモリを使用することに注意してください。これは TransformerEngine がウェイトを高精度に保ち、オンザフライでキャストするためです。メモリが制約であり、オプティマイザがサポートしている場合は、quantized_model_init を使用してウェイトを直接 FP8 で保存し、メモリを FP16 レベル以下に削減してください。
大規模モデル（約 14B+ パラメータ）：ここが低精度が最大の効果を発揮する場所です。FP8 は通常、スループットとメモリ効率の強力なバランスを提供します。MXFP8 は理論的にはメモリ効率が高いですが、その転置オーバーヘッドが実際にはその利点を部分的に相殺します。ただし、収束安定性や数値精度がワークロードの優先事項である場合、MXFP8 の方が適している可能性があります。そのより細かい量子化スキームは FP8 よりも確実にモデル精度を維持する傾向があるためです。メモリが主なボトルネックである大規模モデルの場合、NVFP4 は追加のスループット向上をもたらす可能性があります。その行列乗算速度の利点はモデルサイズに応じて拡大するためです。これらの利点を実現するには、意味のあるエンジニアリング投資が必要です。Megatron Core のフレームワークレベルのレシピを使用してください。これらは検証済みの NVFP4 構成を提供し、ゼロから実装するよりも優れています。

NVIDIA の TransformerEngine は、自動混合精度切り替え、融合カーネル、動的損失スケーリングなど、実装の複雑さを処理します。本番環境に移行する前に、フォーマット間で損失曲線を追跡して収束を検証し、選択した精度が精度要件を満たしていることを確認してください。

すべてのワークロードが積極的な最適化から恩恵を受けるわけではありません。モデルがメモリ制限内で快適にトレーニングされ、FP16 でスループット要件を満たしている場合、低精度形式の追加の複雑さはエンジニアリングの労力に見合わない可能性があります。ベースライン測定から始めて、測定可能なボトルネックのみを最適化してください。

Amazon SageMaker AI で Blackwell トレーニングを始める

前のセクションでは、利用可能なメモリ量、アクティベーションチェックポイントがモデルサイズに適しているかどうか、どの精度形式がワークロードに適しているかなど、主要な決定事項を説明しました。次のセクションでは、Amazon SageMaker AI トレーニングジョブを使用して、これらの決定を実践に移します。

Amazon SageMaker AI は、Blackwell インスタンスでの分散トレーニングのための完全マネージド環境を提供し、インスタンスのプロビジョニング、コンテナオーケストレーション、および Amazon Simple Storage Service (Amazon S3)、Amazon CloudWatch、Amazon Elastic Container Registry (Amazon ECR)、AWS Identity and Access Management (AWS IAM) などの AWS サービスとの統合を処理します。

前提条件

始める前に、以下があることを確認してください。

Amazon SageMaker AI トレーニングジョブを作成し、Amazon ECR にアクセスし、Amazon SageMaker AI 実行用の IAM ロールを作成する権限を持つ AWS アカウント。
Flexible Training Plan またはマネージドスポットトレーニングを通じて ml.p6-b200.48xlarge インスタンスにアクセスできること。開始する前にサービス割り当てを確認してください。
ローカルにインストールされた Docker、Python 3.9 以降、および SageMaker Python SDK。
PyTorch と FSDP に関する知識。FSDP が初めての場合は、分散トレーニングの開始方法を参照してください。

トレーニングジョブを起動する

トレーニングジョブを起動するには、次の手順を実行します。

ステップ 1: スクリプトを作成する

NVIDIA TransformerEngine リポジトリの FSDP 例から fsdp.py ファイルをダウンロードします。このスクリプトは FSDP トレーニングを実装し、コマンドライン引数としてハイパーパラメータを受け入れます。

ステップ 2: エントリポイントスクリプトを作成する

torchrun を設定し、トレーニングスクリプトを起動するための train.sh ファイルを準備します。

#!/bin/bash

# SageMaker passes hyperparameters as environment variables (SM_HP_)
PRECISION=${SM_HP_PRECISION:-"mxfp8"}
NUM_LAYERS=${SM_HP_NUM_LAYERS:-10}
BATCH_SIZE=${SM_HP_BATCH_SIZE:-8}
SEQ_LENGTH=${SM_HP_SEQ_LENGTH:-2048}

NUM_GPUS=8

# Launch training with torchrun
torchrun --nproc_per_node=$NUM_GPUS fsdp.py \
  --precision $PRECISION \
  --num-layers $NUM_LAYERS \
  --batch-size $BATCH_SIZE \
  --seq-length $SEQ_LENGTH

ステップ 3: Docker イメージをビルドして Amazon ECR にプッシュする

PyTorch イメージに基づき TransformerEngine を含む Dockerfile を作成し、イメージを Amazon ECR リポジトリにプッシュします。

ステップ 4: SageMaker Python SDK を使用してトレーニングジョブを起動する

SageMaker Python SDK の Estimator クラスを使用して、イメージ、インスタンスタイプ、ハイパーパラメータを指定します。

結論

上記の戦略（より大きなバッチサイズ、より長いシーケンス長、適切な精度形式、アクティベーションチェックポイントの活用）を採用することで、トレーニングスループットを大幅に向上させ、コストを削減できます。Blackwell アーキテクチャは SageMaker AI 上で強力な機能を提供し、注意深い構成により、AI モデルトレーニングの新たなパフォーマンス基準を打ち立てることができます。この記事で提供するフレームワークは、P6-B200 インスタンスで最適化を開始するための実践的なステップを提供します。