Hugging Face Jobs環境でLoRAを用いた非同期GRPOの実現：分散学習の新たな形

Hugging Faceは、TRLのAsyncGRPOTrainerにLoRAアダプターのサポートを導入し、TRL v1.14で提供を開始しました。

この機能により、トレーナーがモデル全体ではなく、数メガバイト程度の軽量なLoRAアダプターのみを学習し、それをvLLMに同期することが可能になりました。

この進展は、トレーニングと推論プロセスを物理的に分離した環境での運用を可能にします。

従来のシステムでは、トレーナーとvLLMのレプリカが同じノードを共有するか、NCCLグループを形成して通信する必要がありましたが、Hugging Face Jobs環境では、トレーナーとvLLMのレプリカが別々のマシンで独立して動作します。

この分離を実現するため、各レプリカの前に小さなプロキシが配置され、認証ヘッダーの付与、ロールアウトの適切なレプリカへのルーティング、およびアダプターのロード通知といった役割を担います。

LoRAの採用は、特に強化学習（RL）において有効性が示されています。

Thinking Machinesのブログによれば、ランク1のアダプターであっても、ポリシー勾配RLにおいて完全なファインチューニングと同等の性能を達成できることが示されています。

これは、アドバンテージ関数がエピソードあたりに提供する情報量が限定的であるため、ランク1アダプターの容量で十分吸収できるためと説明されています。

この新しいアーキテクチャの利点は、システム的な効率性の向上にあります。

例えば、1.5Bモデルの場合、フルモデルは約3GBであるのに対し、ランク1アダプターは数メガバイトに過ぎません。

これにより、更新のたびに完全なポリシーを推論ワーカーに送信する代わりに、アダプターのみを送信することができ、vLLMは複数のアダプターを同時にロードして保持することが可能です。

この実装により、トレーニングと推論が異なるマシンで独立して動作する、よりスケーラブルなセットアップがHugging Face Jobs上で実現しました。