LazyTrain、混合整数スケジューリングでメモリ転送をオーバーラップし、H800単体での27Bモデル学習を1.24倍高速化

メモリが不足する場合、単一GPUでの大規模モデル学習では通常、activation checkpointing、CPU offload、再計算の間で固定ルールが採用される。LazyTrainは、これらの選択を混合整数線形計画問題として定式化する。GPU、メインメモリ、NVMeの容量に加え、PCIeとSSDの帯域幅、各層の計算時間を入力し、どの境界を保存するか、どのストレージ階層に配置するか、どのブロックを再計算するか、さらに転送をどの計算ウィンドウに割り当てるかを一括して求める。

システムは、バックプロパゲーションの開始点から終了点までの経路にフロー制約を設定すると同時に、HBM、DRAM、NVMeの容量、およびパラメータ、勾配、activationの双方向トラフィックを制約する。SCIPソルバーが学習前にスケジュールを生成するため、runtimeではリアルタイム探索を行わない。主要なQwen3.6-27B構成は、64個のTransformer block、シーケンス長1,024、batch 72。実体化された52個のcheckpointのうち、30個をGPU、21個をCPU、1個をNVMeに配置し、パラメータとoptimizer stateは主にCPUに置く。

同一のH800、同一のデータ分割およびbatchという条件で、LazyTrainはMegaTrainと比べ、持続演算性能を176.90 TFLOPSから219.95 TFLOPSへ引き上げ、約1.24倍を達成した。スループットは毎秒1,361 token、GPUメモリ使用量のピークは約68.84GBだった。混合整数スケジューリングを除くと性能は193.17 TFLOPSまで低下しており、主な性能向上が転送と再計算の共同最適化によることを示している。一方、いわゆるHybrid 8-bit経路を除いた場合の低下は219.29 TFLOPSにとどまり、量子化されたstateが使用されるのも全パラメータの約2%にすぎない。

オープンソースのalpha版は、MegaTrainのCPU-masterおよび層ごとのストリーミングアーキテクチャを基盤とし、Qwen3.6-27B、MetaMathQA、NVMe向けのスケジューリング設定を提供する。ただし、論文で実施された実験は単一のH800およびRTX 3090に限られ、表の各行も通常は1回の実行結果にすぎない。実行中に帯域幅が変動してもスケジュールは再構成されない。また、追加で露出する転送時間がゼロという主張もソルバーによる推定値であり、ステップ単位のstall instrumentationでは検証されていない。注目すべきは、単一の27Bモデル事例におけるピーク値ではなく、マルチGPU、クロスノード環境、異なるMoEアーキテクチャでも性能向上を維持できるかどうかだ。