This vLLM blog post explains weight updates in RL + KV cache recompute in a very clear and illustrated way, and it also covers the API for weight sync that was recently introduced in vLLM, so I think it’s a must-see, but within my own observations, it doesn’t seem to be getting that much attention…

For reference, I’ve also written a clear Japanese-language article about the first half, so please check it out by all means (I had planned to write about the second half as well, but I couldn’t find the time…) https://zenn.dev/kaz20/articles/e3c5dfc5111f1d…

vLLM Blog:

RLVR時代におけるInference Framework: Weight Syncing編

Source: https://zenn.dev/kaz20/articles/e3c5dfc5111f1d

はじめに

東京科学大学博士課程の藤井です。

本記事は、vLLMやSGLangなどのInference FrameworkがRL(Reinforcement Learning)全盛の現代(2026年)において、RLのためにどのような機能を有しているのか、またどのような意図/背景でその機能を拡張してきたのかについて解説を行うシリーズ記事の一環として執筆されています。今回は、Claude CodeやCodexなどのCoding Agentを実現するために必要なReasoning LLMの開発において非常に重要な役割を果たすRLVR(Reinforcement Learning with Verifiable Reward)において重要なWeight Syncingに焦点を当ててInference Framework側がどのような背景のもと、どのような機能をサポートしているのか解説を行います。

vLLM ロゴ

Weights Syncing

本記事では、主要なInference Engineの１つであるvLLMにおけるWeight Syncing機能の実装背景とvLLMが機能としてこれをサポートすることでどのような恩恵があるのかについて解説を行います。

Weight Syncing Background

On-policy / synchronous RLの設定では、Trainer側で更新されるpolicyとGenerator側でrolloutを生成するpolicyの乖離をできるだけ小さく保つ必要があります。これは、rollout generationとtrainingを非同期に進め、一定のpolicy lagやstale rolloutsを許容するasynchronous RL設定でも同様です。Generatorで利用されるmodel weightsを最新、または十分に直近のものに更新することで、現在のTrainer policyに対して極端に古いrolloutsが混ざることを避け、学習信号(training feedback)の有用性を保ちやすくします。

前提の説明が長くなりましたが、RLにおいてはweight syncが必要という訳です。そしてこのweight syncは、伝統的にInference Engineではなく、それぞれのRL Training Frameworkが別々に担ってきました。vLLMを例に取ると、多くの場合、vLLM workersがweightを受け取り、そのweightをloadするという拡張を行うことで実現されてきました。もちろん、そのようなvLLM拡張は機能しますが、以下のような問題も存在します。

• Training Framework側の複雑性の増加: Training Frameworkの開発者にvLLM/SGLangのcustom worker extensionの実装/メンテナンスコストを押し付けることになります。仮にWeight syncのための機能をInference Engine側が有していればこれを解決できます。
• Training Frameworkごとに重複する作業が発生: 多くのRL frameworkでpacked tensor transferやRPC endpointsなど似通った実装が採用されています。NeMo-RL、slime、verlなど有名なRL Frameworkだけでも似たような実装が重複しているのは非常に無駄が多いと言えます。
• 特定のversionにlockingされる: 多くのFrameworkで、特定のvLLM/SGLang versionに対してpatch処理としてweightsを受信(receiving)した後にpre/post-processingする処理をad-hocに実装しています。そのため、Inference Engine versionが異なると上手く動作しなくなるおそれがあり、利用者、開発者双方にコストが生じています。

以上の背景から、vLLMはweight syncingのためのAPIをnative supportするようになりました。次のsectionでどのようなAPIがサポートされたのか、そしてその概要について解説を行います。

Native weight transfer

New APIs in vLLM

vLLMがサポートしたweight syncingのためのAPIは主に以下の4つに分類されます。

