
2026/01/20 4:53
**RLのポスト・トレーニング時に、重みの転送を2 秒未満で行う**
RSS: https://news.ycombinator.com/rss
要約▶
Japanese Translation:
要約
Kimi‑K2 は、256 個の BF16 訓練 GPU から 128 個の FP8 推論 GPU に重みを転送することで、1 兆パラメータモデルに対して 1.3 秒 のクロスマシンパラメータ更新を実現します。転送は一方向の RDMA WRITE を使用しています。
更新ワークフローはコントローラーによって駆動され、以下の手順で構成されています:
- すべての訓練 GPU からパラメータメタデータを収集します。
- 各ソース GPU をターゲット推論 GPU にマッピングし、転送順序を定義する静的重み転送スケジュールを計算します。
- このスケジュールをすべての訓練 GPU に配布します。
- 各訓練ステップ後に GPU を通知して予定された転送を開始させます。
訓練パラメータは Fully Sharded Data Parallel (FSDP) でシャーディングされ、
full_tensor() は任意の GPU がフルパラメータを再構築できるようにします。複数のDeviceMesh グループが並列に動作し、グループ内ではポイント・ツー・ポイント通信、グローバル同期は Ethernet 上の GLOO を介して行われます。
転送タスクは 4 段階にパイプライン化されます:
- ホストからデバイスへの memcpy(オフロード時)。
- 投影融合 / 量子化。
- RDMA WRITE による宛先 GPU メモリへの書き込み。
- グローバルバーリア同期。
各段階のタスクはキューに入れられ重ね合わせて、ハードウェア資源を最大限活用します。GPU メモリ使用量は設定可能な ウォーターマーク により制御され、新しいタスクはオンザフライ一時メモリがこの上限未満のときのみ開始されます。これにより過負荷を防ぎつつネットワークを飽和状態に保ちます。
1.3 秒更新時間を実現する主な設計選択:
- ポイント・ツー・ポイント通信は、単一の rank‑0 ボトルネックではなくネットワーク全体を飽和させる。
- 一方向 RDMA は RPC/シリアライズオーバーヘッドを排除し、推論エンジンに変更を加える必要がない。
- 静的スケジュールはステップごとの制御プレーンレイテンシーを削減する。
- コンポーネントの明確な分離は、設計・テスト・最適化を容易にする。
結果として得られる 1.3 秒更新ウィンドウ は、大規模強化学習微調整で Kimi‑K2 上の 1 兆パラメータシステムを迅速かつ信頼性高く転送できる基盤となり、研究者や企業がダウンタイムを最小限に抑えながら生産性向上を図れるようにします。
本文
最近、Kimi‑K2(1 Tパラメータ)に対して 1.3秒でクロスマシンのパラメータ更新 を実現しました。これは256台のトレーニングGPU(BF16)から128台の推論GPU(FP8)へ重みを転送したものです。
非同期強化学習ファインチューニングでは、トレーニングと推論が別々のGPUで実行されます。各トレーニングステップ後に新しい重みを推論ノードへプッシュする必要がありますが、従来のフレームワークだと数秒、あるいは数分かかることもあります。
RDMA ポイント‑ツー‑ポイント通信 を活用することで、推論エンジンを変更せずに重み転送を爆速化し、コードを書きやすく保守もしやすくしています。
RDMA WRITE:一方通行転送
私たちのソリューションは
RDMA_WRITE に基づいています。これは一方向性のプリミティブで、ソースが直接宛先GPUメモリに書き込みます:
def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr): ...
宛先側は転送を検知しません。
これにより、トレーニングノードから推論ノードへ低遅延・高スループット・ゼロコピーの転送が可能になり、推論ノード側には制御ロジックが不要です。
高水準ワークフロー
- メタデータ収集 – コントローラが全トレーニングおよび推論GPUからパラメータメタデータを取得
- スケジュール計算 – スタティックな重み転送スケジュールを算出し、どのトレーニングGPUがどのパラメータをどの推論GPUへ何順で送るか決定
- スケジュール配布 – コントローラが全トレーニングGPUにスケジュールを送信
- 実行 – 各トレーニングステップ後、コントローラが転送開始をトレーニングGPUへシグナル
重み転送の実行
DeviceMesh と Mesh Groups
- トレーニング中のパラメータは FSDP 配置に従って分散
を使うと、DeviceMesh の全GPUが完全なパラメータを再構築できるため、任意の GPU が転送元になり得るfull_tensor()- 複数の disjoint DeviceMeshes は mesh group を形成。同一グループ内は分離されているので転送が干渉せず完全に並列で実行可能
- Mesh グループ間ではグローバルバリアを挿入し順序を保証
タスクパイプライン
各パラメータテンソルの転送をタスクとみなし、重み転送プロセスは複数種類のハードウェア源を利用するため、時間的に重なるように異なるパイプラインステージへ分割します:
| ステージ | 操作 |
|---|---|
| 1. Host → Device memcpy | FSDP が CPU にオフロードした場合 |
| 2. パラメータ準備 | で完全重みを再構築、プロジェクション融合適用、必要なら量子化 |
| 3. RDMA 転送 | リモート推論GPUメモリへゼロコピー書き込み |
| 4. グローバルバリア | 全 が完了した後、Ethernet 経由で GLOO を使って Mesh グループ間を同期 |
実装では各ステージの FIFO キューを保持し、キュー先頭が完了すると次のステージの末尾へ移動させます。
GPU メモリ使用量制御
full_tensor() やその他 GPU 操作は一時的なメモリ消費を増加させるため、タスク実行は現在のオンザフライタスクが占有する一時 GPU メモリが設定可能ウォーターマーク未満の場合にのみ開始します。
なぜ高速でシンプルなのか
-
ポイント‑ツー‑ポイント通信
- 従来は全パラメータを rank‑0 GPU 経由で集約→送信→再分散するため、PCIe バンド幅と NIC(≈400 Gbps ≈ 50 GB/s)に限定されるボトルネックが発生
- ポイント‑ツー‑ポイント構成は各トレーニング GPU が直接推論 GPU に送信し、ネットワーク全体をフルに活用
-
一方通行データ転送
- いくつかのシステムでは
を推論エンジンで呼び出すたびに RPC、シリアライズ、制御平面調整が発生update_weight()
により推論 GPU メモリを静かに更新(コピー不要、CPU 制御ロジック不要、推論エンジン変更不要)RDMA_WRITE
- いくつかのシステムでは
-
パイプライニング
- 重み転送はホスト‑デバイスのデータ移動、GPU でのプロジェクション融合/量子化計算、RDMA ネットワーク、Ethernet の制御平面という四つのハードウェアリソースを活用
- タスクをパイプラインステージに分割し、これらリソース間で時間的オーバーラップを実現
-
スタティックスケジュール
- 一部実装は各トレーニングステップごとに転送スケジュールを再計算し、余分な制御平面遅延を招く
- 当社のスケジュールは初期化時に一度だけ計算され、以降はそのプランを再生するだけで実行が予測可能かつ軽量
-
クリーンな分離
- メタデータ収集、ノード内集約、プロジェクション融合、量子化、ネットワーク転送などを個別コンポーネントに分割
- 各コンポーネントは単体テスト可能で、独立して最適化が行える
結論
高速かつ信頼性の高い重み転送は、大規模 RL ファインチューニングに不可欠な構成要素です。
RDMA_WRITE プリミティブ、スタティック転送スケジュール、およびパイプライン実行を組み合わせることで、Kimi‑K2 で 1.3 秒 のトランザクション時間にまで縮小しました。この手法は直感的に理解でき、保守が容易であり、従来設計のボトルネックを回避します。