Hugging Face Jobs 間で LoRA を用いた Async GRPO
Hugging Face は AsyncGRPOTrainer（TRL v1.14）に LoRA サポートを実装しました。これにより、トレーナーは完全なモデル重みではなく、小さな LoRA アダプターのみを vLLM に同期できます。Hugging Face Storage Buckets を共有ファイルシステムとして活用し、リクエストルーティング用のカスタムプロキシを使用することで、このアーキテクチャは NCCL や共有ローカルディスクを必要とせずに、トレーニングと推論を別々のマシン（Hugging Face Jobs）上で実行することを可能にします。
アーキテクチャ: Storage Buckets による分散同期
AsyncGRPOTrainer は、テンソルを直接 vLLM に送信する必要性を排除するアダプター専用の同期パスをサポートするようになりました。代わりに、トレーナーは LoRA アダプターを Storage Bucket の特定のディレクトリに保存し、アトミックなリネームを実行し、/v1/load_lora_adapter エンドポイント経由で vLLM に通知します。
Hugging Face Jobs は hf-mount を使用して Storage Buckets を FUSE ファイルシステムとしてマウントできるため、トレーナーと vLLM レプリカは異なる VM 間で同じ絶対パスを共有できます。これにより、トレーナーと推論サーバーが物理ノードや高密度クラスタネットワークを共有する必要がなくなります。
ジョブの構成
システムは主に3つのコンポーネントで構成されています：
- トレーナージョブ: LoRA と FSDP を使用して
AsyncGRPOTrainerを実行します。 - vLLM ジョブ: ベースモデルを提供し、バケットから最新のアダプターを読み込む複数のレプリカ。
- プロキシサーバー: トレーナージョブ上で動作し、認証ヘッダーを処理してリクエストを vLLM レプリカにルーティングする、小さな asyncio ベースのプロキシ。
プロキシ: KV プレフィックスルーティングとブロードキャスト
複数の vLLM レプリカ間で効率を最大化するために、リクエストの分配と状態の同期を管理するカスタムプロキシが使用されます。
KV プレフィックスによるルーティング
冗長なプリフィル計算を避けるために、プロキシは KV キャッシュのプレフィックスに基づいてリクエストをルーティングします。プロンプトを 16 トークンのブロックに分割し、アダプター名をシードとしてチェーンハッシュを計算します。
ルーターは、どのレプリカがどのブロックハッシュを提供したかを追跡します。リクエストのプロンプトが特定のレプリカにすでにキャッシュされているプレフィックスと一致し（かつそのレプリカが過負荷になっていない場合）、リクエストはそこにルーティングされます（「アフィニティヒット」）。これにより、同じプロンプトのプリフィルを複数のロールアウト間で再計算することが防止されます。これは、1つのプロンプトに対して複数の補完が生成される GRPO において重要であり、
状態のブロードキャスト
各 vLLM レプリカは独立したジョブであるため、プロキシはアダプターの読み込み、一時停止、再開などの状態変更リクエストをすべてのレプリカにブロードキャストすることで整合性を確保します。これにより、特定のポリシーバージョン名がフルートの全体で同じ重みを指すことが保証されます。
パフォーマンス最適化とボトルネック分析
sail/Sanity-Test-R1D-1.5B データセットと Qwen/Qwen2.5-Math-1.5B モデルを使用して、Hugging Face はパイプラインを最適化するために5回の実験実行を行いました。結果は、Async RL におけるボトルネックがトレーニングと生成の間で移動し得ることを示しています。
主要な最適化
- トークンバジェットバッチング: デバイスごとのトレーニングバッチサイズを 1 からトークンバジェットバッチング（例：
token_budget=16384）に変更することで、各行に複数のシーケンスをパックし、マイクロバッチの数を減らすことで、MFU が 3.9% から 19% に向上しました。 - 勾配チェックポイントの無効化: 小さなモデル（1.5B）の場合、勾配チェックポイントを無効にすることで冗長なフォワードパスを排除し、フォワード+バックワードの時間を短縮しました。これにより、ボトルネックがトレーナーから生成レプリカに移行しました。
- インフライトリクエストの増加:
max_inflight_tasksを増加させる（例：384 に設定）ことで、システムが複数の vLLM レプリカを完全に活用できるようになり、クライアント側の並行性制限によるスループットのスロットリングを防止しました。
最終結果
これらの最適化を組み合わせることで、500 ステップを完了するまでの総時間は 3 時間 27 分から 53 分（3.9 倍の高速化）に短縮されました。
| 指標 | Run 1 (ベースライン) | Run 5 (最適化済み) |
|---|---|---|
| ウォールクロック時間 | 3 時間 27 分 | 53 分 |
| 中央値ステップ時間 | 22.9 秒 | 4.8 秒 |
| MFU (Fwd/Bwd) | 3.9% | 23.5% |
| トレーニング済みサンプル数 | 64,000 | 84,078 |
| 平均鮮度 (Staleness) | 1.5 バージョン | 2.0 バージョン |
技術的な実装詳細
- vLLM バージョン: ランタイム LoRA エンドポイントとの互換性のために
v0.27.1に固定されています。 - アダプタースロット:
max_staleness=4をサポートするために、スワップ中に現在のポリシーと以前のバージョンが読み込まれたままになるよう、vLLM は--max-loras 6で設定されています。 - 整合性: KV キャッシュが古いポリシーバージョンによって生成されたプレフィックスと誤って一致することを防ぐために、バージョン付きのアダプター名が使用されます。