feat: dynamo prime-rl integration - #3062
Conversation
16db6e6 to
ae95617
Compare
| """Narrow compatibility surface for vLLM's token-in/token-out API move. | ||
|
|
||
| vLLM moved these types from ``entrypoints.serve.disagg`` to | ||
| ``entrypoints.scale_out.token_in_token_out`` without a compatibility alias. | ||
| Keep the version branch in one module so Prime's serving and client code do not | ||
| grow parallel implementations. | ||
| """ |
There was a problem hiding this comment.
I didn't follow why we need this compat layers, is this related to a breaking change in upstream vllm that we used for dynamo ?
There was a problem hiding this comment.
yes there are upstream vLLM breaking changes and this adds the compatibility layer
| Everything else (request/response schema, sampling params, error handling) | ||
| delegates to upstream so we track future vLLM changes for free. | ||
| """ | ||
| """Small Prime extensions to vLLM's canonical token-in/token-out handler.""" |
There was a problem hiding this comment.
same, I didn't follow here. Maybe would be cleaner to do a seperate for this vllm upstream so that dynamo change are clear
There was a problem hiding this comment.
sounds good. I'm rebasing this and sibling PRs for vllm/dynamo/prime-rl on top of respective main.
I'll address it together
17f67eb to
f54fbb2
Compare
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit f54fbb2. Configure here.
…adcast and guard optional world size
…rop unused import
f54fbb2 to
7a8f90b
Compare

Summary
This PR has been broken into smaller PRs below:
This PR lets Prime-RL use an externally deployed Dynamo inference stack while Prime continues to own the trainer and orchestrator. Rollouts use Dynamo's OpenAI-compatible frontend, and Prime discovers each engine's existing vLLM HTTP control endpoint through the opt-in
/v1/rl/workersAPI.Companion PRs:
Exact tested heads:
7a8f90b2495e344f5b9198e20f418c472018b1adfc556d992cd2525eaf6ae0f612f31a46db2e5a7be74fc3f1b06258e25519a5d5ed4d9a1b05d9a1cbvLLM 0.26.0 alignment
The branch is based on Prime main commit
2ffe374e020a64b046f8fc3e7dcbb5278788d83e, which includes the upstream vLLM 0.26.0 dependency update in9f0d41e23. Prime now imports the canonicalvllm.entrypoints.scale_out.token_in_token_outpackage directly. The previous import-layout probing and compatibility shim have been removed.PrimeRlServingTokensis a 75-line adapter over upstream vLLM 0.26.0 serving. It preserves Prime's two required behaviors: forwarding P/Dkv_transfer_paramsand compacting routed-expert metadata. Validation, sampling, rendering, multimodal preprocessing, streaming, usage, logprobs, cancellation, and response metadata remain owned by vLLM.Minimal configuration
Prime launches no local inference process for this topology, so
deployment.num_infer_gpusmay be zero and no[inference]block is required. Existingadmin_base_urlbehavior remains unchanged for native-vLLM users and debugging.Discovery and control lifecycle
<dynamo_discovery_url>/v1/rl/workerswithinwait_for_ready_timeout.world_sizevalues must match the explicitly configuredinference_world_sizebefore collective initialization.Each discovery record represents one engine endpoint, not one GPU. For example, engine sizes
[2, 2, 2, 2]produce rank offsets[0, 2, 4, 6]and ranks0..7.Compatibility and cleanup
weight_broadcastresolution, not by runtime-only branches.openengine.v1,prime_rl.engine.v2, orVLLM_PLUGINSintegration is added.Examples
examples/dynamo/qwen3_06b_math/examples/dynamo/qwen3_30b_Thinking/examples/dynamo/glm52_fp8_r2e/Validation
The exact final ARM64 image passed 35 focused configuration, discovery, broadcaster, serving-adapter, and rank tests. Additional explicit compatibility gates confirmed that non-Dynamo requests omit Dynamo topology metadata and retain the original
rank_offset + device.indexcalculation.The exact tested stack was built with DIND and deployed as
q06v4i-0803inbis-rl-3:/v1/modelsand/v1/rl/workersreturned HTTP 200, discovery reportedworld_size=1, all three policy versions were published and reloaded, and post-update generation returned HTTP 200. The trainer exited successfully with losses0.0000,0.0066, and0.0069; inference generated 14,510 tokens at 172.21 tokens/s aggregate. The step-3 displayed error percentage came from intentional stale-rollout cancellation withmax_off_policy_steps=0; dispatcher errored counters remained zero and the pipeline drained cleanly.Images:
The checkpoint reload path emitted a
RotaryEmbedding: Failed to load weightswarning, but all three reload calls completed, policy versions advanced, and the post-update generation gate passed. This warning remains worth following up independently.Scope
Active NCCL/NIXL communicators are not elastic. Topology epochs, worker replacement, autoscaling, and communicator reconstruction remain follow-up work.
Important
The automated block below is retained from the previous
f54fbb267head. The current7a8f90b24head removescompat.py; the manually maintained sections above are canonical until the bot refreshes this block.Note
High Risk
Touches weight broadcast, NCCL rank assignment, and external inference discovery—errors can mis-size communicators or push updates to wrong ranks; vLLM 0.26 upgrade affects serving compatibility.
Overview
Prime can now train against an externally deployed Dynamo stack: rollouts use the OpenAI frontend while the orchestrator discovers per-engine admin endpoints from
dynamo_discovery_url(/v1/rl/workers) instead of staticadmin_base_url. A newDynamoInferencePoolvalidates protocol-v1 snapshots, requiresweight_broadcast.inference_world_sizeto match the sum of discovered ranks, and drives pause/resume, NCCL init, and weight updates via vLLM/collective_rpc(plus optional Dynamo LoRA system routes).Config and RL entrypoint gain
dynamo_discovery_url,is_dynamo, and optionalinference_world_sizeon NCCL/NIXL/filesystem broadcast types. Validation enforces Dynamo + explicit world size, NCCL for full-weight Dynamo (filesystem only with LoRA), and allows single-GPU trainers when externalinference_world_sizeis set.examples/dynamo/documents deployment overlays and adds Qwen smoke and GLM R2E TOMLs.NCCL path adds per-engine
engine_world_size/ cumulative rank offsets,global_inference_rankfor managed-DP layouts, layered checkpoint reload with empty-broadcast guards, and completion logging on workers. vLLM is bumped to 0.26.0 with a narrowcompat.pyimport shim;PrimeRlServingTokensshrinks to KV handoff plus compact routed-experts on top of upstream serving.Reviewed by Cursor Bugbot for commit f54fbb2. Bugbot is set up for automated code reviews on this repo. Configure here.