Add PPO and SAO value-model training

packages/prime-rl-configs/src/prime_rl/configs/algorithm.py

@@ -1,7 +1,7 @@
 """Algorithm abstraction: sampling and the per-token training signal.
 
 An algorithm is a named, self-contained config — a discriminated union keyed
-on ``type`` (``grpo``, ``max_rl``, ``rae``, ``hierarchical_grpo``, ``opd``,
+on ``type`` (``grpo``, ``ppo``, ``sao``, ``max_rl``, ``rae``, ``hierarchical_grpo``, ``opd``,
 ``opsd``, ``sft``, ``echo``, ``debug``).
 The bundle *is* the algorithm: each variant carries
 its sampling component and its credit-assignment / loss-routing parameters,
@@ -24,9 +24,9 @@
    to env-provided observation tokens (masked out by default; ``echo`` trains on
    them with weighted CE).
 
-prime-rl only ever hosts the trainable policy. Every other model an algorithm
-uses is an external OpenAI-compatible endpoint, declared inline on the
-algorithm that uses it (a :class:`FrozenModelConfig`). Model roles like
+Frozen model references use external OpenAI-compatible endpoints, declared inline on the
+algorithm that uses them (a :class:`FrozenModelConfig`). PPO and SAO use the
+separate value trainer configured under ``value``. Model roles like
 "teacher" are algorithm-local vocabulary over these references; the pipeline
 branches on liveness alone. The trainer is algorithm-blind: the loss is a sum
 of three components (rl, ce, ref_kl), each normalized by its own global token
@@ -48,8 +48,7 @@ class FrozenModelConfig(ClientConfig):
     """An externally hosted model behind an OpenAI-compatible endpoint: the
     client config plus the served model's ``name``.
 
-    prime-rl never launches or updates these — only the trainable policy is
-    ever hosted by prime-rl itself. Frozen models are reachable-but-unmanaged:
+    prime-rl never launches or updates these. Frozen models are reachable-but-unmanaged:
     ``base_url`` is required, their weights never change, and rollouts or
     scores from them never go stale (stable prefix cache, no off-policy
     aging)."""
@@ -209,6 +208,22 @@ class GRPOAlgoConfig(BaseAlgoConfig):
     """Linear length penalty subtracted from each reward before the GRPO baseline (see ``LinearLengthPenaltyConfig``): a ``pass_rate``-scaled sum of output-token, input-token, and turns terms, each normalized by the group's own max for that quantity. None disables it."""
 
 
