大規模言語モデル（LLM）のトレーニングワークロードでは、計算リソースが完全に活用される前にGPUメモリ制限に達することが増えています。モデル重み、勾配、オプティマイザの状態、通信バッファ、中間活性化はすべて、限られたGPU高帯域幅メモリ（HBM）を競合しています。モデルサイズ、シーケンス長、バッチサイズが増大するにつれて、HBM容量が主要なスケーリングボトルネックとなることがよくあります。

この記事では、JAXベースのトレーニングワークフローにおいてこれらの制約を緩和するためにホストオフロードをどのように使用できるかを解説します。