再現可能なPyTorch実行を計画する

スケーリング前に1つの完全なステップを実証する

モデルとデータセットのリビジョン、依存関係のバージョン、トレーニング構成、乱数シード、コンテナダイジェストを固定してください。1つのアクセラレータでモデルの読み込み、完全なトレーニングステップ、評価、チェックポイント保存、チェックポイント復元を実行してください。

長期実行で意図しているものと同じシーケンス長とマイクロバッチ設定を使用してください。オプティマイザ状態を省略したり、短いシーケンスを使用したりする小さなパイロットでは、測定しようとしているメモリ問題を隠してしまう可能性があります。

代表的な最適化ステップを測定する

既存のトレーニングプログラム内で、代表的なステップをピークメモリ測定で囲みます。この例では、モデル、バッチ、オプティマイザがすでにGPU上で初期化されていることを前提としています。一部の状態は遅延作成されるため、オプティマイザステップを含めてください。

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

これはピーク時のテンソル割り当てを報告するものであり、デバイス全体の使用量ではありません。予約済みメモリ、ランタイムのオーバーヘッド、システムのデバイスメモリレポートも確認してください。

正しい理由でDDPまたはシャーディングを選択する

DistributedDataParallel は各プロセスにモデルのレプリカを保持し、勾配を同期します。これは、完全な学習状態が各 GPU に収まり、より多くのデータを並列処理したい場合に有用です。VRAM を1つの連続したプールに統合するものではありません。

Fully Sharded Data Parallelは、パラメーター、勾配、オプティマイザー状態を分散できます。メモリ圧力がシャーディングを正当化する場合に使用し、追加の通信と分散チェックポイントの複雑さを考慮してください。

混合精度とアクティベーションチェックポイントは別々の選択です。低精度の数値的挙動を検証し、チェックポイントがメモリを節約する場合のアクティベーション再計算の計算コストを測定してください。

単一ノードのパイロットを開始する

トレーニングスクリプトは分散実行を初期化し、ローカルランク情報を使用して各GPUに1つのプロセスを割り当てる必要があります。プロセス数を可視GPU数に一致させてください。以下は分散対応スクリプトの起動例であり、完全なトレーニング実装ではありません。

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

これはプロビジョニング済みノードでのみ実行してください。ローカルのAnchorGPUデモでは学習を実行したり、リモートGPUを割り当てたりしません。

パイロット結果に基づいてハードウェアを選択

A100を低価格の80 GB CUDAベースラインとして使用し、ワークロードに有用なHopper固有のパスがある場合はH100をテストしてください。H200はより大きなCUDAメモリ構成を提供します。MI300Xは、完全なアプリケーションがROCmで認定されている場合、このカタログでアクセラレータあたりの容量がより大きくなります。

AMDについては、カスタム拡張機能、カーネル、パッケージバージョン、ROCmでテスト済みのPyTorchイメージを一式として確認してください。フレームワークレベルの互換性は、すべてのオプション操作を検証するものではありません。

保存だけでなく復元もテストする

モデル、オプティマイザ、スケジューラ、トレーニングステップ、構成、および該当する場合は混合精度スケーラの状態を記録します。チェックポイントを保存し、プロセスを終了してから、長期実行にコミットする前に復元して再開します。

分散状態に適したチェックポイント方式を使用してください。保存されたアーティファクトが元のプロセスから独立して読み取れること、評価出力が期待される許容範囲内に収まることをテストしてください。

パイプライン全体の予算を立てる

ダウンロード、前処理、データロード、検証、チェックポイント、エクスポート、ノードからの出力のコピーを含めてください。メモリ、データロードの停止、通信時間、チェックポイント所要時間、検証品質と合わせてスループットを測定してください。

7日間の期間は、限定的な互換性およびプロファイリングの実施に有用です。30日間の予約は、有用な作業量と運用上のバッファを見積もれる場合にのみ選択してください。GPUの追加は、追加の計算が通信および入力パイプラインの制限を上回る場合にのみ有効です。