JAXホストオフロードでLLM学習が57%高速化

エヌビディア（NVIDIA）は、オープンソースのPython数値計算ライブラリ「JAX」において、GPU（グラフィックス処理ユニット）のメモリ不足問題を解消する新技術「ホストオフロード（Host Offloading）」の検証結果を公開しました。この技術は、AIのトレーニング時にGPU内の超高速メモリ（HBM）を節約するため、一部の計算過程データをCPU側のシステムメモリ（ホストメモリ）に一時退避させるものです。これにより、巨大な大規模言語モデル（LLM）の学習時に、メモリ容量不足（OOMエラー）によるクラッシュを防ぎ、効率的な実行を可能にします。

1. 再計算を防ぐ「ホストオフロード」のメリット

従来の学習では、GPUメモリを節約するために「アクティベーション再マテリアライズ（順伝播データを破棄し、逆伝播時に再計算する手法）」が用いられていました。これに対しホストオフロードは、順伝播時の大容量データをCPU側に転送し、逆伝播時に順次GPUへ送り返すことで、再計算の計算負荷を取り除きます。これは、CPUとGPUを「NVLink-C2C（双方向で900 GB/s〜1.8 TB/sという超高速接続）」で結ぶ「NVIDIA Grace Blackwell (GB200 NVL72)」システムなどで特に威力を発揮します。

2. MaxTextでの検証とDeepSeek-V3の成果

JAXをベースにした大規模学習フレームワーク「MaxText」を用い、128基のGPU（GB200 NVL72）環境で測定が実施されました。

• DeepSeek-V3 671B（MoEモデル）：​ ホストオフロードにレイテンシ隠蔽スケジューラ（LHS）や転送のパイプライン化を適用した結果、スループットは 908.2 TFLOPs/s/device を記録。従来の再計算手法と比べて 57%高速化​ しました。
• Llama 3.1 405B：​ QKV（Query/Key/Value）アクティベーションオフロードにより、スループットが2.9%向上（2,746 TFLOPs/s/device）。
• ​バッチサイズの拡大：​ 通常はメモリ制限によりグローバルバッチサイズが256まででしたが、オフロードを適用することで1024へと 4倍に拡大​ でき、学習効率が大幅に向上しました。

本技術は、AI半導体の使用量を物理的に増やさずとも、コンパイラ（プログラム翻訳ソフト）とチップ間インターコネクト（接続技術）の密な連携設計によってソフトウェア的にメモリ上限の壁を突破した点が非常に画期的です。DeepSeek-V3やLlama 3.1といった超巨大モデルのトレーニングにかかる期間とインフラコストを大きく圧縮し、今後の超巨大AI開発のあり方に影響を与えると見られます。