LLM学習のHBM不足をどう超えるか──NVIDIAがJAXのホストオフロード性能を解説

NVIDIAは、JAXベースの大規模言語モデル学習でGPUの高帯域メモリ（HBM）不足を緩和する「ホストオフロード」の性能を解説した。モデル重み、勾配、オプティマイザ状態、通信バッファ、中間活性化が同じHBMを取り合うため、モデル規模やシーケンス長、バッチサイズが伸びるほどメモリ容量が先に限界になる。

ホストオフロードは、順伝播時に一部の活性化をピン留めされたホストメモリへ移し、逆伝播時に必要なタイミングで戻す手法だ。再計算でメモリを節約する activation rematerialization とは異なり、再計算ではなく転送と重ね合わせでボトルネックを避ける。

Grace Blackwellでなぜ効くのか

NVIDIAは、Grace CPUとBlackwell GPUを結ぶNVLink-C2Cが双方向900GB/sの帯域を持つため、ホストメモリを活性化の退避先として使いやすいと説明している。さらにVera CPUとRubin GPUでは、双方向1.8TB/sのコヒーレント帯域へ拡張される。

ただし、帯域だけでは十分ではない。転送がGPU計算を止めないよう、XLAのスケジューリング、Latency Hiding Scheduler、パイプライン転送と組み合わせ、コピーと計算を重ねる必要がある。

メモリ節約だけでなく、バッチ構成を変える

DeepSeek-V3 671Bの例では、ホストオフロードによって大きな活性化をGPU外に逃がし、より大きなバッチ構成を実行可能にした。単にピークメモリを下げるだけでなく、これまでOOMになっていた設定を候補に戻せる点が大きい。

一方で、転送バッファやプリフェッチ済み活性化をGPU側に持つため、最適化したオフロードではGPU使用量が増える場合もある。メモリを減らす技術というより、計算、通信、転送、容量の全体最適化と見るべきだ。

AIインフラ投資への示唆

企業がLLM学習や大規模ファインチューニングを検討する際、GPU単体の演算性能だけを見ても不十分だ。CPU-GPU間の帯域、コンパイラ、フレームワーク、スケジューラが揃って初めて、巨大モデルの訓練効率が上がる。

特にMoEや長文コンテキストでは、活性化のサイズが大きく、メモリ配置の戦略が成果を左右する。クラウドやオンプレのAI基盤を選ぶ際は、モデルが動くかだけでなく、どのバッチサイズで、どのスループットで、どれだけ再現性高く動くかを検証すべきだ。

注意点はワークロード依存性

NVIDIAの結果は、GB200 NVL72上でのMaxTextワークロードに基づく。効果はモデル構造、シーケンス長、バッチサイズ、接続帯域、JAX/XLA設定によって変わる。

導入企業は、代表的な自社ワークロードで、再計算、オンデバイス保存、ホストオフロードを比較し、ピークメモリ、ホストメモリ、スループット、安定性を同時に見る必要がある。LLM学習の競争力は、モデル設計だけでなく、メモリ階層を使い切る運用技術にも移っている。

参考：NVIDIA Technical Blog