TRL 1.14、LoRA重みのみを同期し、非同期GRPOのマルチノード学習をNCCL非依存に

Hugging FaceはTRL 1.14で、実験的な`AsyncGRPOTrainer`にadapter-only同期を追加した。trainerはFSDP2でLoRAを更新し、数optimizer stepごとに新しいバージョンをStorage Bucketへ書き込んだうえで、独立したvLLMジョブに`/v1/load_lora_adapter`経由での読み込みを要求する。Qwen2.5-Math-1.5Bとrank-1 LoRAの例では、転送されるのは約3 GBの完全なモデルではなく数MBのadapterであるため、学習プロセスと生成プロセスを同一ノードに配置する必要がなく、ノード間NCCL groupを構築する必要もない。

この設計の難しさは、単なるファイル共有ではない。フロントエンドのproxyはHF Jobsの認証ヘッダーを付与し、同じrolloutを、そのKV prefixを保持している可能性が高いreplicaへルーティングするとともに、adapterの読み込み、一時停止、再開をすべてのreplicaへブロードキャストする。更新のたびに新しいpolicy名を使用し、古い重みで作成されたKV blockがhot update後に誤って再利用されるのを防ぐ。サンプルに最大4世代のpolicy version遅延を許容する場合、vLLMは現在のversion、4つの旧version、さらに切り替え中に使用する追加slotを保持する必要があるため、`max_loras=6`に設定される。slotが不足すると、rolloutの実行中にもかかわらず古いpolicyがevictされる可能性がある。

チームはtoken-budget packing、concurrency、gradient checkpointing、adapter読み込みのretryを段階的に調整し、同一の500-stepレシピを3時間27分から約53分へ短縮したほか、実行scriptとproxy testも公開した。これは、低ランク更新によって同期問題をversioned artifactの受け渡しへ置き換えられることを示しており、高速なcluster interconnectを利用できない作業環境に適している。ただし、「ネットワーク不要」という意味ではなく、直接的なtensor collectiveをFUSE-backed bucketとHTTPS control flowで置き換えたものだ。

エンジニアリングチームは次に、object storageのvisibility latency、adapterのatomic publish、policy staleness、破棄されたrolloutの割合を検証し、GPUコスト、最終reward、収束安定性も個別に測定する必要がある。公式の数値は1.5Bモデルと特定のH200構成だけに基づいており、成熟したNCCL／Slurm clusterとの同等品質・同等コストでの比較はまだ行われていない。また、DoRA、追加のtrainable module、またはvLLMのrank上限を超える構成では、完全にmergeされた重みの同期へフォールバックする。