TRLのAsyncGRPOがLoRAアダプタ同期に対応、HF Jobsで学習を短縮

何が変わったのか

Hugging Faceのブログによると、TRLのAsyncGRPOTrainerがLoRAアダプタを学習し、そのアダプタだけをvLLMへ同期できるようになった。この変更はPR #7017としてマージされ、TRL v1.14に含まれる。従来はモデル全体の重みを推論側へ渡す必要があり、その転送量が構成上の制約になっていた。

AsyncGRPOTrainerはもともと学習と生成を分離できる設計で、トレーナーとvLLMを別々のマシンで動かせた。ただし、両者がファイルシステムを共有できるか、NCCLグループを組めることが前提だった。

今回の検証では、トレーナーとvLLMレプリカを別々のHugging Face Jobsとして動かす構成が試された。同じレシピで500ステップを回した5回の実行で、所要時間が3時間27分から53分に短縮されたと報告されている。

アダプタだけを同期する仕組み

トレーナーはテンソルをvLLMへ送らない。数オプティマイザステップごとにアダプタを <output_dir>/.vllm_lora/trl-policy-v{N} に保存し、原子的なリネームで公開したうえで、そのパスをvLLMの /v1/load_lora_adapter エンドポイントへ送る。vLLMはディスクからファイルを読み込み、ロールアウト側は model="trl-policy-v{N}" を指定して呼び出す。

Job間のファイル共有にはStorage Bucketを使う。バケットをFUSE(hf-mount)で各Jobの同じパスにマウントすれば、ネットワークファイルシステムと同じように扱える。Job間にネットワーク経路は不要で、NCCLも使わない。

LoRAがRLに向いている理由も説明されている。Thinking Machinesのブログによれば、アドバンテージ関数が1エピソードあたり約O(1)ビットの情報しか与えないため、rank 1でもポリシー勾配RLで全微調整に匹敵しうる。1.5Bモデルのrank-1アダプタは数メガバイトで、モデル全体の約3GBと比べて転送量が大きく減る。

3つのJobとプロキシ

構成は小さい。トレーナーJob、vLLMを動かすJob 2つ、3つすべてに同じパスでマウントするStorage Bucket、そしてvLLMレプリカの前段に置くプロキシだ。vLLMはv0.27.1に固定されている。vLLMの更新が速いため、使うフラグとランタイムLoRAエンドポイントはこのバージョンのものとして扱う必要がある。

プロキシには複数の役割がある。公開されたJobのポートは Authorization: Bearer <HF token> ヘッダを要求するため、それを付与する。加えて、各ロールアウトをKVプレフィックスを保持している可能性の高いレプリカへ振り分け、アダプタ更新を全レプリカへブロードキャストする。

アダプタのスロット数はmax_stalenessから決まる。max_staleness=4なら現在のポリシーと過去4つを同時に扱う必要があり、切り替え中の1枠を足して --max-loras 6 になる。5では同期のたびに、まだ実行中のロールアウトを持つポリシーが黙って追い出されてしまうという。

検証データと注意点

検証には sail/Sanity-Test-R1D-1.5B が使われた。論文「Defeating the Training-Inference Mismatch via FP16」(Qi et al., 2025)で使われたデータセットで、DeepSeek-R1-Distill-Qwen-1.5BでMATHの問題ごとに40の回答を生成し、成功率20%〜80%の問題を残して1,460問に絞ったものだという。すぐ解ける問題でも絶望的な問題でもないため、初期の学習信号が得やすいと説明されている。

ハイパーパラメータは Qwen/Qwen2.5-Math-1.5B、LoRA rank 1、alpha 2、学習率4e-5、プロンプトあたり8サンプル、1ステップあたり128コンプリーション、最大3,000生成トークン、4,096トークン文脈となっている。重み同期は4オプティマイザステップごと、チェックポイントは50ステップごとにバケットへ保存される。

注意点も示されている。TRLは --data-parallel-size を1より大きくした構成ではアダプタのみの同期を拒否する。DoRAやmodules_to_save、--max-lora-rankを超えるrankはマージ済み重み同期にフォールバックし、警告が出る。プロキシの実装詳細やスループットの数値は、本記事で確認できる範囲では示されていない。