feat(rollout): add VllmSamplerAdapter with Raiden weight sync support - #2012
Open
SurbhiJainUSC wants to merge 2 commits into
Open
feat(rollout): add VllmSamplerAdapter with Raiden weight sync support#2012SurbhiJainUSC wants to merge 2 commits into
SurbhiJainUSC wants to merge 2 commits into
Conversation
SurbhiJainUSC
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
August 27, 2026 19:43
SurbhiJainUSC
changed the title
wang2yn84
reviewed
Aug 28, 2026
|
|
||
| _ws_lib: Any = None | ||
| try: | ||
| from tpu_sync.api.jax import weight_synchronizer as _ws_lib # pytype: disable=import-error pylint: disable=g-import-not-at-top |
Collaborator
There was a problem hiding this comment.
This might break on g3, please fix in your patch cl.
Comment on lines
+77
to
+78
| del leaf, arr | ||
| # Periodically run GC to release Pathways proxy transit buffers incrementally |
Collaborator
There was a problem hiding this comment.
gc is not cheap, please use "jax.tree.map(lambda x: x.delete(), new_weights)
", also can you add TODO to remove it in the FFI path when FFI is enabled?
| from tunix.experimental.weight_sync import weight_sync | ||
|
|
||
|
|
||
| def _log_rss(tag: str) -> None: |
Collaborator
There was a problem hiding this comment.
Make this logging optional
| Attributes: | ||
| sampler_type: Type of sampler adapter to construct ("vanilla", | ||
| "inprocess_vllm", "vllm"). | ||
| weight_sync_mode: Mode of weight synchronization ("default", "fallback", |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
VllmSamplerAdapterto integrate vLLM/TPU inference backends into Tunix distributed RL with native Raiden weight synchronization support.Key Changes
DEPENDS ON #1990
Checklist