Abstract
本稿では、MegaTrainというメモリ中心のシステムを提案する。単一GPU上で、フル精度により100B+パラメータ級の大規模言語モデルを効率的に学習する。従来のGPU中心のシステムとは異なり、MegaTrainはパラメータとオプティマイザ状態をホストメモリ（CPUメモリ）に保持し、GPUを一時的な計算エンジンとして扱う。各層について、パラメータをストリーミングで読み込み、勾配を計算して外部へ出力することで、デバイス上に保持され続ける状態を最小化する。CPU-GPU帯域幅ボトルネックに対処するため、我々は2つの主要な最適化を採用する。1) 複数のCUDAストリームにまたがって、パラメータのプリフェッチ、計算、勾配オフロードをオーバーラップさせるパイプライン化された二重バッファ付き実行エンジンを導入し、GPUの継続的な実行を可能にする。2) 恒常的（persistent）な自動微分グラフを、ステートレスな層テンプレートに置き換える。ストリーミングで読み込まれる際に重みを動的にバインドすることで、恒常的なグラフのメタデータを排除しつつ、スケジューリングの柔軟性を確保する。ホストメモリ1.5TBを搭載した単一H200 GPUにおいて、MegaTrainは最大120Bパラメータまでのモデルを確実に学習できる。また、14Bモデルを学習する際、CPUオフロードを伴うDeepSpeed ZeRO-3に対して学習スループットを1.84\times向上させる。さらにMegaTrainは、単一のGH200上で512kトークンのコンテキストを持つ7Bモデルの学習も可能にする。
