TRLにおけるDelta Weight Syncにより、最小限の帯域幅で兆単位のパラメータを持つモデルの学習が可能に

1テラバイトの問題

非同期RL（Async RL）トレーニングでは、ポリシーの同期を維持するために、ステップごとにトレーナーから推論エンジンへモデル全体を送信する必要があります。bf16形式の7Bモデルの場合、これは1ステップあたり14 GBです。最先端の1Tパラメータモデルの場合、1ステップあたり約1 TBに達します。この転送はクリティカルパス上に位置しており、GPUがトークンを生成できないアイドル時間（計算の空き時間）を引き起こします。

なぜbf16のRL重みはほぼ常にスパース（疎）なのか

連続するRLオプティマイザのステップ間で、bf16の重みの約99%はビットレベルで同一のままです（最悪の場合でも98%を下回ることはありません）。これは、bf16の精度が限られているために起こります。更新量が、重みの周囲の表現可能な値の間隔の半分を下回る場合、更新は丸めによって吸収されてしまいます。典型的なRLの学習率（例：3×10⁻⁶）では、更新サイズはほとんどの重みにおいてこの閾値よりも小さいため、bf16の表現は変化しません。このスパース性は、運による測定ではなく、算術的に保証されています。

HF Bucketsとアーキテクチャ

Bucketとは何か？

Bucketは、高頻度なオブジェクトストレージのためのHugging Face Hub上のリポジトリタイプです。コミットの手順やPRワークフローを必要としません。ファイルは、アップロード用の batch_bucket_files とダウンロード用の download_bucket_files という2つの関数を介して追加、リスト表示、またはダウンロードされます。内部的には、コンテンツに基づいてチャンクを重複排除するHubのコンテンツ定義チャンキングストレージ層であるXetを使用しています。

3つのボックス

アーキテクチャは、3つのコンポーネントと1つの共有基盤で構成されています：

• Trainer: モデルの重みを所有し、オプティマイザを実行し、スパースなデルタ（sparse deltas）を放出します（単一のGPU、複数のGPU、またはノートPCなど、どこでも可能です）。
• HF Bucket: フルスナップショット用の anchors/ とスパースなパッチ用の deltas/ を持つ単一のリポジトリであり、両者が合意する唯一の要素です。
• vLLM rollout server: バケットからデータを取得し、デルタを適用し、ロールアウトを提供します（必ずしもトレーナーと同じ場所に位置している必要はありません）。
• Environment: HTTPまたは関数呼び出しを介してロールアウトサーバーに接続します。

トレーナーとロールアウトサーバーが重みデータを直接交換することはありません。両者が共有するのは、バケットの座標を含む極めて小さなPOSTリクエストのみです。すべてのバイト転送は、各サイドとバケットの間で並列に行われます。

プロトコル

ワイヤフォーマットとしてのSafetensors

ディスク上および転送用のフォーマットとしてsafetensorsを使用しています。バケットには2種類のファイルが存在します：

• Anchors: フルbf16重みを持つ通常のチェックポイント（Nステップごと、デフォルト N=10 に書き込み）。
• Deltas: 変更された各パラメータについて、要素インデックスのint32テンソルと、それらのインデックスにおける値のbf16テンソルを保存します。

メタデータは、ファイルがスパースかアンカーかを示し、受信者がそれに応じて処理を分岐できるようにします。

トレーナー側：オプティマイザフックからのブーリアンマスク

BF16ChangeDetector は、ステップの前後でbf16の重みをスナップショットするために、オプティマイザにpre-stepおよびpost-stepフックを登録します。変更された要素のブーリアンマスクは、これらのスナップショットを比較することで計算されます。Adamの統計量からマスクを予測する手法は再現率（recall）が低かった（約30%）ため、このグラウンドトゥルース（正解）に基づくアプローチが採用されています。

vLLM側：30行の拡張機能

--worker-extension-cls フラグを介してvLLMにプラグインする DeltaWeightTransferEngine を実装しています（フォークは不要です）。重みの更新を受け取ると：

1. バケットからデルタのsafetensorsファイルをダウンロードします。
2. アンカーの場合：すべてのテンソルをロードし、将来のデルタのためにスナップショットを作成します。
3. デルタの場合：変更された各パラメータについて、インデックスと値を取得し、ローカルのbf16スナップショットに適用し、再構成されたフルテンソルをvLLMの load_weights に渡します。

Spacesでの実運用

共有ネットワークのない、完全に分散されたトレーニングを実行しました：

• Trainer: 単一のGPUボックス。
• vLLM rollout server: 拡張機能がインストールされたHugging Face Space (Docker SDK, L4 GPU)。
• Wordle environment: 256の同時セッション容量を持つ2つ目のHugging Face Space (CPU)。
• Hub bucket: 重みのデルタとアンカーのためのセントラルリポジトリ。

セットアップには数回の hf CLI呼び出しが必要でした。vLLM SpaceのDockerfileは、delta-weight-sync ブランチからTRLをインストールし、ワーカー拡張クラスを設定します。トレーニングは、SpacesとバケットへのHTTPSアクセスがあればどこからでも開始できます。

これによって何が可能になるのか？

• クラスターなしでの非同期RLトレーニング: 単一のGPUトレーナーが、バケットを介して重みを移動させながら、Spacesをロールアウトと環境として使用できます。
• 無料のマルチレプリカ推論: 複数のvLLM Spacesが同じバケットからデータを取得します。Xetが保存されたチャンクを重複排除し、Hubのエッジキャッシュが繰り返されるダウンロードを安価に提供します。
• デバッグ可能なワイヤフォーマット: デルタはPythonの safe_open で検査可能なsafetensorsファイルです。
• フロンティアスケールへの道: Qwen3-0.6Bモデルの場合、1ステップあたりのペイロードは1.2 GBから20–35 MBに減少します。Llama-3.1-405Bモデル（bf16で810 GB）の場合、計算上、1ステップあたりのデルタは約6 GBとなり（フルの場合は810 GB）、推論の停止時間は約8秒（100 GB/sのNCCLを使用した場合）から数秒に短縮されます。1 GB/sの帯域幅を持つクラウド間通信では、フルブロードキャストには13分かかりますが、デルタなら6秒で済みます。

今後の課題

• 2つのCPU bf16スナップショット: トレーナーは変更検出用に1つ、ロールアウトサーバーはvLLMの load_weights 用にフルテンソルを再構成するために1つを保持します。後者は、vLLMがスパースな load_weights APIを獲得した際に削除される予定です。
• 固定されたアンカー間隔: 現在はNステップごとにアンカーを作成していますが、適応的なポリシー（累積ドリフトが閾値を超えたときにアンカーを作成する）により、コストを削減できる可能性があります。
• マルチノードFSDP2トレーナー: BF16ChangeDetector は単一プロセスのオプティマイザフック向けに構築されており、マルチノードFSDP2のサポートはまだ検証されていません。
• オプティマイザへのフック: Adamの統計量からマスクを予測することは、複雑な相互作用のため、依然として困難な課題です。
• 転送時の圧縮とのスタッキング: スパースなsafetensorsとチャンクごとのgzipは直交していますが、まだ組み合わされていません。

試してみる

• PR: huggingface/trl#5417 (branch: delta-weight-sync).
• 完全なWordleの例: examples/scripts/openenv/async_wordle.py.
• SpacesのDockerfile: examples/scripts/openenv/vllm_space/ および examples/scripts/openenv/wordle_space/.
• 背景知識: 私たちのasync RL landscape post、Fireworks 1 TB post、Cursor Composer 2 report。