1. initialization:init\_weight\_transfer\_engineTrainer workers(processes)とInference workersの間に通信チャンネル(communication channel)をTraining Loopを始める前に確立する役割を担います。
2. Start weight update:start\_weight\_updateTraining Loop内で各training step後、または複数stepごとに呼び出され、weight update処理の開始をvLLM workersに通知します。このAPIを呼び出すことでvLLM workersは新しいweightsを受け取る準備を行います。
3. Update weights:update\_weightsは、Trainerから受け取ったweightsでInference Engine側のweightsを更新するAPIです。更新対象はmodel全体でも、一部のweightsでも構いません。また、大きなmodelでは一度にすべてのweightsを転送せず、chunked weightsごとに複数回update\_weightsを呼び出すことで、段階的にweightsをloadできます。
4. Finish weight update:finish\_weight\_updateweight updateを終了するためのAPIです。Quantization(量子化)などの新たに受け取ったweightに基づき、post-processを行う必要がある場合は、それを行うための起点になります。

なお、これらのAPIはvLLMのAPI server mode及びengine modeのいずれでも実装されています。 また、weight syncing(weight transfer)におけるcommunication backendとして現状では以下の2つをサポートしています。

• NCCL: Trainer workersからのweight transferにNCCL broadcast operationsを利用するbackendです。Trainer workersとInference workersは一般に別のGPU上に存在するためNCCLによる通信を行うことでweight transferを実現しています。
• IPC: 同じdevice内のweight transferには、shared memory handlesを利用したCUDA IPCを利用してweight transferを行います。

以上のようなweight transferを実現するためのcore transport logicは、プラグイン可能な抽象化実装であるWeightTransferEngineの形で実装されています。(WeightTransferEngine classを継承してNCCLWeightTransferEngineやIPCWeightTransferEngineを実装するという形です)このようにworker実装(worker implementation)とweight transport(重み転送)の実装を分離することで、vLLMを利用するユーザーが独自の実装を簡単に組み込めるようにしています。

NCCL Weight Transfer Engine

IPC Weight Transfer Engine

Initialization(初期化)とUpdate Weights(重みの更新)のphaseは一般的にRL Frameworkの開発者によってカスタマイズされ、transport logicについても同様にその一環として利用されるという想定になっています。しかし、従来のようにすべてをRL Framework側が担当するのではなく、vLLMのAPIを利用して実現することでRL Framework間で重複していた実装を簡略化できる(=統一的なAPI利用に置き換えられる)ということです。

次のsectionでは、NCCLを利用したweight transferとFP8 Quantizationをpost processingで行う場合について解説します。

Weight transfer Example

vLLM Blog: Native RL APIs in vLLMより

上図を参考にvLLMのAPIをどのように呼ぶのかと照らし合わせながら解説を行います。

1. weight transfer engine の設定

以下のように通信で利用するbackendに今回利用するNCCLを指定して、Weight Transferの準備を行います。

from vllm import LLM
from vllm.config import WeightTransferConfig
 
llm = LLM(
    model="my-model",
    weight_transfer_config=WeightTransferConfig(backend="nccl"),
)

1. communication stateの初期化

上図では、/init\_weight\_transfer\_engineがそれに該当します。通信で利用するMASTER\_PORT、MASTER\_ADDRESS、NCCL process groupに参加するrank数を表すworld\_sizeを設定します。この設定により、Trainer workerとInference worker間のcommunication stateを初期化します。上図のinit NCCL Process Groupはまさに通信 Process Groupの初期化に該当します。

from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
 
# Initialization for inference
llm.init_weight_transfer_engine(
    WeightTransferInitRequest(      # <--- initialization parameters
        init_info=dict(
            master_address=master_address,
            master_port=master_port,
            rank_offset=1,          # <--- offset accounts for trainer rank 0
            world_size=world_size,  # <--- trainer + all inference workers
        )
    )
)
 
# Initialization for training
from vllm.distributed.weight_transfer.nccl_engine import (
    NCCLWeightTransferEngine,
)
 
group = NCCLWeightTransferEngine.trainer_init(
    dict(
        master_address=master_address,
        master_port=master_port,
        world_size=world_size,
    )
)

1. Trainer workerからweightを送信

以下でweight transferを開始します。trainer\_send\_weightsメソッドを利用してiterator=model\.named\_parameters\(\)で受け取ったparametersを送信します。なおNCCLTrainerSendWeightsArgsは、効率化のためにpacked tensor broadcastingを利用するかどうかなどオプションが指定可能です。

