Skip to content

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

Merged
copybara-service[bot] merged 4 commits into
mainfrom
feat/vllm-sampler-adapter
Sep 1, 2026
Merged

feat(rollout): add VllmSamplerAdapter with Raiden weight sync support#2012
copybara-service[bot] merged 4 commits into
mainfrom
feat/vllm-sampler-adapter

Conversation

@SurbhiJainUSC

@SurbhiJainUSC SurbhiJainUSC commented Aug 27, 2026 •
edited
Loading

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".

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
wang2yn84
wang2yn84 reviewed Aug 28, 2026
View reviewed changes
Comment thread tunix/experimental/weight_sync/raiden_synchronizer.py
Comment thread tunix/experimental/weight_sync/raiden_synchronizer.py Outdated
Comment thread tunix/experimental/weight_sync/raiden_synchronizer.py
Comment thread tunix/experimental/worker/rollout_worker.py Outdated
@khatwanimohit
khatwanimohit force-pushed the feat/vllm-sampler-adapter branch from 88698e9 to e317910 Compare August 31, 2026 15:52
@khatwanimohit
khatwanimohit force-pushed the feat/vllm-sampler-adapter branch from e317910 to 11a2c5f Compare August 31, 2026 17:22
tianshub
tianshub reviewed Aug 31, 2026
View reviewed changes
Comment thread tunix/experimental/worker/rollout_worker.py Outdated
@khatwanimohit
khatwanimohit force-pushed the feat/vllm-sampler-adapter branch from 20bac7f to 6e4e701 Compare August 31, 2026 20:49
@khatwanimohit
khatwanimohit force-pushed the feat/vllm-sampler-adapter branch from 6e4e701 to 06e0682 Compare August 31, 2026 21:06
wang2yn84
wang2yn84 reviewed Aug 31, 2026
View reviewed changes
Comment thread tunix/experimental/worker/rollout_worker.py Outdated
Comment thread tunix/experimental/worker/rollout_worker.py Outdated
Comment thread tunix/experimental/rollout/vllm_sampler_adapter.py Outdated
Comment thread tunix/experimental/rollout/vllm_sampler_adapter.py Outdated
Comment thread tunix/experimental/rollout/vllm_sampler_adapter.py Outdated
Comment thread tunix/experimental/rollout/vllm_sampler_adapter.py Outdated
@khatwanimohit
fix(rollout): align VllmSamplerAdapter signatures and import for pytype
33a1af7 
Resolve RLVllmSampler via importlib so static analyzers do not follow the
optional tpu-inference dependency. Match the in-process adapter's
bind_weight_sync/get_weight_sync_metadata signatures and return the bind
result, and drop the lazy initialize() in start().
@copybara-service
copybara-service Bot merged commit d2be75a into main Sep 1, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@khatwanimohit khatwanimohit khatwanimohit left review comments
@wang2yn84 wang2yn84 wang2yn84 approved these changes
@lc5211 lc5211 Awaiting requested review from lc5211 lc5211 is a code owner
@hgao327 hgao327 Awaiting requested review from hgao327 hgao327 is a code owner
@sizhit2 sizhit2 Awaiting requested review from sizhit2 sizhit2 is a code owner
@abheesht17 abheesht17 Awaiting requested review from abheesht17 abheesht17 is a code owner
@jiangyangmu jiangyangmu Awaiting requested review from jiangyangmu jiangyangmu is a code owner
@s-noghabi s-noghabi Awaiting requested review from s-noghabi s-noghabi is a code owner
@tianshub tianshub Awaiting requested review from tianshub tianshub is a code owner

Assignees

@lc5211 lc5211

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@SurbhiJainUSC @tianshub @wang2yn84 @khatwanimohit @lc5211