MegaTrain: 単一GPUでの100B+パラメータ大規模言語モデルの完全精度学習

MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU

April 6, 2026

著者: Zhengqing Yuan, Hanchi Sun, Lichao Sun, Yanfang Ye

cs.AI

要旨

我々はMegaTrainを提案する。これは、単一GPU上で100B+パラメータ大規模言語モデルの完全精度訓練を効率的に実現する、メモリ中心のシステムである。従来のGPU中心システムとは異なり、MegaTrainはパラメータとオプティマイザ状態をホストメモリ（CPUメモリ）に格納し、GPUを一時的な計算エンジンとして扱う。各レイヤーごとに、パラメータをストリーミングで読み込み、勾配を計算して出力することで、デバイス上に永続的に保持される状態を最小化する。CPU-GPU間の帯域幅ボトルネックに対処するため、二つの主要な最適化を採用する。1) パイプライン化されたダブルバッファ実行エンジンを導入し、複数のCUDAストリーム間でパラメータのプリフェッチ、計算、勾配のオフロードを重畳させ、GPUの連続的な実行を実現する。2) 永続的な自動微分グラフを、状態を持たないレイヤーテンプレートで置き換え、重みをストリーミング入力時に動的にバインドすることで、永続的なグラフメタデータを排除し、スケジューリングの柔軟性を提供する。1.5TBのホストメモリを搭載した単一のH200 GPU上で、MegaTrainは120Bパラメータまでのモデルを確実に訓練する。また、14Bモデルの訓練において、CPUオフロードを行うDeepSpeed ZeRO-3と比較して1.84倍の訓練スループットを達成する。MegaTrainはさらに、単一のGH200上で512kトークンのコンテキスト長による7Bモデルの訓練も可能にする。

English

We present MegaTrain, a memory-centric system that efficiently trains 100B+ parameter large language models at full precision on a single GPU. Unlike traditional GPU-centric systems, MegaTrain stores parameters and optimizer states in host memory (CPU memory) and treats GPUs as transient compute engines. For each layer, we stream parameters in and compute gradients out, minimizing persistent device state. To battle the CPU-GPU bandwidth bottleneck, we adopt two key optimizations. 1) We introduce a pipelined double-buffered execution engine that overlaps parameter prefetching, computation, and gradient offloading across multiple CUDA streams, enabling continuous GPU execution. 2) We replace persistent autograd graphs with stateless layer templates, binding weights dynamically as they stream in, eliminating persistent graph metadata while providing flexibility in scheduling. On a single H200 GPU with 1.5TB host memory, MegaTrain reliably trains models up to 120B parameters. It also achieves 1.84times the training throughput of DeepSpeed ZeRO-3 with CPU offloading when training 14B models. MegaTrain also enables 7B model training with 512k token context on a single GH200.