+class PPOAlgoConfig(BaseAlgoConfig):
+    type: Literal["ppo"] = "ppo"
+    action_loss_type: ClassVar[ActionLossType] = "rl"
+    value_url: str = "http://127.0.0.1:8123"
+    value_seq_len: int | None = Field(None, ge=1)
+    gamma: float = Field(1.0, ge=0, le=1)
+    policy_lambda: float = Field(0.95, ge=0, le=1)
+    value_lambda: float = Field(1.0, ge=0, le=1)
+    length_adaptive_alpha: float | None = Field(None, gt=0)
+
+
+class SAOAlgoConfig(PPOAlgoConfig):
+    type: Literal["sao"] = "sao"  # type: ignore[assignment]
+    length_adaptive_alpha: float | None = Field(1.5, gt=0)
+
+
 class EchoAlgoConfig(GRPOAlgoConfig):
     type: Literal["echo"] = "echo"  # type: ignore[assignment]
     """ECHO: group-relative advantage on action tokens (GRPO), plus weighted
@@ -393,6 +408,8 @@ def require_nonzero_advantage(self):
 
 AlgoConfig: TypeAlias = Annotated[
     GRPOAlgoConfig
+    | PPOAlgoConfig
+    | SAOAlgoConfig
     | EchoAlgoConfig
     | MaxRLAlgoConfig
     | RAEAlgoConfig
@@ -408,6 +425,8 @@ def require_nonzero_advantage(self):
 its class defaults are the vetted setting.
 
 - ``grpo`` — policy group sampling, group-relative advantage, RL loss (the default).
+- ``ppo`` — grouped sampling with token-value GAE and a clipped policy surrogate.
+- ``sao`` — grouped sampling with token-value GAE and a direct ratio interval mask.
 - ``max_rl`` — GRPO with mean-normalized advantages (maximum-likelihood RL).
 - ``rae`` — reward minus a per-agent EMA baseline (SPIRAL), for multi-agent self-play envs.
 - ``hierarchical_grpo`` — GRPO for proposer-solver envs: solvers are compared within one proposed problem and proposers across proposals. Needs ``episode_agents``.

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -631,6 +631,17 @@ class OrchestratorConfig(BaseConfig):
     rollout_transport: TransportConfig = ZMQTransportConfig()
     """Transport used to ship rollouts from orchestrator to trainer."""
 
+    value_rollout_transport: TransportConfig | None = None
+    """Transport used to ship PPO batches to the separate value trainer."""
+    value_service_url: str | None = None
+    """HTTP endpoint for PPO value scoring and update status."""
+    value_seq_len: int | None = Field(None, ge=1)
+    """Critic sequence length for independent value batch packing."""
+    value_num_train_workers: int | None = Field(None, ge=1)
+    """Critic DP rank count; one ZMQ stream is sent per rank."""
+    value_pad_to_multiple_of: int | None = Field(None, ge=1)
+    """Critic CP degree for padding value batches."""
+
     output_dir: Path = Field(default_factory=default_output_dir)
     """Directory to write outputs to — checkpoints, weights, rollouts, and logs are written as subdirectories. Shared with the trainer; should be a persistent directory with enough disk space and unique per experiment running on a single node. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``."""
 
@@ -697,6 +708,14 @@ def inherit_env_algorithms(self):
                 env_cfg.algo = self.algo.model_copy(deep=True)
         return self
 
+    @model_validator(mode="after")
+    def sync_value_service_url(self):
+        if self.value_service_url is not None:
+            for source in self.train.source:
+                if source.algo is not None and source.algo.type in ("ppo", "sao"):
+                    source.algo.value_url = self.value_service_url
+        return self
+
     @model_validator(mode="after")
     def validate_env_algorithms(self):
         """Let each algorithm reject environments it cannot score correctly."""

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -32,15 +32,19 @@
     FileSystemWeightBroadcastConfig as TrainerFileSystemWeightBroadcastConfig,
 )
 from prime_rl.configs.trainer import (
-    NCCLWeightBroadcastConfig as TrainerNCCLWeightBroadcastConfig,
+    ModelConfig,
+    PPOLossConfig,
+    SAOLossConfig,
+    TokenizerConfig,
+    TrainerConfig,
 )
 from prime_rl.configs.trainer import (
-    NIXLWeightBroadcastConfig as TrainerNIXLWeightBroadcastConfig,
+    NCCLWeightBroadcastConfig as TrainerNCCLWeightBroadcastConfig,
 )
 from prime_rl.configs.trainer import (
-    TokenizerConfig,
-    TrainerConfig,
+    NIXLWeightBroadcastConfig as TrainerNIXLWeightBroadcastConfig,
 )
+from prime_rl.configs.value import ValueConfig
 from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource
 from prime_rl.utils.validation import (
     propagate_shared_fields,
@@ -186,12 +190,16 @@ class SingleNodeDeploymentConfig(BaseDeploymentConfig):
     num_infer_gpus: int = 1
     """GPUs allocated to inference."""
 
+    num_value_gpus: int = Field(0, ge=0)
+    """GPUs allocated to a separate PPO value trainer."""
+
     @model_validator(mode="after")
     def validate_gpu_count(self):
-        total = self.num_train_gpus + self.num_infer_gpus
+        total = self.num_train_gpus + self.num_infer_gpus + self.num_value_gpus
         if total > self.gpus_per_node:
             raise ValueError(
-                f"Total GPU count ({total} = {self.num_train_gpus} train + {self.num_infer_gpus} infer)"
+                f"Total GPU count ({total} = {self.num_train_gpus} train + {self.num_infer_gpus} infer + "
+                f"{self.num_value_gpus} value)"
                 f" exceeds gpus_per_node ({self.gpus_per_node})."
             )
         return self
@@ -203,6 +211,12 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig):
     num_train_nodes: int
     """Training nodes."""
 
+    num_value_nodes: int = Field(0, ge=0)
+    """Dedicated value-trainer nodes for PPO/SAO."""
+
+    num_value_gpus_per_train_node: int = Field(0, ge=0)
+    """GPUs reserved for the critic on each training node instead of dedicated value nodes."""
+
     num_infer_nodes: int | None = Field(None, ge=0)
     """Inference nodes per replica. If unset, inferred from ``inference.deployment``. Set to 0 to skip inference and orchestrator (requires fake data)."""
 
@@ -215,6 +229,26 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig):
     orchestrator_on_inference: bool = False
     """Run the orchestrator on the last inference node instead of trainer rank 0 (frees host RAM on the trainer node)."""
 
+    @model_validator(mode="after")
+    def validate_value_gpu_layout(self):
+        if self.num_value_gpus_per_train_node >= self.gpus_per_node:
+            raise ValueError("num_value_gpus_per_train_node must leave at least one policy GPU per training node")
+        if self.num_value_nodes and self.num_value_gpus_per_train_node:
+            raise ValueError("Use dedicated value nodes or shared training-node GPUs, not both")
+        return self
+
+    @property
+    def train_gpus_per_node(self) -> int:
+        return self.gpus_per_node - self.num_value_gpus_per_train_node
+
+    @property
+    def value_gpus_per_node(self) -> int:
+        return self.num_value_gpus_per_train_node or self.gpus_per_node
+
+    @property
+    def value_trainer_nodes(self) -> int:
+        return self.num_train_nodes if self.num_value_gpus_per_train_node else self.num_value_nodes
+
     @property
     def infer_nodes_per_replica(self) -> int:
         return self.num_infer_nodes or 0
@@ -232,6 +266,9 @@ def total_infer_nodes(self) -> int:
 class RLConfig(BaseConfig):
     trainer: TrainerConfig
 
+    value: ValueConfig | None = None
+    """Separate PPO/SAO value trainer."""
+
     orchestrator: OrchestratorConfig
 
     inference: InferenceConfig | None = None
@@ -337,6 +374,20 @@ def auto_setup_infer_nodes(self):
 
     @model_validator(mode="after")
     def validate_deployment(self):
+        if self.value is not None:
+            if self.deployment.type == "single_node" and self.deployment.num_value_gpus == 0:
+                raise ValueError("A value trainer requires num_value_gpus > 0")
+            if self.deployment.type == "multi_node":
+                if self.deployment.value_trainer_nodes == 0:
+                    raise ValueError("A multi-node value trainer requires value nodes or GPUs on training nodes")
+                if self.deployment.infer_nodes_per_replica == 0:
+                    raise ValueError("A multi-node value trainer requires inference nodes for the orchestrator")
+        elif self.deployment.type == "single_node" and self.deployment.num_value_gpus:
+            raise ValueError("num_value_gpus requires a [value] configuration")
+        elif self.deployment.type == "multi_node" and (
+            self.deployment.num_value_nodes or self.deployment.num_value_gpus_per_train_node
+        ):
+            raise ValueError("Value GPUs require a [value] configuration")
         if self.deployment.type == "multi_node":
             if self.slurm is None:
                 raise ValueError("Must use SLURM for multi-node deployment.")
@@ -393,6 +444,9 @@ def auto_setup_run_dir(self):
                     "output_dir / run.name — set those instead."
                 )
             sub.output_dir = run_dir
+        if self.value is not None:
+            self.value.output_dir = run_dir / "value"
+            self.value.rollout_dir = run_dir
         return self
 
     @model_validator(mode="after")
@@ -402,6 +456,8 @@ def auto_setup_resume(self):
             return self
         self.trainer.resume = self.resume.model_copy()
         self.orchestrator.resume = self.resume.model_copy()
+        if self.value is not None:
+            self.value.resume = self.resume.model_copy()
         return self
 
     @model_validator(mode="after")
@@ -516,6 +572,84 @@ def auto_setup_rollout_transport(self):
             self.rollout_transport = self.trainer.rollout_transport
         return self
 
+    @model_validator(mode="after")
+    def auto_setup_value(self):
+        if self.value is None:
+            if any(
+                source.algo is not None and source.algo.type in ("ppo", "sao")
+                for source in self.orchestrator.train.source
+            ):
+                raise ValueError("PPO/SAO requires a [value] configuration")
+            return self
+        algo_types = {source.algo.type for source in self.orchestrator.train.source if source.algo is not None}
+        if len(algo_types) != 1 or not algo_types <= {"ppo", "sao"}:
+            raise ValueError("A value trainer requires one PPO or SAO algorithm across all training sources")
+        algorithm_type = next(iter(algo_types))
+        if algorithm_type == "sao" and "freeze_attention" not in self.value.model_fields_set:
+            self.value.freeze_attention = True
+        if "loss" not in self.trainer.model_fields_set:
+            self.trainer.loss = PPOLossConfig() if algorithm_type == "ppo" else SAOLossConfig()
+        elif self.trainer.loss.type != algorithm_type:
+            raise ValueError("trainer.loss.type must match the PPO or SAO algorithm")
+        model_overrides = self.value.model.model_dump(exclude_unset=True)
+        self.value.model = ModelConfig.model_validate({**self.trainer.model.model_dump(), **model_overrides})
+        if self.value.model.cp > 1 and self.value.model.seq_len % (2 * self.value.model.cp):
+            raise ValueError("value.model.seq_len must be divisible by 2 * value.model.cp")
+        if self.value.model.name != self.trainer.model.name:
+            raise ValueError("value.model.name must match trainer.model.name")
+        if self.deployment.type == "single_node":
+            num_train_gpus = self.deployment.num_train_gpus
+            num_value_gpus = self.deployment.num_value_gpus
+        else:
+            num_train_gpus = self.deployment.num_train_nodes * self.deployment.train_gpus_per_node
+            num_value_gpus = self.deployment.value_trainer_nodes * self.deployment.value_gpus_per_node
+        if num_train_gpus % self.trainer.model.cp:
+            raise ValueError("Policy GPU count must be divisible by trainer.model.cp")
+        if num_value_gpus % self.value.model.cp:
+            raise ValueError("Value GPU count must be divisible by value.model.cp")
+        if (
+            self.value.model.lora is not None
+            and "experts" in self.value.model.lora.target_modules
+            and self.value.model.lora.rank % 8
+        ):
+            raise ValueError("Expert LoRA requires a rank divisible by 8")
+        if self.value.policy_sync_interval is not None:
+            if self.value.model.lora is None:
+                raise ValueError("value.policy_sync_interval requires value.model.lora")
+            if self.trainer.model.lora is not None:
+                raise ValueError("Policy backbone sync requires a full-weight policy trainer")
+            if self.value.model.cp != self.trainer.model.cp or self.value.model.ep != self.trainer.model.ep:
+                raise ValueError("Policy and value models must use the same CP and EP for backbone sync")
+            if num_train_gpus != num_value_gpus:
+                raise ValueError("Policy and value models must use the same GPU count for backbone sync")
+            self.value.policy_sync_dir = self.run_dir / "policy_sync"
+            self.trainer.ppo_policy_sync_dir = self.value.policy_sync_dir
+            self.trainer.ppo_policy_sync_interval = self.value.policy_sync_interval
+        if self.trainer.rollout_transport.type != "zmq":
+            raise ValueError("The separate value trainer currently requires ZMQ rollout transport")
+        self.value.rollout_transport = self.trainer.rollout_transport.model_copy(deep=True)
+        self.value.rollout_transport.port += 2
+        self.orchestrator.value_rollout_transport = self.value.rollout_transport.model_copy(deep=True)
+        self.orchestrator.value_service_url = f"http://{self.value.service_host}:{self.value.service_port}"
+        self.orchestrator.value_seq_len = self.value.model.seq_len
+        self.orchestrator.value_num_train_workers = num_value_gpus // self.value.model.cp
+        self.orchestrator.value_pad_to_multiple_of = self.value.model.cp
+        self.value.max_steps = self.trainer.max_steps
+        self.value.ckpt = self.value.ckpt or self.trainer.ckpt
+        if self.value.resume is not None and self.value.resume.dir is not None:
+            source_run_dir = self.value.resume.dir.parent.parent
+            self.value.resume.dir = source_run_dir / "value" / "checkpoints" / self.value.resume.dir.name
+        self.trainer.ppo_head_warmup_steps = max(0, self.value.head_warmup_steps - self.value.pretrain_steps)
+        for source in self.orchestrator.train.source:
+            algo = source.algo
+            if algo is not None and "value_url" not in algo.model_fields_set:
+                algo.value_url = f"http://{self.value.service_host}:{self.value.service_port}"
+            if algo is not None and algo.value_seq_len is None:
+                algo.value_seq_len = min(self.value.model.seq_len, self.orchestrator.seq_len)
+            if algo is not None and algo.value_seq_len > self.value.model.seq_len:
+                raise ValueError("algo.value_seq_len cannot exceed value.model.seq_len")
+        return self
+
     @model_validator(mode="after")
     def validate_eplb(self):
         if self.inference is not None and self.inference.vllm.enable_eplb:
@@ -710,7 +844,7 @@ def auto_setup_deployment(self):
 
         elif self.deployment.type == "multi_node":  # multi-node
             self.orchestrator.num_train_workers = (
-                self.deployment.num_train_nodes * self.deployment.gpus_per_node // self.trainer.model.cp
+                self.deployment.num_train_nodes * self.deployment.train_gpus_per_node // self.trainer.model.cp
             )
 
             if self.deployment.nodes_per_fsdp_group is not None:

packages/prime-rl-configs/src/prime_rl/configs/shared.py

@@ -91,6 +91,9 @@ class SlurmConfig(BaseConfig):
     partition: str = "cluster"
     """SLURM partition (#SBATCH --partition)."""
 
+    cpus_per_task: int | None = Field(None, ge=1)
+    """CPUs allocated to each node's launcher task (#SBATCH --cpus-per-task)."""
+
     nodelist: str | None = None
     """Comma-separated list of specific nodes to run on (#SBATCH --nodelist)."""
 
@@ -122,6 +125,7 @@ def template_vars(self) -> dict:
             "job_name": self.job_name,
             "project_dir": self.project_dir,
             "partition": self.partition,
+            "cpus_per_task": self.cpus_per_task,
             "nodelist": self.nodelist,
             "exclude": self.exclude,
             "account": self.account,

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -67,6 +67,9 @@ class OptimizerInBackwardOffloadConfig(BaseConfig):
     numa_bind: bool = True
     """Pin each rank's CPUs to its GPU's NUMA node. Disable when the launcher already manages CPU affinity or GPU sysfs topology is unavailable."""
 
+    timeout_seconds: float = Field(120.0, gt=0)
+    """Maximum wait for a native CPU optimizer transfer or update to complete."""
+
 
 def _normalize_optimizer_in_backward_offload(value: Any) -> Any:
     if value is True:
@@ -638,6 +641,12 @@ class CISPOLossConfig(BaseConfig):
     """Temperature for the advantage term."""
 
 
+class SAOLossConfig(IcePopLossConfig):
+    type: Literal["sao"] = "sao"  # type: ignore[assignment]
+    ratio_low: float = Field(0.7, gt=0)
+    ratio_high: float = Field(6.0, gt=0)
+
+
 class CustomLossConfig(BaseConfig):
     type: Literal["custom"] = "custom"
 
@@ -649,7 +658,8 @@ class CustomLossConfig(BaseConfig):
 
 
 LossConfig: TypeAlias = Annotated[
-    IPOLossConfig | IcePopLossConfig | PPOLossConfig | CISPOLossConfig | CustomLossConfig, Field(discriminator="type")
+    IPOLossConfig | IcePopLossConfig | PPOLossConfig | CISPOLossConfig | SAOLossConfig | CustomLossConfig,
+    Field(discriminator="type"),
 ]
 
 
@@ -718,6 +728,11 @@ class TrainerConfig(BaseConfig):
     loss: LossConfig = IPOLossConfig()
     """Loss config for the rl loss component (see ``setup_rl_loss_fn``). The ce / ref_kl components are fixed and do not read this."""
 
+    ppo_head_warmup_steps: int = Field(0, ge=0)
+    """Live critic head warmup steps before policy-backbone snapshots begin."""
+    ppo_policy_sync_interval: int | None = Field(None, ge=1)
+    ppo_policy_sync_dir: Path | None = None
+
     optim: OptimizerConfig = AdamWConfig()
 
     scheduler: SchedulerConfig = ConstantSchedulerConfig()