Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
3690dc0
feat(inference): add Dynamo backend configuration
biswapanda Jul 3, 2026
c58de96
feat(inference): launch Dynamo disaggregated workers
biswapanda Jul 3, 2026
7c1a61d
feat(orchestrator): administer Dynamo inference workers
biswapanda Jul 3, 2026
908705d
fix(config): stop injecting response-only token flag
biswapanda Jul 3, 2026
770b568
fix(client): preserve native admin routing
biswapanda Jul 3, 2026
7510010
fix(dynamo): preserve weight update failures
biswapanda Jul 3, 2026
beaec34
refactor(inference): isolate Dynamo runtime boundaries
biswapanda Jul 5, 2026
c6fab8f
refactor(inference): keep process environment canonical
biswapanda Jul 5, 2026
5ce8959
refactor(inference): remove redundant process wrappers
biswapanda Jul 5, 2026
69986cc
refactor(inference): derive worker port evidence
biswapanda Jul 5, 2026
1ac4298
test: update Dynamo topology fixtures
biswapanda Jul 9, 2026
04eeceb
fix(dynamo): enforce runtime topology contracts
biswapanda Jul 9, 2026
0d22c2b
fix(dynamo): close topology and readiness gaps
biswapanda Jul 9, 2026
4a8a185
fix(dynamo): classify readiness and fixed pools
biswapanda Jul 9, 2026
f5a9d56
fix(dynamo): fail fast on discovery contracts
biswapanda Jul 10, 2026
66c40e4
fix(dynamo): make policy updates fail closed
biswapanda Jul 10, 2026
3ece154
fix(dynamo): harden policy update transactions
biswapanda Jul 10, 2026
f077fbf
fix(orchestrator): adapt policy markers to TraceTask
biswapanda Jul 11, 2026
2bc3754
feat(slurm): support single-node Dynamo RL jobs
furionw Jul 17, 2026
61484f5
fix(rl): preserve Dynamo deployment topology
furionw Jul 17, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions docs/scaling.md
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,26 @@ uv run rl @ base_rl.toml @ my_slurm.toml # submits via sbatch
uv run rl @ base_rl.toml @ my_slurm.toml --dry-run # writes the sbatch script + resolved config, exits
```

Single-node RL jobs may use the Dynamo inference backend. The generated job
enters the same local Dynamo launcher used without SLURM, after the scheduler
assigns the node and GPUs:

```bash
# One aggregated Dynamo worker and one trainer GPU.
uv run rl @ examples/reverse_text/rl.toml \
@ examples/reverse_text/dynamo_aggregated_rl.toml \
@ examples/reverse_text/slurm_rl.toml

