7月11日、NVIDIAが「Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading」と題した記事を公開した。大規模LLMの訓練においてGPUのHBM（高帯域メモリ）容量は依然として深刻なボトルネックであり、モデルを再計算なしに大きなバッチで走らせるための実践的な打ち手として、CPU側ホストメモリへのオフローディングが有力な選択肢として浮上している。この記事では、JAXベースの訓練環境でホストオフローディングを活用してHBMボトルネックを解消する技術手法が詳しく紹介されている。

HBMが詰まる前に「ホストメモリへ逃がす」という発想

大規模なLLM訓練では、モデルの重み、勾配、オプティマイザの状態、通信バッファ、中間活性化テンソルがすべてGPUのHBMを奪い合う。モデルサイズやシーケンス長、バッチサイズが大きくなるにつれ、HBM容量がスケールアップの主要なボトルネックになる。

その解決策として記事が取り上げているのがホストオフローディングだ。フォワードパスで選択した活性化テンソルをピン留めされたホストメモリ（CPU側）に移動し、バックワードパスで必要になったタイミングでGPUに戻す仕組みである。

比較対象となるのがアクティベーション再マテリアライゼーション（activation rematerialization）——バックワードパスで活性化テンソルを保持せず、必要なときに再計算する手法だ。ホストオフローディングは再計算の代わりにホストメモリからのリロードを選ぶ。どちらが有利かはワークロードによるが、今回の実験では大規模MoEモデルで大差がついた。

Blackwellが有利な理由：NVLink-C2Cの存在

ホストオフローディングはどのシステムでも有効なわけではない。CPU-GPU間の帯域が細いと、転送がボトルネックになる。

NVIDIA Grace Blackwellでは、Grace CPUとBlackwell GPUをNVLink-C2Cで接続しており、帯域幅は900 GB/sに達する（※元記事では双方向合計か片方向かの明示がないため、正確な解釈は元記事を直接参照されたい）。次世代のVera CPU＋Rubin GPU構成ではこれが1.8 TB/sに倍増する。PCIeの帯域（数十GB/s程度）とは桁が異なる。

ただし帯域幅だけでは不十分で、転送とGPU演算を重ね合わせる（オーバーラップ） ことが性能の鍵になる。そのためにXLAコンパイラ側の2つの機能が重要な役割を担う。

• LHS（Latency Hiding Scheduler）：転送スケジュールを調整し、コピーをコンピューティングや通信と並列実行させる
• パイプライン化されたホストオフローディング：転送と演算のオーバーラップをさらに強化する

実測値：DeepSeek-V3 671Bでリマテリアライゼーションより57%高速

実験はMaxText（JAXベースのLLM訓練フレームワーク）を用い、NVIDIA GB200 NVL72システム上の128 GPU構成で実施した。

最も大きな効果が出たのがDeepSeek-V3 671B（疎なMixture-of-Expertsモデル）だ。

ホストオフローディング＋LHS＋パイプラインの組み合わせで908.2 TFLOPs/s/deviceを達成。リマテリアライゼーションと同一バッチ構成での比較で57%高速、LHS・パイプラインなしのオフローディング単体と比べても67.7%高速だった。

重要なのはバッチサイズの話でもある。オフロードなしでマイクロバッチ8・グローバルバッチ1024を試みるとOOMになる。ホストオフローディングによってこの構成が実現可能になり、MLA・MoE・MLPブロックの大きな中間活性化をGPUメモリから追い出すことで、モデルの状態や通信バッファに使えるHBMが確保された。

Llama 3.1 405Bでは2.9%の改善

密なデコーダモデルであるLlama 3.1 405Bでは、QKV（Query/Key/Value）活性化のオフローディングを検証した。バッチサイズ2、シーケンス長8,192、128 GPU FSDP、bfloat16活性化＋NVFP4 4ビット重み量子化という構成だ。

※ここで用いられているNVFP4はNVIDIAが導入した比較的新しい4ビット浮動小数点量子化形式であり、Blackwell世代のGPUで対応が進んでいる。bfloat16で活性化を保ちつつ重みだけをNVFP4に落とす混在構成により、精度と省メモリのバランスを取っている。

LHSを有効にしたQKVオフローディングで2,746 TFLOPs/s/device（ベースライン2,669から2.9%向上）。LHSをオフにすると2,569に落ち、転送のオーバーラップが性能に直結することが示された。

この構成ではパイプラインを追加しても2,718に微減した。LHSだけで転送レイテンシの大半が隠蔽されており、パイプライン化のメリットが薄い状況だ。DeepSeek-V3のような巨大なMoEモデルとは対照的な結果である。

使いどころと設定方法

ホストオフローディングが効果的な条件を整理すると：

• GPU メモリがモデルサイズ・シーケンス長・バッチサイズを制約している
• オフロード対象のテンソルが十分に大きい
• 転送と並行して実行できる独立した演算・通信がある

逆に、テンソルが小さい場合、独立した並行作業が少ない場合、ボトルネックがメモリ以外にある場合は効果が出にくい。

実際に試す際の出発点として、NVIDIAは以下のXLAフラグを推奨している：

--xla_gpu_enable_latency_hiding_scheduler=true
--xla_gpu_enable_pipelined_host_offloading=true
--xla_gpu_experimental_parallel_async_compute_limit=8

3つ目のフラグは非同期処理のキュー深さを増やし、LHSがオーバーラップを最大化できる余地を広げる。プロファイリングにはNVIDIA Nsight Systemsを使い、device-to-hostおよびhost-to-deviceのコピーが演算・NCCL通信と実際に重なっているか確認することが推奨されている。

JAX APIとしてはJAX host offloading tutorialが参考になる。jax.rematによるアクティベーションオフロード、memory_kind="pinned_host"、jax.device_put()によるパラメータ・オプティマイザ状態のオフロードがカバーされている。

詳細はReducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloadingを参照していただきたい。