大規模言語モデル（LLM）の訓練現場では、GPUの演算能力よりも容量が先に限界を迎える「HBM（高帯域幅メモリ）の壁」が深刻化している。NVIDIAはJAXフレームワーク上で、活性化関数のデータをGPUからCPUのメインメモリへ退避させる「ホストオフロード」手法を発表。Grace Blackwellプラットフォームの高速インターコネクトを活用し、DeepSeek-V3やLlama 3.1で再計算手法を上回る性能を実証した。演算効率を追求する次世代AIハードウェアの競争が、メモリ帯域とデータ移動の最適化へと重心を移しつつある。

GPU演算よりも先にメモリが枯渇する構造問題

LLM訓練では、モデルの重みや勾配、オプティマイザの状態に加え、順伝播で生じる中間的な活性化関数の保存がHBMを圧迫する。特にバッチサイズやシーケンス長の拡大に伴い、このHBM消費はGPUの理論演算性能を十分に引き出す前の制約要因となってきた。NVIDIAの今回の発表は、この課題に対し活性化関数を計算し直す「リマテリアライゼーション」から、高性能CPU-GPU接続を活用した「ホストオフロード」への転換を提案するものだ。従来のPCIe接続では帯域不足で非現実的だった手法が、NVLink-C2Cの900GB/s超の双方向帯域によって実用水準に達したことを示している。

Grace Blackwellで活性化関数がCPUへ退避可能に

ホストオフロードの成否は、CPU-GPU間のデータ転送をいかにGPUの計算処理と並行して隠蔽できるかにかかっている。NVIDIA Grace Blackwellアーキテクチャでは、Grace CPUとBlackwell GPUをNVLink-C2Cで直結し、専用のコピーストリームとXLAコンパイラのスケジューリングフラグにより転送レイテンシを隠蔽する機構を実装した。これにより、前方の計算で生成された活性化関数をピン留めされたホストメモリへ退避し、逆伝播時に計算と重ねてストリーミングで再読み込みできる。MaxTextフレームワークを用いた検証では、128基のGB200で最大57%のスループット向上を記録している。

DeepSeek-V3の大規模MoEモデルで顕著な効果

この手法が特に有効性を発揮したのは、DeepSeek-V3 671Bに代表されるスパースMoE（Mixture of Experts）モデルでの訓練だ。大規模MoEは専門家ネットワークごとに活性化関数を保持する必要性からメモリ消費が顕著で、従来のHBMだけではバッチサイズが制限されていた。ホストオフロードの導入により、GPUメモリ上には収まらなかったバッチサイズでの訓練が可能となり、メモリ容量の制約そのものが解除されるシナリオが実証された。GPUの演算リソースを限界まで使うための手段が、従来の「計算を節約」から「データ配置を最適化」へと変化している。

Vera Rubinでさらに倍増するCPU-GPU間帯域の意味

今回のアプローチは単に現行ハードウェアの最適化手法にとどまらず、将来のプラットフォーム設計への布石でもある。NVIDIAが発表した次世代アーキテクチャVera Rubinでは、NVLink-C2Cの帯域が1.8TB/sへ倍増される計画だ。帯域の拡大は、より多くのデータをGPU外部へオフロードできる余地を広げ、モノリシックなGPUメモリ容量への依存度を相対的に下げることを示唆する。AIインフラ設計においては、GPUの演算ピーク性能だけでなく、システム全体でのデータ移動帯域とそのスケジューリング精度が次の主要評価指標となる。