JAXでの「ホストオフロード」がLLMトレーニングのメモリ制限を緩和、NVIDIA Grace Blackwell環境で最大57%のスループット向上を実証

要点

• NVIDIAは、オープンソースの機械学習ライブラリ「JAX」において、大規模言語モデル（LLM）のトレーニング時のメモリ制限を緩和する「ホストオフロード」技術を解説した。

• 本技術は、トレーニング中の特定の中間データをGPUからCPU側のホストメモリに一時退避させることで、GPUの高帯域幅メモリ（HBM）の容量不足を解消する。

• 900 GB/sの双方向帯域幅を持つNVLink-C2Cで接続された「NVIDIA Grace Blackwell」や、1.8 TB/sに高速化された「Vera Rubin」プラットフォームで特に高い効果を発揮する。

• MaxTextを用いた実験では、DeepSeek-V3 671BやLlama 3.1 405Bのトレーニングにおいて、従来の再計算手法と比べて最大57%のスループット向上を達成した。

• NVIDIAは2026年7月10日、オープンソースのPythonライブラリである「JAX」（数値計算ライブラリ）を用いた大規模言語モデル（LLM）のトレーニングにおいて、GPUのメモリボトルネックを軽減する「ホストオフロード（Host Offloading）」技術に関する検証結果を公式技術ブログで公開した。同社によると、この技術は最新の「NVIDIA Grace Blackwell」や将来の「Vera Rubin」プラットフォームに搭載される高速インターコネクト（接続用半導体やネットワーク技術）を活かすことで、モデルのサイズやバッチサイズ（一度に処理するデータ量）の制限を大幅に緩和できるという。

• 大規模言語モデルのトレーニングでは、モデルのパラメータ数や処理する文章の長さ（シーケンス長）、バッチサイズが大きくなるにつれて、GPUが搭載する高帯域幅メモリ（HBM：GPUに直結された極めて高速なメモリ）の容量制限に達してしまう課題が存在する。モデルの重み、勾配、オプティマイザ（最適化手法）の状態、通信バッファ、そして処理の過程で生じる「アクティベーション（中間データ）」が限られたメモリ容量を奪い合うため、計算性能を使い切る前にメモリ不足で動作しなくなることが増えている。

• 今回発表されたホストオフロード技術は、このHBM의容量圧迫を避けるためのアプローチである。ニューラルネットワークの順伝播（データを前方に流す計算）の際に、選択した一部のアクティベーションをシステム側のピン留めされたホストメモリ（CPUが管理するメインメモリ）へと移動させ、逆伝播（誤差を逆方向に流して学習する計算）のタイミングで再びGPUにストリームで読み戻す。これは、必要なときにアクティベーションを再計算する従来の「アクティベーション再マテリアライズ（再計算）」と呼ばれる手法に代わる選択肢となる。

• この処理が実用的になった背景には、CPUとGPUを結ぶ接続技術の劇的な高速化がある。NVIDIA Grace CPUとBlackwell GPUの間は「NVLink-C2C」と呼ばれる接続技術を介して双方向900 GB/sという極めて高い帯域幅で結ばれており、ホストメモリをアクティベーションの一時退避場所として十分に活用できる性能を備えている。さらに、将来のVera CPUとRubin GPUの組み合わせでは、この双方向速度が2倍となる1.8 TB/s（コヒーレント帯域幅：CPUとGPUの間で一貫性を保った高速な通信路）へと向上する。

• ただし、単に接続帯域幅が広いだけでは性能向上には不十分である。ホストメモリとGPU間のデータ転送を、GPUが行う本来の計算処理と時間的に重ね合わせる（オーバーラップさせる）必要がある。

• NVIDIAは、JAXベースのLLMトレーニングフレームワーク「MaxText」と、最適化コンパイラ「XLA（Accelerated Linear Algebra：機械学習モデルの計算を最適化するコンパイラ）」を使用し、128基のGPUからなるシステム「NVIDIA GB200 NVL72」上で実験を行った。検証には、Metaのdense（密）なモデル「Llama 3.1 405B」と、巨大なsparse MoE（まばらな構造を持つ混合専門家）モデルである「DeepSeek-V3 671B」の2つのワークロード（処理負荷）が用いられた。

• DeepSeek-V3 671Bは、合計61のデコーダレイヤー（モデルを構成する基本レイヤー）のうち、最初の3つが密なMLP（多層パーセプトロン）ブロックで、残りの大部分はMoE（混合専門家）ブロックで構成されている。この実験において、MoEレイヤー内の特定の計算ステップ（MLAクエリ、キー/値プロジェクションの中間体、MoEアッププロジェクションの中間体など）で生じる巨大なアクティベーションを選択的にホストメモリへ退避させるオフロードポリシー（退避基準）が適用された。これらのデータは非常に大きく、メモリに収まる最大バッチサイズに大きな影響を与える。

• 実験の結果、ホストオフロードを「Latency Hiding Scheduler（遅延を隠蔽するスケジューラ）」やパイプライン転送と組み合わせることで、従来のアクティベーション再計算手法と比較して最大57%のスループット（単位時間あたりの処理能力）向上が確認された。また、本来であればGPUのメモリ制限により実行不可能であった大きなバッチサイズでのトレーニングも実行可能になるなど、メモリ容量の壁を突破できることが実証された。特にこの恩恵は、スパースMoEモデルにおいて顕著に現れたという。

• 最適なパフォーマンスを引き出すためには、データ転送と計算・通信の緻密な重ね合わせが不可欠となる。これはXLAのカスタムスケジューリングフラグや、データコピー専用のストリームを使用することで管理される。さらにNVIDIAは、これらの非同期的なデータ移動やメモリの利用状況が期待通りに動作しているかを検証するために、同社のプロファイリングツール（性能分析ソフト）「NVIDIA Nsight Systems」を活用することを推奨している。