Skip to content

feat: dynamo prime-rl integration - #3062

Open
biswapanda wants to merge 11 commits into
PrimeIntellect-ai:mainfrom
biswapanda:feat/dyn-pi-sidecar-v2
Open

feat: dynamo prime-rl integration#3062
biswapanda wants to merge 11 commits into
PrimeIntellect-ai:mainfrom
biswapanda:feat/dyn-pi-sidecar-v2

Conversation

@biswapanda

@biswapanda biswapanda commented Jul 16, 2026

Copy link
Copy Markdown

Summary

This PR has been broken into smaller PRs below:

  1. feat(inference): Dynamo integration with prime-rl #3176
  2. fix(inference): map external engine ranks #3177
  3. feat(dynamo): support native weight transfer #3178

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/workers API.

Companion PRs:

Exact tested heads:

Repository Commit
Prime-RL 7a8f90b2495e344f5b9198e20f418c472018b1ad
Dynamo fc556d992cd2525eaf6ae0f612f31a46db2e5a7b
vLLM e74fc3f1b06258e25519a5d5ed4d9a1b05d9a1cb

vLLM 0.26.0 alignment

The branch is based on Prime main commit 2ffe374e020a64b046f8fc3e7dcbb5278788d83e, which includes the upstream vLLM 0.26.0 dependency update in 9f0d41e23. Prime now imports the canonical vllm.entrypoints.scale_out.token_in_token_out package directly. The previous import-layout probing and compatibility shim have been removed.

PrimeRlServingTokens is a 75-line adapter over upstream vLLM 0.26.0 serving. It preserves Prime's two required behaviors: forwarding P/D kv_transfer_params and compacting routed-expert metadata. Validation, sampling, rendering, multimodal preprocessing, streaming, usage, logprobs, cancellation, and response metadata remain owned by vLLM.

Minimal configuration

[orchestrator.model.client]
base_url = ["http://dynamo-frontend:8000/v1"]
dynamo_discovery_url = "http://dynamo-frontend:8001"

[weight_broadcast]
type = "nccl"
host = "trainer-rank-zero"
inference_world_size = 8

Prime launches no local inference process for this topology, so deployment.num_infer_gpus may be zero and no [inference] block is required. Existing admin_base_url behavior remains unchanged for native-vLLM users and debugging.

Discovery and control lifecycle

  1. Prime polls <dynamo_discovery_url>/v1/rl/workers within wait_for_ready_timeout.
  2. Pydantic protocol models validate the complete snapshot and reject empty, errored, duplicate, or mixed-capability worker records.
  3. The sum of discovered engine world_size values must match the explicitly configured inference_world_size before collective initialization.
  4. Prime constructs its existing vLLM admin clients, computes cumulative per-engine rank offsets, and waits for both direct engine health and Dynamo frontend model publication.
  5. Pause, resume, full-weight updates, filesystem updates, and optional LoRA mutations use the existing vLLM/Dynamo control surfaces.

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 ranks 0..7.

Compatibility and cleanup

  • Existing non-Dynamo static and elastic inference paths retain their previous payloads and rank calculation. Engine topology fields are sent only when Dynamo supplies explicit per-engine sizes.
  • Dynamo topology validation is performed by configuration models after shared weight_broadcast resolution, not by runtime-only branches.
  • The pool API carries only the operations required by all inference pool implementations; the unused transport argument and one-use helper methods were removed.
  • Runtime debug/evidence branches, exact-log mocks, speculative import compatibility, and one-use private adapters were removed.
  • Topology-aware rank logic is isolated as a pure utility with focused cases for DP/TP/PP/PCP layouts.
  • No Helm, DGD, Kubernetes chart, openengine.v1, prime_rl.engine.v2, or VLLM_PLUGINS integration 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.index calculation.

The exact tested stack was built with DIND and deployed as q06v4i-0803 in bis-rl-3:

  • one CPU Dynamo frontend;
  • one aggregated inference GPU running vLLM-rs/Python EngineCore plus a separate Dynamo sidecar;
  • one external Prime trainer GPU;
  • Qwen3-0.6B Math for three orchestrator and trainer steps.

/v1/models and /v1/rl/workers returned HTTP 200, discovery reported world_size=1, all three policy versions were published and reloaded, and post-update generation returned HTTP 200. The trainer exited successfully with losses 0.0000, 0.0066, and 0.0069; inference generated 14,510 tokens at 172.21 tokens/s aggregate. The step-3 displayed error percentage came from intentional stale-rollout cancellation with max_off_policy_steps=0; dispatcher errored counters remained zero and the pipeline drained cleanly.

Images:

nvcr.io/nvidian/dynamo-dev/biswa@sha256:be98b7cbbea4f3053eb20ad0d0596beab6644f02adc4d01c94ec845603c225dc
nvcr.io/nvidian/dynamo-dev/biswa@sha256:6740b20164a7091b252046b98786fbf39a18902b463eeb5b480384de4e9f4cfb

The checkpoint reload path emitted a RotaryEmbedding: Failed to load weights warning, 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 f54fbb267 head. The current 7a8f90b24 head removes compat.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 static admin_base_url. A new DynamoInferencePool validates protocol-v1 snapshots, requires weight_broadcast.inference_world_size to 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 optional inference_world_size on 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 external inference_world_size is 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_rank for managed-DP layouts, layered checkpoint reload with empty-broadcast guards, and completion logging on workers. vLLM is bumped to 0.26.0 with a narrow compat.py import shim; PrimeRlServingTokens shrinks 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.

@biswapanda biswapanda changed the title feat: discover Dynamo inference workers feat: dynamo integration Jul 16, 2026
@biswapanda biswapanda changed the title feat: dynamo integration feat: dynamo prime-rl integration Jul 16, 2026
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
@biswapanda
biswapanda force-pushed the feat/dyn-pi-sidecar-v2 branch from 16db6e6 to ae95617 Compare July 17, 2026 19:56
Comment thread src/prime_rl/orchestrator/utils.py
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/utils/client.py Outdated
Comment thread src/prime_rl/orchestrator/utils.py Outdated
Comment thread packages/prime-rl-configs/src/prime_rl/configs/rl.py Outdated
Comment thread src/prime_rl/inference/vllm/routed_experts.py Outdated
Comment thread src/prime_rl/inference/vllm/compat.py Outdated
Comment on lines +1 to +7
"""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.
"""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same, I didn't follow here. Maybe would be cleaner to do a seperate for this vllm upstream so that dynamo change are clear

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good. I'm rebasing this and sibling PRs for vllm/dynamo/prime-rl on top of respective main.
I'll address it together

Comment thread src/prime_rl/utils/dynamo.py Outdated
Comment thread src/prime_rl/utils/dynamo.py Outdated
Comment thread src/prime_rl/utils/dynamo.py
Comment thread src/prime_rl/utils/dynamo.py
@biswapanda
biswapanda force-pushed the feat/dyn-pi-sidecar-v2 branch from 17f67eb to f54fbb2 Compare July 30, 2026 21:47

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

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

Comment thread packages/prime-rl-configs/src/prime_rl/configs/rl.py Outdated
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.

2 participants