テンソルを飛翔させる：R-Forkが大規模モデルの重み読み込みを加速

TL;DR

私たちはTensor R-Fork（Tensor Remote Fork）を発表します。これは、高効率なノード間デバイス間インターコネクトを活用し、実行中のSGLangインスタンスから新しいインスタンスへゼロコピーでテンソルを読み込む新しい重み読み込み手法です。

この手法には3つの主要な利点があります：

• 重み読み込みパフォーマンスの大幅な向上
• ローカルディスク/DRAMでの冗長なモデル重みストレージの削除
• 推論サービスが干渉を受けないことの保証

Deepseek-R1モデルを例にとると、読み込み時間は数分からわずか数秒に短縮され、ローカルディスク/DRAMのストレージは約600GB節約され、同時に推論サービスの品質は安定したままです。

背景

LLMサービスの規模とモデル重みのサイズが継続的に拡大する中、SGLangインスタンスのコールドスタート時間は生産効率の重要なボトルネックとなっています。その中でも、重み読み込みは最も時間のかかる段階です。

Deepseek-R1を例にとると、ローカルディスクからの読み込みには数分かかり、リモートストレージからは数十分に達することもあります。モデルサイズが指数関数的に増大するにつれて、初期化とデータ転送時間はさらに悪化します。

重み読み込みを最適化するにはどうすればよいでしょうか？最も直接的な方法は、重みデータフローのボトルネック帯域幅を最大化することです。業界で一般的なモデル読み込み方法とそのボトルネックは以下の表のとおりです：

より高い帯域幅のデータフローを利用してテンソルを転送できるでしょうか？答えはYESです——ノード間デバイス間インターコネクトは数百GB/sのスループットを提供できます。しかし重要な問題は：SGLangでこのインターコネクト帯域幅を効率的な重み読み込みに十分活用するにはどうすればよいかということです。

このため、私たちは新しい重み読み込みフレームワークTensor R-Forkを開発し、Deepseek-R1の読み込み時間を数秒に短縮し、すでに本番環境で利用可能になっています。

設計

Tensor R-Forkの核心的な理念はGPU-Direct RDMAを利用してポイントツーポイント（P2P）重みストレージアーキテクチャを構築することです。

従来のデータ転送性能が低い理由は、パスに常にデバイス間インターコネクトよりもはるかに低い帯域幅のボトルネックが存在するからです。データフロー分析から、重みテンソルは各GPUに格納され、GPU-Direct RDMAを介して直接ノード間で転送できることがわかります。

RDMA NIC帯域幅の利用を最大化するため、私たちはGPUペアごとのデータ転送戦略を設計しました：ローカルGPUは、ペアになったリモートGPUとの間で直接データを転送します。この設計はGPUとCPU間のPCIeボトルネックを回避し、CPU/ホストメモリに依存しない高スループット通信を実現します。

リモートSGLangインスタンスから重みを読み込むフローは以下のとおりです：

実装

各実行インスタンスが同じモデルの新しいインスタンスの重みソースとして機能し、同時に実行インスタンスの推論サービスへの干渉を最小化（または排除）するため、私たちはNCCLとTransferEngineの2つのバックエンドを実装しました。実行インスタンスA（ソースインスタンス）と新しいインスタンスB（ターゲットインスタンス）を例に、重み転送メカニズムを詳しく説明します。

NCCLバックエンド

NCCLをバックエンド[1]として使用し、プロセスは2つの段階に分かれます：

1. ソースとターゲットインスタンス間で通信グループを確立。
2. 通信グループを介して重みを転送。

ターゲットインスタンスの初期化時に、指定されたソースインスタンスにHTTPリクエストを送信して通信グループの作成を開始します。各ターゲットTPWorkerは対応するソースTPWorkerとNCCL通信グループを確立し（ソースランク0とターゲットランク0がペアになるなど）、各グループにはメンバーが2つだけです。

通信グループが確立されると、各ソースTPWorkerはNCCLブロードキャストを介してGPUメモリの重みテンソルをブロードキャストし、ターゲットTPWorkerは中間コピーなしで直接そのGPUメモリに受信します。

NCCLはGPU-Direct RDMAを利用していますが、重要な制限があります：転送がソースインスタンスの推論サービスに干渉することです。理由は：

1. 通信グループの確立：ソースインスタンスが能動的に参加する必要があります。
2. CUDAカーネルの干渉：NCCLブロードキャストがCUDAカーネルをトリガーし、GPUリソースを競合して生成タスクの遅延ピークを引き起こします。

TransferEngineバックエンド

干渉のない転送を実現するため、私たちはTransferEngineバックエンド[2]を導入しました。これはRDMAベースの軽量転送ランタイムで、ソースインスタンスの各TPWorkerと並行して実行され、ソース側のCUDAカーネルを必要とせずに、GPU常駐重みテンソルをリモートリーダーに公開します。

ソースSGLangインスタンスの初期化時：

1. 各TPWorkerがTransferEngineインスタンスを起動。
2. TransferEngineが重みGPUメモリアドレスをRDMAチャネルに登録。

ターゲットインスタンスの初期化時：

1. HTTPリクエストを送信してソースTransferEngineメタデータを取得（GPUメモリアドレスにマップされたRDMAキーを含む）。
2. RDMAキーを利用して、ソースサービスを中断することなく、ソースGPUメモリから直接重みを読み込み。

TransferEngineについてもっと知りたいですか？付録をご覧ください 🚀

NCCL vs. TransferEngine

使用方法

詳細はR-Forkドキュメントを参照してください

NCCLバックエンドの使用

シードインスタンス：

python -m sglang.launch_server [args]

クライアントインスタンス：

python -m sglang.launch_server [args] \
  --load-format remote_instance  \
  --remote-instance-weight-loader-seed-instance-ip [seed_instance_ip] \
  --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \
  --remote-instance-weight-loader-send-weights-group-ports [send_weights_nccl_group_ports_list]  \
  --remote-instance-weight-loader-backend nccl

TransferEngineバックエンドの使用

シードインスタンス：

python -m sglang.launch_server [args] \
  --remote-instance-weight-loader-start-seed-via-transfer-engine

クライアントインスタンス：

python -m sglang.launch_server [args] \
  --load-format remote_instance  \
  --remote-instance-weight-loader-seed-instance-ip [seed_instance_ip] \
  --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \
  --remote-instance-weight-loader-backend transfer_engine

パフォーマンス

私たちは、8枚のNVIDIA H20 GPUを搭載した新しいSGLangインスタンスで、異なるソースからDeepSeek-R1モデルを読み込むパフォーマンスをテストしました。

メモリ領域の登録は他の初期化段階と重複でき、総起動時間をさらに最適化できます。

産業実践

前述の手動でのシードインスタンス設定はテストには適していますが、産業デプロイメントでは利用可能なシードインスタンスの識別に大量の運用オーバーヘッドが必要です。

この課題を解決するため、私たちはTenソリューションを提案しています（原文はここまで）。