# One prefill worker, one decode worker, and one trainer GPU.
uv run rl @ examples/reverse_text/rl.toml \
@ examples/reverse_text/dynamo_disaggregated_rl.toml \
@ examples/reverse_text/slurm_rl.toml
```

Standalone Dynamo inference with `[slurm]` and multi-node Dynamo RL are not
supported. Those deployment shapes require a DynamoGraphDeployment or an
external launcher.

### `[deployment]` Block

`[deployment]` is a discriminated union picked by `type` — `single_node` or `multi_node` for RL/SFT, with an extra disaggregated variant for inference. RL multi-node:
Expand Down
2 changes: 1 addition & 1 deletion docs/training.md
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli
|---|---|
| `orchestrator.batch_size` | Tasks per trainer step. |
| `orchestrator.group_size` | Rollouts generated per task. |
| `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. |
| `orchestrator.max_off_policy_steps` | On the vLLM admin backend, how many distinct policies may contribute to one rollout before it is discarded (default 8). Dynamo instead drains all live-policy/eval requests before mutating weights, so this setting does not apply to Dynamo runs. |
| `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). |
| `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). |
| `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). |
Expand Down
25 changes: 25 additions & 0 deletions examples/reverse_text/dynamo_aggregated_rl.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
# Usage: uv run rl @ examples/reverse_text/rl.toml @ examples/reverse_text/dynamo_aggregated_rl.toml @ examples/reverse_text/slurm_rl.toml

[weight_broadcast]
type = "nccl"

[trainer]

[orchestrator]

[deployment]
type = "single_node"
gpus_per_node = 2
num_train_gpus = 1
num_infer_gpus = 1

[inference.backend]
type = "dynamo"

[inference.parallel]
tp = 1
dp = 1

[inference.deployment]
type = "single_node"
gpus_per_node = 1
32 changes: 32 additions & 0 deletions examples/reverse_text/dynamo_disaggregated_rl.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# Usage: uv run rl @ examples/reverse_text/rl.toml @ examples/reverse_text/dynamo_disaggregated_rl.toml @ examples/reverse_text/slurm_rl.toml

[weight_broadcast]
type = "nccl"

[trainer]

[orchestrator]

[deployment]
type = "single_node"
gpus_per_node = 3
num_train_gpus = 1
num_infer_gpus = 2

[inference]
enable_expert_parallel = false

[inference.backend]
type = "dynamo"

[inference.parallel]
tp = 1
dp = 1

[inference.deployment]
type = "disaggregated"
gpus_per_node = 1
prefill_nodes_per_replica = 1
decode_nodes_per_replica = 1
num_prefill_replicas = 1
num_decode_replicas = 1
100 changes: 92 additions & 8 deletions packages/prime-rl-configs/src/prime_rl/configs/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ class ServerConfig(BaseConfig):


class ParallelConfig(BaseConfig):
tp: int = 1
tp: int = Field(1, ge=1)
"""Tensor parallel size. Forwarded to vLLM as ``--tensor-parallel-size``."""

dp: int = Field(1, ge=1)
Expand Down Expand Up @@ -251,7 +251,7 @@ def validate_scorers(self):


class BaseInferenceDeploymentConfig(BaseConfig):
gpus_per_node: int = 8
gpus_per_node: int = Field(8, ge=1)
"""GPUs per node."""

backend_port: int = 8100
Expand Down Expand Up @@ -304,10 +304,10 @@ class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig):
"""Extra environment variables exported only on decode nodes."""

prefill_vllm_overrides: dict[str, Any] = {}
"""Extra vLLM config options merged into --vllm-extra only for prefill ranks (SLURM only)."""
"""Extra vLLM config options merged into the resolved config only for prefill workers."""

decode_vllm_overrides: dict[str, Any] = {}
"""Extra vLLM config options merged into --vllm-extra only for decode ranks (SLURM only)."""
"""Extra vLLM config options merged into the resolved config only for decode workers."""

@property
def num_prefill_nodes(self) -> int:
Expand All @@ -328,7 +328,24 @@ def num_nodes(self) -> int:
]


class VllmInferenceBackendConfig(BaseConfig):
type: Literal["vllm"] = "vllm"


class DynamoInferenceBackendConfig(BaseConfig):
type: Literal["dynamo"] = "dynamo"


InferenceBackendConfig: TypeAlias = Annotated[
VllmInferenceBackendConfig | DynamoInferenceBackendConfig,
Field(discriminator="type"),
]


class InferenceConfig(BaseConfig):
backend: InferenceBackendConfig = VllmInferenceBackendConfig()
"""Serving backend. Existing configs default to Prime's native vLLM launcher."""

server: ServerConfig = ServerConfig()

model: ModelConfig = Field(default_factory=ModelConfig)
Expand Down Expand Up @@ -426,12 +443,81 @@ class InferenceConfig(BaseConfig):
dry_run: bool = False
"""Only validate and dump resolved configs, then exit early."""

@property
def dynamo_worker_roles(self) -> tuple[Literal["agg", "prefill", "decode"], ...]:
"""Canonical admin and launch order for Dynamo worker groups."""
if self.deployment.type == "disaggregated":
return ("prefill",) * self.deployment.num_prefill_replicas + (
"decode",
) * self.deployment.num_decode_replicas
return ("agg",)

@property
def dynamo_gpus_per_worker(self) -> int:
"""GPU allocation owned by one independently administered Dynamo worker."""
if self.deployment.type == "disaggregated":
return self.deployment.gpus_per_node
return self.parallel.tp * self.parallel.dp

@property
def dynamo_local_dp(self) -> int:
"""vLLM data-parallel ranks inside one Dynamo worker."""
if self.deployment.type == "disaggregated":
return self.deployment.gpus_per_node // self.parallel.tp
return self.parallel.dp

@model_validator(mode="after")
def validate_multi_node_requires_slurm(self):
if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None:
if self.deployment.type == "multi_node" and self.slurm is None:
raise ValueError("Must use SLURM for multi-node deployment.")
if self.deployment.type == "disaggregated" and self.slurm is None and self.backend.type != "dynamo":
raise ValueError("Must use SLURM for multi-node / disaggregated deployment.")
return self

@model_validator(mode="after")
def validate_disaggregated_topology(self):
if self.deployment.type != "disaggregated":
return self
if self.deployment.gpus_per_node % self.parallel.tp != 0:
raise ValueError(
"inference.deployment.gpus_per_node must be divisible by inference.parallel.tp "
"so every worker contains whole tensor-parallel groups."
)
if self.backend.type == "dynamo":
local_dp = self.deployment.gpus_per_node // self.parallel.tp
if "dp" in self.parallel.model_fields_set and self.parallel.dp != local_dp:
raise ValueError(
"inference.parallel.dp must equal inference.deployment.gpus_per_node / "
"inference.parallel.tp for a Dynamo disaggregated worker."
)
if self.data_parallel_size_local is not None and self.data_parallel_size_local != local_dp:
raise ValueError(
"inference.data_parallel_size_local must equal inference.deployment.gpus_per_node / "
"inference.parallel.tp for a Dynamo disaggregated worker."
)
return self

@model_validator(mode="after")
def validate_dynamo_backend(self):
if self.backend.type != "dynamo":
return self
if self.slurm is not None:
raise ValueError(
"Dynamo is launched locally or through a DynamoGraphDeployment, not Prime's SLURM template."
)
if self.deployment.type == "multi_node":
raise ValueError("Dynamo multi-node inference must use a DynamoGraphDeployment.")
if self.enable_lora:
raise ValueError("The Dynamo backend does not support LoRA weight updates.")
router = getattr(self.deployment, "router", None)
if router is not None and router.type == "llm-d":
raise ValueError("The Dynamo backend owns request routing and cannot use the llm-d router.")
if self.deployment.type == "disaggregated":
if self.enable_prefix_caching is False:
raise ValueError("Dynamo disaggregated inference requires prefix caching for exact KV-aware routing.")
self.enable_prefix_caching = True
return self

@model_validator(mode="after")
def validate_llmd_no_routed_experts(self):
"""Reject routed-expert return with the llm-d router (breaks P/D, unverified for multi-node)."""
Expand Down Expand Up @@ -463,9 +549,7 @@ def auto_setup_disaggregated(self):
self.enable_expert_parallel = True
if "enable_eplb" not in self.model_fields_set:
self.enable_eplb = False
gpus_per_node = self.deployment.gpus_per_node
tp = self.parallel.tp
dp_per_node = gpus_per_node // tp
dp_per_node = self.dynamo_local_dp
if self.data_parallel_size_local is None:
self.data_parallel_size_local = dp_per_node
if self.parallel.dp == 1:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -529,7 +529,9 @@ class OrchestratorConfig(BaseConfig):
"""Maximum training steps. If None, runs indefinitely."""

max_off_policy_steps: int = Field(8, ge=0)
"""Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise."""
"""Maximum policies allowed to generate one rollout on the vLLM admin backend. Dynamo uses a strict
application drain before every weight mutation, so live-policy and eval requests never span versions and
this tolerance does not apply there."""

bench: bool = False
"""Benchmark mode. Sets ``max_steps`` to 5 and disables W&B."""
Expand Down Expand Up @@ -728,7 +730,6 @@ def resolve_env_config(self):
if env.algo.sampling.source == "policy":
env.sampling.extra_body.setdefault("top_k", -1)
env.sampling.extra_body.setdefault("min_p", 0.0)
env.sampling.extra_body.setdefault("return_token_ids", True)
if env.is_legacy:
# v0 env: cap per-turn response tokens to the training budget (the legacy
# bridge applies extra_env_kwargs via env.set_kwargs).
Expand Down
55 changes: 50 additions & 5 deletions packages/prime-rl-configs/src/prime_rl/configs/rl.py
Original file line number Diff line number Diff line change
Expand Up @@ -460,6 +460,12 @@ def auto_setup_lora(self):

return self

@model_validator(mode="after")
def validate_auto_setup_does_not_enable_dynamo_lora(self):
if self.inference is not None and self.inference.backend.type == "dynamo" and self.inference.enable_lora:
raise ValueError("The Dynamo backend does not support LoRA weight updates.")
return self

@model_validator(mode="after")
def auto_setup_router_replay(self):
if self.trainer.enable_router_replay:
Expand Down Expand Up @@ -534,7 +540,23 @@ def auto_setup_deployment(self):
# fill up inference capacity with dp ranks
if self.inference is not None:
num_infer_gpus = self.deployment.num_infer_gpus
if num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp:
is_dynamo_disaggregated = (
self.inference.backend.type == "dynamo" and self.inference.deployment.type == "disaggregated"
)
if is_dynamo_disaggregated:
infer_deploy = self.inference.deployment
expected_infer_gpus = infer_deploy.num_nodes * infer_deploy.gpus_per_node
if num_infer_gpus != expected_infer_gpus:
raise ValueError(
"deployment.num_infer_gpus must equal the Dynamo prefill/decode topology GPU count "
f"({expected_infer_gpus}), got {num_infer_gpus}."
)
if self.weight_broadcast is not None and self.weight_broadcast.type == "nccl":
assert self.trainer.weight_broadcast.type == "nccl"
self.trainer.weight_broadcast.inference_world_size = expected_infer_gpus
assert self.orchestrator.weight_broadcast.type == "nccl"
self.orchestrator.weight_broadcast.inference_world_size = expected_infer_gpus
elif num_infer_gpus != self.inference.parallel.dp * self.inference.parallel.tp:
assert num_infer_gpus % self.inference.parallel.tp == 0, (
"Number of inference GPUs must be divisible by the tensor parallel size"
)
Expand Down Expand Up @@ -675,15 +697,38 @@ def auto_setup_inference_client(self):
if self.inference is None:
return self
client = self.orchestrator.model.client
client_updates: dict[str, Any] = {}
if "admin_api" in client.model_fields_set and client.admin_api != self.inference.backend.type:
raise ValueError(
"orchestrator.model.client.admin_api conflicts with inference.backend.type; "
"configure the backend only under inference."
)
client_updates["admin_api"] = self.inference.backend.type
if "dp_rank_count" not in client.model_fields_set:
if self.deployment.type == "multi_node":
client.dp_rank_count = 1
if self.inference.backend.type == "dynamo" or self.deployment.type == "multi_node":
client_updates["dp_rank_count"] = 1
else:
client.dp_rank_count = self.inference.data_parallel_size_local or self.inference.parallel.dp
client_updates["dp_rank_count"] = self.inference.data_parallel_size_local or self.inference.parallel.dp
if self.inference.backend.type == "dynamo":
expected_topology = {
"dynamo_worker_roles": self.inference.dynamo_worker_roles,
"dynamo_gpus_per_worker": self.inference.dynamo_gpus_per_worker,
}
for field, expected in expected_topology.items():
if field in client.model_fields_set and getattr(client, field) != expected:
raise ValueError(
f"orchestrator.model.client.{field} conflicts with the inference topology; "
"configure the topology only under inference."
)
client_updates.update(expected_topology)
if not self.orchestrator.any_policy_sourced and "base_url" not in client.model_fields_set:
host = self.inference.server.host or "localhost"
port = self.inference.server.port
client.base_url = [f"http://{host}:{port}/v1"]
client_updates["base_url"] = [f"http://{host}:{port}/v1"]

updated_client = client.model_copy(update=client_updates)
updated_model = self.orchestrator.model.model_copy(update={"client": updated_client})
self.orchestrator = self.orchestrator.model_copy(update={"model": updated_model})
return self

@model_validator(mode="after")
Expand Down
15 changes: 14 additions & 1 deletion packages/prime-rl-configs/src/prime_rl/configs/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,18 @@ 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."""

admin_api: Literal["vllm", "dynamo"] = "vllm"
"""Admin protocol used for health and weight updates. Auto-derived from the local inference backend."""

rl_base_url: list[str] | None = None
"""Dynamo RL worker-discovery URLs. When omitted, they are derived from the frontend URL or ``DYN_RL_DISCOVERY_URL``."""

dynamo_worker_roles: tuple[Literal["agg", "prefill", "decode"], ...] | None = None
"""Exact Dynamo worker roles expected during readiness. Auto-derived from the local inference topology."""

dynamo_gpus_per_worker: int | None = Field(None, ge=1)
"""GPUs owned by each discovered Dynamo worker. Auto-derived from the local inference topology."""

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

Expand Down Expand Up @@ -255,7 +267,8 @@ class MetricsServerConfig(BaseConfig):


class BaseTransportConfig(BaseConfig):
pass
send_timeout_seconds: float = Field(300.0, gt=0, allow_inf_nan=False)
"""Maximum time to ship one rollout batch before the orchestrator fails closed."""


class FileSystemTransportConfig(BaseTransportConfig):
Expand Down
7 changes: 6 additions & 1 deletion packages/prime-rl-configs/src/prime_rl/utils/validation.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,12 @@ def propagate(shared_path: str, *targets: str) -> None:
# the same SLURM allocation, so the nested inference inherits [slurm]. This is
# what lets the nested InferenceConfig's multi-node / disaggregated SLURM check
# pass (the per-rank inference.toml drops slurm, so each rank still runs locally).
propagate("slurm", "inference.slurm")
# A single-node Dynamo job instead enters the local RL launcher inside its
# allocation, where the existing local Dynamo supervisor owns the topology.
deployment_type = get("deployment.type") or "single_node"
is_single_node_dynamo = deployment_type == "single_node" and get("inference.backend.type") == "dynamo"
if not is_single_node_dynamo:
propagate("slurm", "inference.slurm")

# output_dir: orchestrator gets a ``/run_default`` subdir so trainer +
# orchestrator nest under the same experiment root without colliding.
Expand Down
Loading