VLLMDoubleBufferSyncScheme

class torchrl.weight_update.llm.VLLMDoubleBufferSyncScheme(remote_addr: str, local_addr: str | None = None, num_threads: int = 1, strategy: Literal['tensordict', 'state_dict'] = 'tensordict')

Weight synchronization scheme for vLLM using double-buffered storage.

This scheme uses memory-mapped TensorDict storage to transfer weights from a trainer to vLLM inference workers. It’s simpler than NCCL-based approaches and doesn’t require process group coordination.

Parameters:

• remote_addr – Directory path where sender writes weights.

• local_addr – Directory path where receiver reads weights. If None, uses same path as remote_addr (for local testing).

• num_threads – Number of threads for memmap operations. Defaults to 1.

• strategy – Weight extraction strategy (“tensordict” or “state_dict”).

Example

>>> # Local testing (same machine)
>>> scheme = VLLMDoubleBufferSyncScheme(
...     remote_addr="/tmp/weights",
...     strategy="tensordict"
... )
>>>
>>> # Distributed setup (different machines)
>>> # On trainer node:
>>> scheme = VLLMDoubleBufferSyncScheme(
...     remote_addr="/mnt/shared/weights",  # NFS mount
...     num_threads=4
... )
>>>
>>> # On vLLM worker node:
>>> scheme = VLLMDoubleBufferSyncScheme(
...     remote_addr="/mnt/shared/weights",  # Same NFS mount
...     num_threads=4
... )

create_receiver(vllm_engine) → VLLMDoubleBufferWeightReceiver

Create a weight receiver for a vLLM worker process.

Parameters:

vllm_engine – The vLLM engine instance (must have .llm_engine.model_executor attribute).

create_sender() → VLLMDoubleBufferWeightSender

Create a weight sender for the trainer process.

create_transport(pipe_or_context: Any = None) → VLLMDoubleBufferTransport

Create transport for double-buffered storage.

Parameters:

pipe_or_context – Not used for file-based transport (kept for API compatibility).

Returns:

A VLLMDoubleBufferTransport instance.