diff --git a/examples/dynamo/README.md b/examples/dynamo/README.md new file mode 100644 index 0000000000..067945e157 --- /dev/null +++ b/examples/dynamo/README.md @@ -0,0 +1,80 @@ +# Dynamo native-gRPC deployment requirements + +The recipes in this directory use Dynamo for inference and Prime only for the +trainer/orchestrator. Start from Dynamo's `sidecar_agg.yaml` or +`sidecar_disagg.yaml`, then apply the following RL overlay. The stock manifests +are serving examples and do not enable Prime worker discovery or vLLM's admin +control routes by themselves. + +## Frontend + +Set these variables on the Dynamo frontend and expose both container ports: + +```yaml +env: + - name: DYN_ENABLE_RL + value: "true" + - name: DYN_RL_PORT + value: "8001" +ports: + - name: http + containerPort: 8000 + - name: rl-discovery + containerPort: 8001 +``` + +The frontend Kubernetes Service must also map ports 8000 and 8001. Prime's +`base_url` targets 8000; `dynamo_discovery_url` targets 8001. + +## Every vLLM engine and sidecar pair + +The engine HTTP address published by discovery must be reachable from the +trainer, so bind vLLM to the pod network rather than loopback: + +```text +vllm-rs serve --host 0.0.0.0 --port 8000 --grpc-port 50051 -- \ + --worker-extension-cls prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker \ + +``` + +Install the matching Prime source in the engine image so Python can import the +worker extension. Set this environment variable on the `vllm-engine` +container to expose `/pause`, `/resume`, and `/collective_rpc`: + +```yaml +env: + - name: VLLM_SERVER_DEV_MODE + value: "1" +``` + +Set the following variables on each `dynamo-vllm-sidecar` container. `POD_IP` +must appear before `VLLM_HTTP_ENDPOINT` so Kubernetes expands it: + +```yaml +env: + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: VLLM_HTTP_ENDPOINT + value: "http://$(POD_IP):8000" + - name: DYN_ENABLE_RL + value: "true" +``` + +Keep the existing `--grpc-endpoint 127.0.0.1:50051`: gRPC stays pod-local, +while discovery publishes the pod-reachable HTTP admin address. No per-worker +Kubernetes Service is required when trainer-to-pod networking is routable. + +After startup, `/v1/rl/workers` must return every expected worker with a +non-null `admin_base_url`, positive `world_size`, and no `error` before Prime is +launched. + +## Recipes + +- [`qwen3_06b_math`](qwen3_06b_math): single-GPU trainer and aggregate Dynamo + inference smoke test. +- [`qwen3_30b_Thinking`](qwen3_30b_Thinking): Qwen3-30B Thinking math with an + external prefill/decode deployment. +- [`glm52_fp8_r2e`](glm52_fp8_r2e): multi-node GLM-5.2 FP8 R2E training with a + separately managed DGD. diff --git a/examples/dynamo/glm52_fp8_r2e/README.md b/examples/dynamo/glm52_fp8_r2e/README.md new file mode 100644 index 0000000000..41932321ef --- /dev/null +++ b/examples/dynamo/glm52_fp8_r2e/README.md @@ -0,0 +1,128 @@ +# GLM-5.2 FP8 R2E with external Dynamo inference + +This three-step smoke recipe runs a distributed Prime trainer and orchestrator +against a separately managed Dynamo deployment serving +`zai-org/GLM-5.2-FP8`. Dynamo owns the frontend, vLLM engines, and one native +gRPC sidecar per engine group; Prime discovers the mutable engine control +surface through `/v1/rl/workers`. + +The recipe is split into trainer and orchestrator files because the external +Dynamo DGD and the multi-node trainer have independent lifecycles. It does not +add a Prime inference configuration or require Prime's launcher to manage the +DGD. + +## Reference topology + +| Component | Shape | GPUs | +|---|---|---:| +| Dynamo prefill | 2 nodes, DP4 x TP2 x PP1 x EP8 | 8 | +| Dynamo decode | 2 nodes, DP4 x TP2 x PP1 x EP8 | 8 | +| Prime trainer | 4 nodes, FSDP16 x CP4 x EP8 | 16 | + +The checked-in configuration targets this topology; it is not a claim that +every cluster can use these parallelism dimensions unchanged. The two +discovery records must report `prefill` and `backend` components with a +combined `world_size` of 16. If the DGD topology changes, update +`weight_broadcast.inference_world_size` in both TOML files to the atomic sum +returned by the same `/v1/rl/workers` response. + +The inference engines must load +`prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker`, expose vLLM's +admin routes, and run the version-matched `vllm-rs` and +`dynamo-vllm-sidecar` binaries described in [`../README.md`](../README.md). +For mutable GLM weight reloads, launch every vLLM rank with `--enforce-eager`. +When serving a snapshot path, set `--served-model-name +zai-org/GLM-5.2-FP8`. The tested GLM entrypoint also uses the `glm47` tool +parser, `glm45` reasoning parser, the model chat template, and complementary +NIXL `kv_producer`/`kv_consumer` roles for prefill/decode. +The trainer and every inference engine must also use a compatible NCCL +transport. Apply cluster-specific settings such as `NCCL_IB_DISABLE`, +`NCCL_SOCKET_IFNAME`, and the NCCL network plugin consistently on both sides; +do not force Socket on the trainer while allowing inference to select IB. +The filesystem rollout transport requires the orchestrator and all trainer +nodes to mount the same read-write shared output root. + +## Configure + +Initialize the R2E environment submodule and install its workspace package: + +```bash +git submodule update --init -- deps/research-environments +uv sync --package prime-rl --package r2e-gym-v1 +``` + +The taskset intentionally uses the current `r2e-gym-v1` default, +`PrimeIntellect/R2E-Gym-Subset-Verified`, instead of pinning an older dataset +override in this recipe. + +Replace the checked-in service names when the DGD and trainer use different +DNS names: + +- `model.client.base_url`: Dynamo OpenAI frontend on port 8000; +- `model.client.dynamo_discovery_url`: Dynamo RL discovery on port 8001; +- `weight_broadcast.host`: trainer rank zero, reachable from every inference + engine. + +The R2E harness uses Prime sandboxes. Configure the normal Prime credentials, +or replace the runtime with the sandbox backend used by your cluster. Model +and dataset caches should be shared across the trainer and inference nodes. + +Multi-turn affinity is not enabled merely by sending a session header. Start +the Dynamo frontend with `--router-session-affinity-ttl-secs ` (or +`DYN_ROUTER_SESSION_AFFINITY_TTL_SECS`) and choose an idle TTL longer than the +longest expected R2E turn gap. Prime maps each trajectory ID to the canonical +`X-Dynamo-Session-ID` header in `orchestrator.toml`. + +Verify both the model and the complete atomic worker snapshot before starting +Prime. A successful HTTP status alone is insufficient: + +```bash +MODEL=zai-org/GLM-5.2-FP8 +curl -fsS http://dynamo-frontend:8000/v1/models | + jq -e --arg model "$MODEL" '.data | any(.id == $model)' +curl -fsS http://dynamo-frontend:8001/v1/rl/workers | + jq -e --arg model "$MODEL" ' + .protocol_version == 1 and + (.workers | length == 2) and + (all(.workers[]; .model == $model and + ((.error // "") == "") and + (.instance_id != null) and + ((.admin_base_url // "") != ""))) and + ([.workers[].instance_id] | unique | length == 2) and + ([.workers[].admin_base_url] | unique | length == 2) and + ([.workers[] | select(.model == $model) | .component] | sort == ["backend", "prefill"]) and + ([.workers[] | select(.model == $model) | .world_size] | add == 16) + ' +``` + +## Run + +Launch the trainer on four 4-GPU nodes with the cluster's distributed runner. +For example, rank zero's rendezvous address can be passed to `torchrun` while +all ranks consume the same trainer file: + +```bash +uv run torchrun \ + --nnodes=4 --nproc-per-node=4 \ + --rdzv-backend=c10d --rdzv-endpoint="$TRAINER_RANK_ZERO:29501" \ + --node-rank="$NODE_RANK" \ + -m prime_rl.trainer.rl.train \ + @ examples/dynamo/glm52_fp8_r2e/trainer.toml \ + --output-dir /shared/glm52-dynamo-r2e/train +``` + +After trainer rank zero opens port 29500, launch the orchestrator once: + +```bash +uv run orchestrator \ + @ examples/dynamo/glm52_fp8_r2e/orchestrator.toml \ + --output-dir /shared/glm52-dynamo-r2e/train/run_0 +``` + +The gate succeeds when the first optimizer step completes, policy version 1 +settles on all 16 inference ranks through NCCL, and a later multi-turn rollout +completes without changing its Dynamo session assignment. Three steps are +required because finite NCCL runs skip broadcasts once +`step >= max_steps - 1`; this leaves step 1 as the first non-final broadcast +slot. The disabled post-batch zero-advantage filter keeps this small smoke run +from stalling on a homogeneous batch; enable it for a production training run. diff --git a/examples/dynamo/glm52_fp8_r2e/orchestrator.toml b/examples/dynamo/glm52_fp8_r2e/orchestrator.toml new file mode 100644 index 0000000000..74a0c0cf60 --- /dev/null +++ b/examples/dynamo/glm52_fp8_r2e/orchestrator.toml @@ -0,0 +1,60 @@ +max_steps = 3 +batch_size = 2 +group_size = 2 +seq_len = 32768 +max_inflight_episodes = 2 +max_off_policy_steps = 1 +tasks_per_minute = 1 + +[model] +name = "zai-org/GLM-5.2-FP8" + +[model.client] +base_url = ["http://dynamo-frontend:8000/v1"] +dynamo_discovery_url = "http://dynamo-frontend:8001" +wait_for_ready_timeout = 7200 + +[model.client.extra_headers_from_state] +X-Dynamo-Session-ID = "trajectory_id" + +[tokenizer] +name = "zai-org/GLM-5.2-FP8" + +[renderer] +name = "glm-5.1" +clear_thinking = false + +[train.sampling] +temperature = 1.0 +max_completion_tokens = 2048 +extra_body = { chat_template_kwargs = { clear_thinking = false } } + +[[train.source]] +name = "r2e" +group_size = 2 +serve.pool = { type = "static", num_workers = 2 } +env.taskset = { id = "r2e-gym-v1" } +env.agent.max_turns = 64 +env.agent.max_input_tokens = 30720 +env.agent.max_output_tokens = 16384 +env.agent.max_total_tokens = 32768 +env.agent.timeout = { setup = 600, rollout = 1800, finalize = 300, scoring = 600 } +env.agent.harness = { id = "bash", edit = true } +env.agent.runtime = { type = "prime", labels = ["glm52-dynamo"], cpu = 4, creates_per_min = 2 } + +[[post_batch_filters]] +type = "zero_advantage" +enforce = false + +[weight_broadcast] +type = "nccl" +host = "trainer-0.trainer-headless" +port = 29500 +timeout = 12000 +inference_world_size = 16 + +[rollout_transport] +type = "filesystem" + +[log] +level = "debug" diff --git a/examples/dynamo/glm52_fp8_r2e/trainer.toml b/examples/dynamo/glm52_fp8_r2e/trainer.toml new file mode 100644 index 0000000000..d7476a734d --- /dev/null +++ b/examples/dynamo/glm52_fp8_r2e/trainer.toml @@ -0,0 +1,46 @@ +max_steps = 3 +dist_timeout_seconds = 12000 + +[model] +name = "zai-org/GLM-5.2-FP8" +seq_len = 32768 +impl = "custom" +attn = "flash_attention_2" +dp_replicate = 1 +cp = 4 +ep = 8 +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" +moe_router_dtype = "float32" +optim_cpu_offload = true +fused_lm_head_token_chunk_size = 1024 + +[model.ac] +freq = 1 + +[model.ac_offloading] +max_inflight_activations = 1 + +[tokenizer] +name = "zai-org/GLM-5.2-FP8" + +[optim] +type = "sign_sgd" +lr = 1e-6 +weight_decay = 0.0 + +[scheduler] +type = "constant" + +[weight_broadcast] +type = "nccl" +host = "0.0.0.0" +port = 29500 +timeout = 12000 +inference_world_size = 16 + +[rollout_transport] +type = "filesystem" + +[log] +level = "debug" diff --git a/examples/dynamo/qwen3_06b_math/README.md b/examples/dynamo/qwen3_06b_math/README.md new file mode 100644 index 0000000000..dcfae8c576 --- /dev/null +++ b/examples/dynamo/qwen3_06b_math/README.md @@ -0,0 +1,45 @@ +# Qwen3 0.6B math with external Dynamo inference + +This four-step smoke recipe runs the Prime trainer and orchestrator locally while generation is served by an already-running Dynamo frontend, vLLM sidecar, and vLLM engine. Prime does not launch local inference, so the configuration intentionally has no `[inference]` block. + +## Prerequisites + +Use this version-matched native-gRPC source set: + +- Dynamo `feat/dyn-pi-sidecar-v2-review-001` at `836fe81012` +- vLLM `feat/dyn-pi-sidecar-v2-review-001` at `e56ee21b2c` +- Prime `feat/dyn-pi-sidecar-v2-review-001` + +Build `vllm-rs` and `dynamo-vllm-sidecar` from those revisions into the same +runtime image. For Kubernetes, start from Dynamo's +`examples/backends/vllm/deploy/sidecar_agg.yaml`; its adjacent `README.md` +documents the paired-binary image. Then apply the required Prime RL discovery +and admin overlay in [`../README.md`](../README.md). This contract requires both +native gRPC and the Dynamo `/v1/rl/workers` endpoint; a standard Python-only +vLLM worker is not compatible. + +Install the math environment: + +```bash +prime env install primeintellect/math-env +``` + +Start an aggregated DP1 Dynamo deployment for `Qwen/Qwen3-0.6B`. The orchestrator waits for both model publication and worker discovery. These requests are useful diagnostics: + +```bash +curl http://127.0.0.1:8000/v1/models +curl http://127.0.0.1:8001/v1/rl/workers +``` + +The checked-in URLs assume Dynamo is reachable from the trainer through localhost, as in a shared dev pod. For a remote DGD, replace both URLs with its frontend services. Also replace `weight_broadcast.host` with a trainer hostname or IP reachable from every sidecar; localhost is not valid across pods or nodes. + +`inference_world_size` must equal the sum of `world_size` in one `/v1/rl/workers` response. This recipe assumes one aggregated DP1 engine and therefore uses `1`. + +## Run + +```bash +uv run rl @ examples/dynamo/qwen3_06b_math/rl.toml \ + --output-dir outputs/dynamo-qwen3-06b-math +``` + +The run is successful when four optimizer steps complete, the verifier reports math rewards, weight versions advance after each update, and the Dynamo workers remain healthy. diff --git a/examples/dynamo/qwen3_06b_math/rl.toml b/examples/dynamo/qwen3_06b_math/rl.toml new file mode 100644 index 0000000000..d26bf43ba2 --- /dev/null +++ b/examples/dynamo/qwen3_06b_math/rl.toml @@ -0,0 +1,43 @@ +max_steps = 4 +seq_len = 2048 + +[model] +name = "Qwen/Qwen3-0.6B" + +[deployment] +type = "single_node" +num_train_gpus = 1 +num_infer_gpus = 0 + +[weight_broadcast] +type = "nccl" +host = "127.0.0.1" +inference_world_size = 1 + +[trainer] + +[orchestrator] +batch_size = 32 +group_size = 4 + +[orchestrator.model.client] +base_url = ["http://127.0.0.1:8000/v1"] +dynamo_discovery_url = "http://127.0.0.1:8001" +wait_for_ready_timeout = 1800 + +[orchestrator.model.client.extra_headers_from_state] +X-Session-ID = "trajectory_id" +X-Dynamo-Session-ID = "trajectory_id" + +[orchestrator.renderer] +name = "auto" +thinking_retention = "all" + +[orchestrator.train.sampling] +max_completion_tokens = 2048 + +[[orchestrator.train.source]] +name = "math" +env.taskset = { id = "math-env-v1", dataset_name = "openai/gsm8k", dataset_subset = "main" } +env.agent.harness = { id = "null" } +env.agent.runtime = { type = "subprocess" } diff --git a/examples/dynamo/qwen3_30b_Thinking/README.md b/examples/dynamo/qwen3_30b_Thinking/README.md new file mode 100644 index 0000000000..d01b182882 --- /dev/null +++ b/examples/dynamo/qwen3_30b_Thinking/README.md @@ -0,0 +1,53 @@ +# Qwen3 30B Thinking math with external Dynamo inference + +This four-step scale recipe runs a two-GPU Prime trainer against an external 1-prefill/1-decode Dynamo deployment serving `Qwen/Qwen3-30B-A3B-Thinking-2507`. Prime launches no inference process; Dynamo owns the frontend, sidecars, and vLLM engines. + +The two-GPU trainer uses BF16 optimization and reduction plus CPU optimizer +offload. Leaving Prime's FP32 optimization default in place exceeds two GB200s +when Adam allocates its first-step state; larger production recipes use the +existing eight-GPU training shape instead. + +The same model is used by the existing public examples `qwen30b_math`, `qwen30b_swe`, `multinode/rl.toml`, and `multinode/sft.toml`. Those examples remain the source of truth for larger training and non-Dynamo deployment settings; this recipe adds only the external-Dynamo client shape. + +## Prerequisites + +Use this version-matched native-gRPC source set: + +- Dynamo `feat/dyn-pi-sidecar-v2-review-001` at `836fe81012` +- vLLM `feat/dyn-pi-sidecar-v2-review-001` at `e56ee21b2c` +- Prime `feat/dyn-pi-sidecar-v2-review-001` + +Build `vllm-rs` and `dynamo-vllm-sidecar` from those revisions into the same +runtime image. For Kubernetes, start from Dynamo's +`examples/backends/vllm/deploy/sidecar_disagg.yaml` and change the model plus +parallelism/resources for the 30B topology; its adjacent `README.md` documents +the paired-binary image. Then apply the required Prime RL discovery and admin +overlay in [`../README.md`](../README.md). This contract requires both native +gRPC and the Dynamo `/v1/rl/workers` endpoint; a standard Python-only vLLM +worker is not compatible. + +Install the math environment: + +```bash +prime env install primeintellect/math-env +``` + +Start a Dynamo 1P/1D deployment and verify its public endpoints: + +```bash +curl http://127.0.0.1:8000/v1/models +curl http://127.0.0.1:8001/v1/rl/workers +``` + +The checked-in localhost URLs are for a colocated dev-pod run. For a DGD, replace them with the frontend services. Replace `weight_broadcast.host` with a trainer address reachable from every sidecar and allow the configured NCCL port through the network policy. + +`inference_world_size` must equal the sum of `world_size` in one `/v1/rl/workers` response. This recipe assumes 1P/1D with one rank per engine, for a total of `2`; TP, PP, or managed-DP topologies require the corresponding larger sum. + +## Run + +```bash +uv run rl @ examples/dynamo/qwen3_30b_Thinking/rl.toml \ + --output-dir outputs/dynamo-qwen3-30b-thinking-math +``` + +The run is successful when four optimizer steps complete, math rewards are emitted, all worker weight versions advance, and both prefill and decode workers remain healthy. diff --git a/examples/dynamo/qwen3_30b_Thinking/rl.toml b/examples/dynamo/qwen3_30b_Thinking/rl.toml new file mode 100644 index 0000000000..12c6995bf6 --- /dev/null +++ b/examples/dynamo/qwen3_30b_Thinking/rl.toml @@ -0,0 +1,56 @@ +max_steps = 4 +seq_len = 2048 + +[model] +name = "Qwen/Qwen3-30B-A3B-Thinking-2507" + +[deployment] +type = "single_node" +num_train_gpus = 2 +num_infer_gpus = 0 + +[weight_broadcast] +type = "nccl" +host = "127.0.0.1" +timeout = 1800 +inference_world_size = 2 + +[trainer.model] +impl = "custom" +attn = "flash_attention_3" +ep = 2 +optim_cpu_offload = true +optimization_dtype = "bfloat16" +reduce_dtype = "bfloat16" + +[trainer.model.ac] +freq = 1 + +[orchestrator] +batch_size = 2 +group_size = 2 +max_inflight_episodes = 2 +max_off_policy_steps = 0 + +[orchestrator.model.client] +base_url = ["http://127.0.0.1:8000/v1"] +dynamo_discovery_url = "http://127.0.0.1:8001" +wait_for_ready_timeout = 3600 + +[orchestrator.model.client.extra_headers_from_state] +X-Session-ID = "trajectory_id" +X-Dynamo-Session-ID = "trajectory_id" + +[orchestrator.renderer] +name = "qwen3" +enable_thinking = true + +[orchestrator.train.sampling] +temperature = 1.0 +max_completion_tokens = 2048 + +[[orchestrator.train.source]] +name = "math" +env.taskset = { id = "math-env-v1", dataset_name = "PrimeIntellect/Hendrycks-Math", dataset_subset = "default", task = { math_verify_timeout = 60 } } +env.agent.harness = { id = "null" } +env.agent.runtime = { type = "subprocess" } diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index e9fdbf2a6e..a4958724e4 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -403,6 +403,9 @@ class ZeroAdvantageFilterConfig(BaseConfig): class FileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks for Dynamo discovery completeness; unused by filesystem transfer itself.""" + class InMemoryWeightBroadcastConfig(BaseConfig): host: str = "localhost" @@ -572,6 +575,17 @@ def auto_setup_session_headers(self): self.model.client.extra_headers_from_state.setdefault("X-Session-ID", "trajectory_id") return self + @model_validator(mode="after") + def validate_dynamo_world_size(self): + if not self.model.client.is_dynamo: + return self + if ( + self.weight_broadcast.inference_world_size is None + or "inference_world_size" not in self.weight_broadcast.model_fields_set + ): + raise ValueError("Dynamo inference requires an explicit weight_broadcast.inference_world_size") + return self + @model_validator(mode="after") def auto_setup_prime_monitor_run_name(self): """Default ``prime_monitor.run_name`` to the W&B run name when monitoring diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 8a4a02ae21..06b408f457 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -138,6 +138,9 @@ class SharedNCCLWeightBroadcastConfig(SharedInMemoryWeightBroadcastConfig): quantize_in_weight_transfer: bool = False """Use kernel-format FP8 quantized NCCL transfer for weight updates. When disabled, uses default HF checkpoint-format transfer.""" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks when inference is managed externally.""" + class SharedNIXLWeightBroadcastConfig(SharedInMemoryWeightBroadcastConfig): type: Literal["nixl"] = "nixl" @@ -148,10 +151,16 @@ class SharedNIXLWeightBroadcastConfig(SharedInMemoryWeightBroadcastConfig): session_id: str = "default" """ModelExpress session ID.""" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks when inference is managed externally.""" + class SharedFileSystemWeightBroadcastConfig(BaseConfig): type: Literal["filesystem"] = "filesystem" + inference_world_size: int | None = Field(None, ge=1) + """Expected inference ranks when inference is managed externally (e.g. Dynamo LoRA over filesystem).""" + SharedWeightBroadcastConfig: TypeAlias = Annotated[ SharedFileSystemWeightBroadcastConfig | SharedNCCLWeightBroadcastConfig | SharedNIXLWeightBroadcastConfig, @@ -323,16 +332,6 @@ def validate_deployment(self): ) return self - @model_validator(mode="after") - def validate_enough_devices_for_nccl(self): - if self.deployment.type == "single_node": - if self.trainer.weight_broadcast.type == "nccl": - if self.deployment.num_train_gpus + self.deployment.num_infer_gpus < 2: - raise ValueError( - "NCCL weight broadcast requires at least 2 GPUs to build the broadcast process group." - ) - return self - @model_validator(mode="after") def validate_quantize_in_weight_transfer(self): if not isinstance(self.weight_broadcast, SharedNCCLWeightBroadcastConfig): @@ -393,13 +392,18 @@ def auto_setup_weight_broadcast(self): "Set weight_broadcast.type = 'filesystem'." ) if self.weight_broadcast.type in ("nccl", "nixl"): - inference_world_size = self.inference.parallel.dp * self.inference.parallel.tp if self.inference else 1 + inference_world_size = ( + self.inference.parallel.dp * self.inference.parallel.tp + if self.inference + else self.weight_broadcast.inference_world_size + ) common_config = dict( host=self.weight_broadcast.host, port=self.weight_broadcast.port, timeout=self.weight_broadcast.timeout, - inference_world_size=inference_world_size, ) + if inference_world_size is not None: + common_config["inference_world_size"] = inference_world_size if self.weight_broadcast.type == "nccl": transport_config = dict( quantize_in_weight_transfer=self.weight_broadcast.quantize_in_weight_transfer, @@ -414,7 +418,9 @@ def auto_setup_weight_broadcast(self): self.orchestrator.weight_broadcast = orchestrator_config_type(**common_config, **transport_config) elif self.weight_broadcast.type == "filesystem": self.trainer.weight_broadcast = TrainerFileSystemWeightBroadcastConfig() - self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig() + self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig( + inference_world_size=self.weight_broadcast.inference_world_size + ) if self.inference is not None: self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type) @@ -444,6 +450,19 @@ def auto_setup_rollout_transport(self): self.rollout_transport = self.trainer.rollout_transport return self + @model_validator(mode="after") + def validate_enough_devices_for_nccl(self): + if self.deployment.type != "single_node" or self.trainer.weight_broadcast.type != "nccl": + return self + if self.inference is None and self.weight_broadcast.inference_world_size is not None: + return self + local_inference_gpus = self.deployment.num_infer_gpus if self.inference is not None else 0 + if self.deployment.num_train_gpus + local_inference_gpus < 2: + raise ValueError( + "NCCL weight broadcast requires at least 2 local GPUs or an explicit external inference_world_size." + ) + return self + @model_validator(mode="after") def validate_eplb_requires_quantized_weight_transfer(self): if self.inference is None or not self.inference.enable_eplb: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 74f2a859b8..4243fa1cc5 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import Annotated, Literal, TypeAlias +from typing import Annotated, Literal, Self, TypeAlias from pydantic import AfterValidator, Field, model_validator @@ -144,17 +144,33 @@ class ClientConfig(BaseConfig): admin_base_url: list[str] | None = None """Separate base URLs for admin operations (weight updates, health checks). When set, admin clients bypass routers and hit each server directly — used in disaggregated P/D deployments where the router must not handle admin traffic.""" + dynamo_discovery_url: str | None = None + """Dynamo discovery URL. When set, Prime discovers vLLM admin endpoints and per-engine world sizes from ``/v1/rl/workers`` instead of requiring ``admin_base_url`` entries.""" + elastic: ElasticConfig | None = None """Elastic inference pool config for DNS-based service discovery. When set, ``base_url`` is ignored and inference servers are discovered dynamically via DNS.""" router_url: str | None = None """vllm-router URL for load-aware inference routing. With elastic mode, inference requests go through the router while admin ops still hit discovered pods directly.""" + @model_validator(mode="after") + def validate_pool_mode(self) -> Self: + if self.dynamo_discovery_url is not None and self.admin_base_url is not None: + raise ValueError("dynamo_discovery_url cannot be combined with admin_base_url") + if self.dynamo_discovery_url is not None and self.elastic is not None: + raise ValueError("dynamo_discovery_url cannot be combined with elastic discovery") + return self + @property def is_elastic(self) -> bool: """Check if elastic mode is enabled.""" return self.elastic is not None + @property + def is_dynamo(self) -> bool: + """Check if Dynamo worker discovery is enabled.""" + return self.dynamo_discovery_url is not None + class LogConfig(BaseConfig): level: str = Field(default_factory=lambda: os.environ.get("PRIME_LOG_LEVEL", "info")) diff --git a/packages/prime-rl-configs/src/prime_rl/utils/validation.py b/packages/prime-rl-configs/src/prime_rl/utils/validation.py index 6dac87c12e..5d507424c8 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -129,6 +129,9 @@ def propagate(shared_path: str, *targets: str) -> None: # [rollout_transport] → both sub-configs (host is launcher-injected for zmq multi-node). propagate("rollout_transport", "trainer.rollout_transport", "orchestrator.rollout_transport") + # The orchestrator validates external inference topology during construction. + propagate("weight_broadcast", "orchestrator.weight_broadcast") + # Top-level scalars. propagate("max_steps", "trainer.max_steps", "orchestrator.max_steps") propagate("seq_len", "trainer.model.seq_len", "orchestrator.seq_len") diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 374a1070f0..4929207c8a 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -45,6 +45,7 @@ async def setup_policy_inference_pool(*, config: OrchestratorConfig, tokenizer): train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=config.renderer, + expected_inference_world_size=config.weight_broadcast.inference_world_size, ) return renderer, inference_pool diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 2d35f122ef..e652f3cc2c 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -5,7 +5,7 @@ from collections.abc import Mapping from itertools import cycle from pathlib import Path -from typing import Protocol, runtime_checkable +from typing import Protocol, cast, runtime_checkable import httpx import verifiers.v1 as vf @@ -122,6 +122,8 @@ def __init__( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, + *, + admin_clients: list[AsyncClient] | None = None, ): renderer_model_name = model_name if train_client_type == "renderer" else None self._train_clients = setup_clients( @@ -131,7 +133,7 @@ def __init__( renderer_model_name=renderer_model_name, ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) - self._admin_clients = setup_admin_clients(client_config) + self._admin_clients = setup_admin_clients(client_config) if admin_clients is None else admin_clients # When admin URLs bypass a router, also health-check the client-facing # (router) endpoint - it only starts serving once its workers are healthy. self._router_clients = ( @@ -193,6 +195,7 @@ async def setup_inference_pool( train_client_type: str = "openai_chat_completions", eval_client_type: str = "openai_chat_completions", renderer_config: RendererConfig | None = None, + expected_inference_world_size: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" if client_config.is_elastic: @@ -206,6 +209,18 @@ async def setup_inference_pool( renderer_config=renderer_config, ) + if client_config.is_dynamo: + from prime_rl.utils.dynamo import DynamoInferencePool + + return await DynamoInferencePool.from_config( + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_config=renderer_config, + expected_inference_world_size=cast(int, expected_inference_world_size), + ) + return StaticInferencePool( client_config, model_name=model_name, diff --git a/src/prime_rl/utils/dynamo.py b/src/prime_rl/utils/dynamo.py new file mode 100644 index 0000000000..6253adee54 --- /dev/null +++ b/src/prime_rl/utils/dynamo.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import asyncio +from typing import Any, cast + +import httpx +from httpx import AsyncClient +from pydantic import BaseModel, ConfigDict, Field +from tenacity import AsyncRetrying, retry_if_exception, stop_after_delay, wait_exponential + +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.client import StaticInferencePool, setup_admin_clients + +DYNAMO_RL_DISCOVERY_PROTOCOL_VERSION = 1 +DYNAMO_READINESS_REQUEST_TIMEOUT_S = 30.0 + + +class DiscoveredDynamoWorker(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + component: str = Field(min_length=1) + instance_id: int = Field(ge=0, strict=True) + model: str + admin_base_url: str = Field(min_length=1) + world_size: int = Field(gt=0, strict=True) + + +class DynamoDiscoverySnapshot(BaseModel): + model_config = ConfigDict(extra="ignore") + + protocol_version: int = Field( + strict=True, + ge=DYNAMO_RL_DISCOVERY_PROTOCOL_VERSION, + le=DYNAMO_RL_DISCOVERY_PROTOCOL_VERSION, + ) + workers: list[dict[str, Any]] + + +class DynamoDiscoveryPending(ValueError): + """A well-formed discovery snapshot that is not ready yet.""" + + +def _is_retryable_dynamo_error(exception: BaseException) -> bool: + if isinstance(exception, httpx.HTTPStatusError): + return exception.response.status_code == 429 or exception.response.status_code >= 500 + return isinstance(exception, (DynamoDiscoveryPending, httpx.TransportError)) + + +def _parse_dynamo_workers(payload: object, model_name: str) -> tuple[DiscoveredDynamoWorker, ...]: + snapshot = DynamoDiscoverySnapshot.model_validate(payload) + workers = [] + for raw_worker in snapshot.workers: + if raw_worker.get("model") not in (None, model_name): + continue + if error := raw_worker.get("error"): + raise DynamoDiscoveryPending(f"Dynamo RL worker probe is not ready: {error}") + workers.append(DiscoveredDynamoWorker.model_validate(raw_worker)) + if not workers: + raise DynamoDiscoveryPending("Dynamo RL discovery returned no workers yet") + + identities = [(worker.component, worker.instance_id) for worker in workers] + admin_urls = [worker.admin_base_url for worker in workers] + if len(set(identities)) != len(identities): + raise ValueError("Dynamo RL discovery returned duplicate worker identities") + if len(set(admin_urls)) != len(admin_urls): + raise ValueError("Dynamo RL discovery returned duplicate admin endpoints") + return tuple(sorted(workers, key=lambda worker: (worker.component, worker.instance_id))) + + +def _setup_control_clients(urls: list[str]) -> list[AsyncClient]: + return [ + AsyncClient( + base_url=url.rstrip("/"), + limits=httpx.Limits(max_connections=4, max_keepalive_connections=1), + timeout=httpx.Timeout(None), + ) + for url in urls + ] + + +async def _wait_for_model(clients: list[AsyncClient], model_name: str, timeout: float) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + async with asyncio.timeout(timeout): + async for attempt in AsyncRetrying( + stop=stop_after_delay(timeout), + wait=wait_exponential(multiplier=0.1, min=0.1, max=1), + retry=retry_if_exception(_is_retryable_dynamo_error), + reraise=True, + ): + with attempt: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError + request_timeout = httpx.Timeout(min(DYNAMO_READINESS_REQUEST_TIMEOUT_S, remaining)) + responses = await asyncio.gather( + *(client.get("/v1/models", timeout=request_timeout) for client in clients) + ) + for response in responses: + response.raise_for_status() + models = response.json().get("data", []) + if not any(model.get("id") == model_name for model in models): + raise DynamoDiscoveryPending(f"Dynamo frontend has not published model {model_name!r}") + + +class DynamoInferencePool(StaticInferencePool): + """Static request pool whose direct admin clients come from Dynamo discovery.""" + + def __init__(self, client_config: ClientConfig, workers: tuple[DiscoveredDynamoWorker, ...], **kwargs): + admin_clients = _setup_control_clients([worker.admin_base_url for worker in workers]) + super().__init__(client_config, admin_clients=admin_clients, **kwargs) + self._admin_world_sizes = [worker.world_size for worker in workers] + self._frontend_model_clients = setup_admin_clients(client_config) + self._readiness_deadline: float | None = None + + async def wait_for_ready(self, model_name: str, timeout: int | None = None) -> None: + effective_timeout = self._wait_for_ready_timeout if timeout is None else timeout + loop = asyncio.get_running_loop() + deadline = ( + self._readiness_deadline + if timeout is None and self._readiness_deadline is not None + else loop.time() + effective_timeout + ) + remaining = max(0.0, deadline - loop.time()) + try: + async with asyncio.timeout(remaining): + await super().wait_for_ready(model_name, timeout=remaining) + if not self._skip_model_check: + await _wait_for_model( + self._frontend_model_clients, + model_name, + timeout=max(0.0, deadline - loop.time()), + ) + finally: + self._readiness_deadline = None + + async def stop(self) -> None: + await super().stop() + await asyncio.gather(*(client.aclose() for client in [*self._admin_clients, *self._frontend_model_clients])) + + @classmethod + async def from_config( + cls, + client_config: ClientConfig, + model_name: str, + expected_inference_world_size: int, + **kwargs, + ) -> DynamoInferencePool: + discovery_url = cast(str, client_config.dynamo_discovery_url).rstrip("/").removesuffix("/v1") + loop = asyncio.get_running_loop() + deadline = loop.time() + client_config.wait_for_ready_timeout + async with asyncio.timeout(client_config.wait_for_ready_timeout): + async with AsyncClient(timeout=httpx.Timeout(None)) as client: + workers: tuple[DiscoveredDynamoWorker, ...] = () + async for attempt in AsyncRetrying( + stop=stop_after_delay(client_config.wait_for_ready_timeout), + wait=wait_exponential(multiplier=0.1, min=0.1, max=1), + retry=retry_if_exception(_is_retryable_dynamo_error), + reraise=True, + ): + with attempt: + remaining = deadline - loop.time() + if remaining <= 0: + raise TimeoutError + response = await client.get( + f"{discovery_url}/v1/rl/workers", + timeout=httpx.Timeout(min(DYNAMO_READINESS_REQUEST_TIMEOUT_S, remaining)), + ) + response.raise_for_status() + workers = _parse_dynamo_workers(response.json(), model_name) + discovered_world_size = sum(worker.world_size for worker in workers) + if discovered_world_size != expected_inference_world_size: + raise DynamoDiscoveryPending( + "Dynamo RL discovery returned " + f"inference_world_size={discovered_world_size}; " + f"waiting for expected inference_world_size={expected_inference_world_size}" + ) + pool = cls(client_config, workers, model_name=model_name, **kwargs) + pool._readiness_deadline = deadline + return pool diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py index 60f13e01aa..48254c2fc7 100644 --- a/tests/unit/orchestrator/test_orchestrator_setup.py +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -18,6 +18,7 @@ async def run() -> None: ), renderer=renderer_settings, any_policy_sourced=True, + weight_broadcast=SimpleNamespace(type="filesystem", inference_world_size=None), ) renderer = object() inference_pool = object() @@ -43,6 +44,7 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, + expected_inference_world_size=None, ) asyncio.run(run()) @@ -64,6 +66,7 @@ async def run() -> None: ), renderer=renderer_settings, any_policy_sourced=False, + weight_broadcast=SimpleNamespace(inference_world_size=8), ) renderer = object() inference_pool = object() @@ -89,6 +92,7 @@ async def run() -> None: train_client_type="renderer", eval_client_type="openai_chat_completions", renderer_config=renderer_settings, + expected_inference_world_size=8, ) asyncio.run(run()) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 53492118fb..0bca37733a 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -215,6 +215,116 @@ def test_trainer_enable_token_export_cli_flag(): assert cli(TrainerConfig, args=["--enable-token-export"]).enable_token_export +def test_external_dynamo_world_size_survives_rl_config_resolution(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + } + }, + "inference": None, + "weight_broadcast": { + "type": "nccl", + "host": "trainer.service", + "inference_world_size": 8, + }, + } + ) + + assert config.trainer.weight_broadcast.inference_world_size == 8 + assert config.trainer.weight_broadcast.host == "trainer.service" + assert config.orchestrator.weight_broadcast.inference_world_size == 8 + assert config.orchestrator.weight_broadcast.host == "trainer.service" + + +def test_external_dynamo_nccl_does_not_require_a_local_inference_gpu(): + config = RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + } + }, + "inference": None, + "deployment": { + "type": "single_node", + "num_train_gpus": 1, + "num_infer_gpus": 0, + }, + "weight_broadcast": { + "type": "nccl", + "host": "trainer.service", + "inference_world_size": 1, + }, + } + ) + + assert config.deployment.num_train_gpus == 1 + assert config.deployment.num_infer_gpus == 0 + assert config.trainer.weight_broadcast.inference_world_size == 1 + + +def test_default_nccl_world_size_does_not_bypass_local_gpu_guard(): + with pytest.raises(ValueError, match="NCCL weight broadcast requires at least 2"): + RLConfig.model_validate( + { + "trainer": {}, + "orchestrator": {}, + "inference": None, + "deployment": { + "type": "single_node", + "num_train_gpus": 1, + "num_infer_gpus": 0, + }, + "weight_broadcast": {"type": "nccl"}, + } + ) + + +def test_external_dynamo_lora_world_size_survives_filesystem_config_resolution(): + config = RLConfig.model_validate( + { + "trainer": {"model": {"lora": {}}}, + "orchestrator": { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + } + }, + "inference": None, + "weight_broadcast": {"type": "filesystem", "inference_world_size": 8}, + } + ) + + assert config.orchestrator.weight_broadcast.inference_world_size == 8 + + +def test_dynamo_orchestrator_requires_explicit_inference_world_size(): + with pytest.raises(ValueError, match="inference_world_size"): + OrchestratorConfig.model_validate( + { + "model": { + "client": { + "base_url": ["http://frontend:8000/v1"], + "dynamo_discovery_url": "http://frontend:8001", + } + }, + "weight_broadcast": {"type": "filesystem"}, + } + ) + + def test_single_node_auto_inference_ports_follow_server_port(): config = RLConfig.model_validate( { diff --git a/tests/unit/utils/test_dynamo.py b/tests/unit/utils/test_dynamo.py new file mode 100644 index 0000000000..d6e0c844eb --- /dev/null +++ b/tests/unit/utils/test_dynamo.py @@ -0,0 +1,119 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from prime_rl.configs.shared import ClientConfig, ElasticConfig +from prime_rl.utils.dynamo import DynamoInferencePool, _parse_dynamo_workers + +MODEL = "Qwen/Qwen3-0.6B" + + +def worker(**updates): + value = { + "component": "backend", + "instance_id": 10, + "model": MODEL, + "admin_base_url": "http://decode:8120", + "world_size": 2, + } + return {**value, **updates} + + +def payload(*workers): + return {"protocol_version": 1, "workers": list(workers)} + + +def response(body): + result = MagicMock() + result.raise_for_status = MagicMock() + result.json.return_value = body + return result + + +def test_parse_workers_orders_identity_and_preserves_topology(): + workers = _parse_dynamo_workers( + payload( + worker(component="prefill", instance_id=20, admin_base_url="http://prefill:8121"), + worker(), + ), + MODEL, + ) + + assert [(item.component, item.instance_id) for item in workers] == [ + ("backend", 10), + ("prefill", 20), + ] + assert [item.world_size for item in workers] == [2, 2] + + +@pytest.mark.parametrize( + "workers", + [ + [], + [worker(error="probe timed out")], + [worker(admin_base_url=None)], + [worker(world_size=0)], + [worker(model="other/model")], + [worker(), worker(instance_id=11)], + [worker(), worker(component="prefill", instance_id=20, admin_base_url="http://decode:8120")], + ], +) +def test_parse_workers_rejects_incomplete_or_duplicate_snapshots(workers): + with pytest.raises(ValueError): + _parse_dynamo_workers(payload(*workers), MODEL) + + +@pytest.mark.parametrize( + "conflict", + [ + {"admin_base_url": ["http://worker:8120"]}, + {"elastic": ElasticConfig(hostname="workers")}, + ], +) +def test_discovery_config_rejects_other_pool_modes(conflict): + with pytest.raises(ValueError, match="dynamo_discovery_url"): + ClientConfig(dynamo_discovery_url="http://frontend:8001", **conflict) + + +def test_discovery_retries_until_expected_world_size_is_complete(): + transient = MagicMock() + transient.raise_for_status.side_effect = httpx.HTTPStatusError( + "Service unavailable", + request=httpx.Request("GET", "http://frontend:8001/v1/rl/workers"), + response=httpx.Response(503), + ) + discovery_client = AsyncMock() + discovery_client.get.side_effect = [ + transient, + response(payload(worker())), + response( + payload( + worker(), + worker(component="prefill", instance_id=20, admin_base_url="http://prefill:8121"), + ) + ), + ] + context = AsyncMock() + context.__aenter__.return_value = discovery_client + + class DiscoveryOnlyPool(DynamoInferencePool): + def __init__(self, _config, workers, **_kwargs): + self.workers = workers + + with patch("prime_rl.utils.dynamo.AsyncClient", return_value=context): + pool = asyncio.run( + DiscoveryOnlyPool.from_config( + ClientConfig( + base_url=["http://frontend:8000/v1"], + dynamo_discovery_url="http://frontend:8001", + wait_for_ready_timeout=1, + ), + model_name=MODEL, + expected_inference_world_size=4, + ) + ) + + assert discovery_client.get.await_count == 3 + assert [item.component for item in pool.workers] == ["backend", "prefill"]