from vllm.distributed.weight_transfer.nccl_engine import (
    NCCLTrainerSendWeightsArgs,
    NCCLWeightTransferEngine,
)
 
trainer_args = NCCLTrainerSendWeightsArgs(
    group=group,
    packed=True,  # use packed broadcasting for efficiency
)
 
# send weights from an `AutoModelForCausalLM` instance
NCCLWeightTransferEngine.trainer_send_weights(
    iterator=model.named_parameters(),
    trainer_args=trainer_args,
)

1. Inference Engineにおいてweightを受信

以下では、update\_weightsが呼ばれてからfinish\_weight\_updateが呼ばれるまでの間でweightを受け取り、parameterの更新を行っています。また、実装では省略していますが、finish\_weight\_updateを呼ぶと、上図のようにQuantizationを実施している場合はFP8 Post-processなどが次に実施されます。

from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
 
# executed asynchronously while trainer sends weights
llm.start_weight_update()
llm.update_weights(
    WeightTransferUpdateRequest(
        update_info=dict(
            names=names,
            dtype_names=dtype_names,
            shapes=shapes,
            packed=True,
        )
    )
)
llm.finish_weight_update()

Customizing Weight Transfer

これまでWeight Transferがなぜ必要なのか、そしてvLLMがどのようにそれをサポートしているのかについて解説を行いました。最後に、custom weight transferを行う場合にどのように拡張を行うのかの概要について解説を行います。

Custom WeightTransferEngineを定義するには、WeightTransferInitInfoとWeightTransferUpdateInfoを継承した独自のdata classを作成する必要があります。なお、WeightTransferInitInfoおよびWeightTransferUpdateInfo自体はAbstract classであり実体はほとんどないため、実装を行う場合はNCCLWeightTransferUpdateInfo,NCCLWeightTransferInitInfoなどを参考にすることをオススメします。

以下の例では、MyInitInfoとMyUpdateInfoを作成し、これを利用してMyWeightTransferEngineを定義しています。weight transfer機能のさらに詳細についてはこちらのREADMEを参照してください。

from dataclasses import dataclass
from typing import Iterator, Callable, Any
from torch import Tensor
from vllm.distributed.weight_transfer.base import (
    WeightTransferEngine,
    WeightTransferInitInfo,
    WeightTransferUpdateInfo,
)
 
 
# define custom dataclasses for initialization and update weights metadata
@dataclass
class MyInitInfo(WeightTransferInitInfo):
    """Custom initialization info."""
    ...
 
 
@dataclass
class MyUpdateInfo(WeightTransferUpdateInfo):
    """Custom update info."""
    ...
 
# custom weight transfer engine
class MyWeightTransferEngine(WeightTransferEngine):
    init_info_cls = MyInitInfo
    update_info_cls = MyUpdateInfo
 
    def init_transfer_engine(self, init_info: MyInitInfo):
        ...
 
    def receive_weights(
        self,
        update_info: MyUpdateInfo,
        load_weights: Callable[[list[tuple[str, Tensor]]], None],
    ):
        ...
 
    @classmethod
    def trainer_send_weights(
        cls,
        iterator: Iterator[tuple[str, Tensor]],
        trainer_args: dict[str, Any] | Any,
    ):
        ...
 
# finally, register the weight transfer engine
from vllm.distributed.weight_transfer import WeightTransferEngineFactory
WeightTransferEngineFactory.register_engine("my_weight_transfer", MyWeightTransferEngine)

おわりに

本記事では、vLLMにおいてサポートがなされたWeight Syncing機能について解説を行いました。RLVRにおいてなぜWeight Syncingが必要なのかという説明から、どのようにWeight Syncingが実現されているのかまで通して解説を行いました。なお、Asynchronous RLにおけるvLLMのweight updateの機能については後続の別の記事にて解説を行う予定です。

私の他の記事では、LLM開発を行った際の知見や、学習フレームワークの裏でどのようなことが行われているのかについて解説を行っていますので、よろしければそちらもご覧ください。