Skip to content

feat(rollout): add VllmSamplerAdapter with Raiden weight sync support - #2012

Open
SurbhiJainUSC wants to merge 2 commits into
mainfrom
feat/vllm-sampler-adapter
Open

feat(rollout): add VllmSamplerAdapter with Raiden weight sync support#2012
SurbhiJainUSC wants to merge 2 commits into
mainfrom
feat/vllm-sampler-adapter

Conversation

@SurbhiJainUSC

Copy link
Copy Markdown
Collaborator

Summary

Adds VllmSamplerAdapter to integrate vLLM/TPU inference backends into Tunix distributed RL with native Raiden weight synchronization support.

Key Changes

  1. VllmSamplerAdapter (vllm_sampler_adapter.py):
    • Implements the Sampler protocol and WeightSyncDestination interface for tpu-inference / RLVllmSampler.
    • Integrates the complete Raiden weight-sync lifecycle (bind_weight_sync, get_weight_sync_metadata, pre_weight_sync, weight_sync, post_weight_sync, abort_weight_sync, get_weight_sync_status).
  2. RolloutManager (manager.py):
    • Instantiates VllmSamplerAdapter when sampler_type == "vllm".

DEPENDS ON #1990

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and all unit tests pass.
  • I have added all appropriate doc-strings/documentation.
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have signed the Contributor License Agreement.
  • I have followed Contribution Guidelines.

@SurbhiJainUSC SurbhiJainUSC changed the title Feat/vllm sampler adapter feat(rollout): add VllmSamplerAdapter with Raiden weight sync support Aug 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants