Hugging Face TRL: 効率的な GRPO トレーニングのための同居 vLLM

Hugging Face は、TRL ライブラリ内で vLLM を同居させることをサポートし、トレーニングと推論が同じ GPU を共有できるようにしました。これにより、生成時にトレーニング GPU がアイドル状態になり、推論時に推論 GPU がアイドルになるという「ピンポン」非効率が解消され、スループットが向上しハードウェア要件が削減されます。

The Inefficiency of Server-Mode Inference

TRL v0.18.0 以前は、vLLM はサーバーモードでのみサポートされていました。この構成では、vLLM は専用 GPU 上で別プロセスとして実行され、トレーニングスクリプトとは HTTP 経由で通信します。このアーキテクチャは、生成が常に行われる Group Relative Policy Optimization (GRPO) のようなオンライン学習アルゴリズムにおいて大きなボトルネックを生み出します。

  • GPU Underutilization: トレーニング GPU は vLLM サーバーが生成を行っている間アイドル状態になります。
  • Resource Waste: 推論専用に GPU を確保する必要があり、コストが増大します。
  • Communication Overhead: REST API 呼び出しに依存するため、ネットワーク遅延が発生します。

Co-located Design and Implementation

同居により、vLLM はトレーニングコードと同じ分散プロセスグループ内で実行されます。外部サーバーの代わりに、トレーナーは外部ランチャーを用いてプロセス内で vLLM を起動します。

Key Technical Capabilities

  • Unified Execution: トレーニングと推論タスクが同じ GPU リソースを交互に使用するため、アイドル時間が削減されます。
  • Direct Memory Communication: HTTP 呼び出しを排除し、vLLM はネイティブな Python 呼び出しでトレーニングループと通信します。
  • Distributed Compatibility: 本システムは torchrun、Tensor Parallelism (TP)、Data Parallelism (DP) と互換性があり、Single Program, Multiple Data (SPMD) 実行パターンを利用します。
  • Simplified Deployment: vLLM はトレーニングジョブ内部で直接制御されるため、別個のサーバースクリプトが不要です。

Configuration

ユーザーは GRPOConfig で vllm_mode="colocate" を設定することでこの機能を有効にできます。

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="colocate",
)

Performance Benchmarks

サーバーモード（8 GPU のうち 1 GPU を vLLM に専念）と同居モード（8 GPU 全てをトレーニングに使用）を比較した実験では、スループットが一貫して向上しました。公平性を保つため、サーバーモードのスループットは 8/7 の係数で正規化しています。

Model-Specific Results

  • 1.5B Model: 最大バッチサイズで最大 1.43 倍の速度向上を達成しました。ただし、Tensor Parallelism (TP) を増やすと通信オーバーヘッドにより性能が低下しました。
  • 7B Model: バッチサイズを変化させた場合に最大 1.35 倍、TP を増やした場合に最大 1.73 倍の速度向上を示し、モデルが大きくなるほどシャーディングの効果が大きいことが分かります。
  • 72B Model (Qwen2.5-Math-72B): 同居構成は、GPU を 4 台少なく使用しながらも、標準 TRL より約 1.26 倍高速でした。

Scaling to Large Models (72B+)

Qwen2.5-Math-72B のような超大規模モデルをトレーニングするには、トレーニングと生成の間でメモリ競合が起きないよう高度なメモリ管理が必要です。

vLLM Sleep Mode

Hugging Face は vLLM の sleep() API を GRPO ループに統合しました。Level 2 スリープを使用すると、モデル重みと KV キャッシュの両方を GPU から完全にアンロードでき、トレーニングステップ用の空きメモリを最大化し、メモリ競合を回避します。

DeepSpeed and Accelerate Optimizations

安定性と効率性を保つため、以下の最適化を採用しています。

  • DeepSpeed ZeRO Stage 3: モデル重み、勾配、オプティマイザ状態を GPU 間で分割します。
  • CPU Offloading: オプティマイザ状態を CPU ("offload_optimizer": {"device": "cpu"}) にオフロードし、GPU メモリを解放します。
  • Memory Management: contiguous_gradients を有効化してメモリ断片化を削減し、overlap_comm を使用してトレーニングを高速化します。
  • Accelerate Integration: マルチ GPU / マルチノードの起動とデータ並列を実現するために利用します。

Challenges and Lessons Learned

同居 vLLM の実装ではいくつかの技術的課題が明らかになりました。

  • Determinism in vLLM ≥ 0.8.0: 外部ランチャーを使用した Tensor Parallelism にバグがあり、乱数シードを明示的に設定する必要がありました。
  • Buffer Restoration: Level 2 スリープのバグ（Issue #16564）により、スリープから復帰した際にモデルバッファ（例: BatchNorm の running mean/var）を明示的に復元する修正が必要でした。
  • Shutdown Stability: vLLM スリープがトレーニングプロセス終了時にセグメンテーションフォルトを引き起こす既知の問題があります（Issue #16993）。ただし、トレーニング結果には影響しません。

Conclusion on Model Quality

Math500 と AIME24 のベンチマーク結果は、同居構成がモデル性能を損なわないことを示しています。同居と標準設定の報酬曲線はほぼ一致しており、同居でトレーニングしたモデルは標準でトレーニングしたモデルと同等の性能を発揮します。これにより、効率向上が精度低下を伴わないことが検証されました。

Sources