feat: Raiden weight sync for GRPO RL training (VllmSamplerAdapter, RaidenSynchronizer, WeightSyncCoordinator) - #1960
Closed
khatwanimohit wants to merge 1 commit into
Closed
feat: Raiden weight sync for GRPO RL training (VllmSamplerAdapter, RaidenSynchronizer, WeightSyncCoordinator)#1960khatwanimohit wants to merge 1 commit into
khatwanimohit wants to merge 1 commit into
Conversation
…idenSynchronizer, WeightSyncCoordinator) Adds a full weight-sync protocol stack so a trainer and a separate rollout/sampler process can exchange policy weights over Raiden (a native TPU-to-TPU weight transfer library) mid-GRPO-run, plus a distributed orchestrator/trainer/rollout example wiring it end to end and a k8s_launcher.sh for running it on GKE. Key pieces: - tunix/experimental/rollout/vllm_sampler_adapter.py (new): VllmSamplerAdapter, the WeightSyncDestination protocol implementer wrapping tpu-inference's RLVllmSampler -- bind_weight_sync/get_weight_sync_metadata/pre_weight_sync/ weight_sync/post_weight_sync/abort_weight_sync, with a WorkerRoundTracker for idempotent phase admission across retried RPCs. - tunix/experimental/worker/raiden_synchronizer.py: RaidenSynchronizer, the shared (trainer- and rollout-side) binding of one process's weights to the Raiden transport -- flattens an nnx.State to (names, arrays), filters to what the native layer can actually bind (RNG-key and non-TPU-resident leaves are undefined behavior to bind, observed as SIGSEGV/RuntimeError), and exposes d2h/h2d/checksums/metrics. - tunix/experimental/orchestrator/weight_sync_coordinator.py, weight_sync.py, raiden_handler.py: the round protocol (bind -> manifest preflight -> quiesce/pre -> transfer -> post/commit, with rollback on any phase failure) and the transport-neutral WorkUnitId/TensorMetadata/ WorkUnitMetadata types RaidenHandler validates registrations against. - tunix/experimental/examples/math_gsm8k_dist/: a 3-process (orchestrator + trainer + rollout) distributed GRPO example over GSM8K, with run_trainer_node.py supporting both Tunix's own PeftTrainer and (via --trainer_backend=maxtext) MaxText's MaxTextTrainingEngine, and run_rollout_node.py loading either tpu-inference's own JAX model reimplementation or (via --maxtext_model_name) MaxText's native model through maxtext_vllm_adapter, so trainer and rollout can share one param tree for Raiden's exact name-matching sync. - k8s_launcher.sh: deploys the three processes as GKE JobSets, each pulling its own repo's feature branch and installing it at pod startup. - Dockerfile.maxtext, requirements/maxtext_requirements.txt, scripts/install_tunix_vllm_requirement.sh, pyproject.toml: build/install plumbing for the MaxText + Raiden + tpu-inference dependency stack. Verified live end-to-end: a 2-step GRPO run on a real 2-slice v5p-8 GKE cluster (MaxTextTrainingEngine trainer + tpu-inference RLVllmSampler rollout, via maxtext_vllm_adapter's native MaxText model) completed both steps with real Raiden weight-sync rounds, checksum-verified between source (trainer) and destination (rollout). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2 tasks
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
2 tasks
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
2 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 Raiden weight-sync support to the GRPO RL example under
tunix/experimental/examples/math_gsm8k_dist/— aVllmSamplerAdapter/RaidenSynchronizer/WeightSyncCoordinatorprotocol stack so a trainer and a separate rollout/sampler process can exchange policy weights over Raiden (a native TPU-to-TPU weight transfer library) mid-GRPO-run, plusk8s_launcher.shfor running it on GKE.Companion PRs (same feature, other two repos):
Test plan
pytest tests/experimental/rollout/vllm_sampler_adapter_test.py tests/experimental/orchestrator/weight_sync_coordinator_test.py(79 passed, 1 pre-existing unrelated failure)k8s_launcher.sh, 2 GRPO steps completed with real Raiden weight-sync rounds, checksum-verified between trainer (source) and rollout (destination)🤖 Generated with Claude Code