From e2230d8431481d260370d5336500e47b1252c661 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 07:53:31 +0000 Subject: [PATCH 1/9] push ppo value-based model changes --- .../src/prime_rl/configs/algorithm.py | 31 +- .../src/prime_rl/configs/orchestrator.py | 5 + .../src/prime_rl/configs/rl.py | 116 ++++- .../src/prime_rl/configs/trainer.py | 32 +- .../src/prime_rl/configs/value.py | 45 ++ src/prime_rl/entrypoints/rl.py | 65 ++- src/prime_rl/orchestrator/algo/__init__.py | 4 + src/prime_rl/orchestrator/algo/base.py | 7 + src/prime_rl/orchestrator/algo/gae.py | 48 ++ src/prime_rl/orchestrator/algo/ppo.py | 89 ++++ src/prime_rl/orchestrator/orchestrator.py | 33 ++ src/prime_rl/orchestrator/train_sink.py | 3 + .../templates/multi_node_rl.sbatch.j2 | 56 ++- src/prime_rl/trainer/batch.py | 36 ++ src/prime_rl/trainer/model.py | 70 ++- src/prime_rl/trainer/models/layers/lm_head.py | 10 + .../trainer/models/layers/lora/multi_moe.py | 85 +--- src/prime_rl/trainer/optim/offload.py | 3 +- src/prime_rl/trainer/rl/data.py | 18 + src/prime_rl/trainer/rl/loss.py | 51 +- src/prime_rl/trainer/rl/policy_sync.py | 62 +++ src/prime_rl/trainer/rl/train.py | 12 +- src/prime_rl/trainer/rl/value.py | 450 ++++++++++++++++++ src/prime_rl/transports/batch/types.py | 10 + src/prime_rl/utils/pathing.py | 1 + tests/unit/orchestrator/test_advantage.py | 25 + tests/unit/orchestrator/test_batch.py | 21 + tests/unit/test_configs.py | 42 ++ tests/unit/train/rl/test_loss.py | 40 +- 29 files changed, 1341 insertions(+), 129 deletions(-) create mode 100644 packages/prime-rl-configs/src/prime_rl/configs/value.py create mode 100644 src/prime_rl/orchestrator/algo/gae.py create mode 100644 src/prime_rl/orchestrator/algo/ppo.py create mode 100644 src/prime_rl/trainer/rl/policy_sync.py create mode 100644 src/prime_rl/trainer/rl/value.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py index dca5242f283..e7be982eead 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py +++ b/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``. 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 1c573cd1825..2a565a70965 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -567,6 +567,11 @@ 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.""" + 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``.""" 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 65e6ed74c92..f65be56cf5c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/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, + PPOClipLossConfig, + 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,9 @@ 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_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).""" @@ -232,6 +243,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 +351,18 @@ 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.num_value_nodes == 0: + raise ValueError("A multi-node value trainer requires num_value_nodes > 0") + 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: + raise ValueError("num_value_nodes requires 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 +419,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 +431,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 +547,79 @@ 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 = PPOClipLossConfig() 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.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.gpus_per_node + num_value_gpus = self.deployment.num_value_nodes * self.deployment.gpus_per_node + 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 on the Qwen3 MoE value model 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 num_value_gpus != self.value.model.cp: + raise ValueError("The value trainer currently requires one data-parallel worker") + if num_train_gpus != self.trainer.model.cp: + raise ValueError("The value trainer currently requires one policy data-parallel worker") + 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.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: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 0bf4704355c..0050ea8961f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/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: @@ -600,6 +603,25 @@ def validate_ratio_bounds(self): return self +class PPOClipLossConfig(BaseConfig): + type: Literal["ppo"] = "ppo" + + eps_low: float = Field(0.2, ge=0, lt=1) + """Lower PPO ratio bound is 1 - eps_low.""" + + eps_high: float = Field(0.2, ge=0) + """Upper PPO ratio bound is 1 + eps_high.""" + + adv_tau: float = Field(1.0, ge=0) + """Scale applied to the advantage.""" + + +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" @@ -610,7 +632,10 @@ class CustomLossConfig(BaseConfig): """Kwargs forwarded to the loss function.""" -LossConfig: TypeAlias = Annotated[IPOLossConfig | IcePopLossConfig | CustomLossConfig, Field(discriminator="type")] +LossConfig: TypeAlias = Annotated[ + IPOLossConfig | IcePopLossConfig | PPOClipLossConfig | SAOLossConfig | CustomLossConfig, + Field(discriminator="type"), +] class FakeDataLoaderConfig(BaseConfig): @@ -678,6 +703,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() diff --git a/packages/prime-rl-configs/src/prime_rl/configs/value.py b/packages/prime-rl-configs/src/prime_rl/configs/value.py new file mode 100644 index 00000000000..34de4efc859 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/value.py @@ -0,0 +1,45 @@ +"""Configuration for the separate PPO value trainer.""" + +from pathlib import Path + +from pydantic import Field, model_validator + +from prime_rl.configs.shared import EnvVars, ResumeConfig, TransportConfig, ZMQTransportConfig +from prime_rl.configs.trainer import ( + AdamWConfig, + CheckpointConfig, + ConstantSchedulerConfig, + ModelConfig, + OptimizerConfig, + SchedulerConfig, +) +from prime_rl.utils.config import BaseConfig, default_output_dir + + +class ValueConfig(BaseConfig): + model: ModelConfig = ModelConfig() + optim: OptimizerConfig = AdamWConfig(lr=5e-6) + scheduler: SchedulerConfig = ConstantSchedulerConfig() + ckpt: CheckpointConfig | None = None + resume: ResumeConfig | None = None + rollout_transport: TransportConfig = ZMQTransportConfig() + output_dir: Path = Field(default_factory=default_output_dir) + rollout_dir: Path = Field(default_factory=default_output_dir) + max_steps: int | None = None + head_warmup_steps: int = Field(10, ge=0) + updates_per_step: int = Field(2, ge=1) + freeze_attention: bool = False + pretrain_data: Path | None = None + pretrain_steps: int = Field(0, ge=0) + service_host: str = "127.0.0.1" + service_port: int = Field(8123, ge=1, le=65535) + policy_sync_interval: int | None = Field(None, ge=1) + policy_sync_retune_updates: int = Field(2, ge=1) + policy_sync_dir: Path | None = None + env_vars: EnvVars = {} + + @model_validator(mode="after") + def validate_pretraining(self): + if self.pretrain_steps and self.pretrain_data is None: + raise ValueError("value.pretrain_steps requires value.pretrain_data") + return self diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 2c9832ddedb..c0e3c308385 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -41,11 +41,13 @@ monitor_process, set_proc_title, ) +from prime_rl.utils.utils import get_free_port RL_CONFIG = "rl.json" RL_SBATCH = "rl.sbatch" TRAINER_CONFIG = "trainer.json" +VALUE_CONFIG = "value.json" ORCHESTRATOR_CONFIG = "orchestrator.json" INFERENCE_CONFIG = "inference.json" @@ -81,6 +83,8 @@ def rl_config_components(config: RLConfig, config_dir: Path) -> list[tuple[str, ] if config.inference is not None: components.append(("Inference", config_dir / INFERENCE_CONFIG)) + if config.value is not None: + components.append(("Value", config_dir / VALUE_CONFIG)) if env_servers(config): components.append(("Envs", f"{config_dir}/{ENVS_DIR}/*/*.json")) return components @@ -93,6 +97,10 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None: with open(output_dir / TRAINER_CONFIG, "w") as f: json.dump(dump_resolved_config(config.trainer), f, indent=2) + if config.value is not None: + with open(output_dir / VALUE_CONFIG, "w") as f: + json.dump(dump_resolved_config(config.value), f, indent=2) + with open(output_dir / ORCHESTRATOR_CONFIG, "w") as f: json.dump(dump_resolved_config(config.orchestrator), f, indent=2) @@ -147,8 +155,10 @@ def rl_local(config: RLConfig): infer_local_gpu_ids = list(range(gpu_offset, gpu_offset + num_infer_gpus)) gpu_offset += num_infer_gpus trainer_local_gpu_ids = list(range(gpu_offset, gpu_offset + config.deployment.num_train_gpus)) + gpu_offset += config.deployment.num_train_gpus + value_local_gpu_ids = list(range(gpu_offset, gpu_offset + config.deployment.num_value_gpus)) - total_requested_gpus = num_infer_gpus + config.deployment.num_train_gpus + total_requested_gpus = num_infer_gpus + config.deployment.num_train_gpus + config.deployment.num_value_gpus physical_gpu_ids = get_physical_gpu_ids() if total_requested_gpus > len(physical_gpu_ids): raise ValueError( @@ -160,6 +170,7 @@ def rl_local(config: RLConfig): infer_gpu_ids = [physical_gpu_mapping[local_gpu_id] for local_gpu_id in infer_local_gpu_ids] trainer_gpu_ids = [physical_gpu_mapping[local_gpu_id] for local_gpu_id in trainer_local_gpu_ids] + value_gpu_ids = [physical_gpu_mapping[local_gpu_id] for local_gpu_id in value_local_gpu_ids] start_command = sys.argv logger.debug(f"RL start command: {' '.join(start_command)}") @@ -241,6 +252,45 @@ def sigterm_handler(signum, frame): "will hang waiting for it." ) + if config.value is not None: + value_cmd = [ + "torchrun", + "--role=value", + f"--rdzv-endpoint=localhost:{get_free_port()}", + f"--rdzv-id={uuid.uuid4().hex}", + f"--log-dir={log_dir / 'value' / 'torchrun'}", + "--redirect=3", + "--tee=3", + f"--nproc-per-node={len(value_gpu_ids)}", + "-m", + "prime_rl.trainer.rl.value", + "@", + (config_dir / VALUE_CONFIG).as_posix(), + ] + logger.info(f"Starting value trainer on GPU(s) {' '.join(map(str, value_gpu_ids))}") + with open(log_dir / "value.log", "w") as log_file: + value_process = Popen( + value_cmd, + env={ + **os.environ, + **DEFAULT_COMMON_ENV_VARS, + **DEFAULT_TRAINER_ENV_VARS, + **config.env_vars, + **config.value.env_vars, + "CUDA_VISIBLE_DEVICES": ",".join(map(str, value_gpu_ids)), + }, + stdout=log_file, + stderr=log_file, + ) + processes.append(value_process) + stop_event = Event() + stop_events["value"] = stop_event + monitor_thread = Thread( + target=monitor_process, args=(value_process, stop_event, error_queue, "value"), daemon=True + ) + monitor_thread.start() + monitor_threads.append(monitor_thread) + frozen_endpoints: list[str] = [] for env in config.orchestrator.train.source: algo = env.algo @@ -325,8 +375,6 @@ def sigterm_handler(signum, frame): monitor_threads.append(monitor_thread) # Start training process - from prime_rl.utils.utils import get_free_port - trainer_cmd = [ "torchrun", "--role=trainer", @@ -451,6 +499,11 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script **config.env_vars, **config.trainer.env_vars, } + value_env_vars = ( + {**DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars, **config.value.env_vars} + if config.value is not None + else {} + ) orchestrator_env_vars = {**DEFAULT_COMMON_ENV_VARS, **config.env_vars, **config.orchestrator.env_vars} inference_env_vars = ( {**DEFAULT_COMMON_ENV_VARS, **DEFAULT_INFERENCE_ENV_VARS, **config.env_vars, **config.inference.env_vars} @@ -501,6 +554,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script launcher_dir=get_launcher_dir(config.run_dir), launcher_log_dir=get_launcher_log_dir(config.run_dir), num_train_nodes=config.deployment.num_train_nodes, + num_value_nodes=config.deployment.num_value_nodes, num_infer_nodes=infer_deploy.num_nodes * config.deployment.num_infer_replicas, nodes_per_infer_replica=infer_deploy.num_nodes, num_infer_replicas=config.deployment.num_infer_replicas, @@ -521,6 +575,8 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script prefill_env_vars=infer_deploy.prefill_env_vars, decode_env_vars=infer_deploy.decode_env_vars, trainer_env_vars=trainer_env_vars, + value_env_vars=value_env_vars, + value_service_port=config.value.service_port if config.value is not None else 0, orchestrator_env_vars=orchestrator_env_vars, inference_env_vars=inference_env_vars, prefill_vllm_extra_json=vllm_overrides_fragment(infer_deploy.prefill_vllm_overrides), @@ -546,6 +602,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script launcher_dir=get_launcher_dir(config.run_dir), launcher_log_dir=get_launcher_log_dir(config.run_dir), num_train_nodes=config.deployment.num_train_nodes, + num_value_nodes=config.deployment.num_value_nodes, num_infer_nodes=config.deployment.total_infer_nodes, nodes_per_infer_replica=config.deployment.infer_nodes_per_replica, num_infer_replicas=config.deployment.num_infer_replicas, @@ -570,6 +627,8 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script ranks_filter=",".join(map(str, config.trainer.log.ranks_filter)), orchestrator_on_inference=config.deployment.orchestrator_on_inference, trainer_env_vars=trainer_env_vars, + value_env_vars=value_env_vars, + value_service_port=config.value.service_port if config.value is not None else 0, orchestrator_env_vars=orchestrator_env_vars, inference_env_vars=inference_env_vars, train_env_names=train_env_names, diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 16c4312b0d3..215846a53d5 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/src/prime_rl/orchestrator/algo/__init__.py @@ -30,6 +30,7 @@ class owns its scoring hooks from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm from prime_rl.orchestrator.algo.opd import OPDAlgorithm from prime_rl.orchestrator.algo.opsd import OPSDAlgorithm +from prime_rl.orchestrator.algo.ppo import PPOAlgorithm from prime_rl.orchestrator.algo.rae import RAEAlgorithm from prime_rl.orchestrator.algo.routing import assign_advantages, stamp_loss_routing from prime_rl.orchestrator.algo.sft import SFTDistillAlgorithm @@ -42,6 +43,8 @@ class owns its scoring hooks # each config class's defaults are its vetted parameterization. ALGORITHM_CLASSES: dict[str, type[Algorithm]] = { "grpo": GRPOAlgorithm, + "ppo": PPOAlgorithm, + "sao": PPOAlgorithm, "echo": EchoAlgorithm, "max_rl": MaxRLAlgorithm, "rae": RAEAlgorithm, @@ -73,6 +76,7 @@ def build_algorithm(config: AlgoConfig, clients: InferenceClient) -> Algorithm: "MaxRLAlgorithm", "OPDAlgorithm", "OPSDAlgorithm", + "PPOAlgorithm", "RAEAlgorithm", "SFTDistillAlgorithm", "build_algorithm", diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index a535acaf34a..b4c57492144 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -7,6 +7,7 @@ import verifiers.v1 as vf from prime_rl.configs.algorithm import ActionLossType, AlgoConfig, FrozenModelConfig +from prime_rl.transports.batch import TrainingSample from prime_rl.utils.logger import get_logger if TYPE_CHECKING: @@ -59,6 +60,9 @@ def __init__(self, config: AlgoConfig, clients: InferenceClient): async def setup(self) -> None: """Connect resources owned by the algorithm.""" + async def aclose(self) -> None: + """Close resources owned by the algorithm.""" + async def connect(self, reference: FrozenModelConfig) -> InferenceClient: """Connect and track the frozen model pool owned by this algorithm.""" self.connected = await connect_frozen_client(reference) @@ -70,6 +74,9 @@ async def score_episode(self, episode: vf.Episode) -> None: async def score_group(self, episodes: list[vf.Episode]) -> None: """Assign group-relative annotations to a finalized cohort.""" + async def score_samples(self, trace: vf.Trace, samples: list[TrainingSample]) -> None: + """Assign token-level training data after branch samples are materialized.""" + async def finalize_episode(self, episode: vf.Episode) -> None: """Run rollout-local scoring when the episode has trainable traces.""" if any(True for _ in iter_trainable_traces([episode])): diff --git a/src/prime_rl/orchestrator/algo/gae.py b/src/prime_rl/orchestrator/algo/gae.py new file mode 100644 index 00000000000..79fe4947a86 --- /dev/null +++ b/src/prime_rl/orchestrator/algo/gae.py @@ -0,0 +1,48 @@ +"""Token-level GAE over model actions in a multi-turn trajectory.""" + +import math +from collections.abc import Sequence + + +def skip_observation_gae( + values: Sequence[float], + action_mask: Sequence[bool], + reward: float, + *, + gamma: float, + policy_lambda: float, + value_lambda: float, + bootstrap_value: float = 0.0, +) -> tuple[list[float], list[float]]: + """Return policy advantages and critic targets, aligned to the full trace. + + Environment observation positions carry zero outputs. The terminal reward + lands on the last model action, and consecutive actions in the recursion + can be separated by any number of observation tokens. + """ + if len(values) != len(action_mask): + raise ValueError("values and action_mask must have the same length") + if ( + any(not math.isfinite(value) for value in values) + or not math.isfinite(reward) + or not math.isfinite(bootstrap_value) + ): + raise ValueError("values, reward, and bootstrap_value must be finite") + if not 0 <= gamma <= 1 or not 0 <= policy_lambda <= 1 or not 0 <= value_lambda <= 1: + raise ValueError("gamma and GAE lambdas must be in [0, 1]") + + advantages = [0.0] * len(values) + returns = [0.0] * len(values) + action_indices = [index for index, sampled in enumerate(action_mask) if sampled] + next_value = bootstrap_value + policy_gae = 0.0 + value_gae = 0.0 + for index in reversed(action_indices): + terminal = index == action_indices[-1] + delta = (reward if terminal else 0.0) + gamma * next_value - values[index] + policy_gae = delta + gamma * policy_lambda * policy_gae + value_gae = delta + gamma * value_lambda * value_gae + advantages[index] = policy_gae + returns[index] = value_gae + values[index] + next_value = values[index] + return advantages, returns diff --git a/src/prime_rl/orchestrator/algo/ppo.py b/src/prime_rl/orchestrator/algo/ppo.py new file mode 100644 index 00000000000..b680e20497e --- /dev/null +++ b/src/prime_rl/orchestrator/algo/ppo.py @@ -0,0 +1,89 @@ +"""PPO credit assignment using a separate token-value service.""" + +from __future__ import annotations + +import asyncio +import time +from typing import TYPE_CHECKING + +import httpx +import verifiers.v1 as vf + +from prime_rl.configs.algorithm import PPOAlgoConfig +from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.orchestrator.algo.gae import skip_observation_gae +from prime_rl.orchestrator.trajectories import iter_trainable_branches +from prime_rl.transports.batch import TrainingSample + +if TYPE_CHECKING: + from prime_rl.orchestrator.clients import InferenceClient + + +class PPOAlgorithm(Algorithm): + def __init__(self, config: PPOAlgoConfig, clients: InferenceClient): + super().__init__(config, clients) + self.config = config + self.value_client = httpx.AsyncClient(base_url=config.value_url, timeout=3600) + + async def setup(self) -> None: + deadline = time.monotonic() + 1800 + while True: + try: + response = await self.value_client.get("/health") + response.raise_for_status() + return + except httpx.ConnectError: + if time.monotonic() >= deadline: + raise TimeoutError(f"Value service did not start at {self.config.value_url}") + await asyncio.sleep(1) + + async def aclose(self) -> None: + await self.value_client.aclose() + + async def score_samples(self, trace: vf.Trace, samples: list[TrainingSample]) -> None: + branches = list(iter_trainable_branches(trace)) + if len(branches) != len(samples): + raise ValueError("PPO branch/sample count mismatch") + for (branch, _), sample in zip(branches, samples, strict=True): + scored_length = min(len(sample.token_ids), self.config.value_seq_len or len(sample.token_ids)) + response = await self.value_client.post("/score", json={"token_ids": sample.token_ids[:scored_length]}) + response.raise_for_status() + score = response.json() + values = score["values"] + if len(values) != scored_length: + raise ValueError("Value service returned a misaligned token stream") + scored_mask = sample.mask[:scored_length] + n_actions = sum(scored_mask) + policy_lambda = self.config.policy_lambda + if self.config.length_adaptive_alpha is not None and n_actions: + policy_lambda = max(0.0, 1 - 1 / (self.config.length_adaptive_alpha * n_actions)) + advantages, targets = skip_observation_gae( + values, + scored_mask, + trace.reward if scored_length == len(sample.token_ids) else 0.0, + gamma=self.config.gamma, + policy_lambda=policy_lambda, + value_lambda=self.config.value_lambda, + bootstrap_value=score["bootstrap_value"] if scored_length < len(sample.token_ids) else 0.0, + ) + unscored = len(sample.token_ids) - scored_length + values += [0.0] * unscored + advantages += [0.0] * unscored + targets += [0.0] * unscored + sample.advantages = advantages + sample.old_values = values + sample.value_targets = targets + sample.value_mask = list(scored_mask) + [False] * unscored + + offset = 0 + for node in branch.nodes: + end = offset + len(node.token_ids) + node_advantages = [ + advantage + for advantage, sampled in zip(advantages[offset:end], sample.mask[offset:end], strict=True) + if sampled + ] + if node_advantages: + node.advantages = node_advantages + offset = end + sample.mask = sample.value_mask.copy() diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 74771cc0663..ee79eb2302b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -25,6 +25,7 @@ import uuid from typing import TYPE_CHECKING +import httpx import verifiers.v1 as vf from verifiers.v1.runtimes import set_base_sandbox_labels @@ -205,6 +206,11 @@ async def setup(self) -> None: ) self.admin_plane = setup_admin_plane(config.model.client, config.model.name) + if config.value_service_url is not None: + for source in config.train.source: + if source.algo is not None and source.algo.type in ("ppo", "sao"): + source.algo.value_url = config.value_service_url + await monitors.setup( producer="orch", wandb=config.monitors.wandb, @@ -264,6 +270,16 @@ async def setup(self) -> None: self.sender = setup_batch_sender( config.output_dir, config.num_train_workers, self.progress.step, config.rollout_transport ) + self.value_sender = ( + setup_batch_sender(config.output_dir, 1, self.progress.step, config.value_rollout_transport) + if config.value_rollout_transport is not None + else None + ) + self.value_client = ( + httpx.AsyncClient(base_url=config.value_service_url, timeout=30) + if config.value_service_url is not None + else None + ) # Wait phase: envs, then inference, then the trainer's startup broadcast — # the last things before the main loop starts. @@ -633,6 +649,18 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: pack_start_time = time.perf_counter() micro_batch_grid = await asyncio.to_thread(self.packer.pack, batch.samples) pack_time = time.perf_counter() - pack_start_time + if self.value_sender is not None: + await self.value_sender.send(micro_batch_grid) + assert self.value_client is not None + value_deadline = time.monotonic() + 3600 + while True: + response = await self.value_client.get("/status") + response.raise_for_status() + if response.json()["completed_step"] >= self.progress.step: + break + if time.monotonic() >= value_deadline: + raise TimeoutError(f"Value trainer did not complete step {self.progress.step}") + await asyncio.sleep(0.1) await self.sender.send(micro_batch_grid) self.progress.step += 1 self.update_dispatch_gate() @@ -1012,6 +1040,10 @@ async def stop(self) -> None: async def teardown() -> None: get_logger().debug("Closing micro batch sender") self.sender.close() + if self.value_sender is not None: + self.value_sender.close() + if self.value_client is not None: + await self.value_client.aclose() if self.dispatcher is not None: get_logger().debug("Stopping dispatcher") await self.dispatcher.stop() @@ -1039,6 +1071,7 @@ async def teardown() -> None: for clients in (env.generation_source.connected, env.algorithm.connected): if clients is not None: await clients.aclose() + await env.algorithm.aclose() get_logger().info("Stopping orchestrator components") t0 = time.perf_counter() diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 50ec66ce30b..80c91c64806 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -33,6 +33,8 @@ def payload_tokens(samples: list[TrainingSample], trace: vf.Trace | None = None) def _prune_zero_advantages(sample: TrainingSample) -> bool: """Remove zero-advantage tokens from the RL component.""" + if sample.value_mask is not None and any(sample.value_mask): + return True if sample.advantages is None: return True @@ -279,6 +281,7 @@ async def process_group(self, group_id: str) -> None: temperature = env.sampling_args["temperature"] for trace in survivors: samples = await asyncio.to_thread(trace_to_samples, trace, env_name=env_name) + await env.algorithm.score_samples(trace, samples) for sample in samples: sample.temperatures = [temperature] * len(sample.token_ids) if env.requires_sampling_masks and sample.sampling_mask is None: diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index b1c225cd995..24dc1ed2c9f 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --job-name={{ job_name }} -#SBATCH --nodes={{ num_train_nodes + num_infer_nodes }} +#SBATCH --nodes={{ num_train_nodes + num_infer_nodes + num_value_nodes }} #SBATCH --ntasks-per-node=1 #SBATCH --gres=gpu:{{ gpus_per_node }} #SBATCH --partition={{ partition }} @@ -28,6 +28,7 @@ set -e # Configs export NUM_TRAIN_NODES={{ num_train_nodes }} export NUM_INFER_NODES={{ num_infer_nodes }} +export NUM_VALUE_NODES={{ num_value_nodes }} export NODES_PER_INFER_REPLICA={{ nodes_per_infer_replica }} export NUM_INFER_REPLICAS={{ num_infer_replicas }} export GPUS_PER_NODE={{ gpus_per_node }} @@ -62,9 +63,12 @@ export LAUNCHER_DIR={{ launcher_dir }} export LOG_DIR={{ log_dir }} export PRL_ATTEMPT_CONFIG_DIR=$CONFIG_DIR export PRL_ATTEMPT_LOG_DIR=$LOG_DIR -mkdir -p $LOG_DIR/trainer $LOG_DIR/inference +mkdir -p $LOG_DIR/trainer $LOG_DIR/inference{% if num_value_nodes > 0 %} $LOG_DIR/value{% endif %} ln -sfn trainer/node_0.log $LOG_DIR/trainer.log ln -sfn inference/node_0.log $LOG_DIR/inference.log +{% if num_value_nodes > 0 -%} +ln -sfn value/node_0.log $LOG_DIR/value.log +{% endif -%} # Clean exits do not trigger srun --kill-on-bad-exit. Each terminal component # publishes its successful exit through the shared launcher directory. Every node @@ -83,7 +87,12 @@ export WANDB_SHARED_MODE=1 # Networking HOSTNAMES=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) -TRAIN_HOSTS=( ${HOSTNAMES[@]:$NUM_INFER_NODES:$SLURM_JOB_NUM_NODES} ) +TRAIN_HOSTS=( ${HOSTNAMES[@]:$NUM_INFER_NODES:$NUM_TRAIN_NODES} ) +{% if num_value_nodes > 0 -%} +VALUE_HOSTS=( ${HOSTNAMES[@]:$((NUM_INFER_NODES + NUM_TRAIN_NODES)):$NUM_VALUE_NODES} ) +export VALUE_ADDR="${VALUE_HOSTS[0]}" +export VALUE_SERVICE_PORT={{ value_service_port }} +{% endif -%} export HOSTNAMES_STR="${HOSTNAMES[*]}" {% if num_infer_nodes > 0 -%} @@ -133,6 +142,9 @@ echo "ADMIN_URLS=${ADMIN_URLS}" echo "HOSTNAMES=${HOSTNAMES[*]}" echo "TRAIN_HOSTS=${TRAIN_HOSTS[*]}" +{% if num_value_nodes > 0 -%} +echo "VALUE_HOSTS=${VALUE_HOSTS[*]}" +{% endif -%} export MASTER_ADDR="${HOSTNAMES[$NUM_INFER_NODES]}" export MASTER_PORT=29500 @@ -493,6 +505,9 @@ else {%- else %} TRAIN_NODE_RANK=$SLURM_PROCID {%- endif %} +{% if num_value_nodes > 0 %} +if [ "$TRAIN_NODE_RANK" -lt "$NUM_TRAIN_NODES" ]; then +{% endif %} # Trainer env vars: launcher defaults + [env_vars] + [trainer.env_vars] {%- for key, value in trainer_env_vars.items() %} @@ -524,7 +539,41 @@ else fi exit "$COMPONENT_CODE" ) & +{% if num_value_nodes > 0 %}fi{% endif %} {% if num_infer_nodes > 0 %} fi{% endif %} +{% if num_value_nodes > 0 %} +if [ "$SLURM_PROCID" -ge "$((NUM_INFER_NODES + NUM_TRAIN_NODES))" ]; then + VALUE_NODE_RANK=$((SLURM_PROCID - NUM_INFER_NODES - NUM_TRAIN_NODES)) +{%- for key, value in value_env_vars.items() %} + export {{ key }}="{{ value }}" +{%- endfor %} + echo $HOSTNAMES_STR | tee $LOG_DIR/value/node_${VALUE_NODE_RANK}.log + + ( + WANDB_SHARED_LABEL=value uv run torchrun \ + --role=value \ + --nnodes=$NUM_VALUE_NODES \ + --nproc-per-node=$GPUS_PER_NODE \ + --node-rank=$VALUE_NODE_RANK \ + --rdzv-endpoint=$VALUE_ADDR:$((MASTER_PORT + 1)) \ + --rdzv-id=value_$SLURM_JOB_ID \ + --log-dir=$LOG_DIR/value/torchrun \ + --tee=3 \ + --redirects=3 \ + -m prime_rl.trainer.rl.value \ + @ $CONFIG_DIR/value.json \ + --service-host 0.0.0.0 \ + --rollout-transport.host $ORCH_ADDR \ + 2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $LOG_DIR/value/node_${VALUE_NODE_RANK}.log + COMPONENT_CODE=$? + if [ "$COMPONENT_CODE" -eq 0 ]; then + echo "Value trainer exited before the run completed" >&2 + exit 1 + fi + exit "$COMPONENT_CODE" + ) & +fi +{% endif %} {% if num_infer_nodes > 0 %} # Launch the orchestrator on the node ORCH_ADDR points at (see the ORCH_PROCID # computation above). Top-level (after the role if/else) so the wait -n teardown @@ -563,6 +612,7 @@ if [ "$SLURM_PROCID" -eq "$ORCH_PROCID" ]; then {% if use_nccl_broadcast %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --weight_broadcast.host $MASTER_ADDR" {% elif use_nixl_broadcast %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --weight_broadcast.host $WEIGHT_BROADCAST_HOST" {% endif %}{% if use_zmq_transport %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --rollout_transport.host $ORCH_ADDR" + {% endif %}{% if num_value_nodes > 0 %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --value-rollout-transport.host $ORCH_ADDR --value-service-url http://$VALUE_ADDR:$VALUE_SERVICE_PORT" {% endif %}WANDB_SHARED_LABEL=orchestrator uv run orchestrator $ORCHESTRATOR_ARGS \ 2>&1 | tee $LOG_DIR/orchestrator.log COMPONENT_CODE=$? diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 63fdea95b1e..16b7a99d6a8 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -393,6 +393,11 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch rl_weights = list(training_example.rl_weights) if training_example.rl_weights is not None else None ce_weights = list(training_example.ce_weights) if training_example.ce_weights is not None else None ref_kl_weights = list(training_example.ref_kl_weights) if training_example.ref_kl_weights is not None else None + old_values = list(training_example.old_values) if training_example.old_values is not None else None + value_targets = list(training_example.value_targets) if training_example.value_targets is not None else None + value_mask = list(training_example.value_mask) if training_example.value_mask is not None else None + if (old_values is None) != (value_targets is None) or (old_values is None) != (value_mask is None): + raise ValueError("old_values, value_targets, and value_mask must be present together") position_ids = list(range(len(input_ids))) mm_token_type_ids = training_example.mm_token_type_ids mm_kwargs = training_example.mm_kwargs @@ -432,6 +437,10 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ce_weights = ce_weights[:cut] if ref_kl_weights is not None: ref_kl_weights = ref_kl_weights[:cut] + if old_values is not None: + old_values = old_values[:cut] + value_targets = value_targets[:cut] + value_mask = value_mask[:cut] if routed_experts is not None: routed_experts = _slice_routed_experts(routed_experts, cut) if sampling_mask is not None: @@ -456,6 +465,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ("rl_weights", rl_weights), ("ce_weights", ce_weights), ("ref_kl_weights", ref_kl_weights), + ("old_values", old_values), + ("value_targets", value_targets), + ("value_mask", value_mask), ): if stream is not None: assert len(stream) == len(input_ids), f"{stream_name}: {len(stream)}" @@ -492,6 +504,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch rl_weights=rl_weights, ce_weights=ce_weights, ref_kl_weights=ref_kl_weights, + old_values=old_values, + value_targets=value_targets, + value_mask=value_mask, seq_lens=[len(input_ids)], trace_ids=[training_example.trace_id or ""], branch_indices=[training_example.branch_index if training_example.branch_index is not None else -1], @@ -578,6 +593,7 @@ def split_by_workload(self, bin_cost: Callable[[Sequence[int]], int]) -> tuple[" def _materialize_bin(bin_content: _MicroBatchBin) -> MicroBatch: has_ref_logprobs = any(sample.ref_logprobs is not None for sample in bin_content.samples) has_mm_token_type_ids = any(sample.mm_token_type_ids is not None for sample in bin_content.samples) + has_values = any(sample.value_targets is not None for sample in bin_content.samples) # A weight stream materializes as soon as one packed sample carries it; the # samples that lack it get the stream's identity fill (STREAM_FILL). has_stream = {name: any(getattr(s, name) is not None for s in bin_content.samples) for name in STREAM_FILL} @@ -593,6 +609,9 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> MicroBatch: temperatures: list[float] = [] env_names: list[str] = [] ref_logprobs: list[float] | None = [] if has_ref_logprobs else None + old_values: list[float] | None = [] if has_values else None + value_targets: list[float] | None = [] if has_values else None + value_mask: list[bool] | None = [] if has_values else None mm_token_type_ids: list[int] | None = [] if has_mm_token_type_ids else None mm_kwargs: dict[str, EncodedTensor] | None = None streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL} @@ -613,6 +632,10 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> MicroBatch: env_names.extend(sample.env_names) if ref_logprobs is not None: ref_logprobs.extend(sample.ref_logprobs if sample.ref_logprobs is not None else [0.0] * sample_len) + if has_values: + old_values.extend(sample.old_values if sample.old_values is not None else [0.0] * sample_len) + value_targets.extend(sample.value_targets if sample.value_targets is not None else [0.0] * sample_len) + value_mask.extend(sample.value_mask if sample.value_mask is not None else [False] * sample_len) for name, fill in STREAM_FILL.items(): stream = streams[name] if stream is not None: @@ -666,6 +689,9 @@ def _materialize_bin(bin_content: _MicroBatchBin) -> MicroBatch: rl_weights=streams["rl_weights"], ce_weights=streams["ce_weights"], ref_kl_weights=streams["ref_kl_weights"], + old_values=old_values, + value_targets=value_targets, + value_mask=value_mask, seq_lens=seq_lens, trace_ids=trace_ids, branch_indices=branch_indices, @@ -769,6 +795,10 @@ def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBa micro_batch.input_ids.extend([1] * padding_size) micro_batch.advantages.extend([0.0] * padding_size) micro_batch.loss_mask.extend([False] * padding_size) + if micro_batch.old_values is not None: + micro_batch.old_values.extend([0.0] * padding_size) + micro_batch.value_targets.extend([0.0] * padding_size) + micro_batch.value_mask.extend([False] * padding_size) micro_batch.position_ids.extend(list(range(padding_size))) micro_batch.sequence_lengths[-1] += padding_size micro_batch.seq_lens[-1] += padding_size @@ -813,6 +843,9 @@ def _assert_token_arrays_aligned(micro_batch: MicroBatch) -> None: "ce_weights", "ref_kl_weights", "mm_token_type_ids", + "old_values", + "value_targets", + "value_mask", ) for name in per_token_fields: values = getattr(micro_batch, name) @@ -854,6 +887,9 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch: dummy.rl_weights = None dummy.ce_weights = None dummy.ref_kl_weights = None + dummy.old_values = None + dummy.value_targets = None + dummy.value_mask = None # Fully loss-masked, so replaying sampling masks would be pure wasted work. dummy.sampling_mask = None # The copied identity would double-annotate the source's traces. diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 7fa2821cc99..39ff3f5f44e 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -1,5 +1,7 @@ +import fcntl import logging import os +import shutil import time from pathlib import Path from typing import cast @@ -699,29 +701,37 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral "Found HF weight format in snapshot state dict and PrimeRL weight format in model state dict. Trying to auto-convert..." ) snapshot_path = convert_dir / "prime" - if not snapshot_path.exists() and get_world().is_master: - logger.debug( - f"Converting snapshot state dict to PrimeRL format and saving to {snapshot_path} on master rank. This is a one-time operation." - ) - snapshot_state_dict = load_state_dict(source_path) - model.convert_to_prime(snapshot_state_dict) - save_state_dict(snapshot_state_dict, snapshot_path) - (snapshot_path / ".prime-v1").touch() - del snapshot_state_dict + if get_world().is_master: + with (convert_dir / ".prime-conversion.lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + if not (snapshot_path / ".prime-v1").is_file(): + if snapshot_path.exists(): + shutil.rmtree(snapshot_path) + logger.debug( + f"Converting snapshot state dict to PrimeRL format and saving to {snapshot_path} on master rank. This is a one-time operation." + ) + snapshot_state_dict = load_state_dict(source_path) + model.convert_to_prime(snapshot_state_dict) + save_state_dict(snapshot_state_dict, snapshot_path) + (snapshot_path / ".prime-v1").touch() + del snapshot_state_dict elif snapshot_is_prime and not snapshot_is_hf and model.is_hf_state_dict(model_keys): logger.warning( "Found PrimeRL weight format in snapshot state dict and HF weight format in model state dict. Trying to auto-convert..." ) snapshot_path = convert_dir / "hf" - if not snapshot_path.exists() and get_world().is_master: - logger.debug( - f"Converting snapshot state dict to HF format and saving to {snapshot_path} on master rank. This is a one-time operation." - ) - snapshot_state_dict = load_state_dict(source_path) - model.convert_to_hf(snapshot_state_dict) - save_state_dict(snapshot_state_dict, snapshot_path) - del snapshot_state_dict + if get_world().is_master: + with (convert_dir / ".hf-conversion.lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + if not snapshot_path.exists(): + logger.debug( + f"Converting snapshot state dict to HF format and saving to {snapshot_path} on master rank. This is a one-time operation." + ) + snapshot_state_dict = load_state_dict(source_path) + model.convert_to_hf(snapshot_state_dict) + save_state_dict(snapshot_state_dict, snapshot_path) + del snapshot_state_dict # All ranks wait for master rank to finish conversion torch.distributed.barrier() @@ -736,6 +746,8 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral load_dcp_start_time = time.perf_counter() state_dict = model.state_dict() state_dict = strip_lora_from_state_dict(state_dict) + state_dict.pop("value_head.weight", None) + state_dict.pop("value_head.bias", None) if model.config.tie_word_embeddings: state_dict.pop("lm_head.weight") dcp_load( @@ -969,7 +981,12 @@ def setup_model( config: ModelConfig, parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False, + *, + value_model: bool = False, + freeze_attention: bool = False, ) -> nn.Module: + if freeze_attention and not value_model: + raise ValueError("freeze_attention is only supported for a value model") resolve_auto_attn(config) if config.attn == "flash_attention_3" and not is_flash_attn_3_available(): @@ -1013,6 +1030,11 @@ def setup_model( frozen_vision_encoder = configure_trainable_parameters(model, config) + if value_model: + model.value_head = nn.Linear( + model.lm_head.in_features, 1, device="meta", dtype=DTYPE_MAP[config.optimization_dtype] + ) + if config.freeze_moe_router: freeze_moe_router(model) @@ -1034,6 +1056,16 @@ def setup_model( if config.lora is not None: freeze_all_except_lora_and_specified(model, config.lora) + if value_model: + for param in model.lm_head.parameters(): + param.requires_grad = False + if freeze_attention: + for layer in model.model.layers: + for param in layer.self_attn.parameters(): + param.requires_grad = False + for param in model.value_head.parameters(): + param.requires_grad = True + if frozen_vision_encoder is not None: freeze_vision_encoder( model, @@ -1075,6 +1107,10 @@ def setup_model( load_dcp_from_hf(model, config, parallel_dims) _reset_runtime_moe_buffers(model) + if value_model: + with torch.no_grad(): + model.value_head.weight.zero_() + model.value_head.bias.zero_() return model diff --git a/src/prime_rl/trainer/models/layers/lm_head.py b/src/prime_rl/trainer/models/layers/lm_head.py index e0ce2ea786e..b27b728d4cc 100644 --- a/src/prime_rl/trainer/models/layers/lm_head.py +++ b/src/prime_rl/trainer/models/layers/lm_head.py @@ -18,6 +18,7 @@ class PrimeLmOutput(TypedDict, total=False): logprobs: Tensor | None entropy: Tensor | None loss: Tensor | None + values: Tensor | None def cast_float_and_contiguous(output: PrimeLmOutput) -> PrimeLmOutput: @@ -31,6 +32,7 @@ def _float_and_contiguous(tensor: Tensor | None) -> Tensor | None: logprobs=_float_and_contiguous(output.get("logprobs")), entropy=_float_and_contiguous(output.get("entropy")), loss=output.get("loss"), + values=_float_and_contiguous(output.get("values")), ) @@ -327,6 +329,8 @@ def new_forward( logits_to_keep: int = 0, temperature: torch.Tensor | None = None, sampling_mask: torch.Tensor | None = None, + return_values: bool = False, + head_only: bool = False, **kwargs: object, ) -> PrimeLmOutput: # For VLM with images, don't create position_ids - let model compute MRoPE internally @@ -340,6 +344,12 @@ def new_forward( outputs = self.model(**model_kwargs) hidden_states = outputs.last_hidden_state + if return_values: + value_head = getattr(self, "value_head", None) + if value_head is None: + raise ValueError("return_values requires a value head") + return PrimeLmOutput(values=value_head(hidden_states.detach() if head_only else hidden_states).squeeze(-1)) + # Slice hidden states for logits_to_keep slice_indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) and logits_to_keep > 0 else slice(None) diff --git a/src/prime_rl/trainer/models/layers/lora/multi_moe.py b/src/prime_rl/trainer/models/layers/lora/multi_moe.py index 9356b67be31..1593a5efe17 100644 --- a/src/prime_rl/trainer/models/layers/lora/multi_moe.py +++ b/src/prime_rl/trainer/models/layers/lora/multi_moe.py @@ -299,9 +299,7 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch up_proj = self.base_layer.up_proj down_proj = self.base_layer.down_proj - # EP handling: convert DTensors to local shards. - # Standard EP also needs token permutation; DeepEP tokens are already dispatched. - permuted_indices = None + # The token dispatcher has already grouped and aligned tokens by local expert. if isinstance(gate_proj, DTensor): gate_proj = gate_proj.to_local() up_proj = up_proj.to_local() @@ -313,27 +311,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch up_proj_lora_a = up_proj_lora_a.to_local() up_proj_lora_b = up_proj_lora_b.to_local() - if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": - from torchtitan.experiments.kernels.moe.indices import generate_permute_indices - - from prime_rl.trainer.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M - - experts_per_ep_rank = gate_proj.shape[0] - num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank - - with torch.no_grad(): - permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( - num_tokens_per_expert, - experts_per_ep_rank, - num_ep_ranks, - x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, - TOKEN_GROUP_ALIGN_SIZE_M, - ) - - x = torch.vstack((x, x.new_zeros((x.shape[-1])))) - input_shape = x.shape - x = x[permuted_indices, :] - # Compute offsets for grouped_mm offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) @@ -354,12 +331,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch down_proj_lora_out = _run_lora_grouped_mm(lora_h, down_proj_lora_a, down_proj_lora_b, offsets) out = (h2_base + scaling * down_proj_lora_out.bfloat16()).type_as(x) - # EP handling: unpermute output back to dispatched token order - if permuted_indices is not None: - out_unpermuted = out.new_zeros(input_shape) - out_unpermuted[permuted_indices, :] = out - out = out_unpermuted[:-1] - return out def __repr__(self) -> str: @@ -530,7 +501,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch up_proj = self.base_layer.up_proj down_proj = self.base_layer.down_proj - permuted_indices = None if isinstance(up_proj, DTensor): up_proj = up_proj.to_local() down_proj = down_proj.to_local() @@ -539,27 +509,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch down_proj_lora_a = down_proj_lora_a.to_local() down_proj_lora_b = down_proj_lora_b.to_local() - if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": - from torchtitan.experiments.kernels.moe.indices import generate_permute_indices - - from prime_rl.trainer.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M - - experts_per_ep_rank = up_proj.shape[0] - num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank - - with torch.no_grad(): - permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( - num_tokens_per_expert, - experts_per_ep_rank, - num_ep_ranks, - x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, - TOKEN_GROUP_ALIGN_SIZE_M, - ) - - x = torch.vstack((x, x.new_zeros((x.shape[-1])))) - input_shape = x.shape - x = x[permuted_indices, :] - offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) lora_x = self.lora_dropout(x) @@ -572,11 +521,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch down_proj_lora_out = _run_lora_grouped_mm(lora_h, down_proj_lora_a, down_proj_lora_b, offsets) out = (out_base + scaling * down_proj_lora_out.bfloat16()).type_as(x) - if permuted_indices is not None: - out_unpermuted = out.new_zeros(input_shape) - out_unpermuted[permuted_indices, :] = out - out = out_unpermuted[:-1] - return out def __repr__(self) -> str: @@ -772,7 +716,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch up_proj_bias = self.base_layer.up_proj_bias down_proj_bias = self.base_layer.down_proj_bias - permuted_indices = None if isinstance(up_proj, DTensor): gate_proj = gate_proj.to_local() up_proj = up_proj.to_local() @@ -785,27 +728,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch d_a = d_a.to_local() d_b = d_b.to_local() - if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": - from torchtitan.experiments.kernels.moe.indices import generate_permute_indices - - from prime_rl.trainer.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M - - experts_per_ep_rank = up_proj.shape[0] - num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank - - with torch.no_grad(): - permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( - num_tokens_per_expert, - experts_per_ep_rank, - num_ep_ranks, - x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, - TOKEN_GROUP_ALIGN_SIZE_M, - ) - - x = torch.vstack((x, x.new_zeros((x.shape[-1])))) - input_shape = x.shape - x = x[permuted_indices, :] - offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) lora_x = self.lora_dropout(x) @@ -826,11 +748,6 @@ def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch out_lora = _run_lora_grouped_mm(lora_h, d_a, d_b, offsets) out = (out_base + scaling * out_lora.bfloat16()).type_as(x) - if permuted_indices is not None: - out_unpermuted = out.new_zeros(input_shape) - out_unpermuted[permuted_indices, :] = out - out = out_unpermuted[:-1] - return out def __repr__(self) -> str: diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py index aeeeb070047..e944b95afdd 100644 --- a/src/prime_rl/trainer/optim/offload.py +++ b/src/prime_rl/trainer/optim/offload.py @@ -1067,7 +1067,6 @@ class FullCPUOffloadOptimizer(OffloadOptimizer): _TARGET_CPU_CHUNK_NUMEL = 16 * 1024**2 _TRANSFER_BUFFER_COUNT = 4 _MAX_INFLIGHT_BACKWARDS = 16 - _TIMEOUT_SECONDS = 120.0 _MASTER_WEIGHT_STATE = "prime_rl_master_weight" def __init__( @@ -1113,7 +1112,7 @@ def __init__( dp_replicate, buffer_count=self._TRANSFER_BUFFER_COUNT, max_inflight_backwards=self._MAX_INFLIGHT_BACKWARDS, - timeout_seconds=self._TIMEOUT_SECONDS, + timeout_seconds=offload_config.timeout_seconds, target_chunk_numel=self._TARGET_CPU_CHUNK_NUMEL, chunk_ready_callback=self._step_cpu_chunk, gradient_dtypes=gradient_dtypes, diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 1d1e9fc28e0..31daaa3168b 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -60,6 +60,9 @@ class TensorMicroBatch(TypedDict): rl_weights: Float[Tensor, "batch seq"] | None ce_weights: Float[Tensor, "batch seq"] | None ref_kl_weights: Float[Tensor, "batch seq"] | None + old_values: Float[Tensor, "batch seq"] | None + value_targets: Float[Tensor, "batch seq"] | None + value_mask: Bool[Tensor, "batch seq"] | None class FakeDataLoader: @@ -140,6 +143,9 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc "rl_weights": None, "ce_weights": None, "ref_kl_weights": None, + "old_values": None, + "value_targets": None, + "value_mask": None, } def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: @@ -172,6 +178,9 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: "rl_weights": None, "ce_weights": None, "ref_kl_weights": None, + "old_values": None, + "value_targets": None, + "value_mask": None, } @@ -264,6 +273,15 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: ref_kl_weights=torch.tensor(micro_batch.ref_kl_weights, dtype=torch.float).unsqueeze(0) if micro_batch.ref_kl_weights is not None else None, + old_values=torch.tensor(micro_batch.old_values, dtype=torch.float).unsqueeze(0) + if micro_batch.old_values is not None + else None, + value_targets=torch.tensor(micro_batch.value_targets, dtype=torch.float).unsqueeze(0) + if micro_batch.value_targets is not None + else None, + value_mask=torch.tensor(micro_batch.value_mask, dtype=torch.bool).unsqueeze(0) + if micro_batch.value_mask is not None + else None, ) diff --git a/src/prime_rl/trainer/rl/loss.py b/src/prime_rl/trainer/rl/loss.py index 9a7d55e6958..a6984b4303d 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -6,7 +6,14 @@ from jaxtyping import Bool, Float, Int, jaxtyped from torch import Tensor -from prime_rl.configs.trainer import CustomLossConfig, IcePopLossConfig, IPOLossConfig, LossConfig +from prime_rl.configs.trainer import ( + CustomLossConfig, + IcePopLossConfig, + IPOLossConfig, + LossConfig, + PPOClipLossConfig, + SAOLossConfig, +) from prime_rl.trainer.models.layers.lm_head import sampling_replay_mask from prime_rl.utils.utils import import_object @@ -39,7 +46,8 @@ class LossOutputs: class Loss(Protocol): """Interface for the config-initialized rl loss objects built by - ``setup_rl_loss_fn``: ``IPOLoss``, ``IcePopLoss`` and ``CustomLoss``.""" + ``setup_rl_loss_fn``: ``IPOLoss``, ``IcePopLoss``, ``PPOClipLoss`` and ``CustomLoss``. + SAO uses ``IcePopLoss`` with its own ratio interval.""" def loss(self, inputs: LossInputs) -> LossOutputs: ... @@ -187,9 +195,7 @@ def __init__(self, config: IcePopLossConfig): def loss(self, inputs: LossInputs) -> LossOutputs: loss_config = self.config - log_importance_ratio, _, mismatch_kl = compute_importance_ratio_and_mismatch_kl( - inputs.trainer_logprobs, inputs.inference_logprobs - ) + log_importance_ratio = inputs.trainer_logprobs.float() - inputs.inference_logprobs.float() log_ratio_low = log_importance_ratio.new_tensor(loss_config.ratio_low).log() log_ratio_high = log_importance_ratio.new_tensor(loss_config.ratio_high).log() @@ -205,6 +211,8 @@ def loss(self, inputs: LossInputs) -> LossOutputs: if inputs.loss_weights is not None: per_token_loss = per_token_loss * inputs.loss_weights + metric_log_ratio = log_importance_ratio.detach().clamp(min=-30, max=30) + mismatch_kl = torch.expm1(metric_log_ratio) - metric_log_ratio metrics = { "masked_mismatch_kl": _safe_mean(mismatch_kl, inputs.loss_mask & is_masked), "unmasked_mismatch_kl": _safe_mean(mismatch_kl, keep_mask), @@ -213,6 +221,35 @@ def loss(self, inputs: LossInputs) -> LossOutputs: return LossOutputs(loss=per_token_loss.sum(), metrics=metrics) +class PPOClipLoss: + """Clipped PPO surrogate using the rollout log probabilities as the old policy.""" + + def __init__(self, config: PPOClipLossConfig): + self.config = config + + def loss(self, inputs: LossInputs) -> LossOutputs: + log_ratio = inputs.trainer_logprobs.float() - inputs.inference_logprobs.float() + log_ratio = torch.where(inputs.loss_mask, log_ratio, torch.zeros_like(log_ratio)) + lower = log_ratio.new_tensor(1 - self.config.eps_low).log() + upper = log_ratio.new_tensor(1 + self.config.eps_high).log() + # A very stale, unfavorably moved token can otherwise overflow exp. + ratio = torch.exp(log_ratio.clamp(max=30)) + clipped_ratio = torch.exp(log_ratio.clamp(min=lower, max=upper)) + advantage = self.config.adv_tau * inputs.advantages + per_token_loss = -torch.minimum(ratio * advantage, clipped_ratio * advantage) + per_token_loss = torch.where(inputs.loss_mask, per_token_loss, 0.0) + if inputs.loss_weights is not None: + per_token_loss = per_token_loss * inputs.loss_weights + + clipped = ((advantage >= 0) & (log_ratio > upper)) | ((advantage < 0) & (log_ratio < lower)) + metric_log_ratio = log_ratio.clamp(min=-30, max=30) + metrics = { + "ppo/clip_fraction": _safe_mean(clipped, inputs.loss_mask), + "ppo/approx_kl": _safe_mean(torch.expm1(metric_log_ratio) - metric_log_ratio, inputs.loss_mask), + } + return LossOutputs(loss=per_token_loss.sum(), metrics=metrics) + + def ref_kl_loss_fn(inputs: LossInputs) -> LossOutputs: """ Ref-KL loss type (on-policy distillation): the reverse KL to the reference @@ -296,6 +333,10 @@ def setup_rl_loss_fn(loss_config: LossConfig) -> Loss: return CustomLoss(loss_config) case IPOLossConfig(): return IPOLoss(loss_config) + case PPOClipLossConfig(): + return PPOClipLoss(loss_config) + case SAOLossConfig(): + return IcePopLoss(loss_config) case IcePopLossConfig(): return IcePopLoss(loss_config) case _: diff --git a/src/prime_rl/trainer/rl/policy_sync.py b/src/prime_rl/trainer/rl/policy_sync.py new file mode 100644 index 00000000000..6e83a168c4b --- /dev/null +++ b/src/prime_rl/trainer/rl/policy_sync.py @@ -0,0 +1,62 @@ +"""Model-only policy backbone snapshots for a LoRA value trainer.""" + +import shutil +from pathlib import Path + +import torch.distributed as dist +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict, set_model_state_dict +from torch.distributed.checkpoint.state_dict_loader import load as dcp_load +from torch.distributed.checkpoint.state_dict_saver import save as dcp_save + +from prime_rl.trainer.world import get_world + + +def _backbone_state(model): + state = { + name: tensor + for name, tensor in get_model_state_dict(model).items() + if name.startswith("model.") and "lora_" not in name + } + if not state: + raise ValueError("Policy backbone snapshot contains no model weights") + return state + + +def snapshot_path(directory: Path, step: int) -> Path: + return directory / f"step_{step}" + + +def save_policy_backbone(model, directory: Path, step: int) -> None: + world = get_world() + path = snapshot_path(directory, step) + if world.is_master: + directory.mkdir(parents=True, exist_ok=True) + if path.exists(): + shutil.rmtree(path) + dist.barrier() + dcp_save({"model": _backbone_state(model)}, checkpoint_id=path) + if world.is_master: + (path / ".ready").touch() + for old_path in directory.glob("step_*"): + if old_path != path and (old_path / ".applied").is_file(): + shutil.rmtree(old_path) + dist.barrier() + + +def latest_ready_step(directory: Path, after_step: int, at_most_step: int) -> int | None: + if not directory.exists(): + return None + steps = [ + int(path.name.removeprefix("step_")) + for path in directory.glob("step_*") + if path.name.removeprefix("step_").isdigit() and (path / ".ready").is_file() + ] + return max((step for step in steps if after_step < step <= at_most_step), default=None) + + +def load_policy_backbone(model, directory: Path, step: int) -> None: + path = snapshot_path(directory, step) + state = _backbone_state(model) + dcp_load({"model": state}, checkpoint_id=path) + set_model_state_dict(model, state, options=StateDictOptions(strict=False)) + dist.barrier() diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index b54126a792c..bb2495a23b1 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -18,6 +18,7 @@ from prime_rl.trainer.scheduler import setup_scheduler from prime_rl.configs.trainer import TrainerConfig from prime_rl.trainer.rl.data import DataLoader, FakeDataLoader +from prime_rl.trainer.rl.policy_sync import save_policy_backbone from prime_rl.utils.cp import ( gather_for_cp, gather_for_cp_wo_grad, @@ -498,7 +499,7 @@ def train(config: TrainerConfig): with maybe_record_function("backward"): begin_backward(gradient_manager, final_backward=micro_step == len(micro_batches) - 1) loss.backward() - finish_backward(gradient_manager) + finish_backward(gradient_manager, wait_for_copies=config.model.full_offload is not None) # Add relevant tensors to tensor dict for logging purposes entropy = out["entropy"][loss_mask].detach().to("cpu") @@ -596,6 +597,15 @@ def train(config: TrainerConfig): weight_sender.broadcast(model, step=progress.step) broadcast_weights_time = time.perf_counter() - broadcast_weights_start_time + if ( + config.ppo_policy_sync_interval is not None + and progress.step > config.ppo_head_warmup_steps + and progress.step % config.ppo_policy_sync_interval == 0 + ): + assert config.ppo_policy_sync_dir is not None + logger.info(f"Saving policy backbone for value trainer at step {progress.step}") + save_policy_backbone(model, config.ppo_policy_sync_dir, progress.step) + # Checkpoint the step we just finished (model = policy v{progress.step}). if ( (config.ckpt and config.ckpt.interval) diff --git a/src/prime_rl/trainer/rl/value.py b/src/prime_rl/trainer/rl/value.py new file mode 100644 index 00000000000..7305e6c9aee --- /dev/null +++ b/src/prime_rl/trainer/rl/value.py @@ -0,0 +1,450 @@ +"""Separate PPO critic: token scoring and online value updates.""" + +import prime_rl._compat # noqa: F401 + +# ruff: noqa: I001 + +import json +import queue +import threading +import time +from dataclasses import dataclass, field +from datetime import timedelta +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import torch +import torch.distributed as dist +import verifiers.v1 as vf + +from prime_rl.configs.value import ValueConfig +from prime_rl.orchestrator.trajectories import trace_to_samples +from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager +from prime_rl.trainer.lora import get_lora_state +from prime_rl.trainer.model import get_full_offload_dtype_policy, setup_model +from prime_rl.trainer.models.layers.lora import set_lora_num_tokens +from prime_rl.trainer.optim import setup_optimizer +from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep +from prime_rl.trainer.rl.data import DataLoader, TensorMicroBatch +from prime_rl.trainer.rl.loss import shift_tensor_right +from prime_rl.trainer.rl.policy_sync import latest_ready_step, load_policy_backbone, snapshot_path +from prime_rl.trainer.scheduler import setup_scheduler +from prime_rl.trainer.utils import ( + begin_backward, + clip_grad_norm_, + finish_backward, + prepare_gradient_offload, + scale_gradients_, + setup_full_cpu_optimizer_offload, + setup_torch_distributed, +) +from prime_rl.trainer.world import get_world +from prime_rl.utils.config import cli +from prime_rl.utils.cp import gather_for_cp, gather_for_cp_wo_grad, setup_context_parallel, setup_cp_params +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.process import set_proc_title +from prime_rl.utils.utils import resolve_latest_ckpt_step + + +@dataclass +class ScoreRequest: + token_ids: list[int] + done: threading.Event = field(default_factory=threading.Event) + values: list[float] | None = None + bootstrap_value: float | None = None + error: BaseException | None = None + + +@dataclass +class ServiceState: + requests: queue.Queue[ScoreRequest] = field(default_factory=queue.Queue) + completed_step: int = 0 + + +def _server_handler(state: ServiceState): + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + if self.path == "/health": + self.send_response(200) + self.end_headers() + return + if self.path != "/status": + self.send_error(404) + return + payload = json.dumps({"completed_step": state.completed_step}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def do_POST(self): + if self.path != "/score": + self.send_error(404) + return + body = self.rfile.read(int(self.headers["Content-Length"])) + token_ids = json.loads(body)["token_ids"] + if ( + not isinstance(token_ids, list) + or not token_ids + or not all(isinstance(token, int) for token in token_ids) + ): + self.send_error(400, "token_ids must be a nonempty list of integers") + return + request = ScoreRequest(token_ids) + state.requests.put(request) + if not request.done.wait(timeout=3600): + self.send_error(504, "value scoring timed out") + return + if request.error is not None: + self.send_error(500, str(request.error)) + return + payload = json.dumps({"values": request.values, "bootstrap_value": request.bootstrap_value}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format, *args): + return + + return Handler + + +def _forward_values( + model, input_ids, position_ids, seq_lens, parallel_dims, cp_style, *, lora_enabled=False, head_only=False +): + cp_enabled = parallel_dims.cp_enabled + if cp_enabled: + cp_mesh = parallel_dims.get_mesh("cp") + input_ids, position_ids = setup_cp_params( + input_ids, + position_ids, + cp_mesh.get_local_rank(), + parallel_dims.cp, + cp_mesh.get_group(), + seq_lens=seq_lens, + cp_style=cp_style, + ) + if lora_enabled: + set_lora_num_tokens(torch.tensor([input_ids.shape[1]], device="cuda", dtype=torch.int32)) + output = model( + input_ids=input_ids, + position_ids=position_ids, + seq_lens=seq_lens, + seq_lens_are_pre_shard=cp_enabled, + return_values=True, + head_only=head_only, + ) + return output["values"] + + +def _score(model, token_ids: list[int], parallel_dims, cp_style: str, lora_enabled: bool) -> tuple[list[float], float]: + original_length = len(token_ids) + padding = (-original_length) % parallel_dims.cp + padded = token_ids + [0] * padding + input_ids = torch.tensor(padded, device="cuda", dtype=torch.long).unsqueeze(0) + position_ids = torch.arange(len(padded), device="cuda", dtype=torch.long).unsqueeze(0) + seq_lens = torch.tensor([len(padded)], device="cuda", dtype=torch.long) + model.eval() + with torch.no_grad(): + values = _forward_values( + model, input_ids, position_ids, seq_lens, parallel_dims, cp_style, lora_enabled=lora_enabled + ) + if parallel_dims.cp_enabled: + values = gather_for_cp_wo_grad(values, parallel_dims.cp, parallel_dims.get_mesh("cp").get_group()) + return ( + shift_tensor_right(values)[0, :original_length].float().cpu().tolist(), + values[0, original_length - 1].float().item(), + ) + + +def _train_batch( + model, optimizer, scheduler, gradient_manager, micro_batches, parallel_dims, config, head_only, updates=None +): + model.train() + dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + total_count = torch.zeros((), device="cuda", dtype=torch.long) + for batch in micro_batches: + mask = batch["value_mask"] + if mask is None: + raise ValueError("PPO critic received a batch without value targets") + mask = mask.to("cuda") + total_count += mask.sum() + dist.all_reduce(total_count, group=dp_cp_group) + denominator = total_count.clamp_min(1) + + last_loss = 0.0 + for _ in range(updates if updates is not None else config.updates_per_step): + update_loss = 0.0 + prepare_gradient_offload(gradient_manager, parallel_dims.fsdp_gradient_divide_factor, overlap_optimizer=True) + for micro_step, batch in enumerate(micro_batches): + input_ids = batch["input_ids"].to("cuda") + position_ids = batch["position_ids"].to("cuda") + seq_lens = batch["seq_lens"].to("cuda") + targets = batch["value_targets"].to("cuda") + mask = batch["value_mask"].to("cuda") + values = _forward_values( + model, + input_ids, + position_ids, + seq_lens, + parallel_dims, + config.model.cp_style, + lora_enabled=config.model.lora is not None, + head_only=head_only, + ) + if parallel_dims.cp_enabled: + values = gather_for_cp(values, parallel_dims.get_mesh("cp").get_group()) + values = shift_tensor_right(values) + loss = (((values.float() - targets.float()) ** 2) * mask).sum() / denominator + begin_backward(gradient_manager, final_backward=micro_step == len(micro_batches) - 1) + loss.backward() + finish_backward(gradient_manager, wait_for_copies=config.model.full_offload is not None) + update_loss += loss.detach().item() + if gradient_manager is None: + scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor) + if config.optim.max_norm is not None: + clip_grad_norm_(gradient_manager, model, config.optim.max_norm, parallel_dims.ep_enabled) + optimizer.step() + optimizer.zero_grad() + scheduler.step() + last_loss = update_loss + return last_loss + + +def _pretrain(model, optimizer, scheduler, gradient_manager, parallel_dims, config): + if config.pretrain_data is None or config.pretrain_steps == 0: + return + with Path(config.pretrain_data).open() as records: + for step in range(config.pretrain_steps): + line = records.readline() + if not line: + raise ValueError(f"value.pretrain_data has fewer than {config.pretrain_steps} records") + item = json.loads(line) + trace = item.get("trace", item) + if "token_ids" in trace: + streams = [(trace["token_ids"], trace["action_mask"])] + reward = float(item["reward"]) + else: + native_trace = vf.Trace.model_validate(trace) + streams = [(sample.token_ids, sample.mask) for sample in trace_to_samples(native_trace)] + reward = float(item.get("reward", native_trace.reward)) + if not streams: + raise ValueError("Pretraining trace has no trainable branches") + batches = [] + for token_ids, action_mask in streams: + if len(token_ids) != len(action_mask): + raise ValueError("Pretraining token_ids and action_mask lengths differ") + if len(token_ids) > config.model.seq_len: + raise ValueError("Pretraining trace exceeds value.model.seq_len") + padding = (-len(token_ids)) % parallel_dims.cp + n = len(token_ids) + padding + mask = action_mask + [False] * padding + target = [reward if sampled else 0.0 for sampled in mask] + batches.append( + { + "input_ids": torch.tensor(token_ids + [0] * padding).unsqueeze(0), + "position_ids": torch.arange(n).unsqueeze(0), + "seq_lens": torch.tensor([n]), + "value_targets": torch.tensor(target, dtype=torch.float).unsqueeze(0), + "value_mask": torch.tensor(mask, dtype=torch.bool).unsqueeze(0), + } + ) + _train_batch( + model, + optimizer, + scheduler, + gradient_manager, + batches, + parallel_dims, + config, + head_only=step < config.head_warmup_steps, + ) + + +def train(config: ValueConfig): + world = get_world() + logger = setup_logger("info") + setup_torch_distributed( + timeout=timedelta(seconds=3600), + enable_gloo=config.model.fsdp_cpu_offload or config.model.full_offload is not None, + ) + if config.model.full_offload is not None: + setup_full_cpu_optimizer_offload(config.model.full_offload) + resolve_ep(config.model) + parallel_dims = get_parallel_dims(config.model) + ckpt_manager = setup_ckpt_manager(config.output_dir, config.ckpt, resume=config.resume) + checkpoint_step = None + if config.resume is not None: + checkpoint_step = ( + config.resume.step + or (config.resume.dir_step if config.resume.dir is not None else None) + or resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) + ) + model = setup_model( + config.model, + parallel_dims, + checkpoint_step is not None, + value_model=True, + freeze_attention=config.freeze_attention, + ) + if parallel_dims.cp_enabled: + setup_context_parallel(model, config.model, parallel_dims) + if config.model.lora is not None: + get_lora_state().reset_adapter_parameters() + optimizer, gradient_manager = setup_optimizer( + config.optim, + list(model.named_parameters()), + parallel_dims, + cpu_offload=config.model.optim_cpu_offload, + full_offload_config=config.model.full_offload, + model=model, + full_offload_dtype_policy=( + get_full_offload_dtype_policy(model, config.model) if config.model.full_offload is not None else None + ), + ) + scheduler_steps = None + if config.max_steps is not None: + scheduler_steps = (config.max_steps + config.pretrain_steps) * config.updates_per_step + if config.policy_sync_interval is not None: + first_actor_step = max(0, config.head_warmup_steps - config.pretrain_steps) + 1 + syncs = config.max_steps // config.policy_sync_interval + syncs -= (first_actor_step - 1) // config.policy_sync_interval + scheduler_steps += syncs * config.policy_sync_retune_updates + scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) + progress = Progress() + if checkpoint_step is not None: + resume_path = config.resume.dir / "trainer" if config.resume and config.resume.dir is not None else None + ckpt_manager.load(checkpoint_step, model, [optimizer], scheduler, progress, path=resume_path) + progress.step += 1 + + if checkpoint_step is None: + _pretrain(model, optimizer, scheduler, gradient_manager, parallel_dims, config) + dataloader = DataLoader(config.rollout_dir, progress.step, 1, config.rollout_transport) + dist.barrier() + + state = ServiceState(completed_step=progress.step - 1) + server = None + server_thread = None + if world.is_master: + server = ThreadingHTTPServer((config.service_host, config.service_port), _server_handler(state)) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + logger.info(f"Value service ready on {config.service_host}:{config.service_port}") + + try: + final_step = None + last_batch = None + last_sync_step = progress.step - 1 + while True: + request = None + if world.is_master: + sync_step = ( + latest_ready_step(config.policy_sync_dir, last_sync_step, final_step or progress.step - 1) + if config.policy_sync_dir is not None and last_batch is not None + else None + ) + if sync_step is not None: + command = ("sync", sync_step) + elif dataloader.receiver.can_receive(): + if final_step is not None: + raise RuntimeError("Value trainer received a batch after max_steps") + command = ("train", None) + else: + try: + request = state.requests.get(timeout=0.1) + except queue.Empty: + continue + command = ("score", request.token_ids) + else: + command = None + objects = [command] + dist.broadcast_object_list(objects, src=0, device=torch.device("cuda")) + action, payload = objects[0] + if action == "sync": + assert config.policy_sync_dir is not None and last_batch is not None + load_policy_backbone(model, config.policy_sync_dir, payload) + _train_batch( + model, + optimizer, + scheduler, + gradient_manager, + last_batch, + parallel_dims, + config, + head_only=False, + updates=config.policy_sync_retune_updates, + ) + last_sync_step = payload + if world.is_master: + (snapshot_path(config.policy_sync_dir, payload) / ".applied").touch() + logger.info(f"Loaded policy backbone at step {payload} and retuned critic LoRA") + continue + if action == "score": + try: + values, bootstrap_value = _score( + model, payload, parallel_dims, config.model.cp_style, config.model.lora is not None + ) + if request is not None: + request.values = values + request.bootstrap_value = bootstrap_value + except BaseException as error: + if request is not None: + request.error = error + raise + finally: + if request is not None: + request.done.set() + continue + + micro_batches: list[TensorMicroBatch] = dataloader.get_batch() + t0 = time.perf_counter() + loss = _train_batch( + model, + optimizer, + scheduler, + gradient_manager, + micro_batches, + parallel_dims, + config, + head_only=progress.step <= max(0, config.head_warmup_steps - config.pretrain_steps), + ) + last_batch = micro_batches + if world.is_master: + logger.info(f"Value step {progress.step} | loss={loss:.5f} | time={time.perf_counter() - t0:.1f}s") + is_last_step = config.max_steps is not None and progress.step >= config.max_steps + if ( + config.ckpt is not None + and config.ckpt.interval + and progress.step % config.ckpt.interval == 0 + and not is_last_step + ): + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) + ckpt_manager.maybe_clean() + if is_last_step: + if config.ckpt is not None: + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress) + final_step = progress.step + else: + progress.step += 1 + state.completed_step = final_step or progress.step - 1 + finally: + if server is not None: + server.shutdown() + server.server_close() + if server_thread is not None: + server_thread.join() + if gradient_manager is not None: + gradient_manager.close() + logger.info("Value trainer finished") + + +def main(): + set_proc_title("Value Trainer") + train(cli(ValueConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/transports/batch/types.py b/src/prime_rl/transports/batch/types.py index f4124ef18ab..11fe8d452a0 100644 --- a/src/prime_rl/transports/batch/types.py +++ b/src/prime_rl/transports/batch/types.py @@ -86,6 +86,12 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr trace_id: str | None = None branch_index: int | None = None + # PPO critic streams. Present together; value_mask selects action tokens + # that contribute to the critic loss in mixed-algorithm batches. + old_values: list[float] | None = None + value_targets: list[float] | None = None + value_mask: list[bool] | None = None + # Orchestrator -> Trainer class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): @@ -123,3 +129,7 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): # (e.g. a dummy micro batch). ``None`` when no packed sample carried one. trace_ids: list[str] | None = None branch_indices: list[int] | None = None + + old_values: list[float] | None = None + value_targets: list[float] | None = None + value_mask: list[bool] | None = None diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index 7dc00a77993..71280bff10c 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -384,6 +384,7 @@ def clean_future_steps(output_dir: Path, resume_step: int) -> None: cleanup_rules = [ (get_batch_dir(output_dir), lambda step: step > resume_step), (get_broadcast_dir(output_dir), lambda step: step >= resume_step), + (output_dir / "policy_sync", lambda step: step > resume_step), ] for directory, should_delete in cleanup_rules: diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index f8aa3a5c75f..a45fbba114f 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -8,6 +8,7 @@ LinearLengthPenaltyConfig, MaxRLAlgoConfig, ) +from prime_rl.orchestrator.algo.gae import skip_observation_gae from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm from prime_rl.orchestrator.algo.routing import assign_advantages @@ -181,6 +182,30 @@ def test_grpo_singleton_group_is_zero(): assert _grpo([_build_episode(0.7, sampled_lengths=[2])]) == pytest.approx([0.0], abs=1e-6) +def test_skip_observation_gae_bridges_tool_output(): + advantages, returns = skip_observation_gae( + [0.0, 0.2, 99.0, 0.4, 0.5], + [False, True, False, True, True], + 1.0, + gamma=1.0, + policy_lambda=0.5, + value_lambda=1.0, + ) + assert advantages == pytest.approx([0.0, 0.375, 0.0, 0.35, 0.5]) + assert returns == pytest.approx([0.0, 1.0, 0.0, 1.0, 1.0]) + partial_advantages, partial_returns = skip_observation_gae( + [0.0, 0.2, 99.0, 0.4, 0.5], + [False, True, False, True, True], + 0.0, + gamma=1.0, + policy_lambda=1.0, + value_lambda=1.0, + bootstrap_value=0.75, + ) + assert partial_advantages == pytest.approx([0.0, 0.55, 0.0, 0.35, 0.25]) + assert partial_returns == pytest.approx([0.0, 0.75, 0.0, 0.75, 0.75]) + + def test_max_rl_mean_normalized(): # mean 0.25: the success gets (1 - 0.25)/0.25 = 3, failures (0 - 0.25)/0.25 = -1 assert _max_rl(_make_group(rewards=[1.0, 0.0, 0.0, 0.0])) == pytest.approx([3.0, -1.0, -1.0, -1.0]) diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 8be69cca458..22127eedf22 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -315,6 +315,27 @@ def test_prepare_sample_uniform_rl_keeps_streams_none(make_training_example): assert micro_batch.ref_kl_weights is None +def test_ppo_value_streams_survive_truncation_and_mixed_packing(make_training_example): + ppo = make_training_example() + ppo.old_values = [0.0, 0.0, 0.2, 0.4] + ppo.value_targets = [0.0, 0.0, 1.0, 1.0] + ppo.value_mask = [False, False, True, True] + truncated = prepare_sample(ppo, seq_len=3) + assert truncated.old_values == [0.0, 0.0, 0.2] + assert truncated.value_targets == [0.0, 0.0, 1.0] + assert truncated.value_mask == [False, False, True] + + mixed = prepare_batch( + rollouts=[ppo, make_training_example()], + seq_len=16, + num_train_workers=1, + bin_cost=build_bin_cost(None), + )[0][0] + assert mixed.old_values == [0.0, 0.0, 0.2, 0.4] + [0.0] * 4 + assert mixed.value_targets == [0.0, 0.0, 1.0, 1.0] + [0.0] * 4 + assert mixed.value_mask == [False, False, True, True] + [False] * 4 + + @pytest.mark.parametrize("streams_on_longer", [True, False]) def test_prepare_batch_packs_mixed_components(make_training_example, streams_on_longer): """Component membership is per token, so samples feeding different diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index a8ce3475598..9a8115c2602 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -516,6 +516,48 @@ def test_multi_node_auto_inference_parallelism(): assert config.inference.vllm.data_parallel_size == 2 +def test_multi_node_value_layout_keeps_policy_and_critic_parallelism_independent(): + payload = { + "trainer": {"model": {"cp": 16, "ep": 8}}, + "orchestrator": { + "algo": {"type": "ppo"}, + "train": {"source": [{"name": "math", "env": {"taskset": {"id": "reverse-text"}}}]}, + }, + "inference": {"vllm": {"tensor_parallel_size": 8}}, + "value": {"model": {"cp": 8, "ep": 8}}, + "deployment": { + "type": "multi_node", + "gpus_per_node": 8, + "num_train_nodes": 2, + "num_infer_nodes": 1, + "num_value_nodes": 1, + }, + "slurm": {}, + "rollout_transport": {"type": "zmq"}, + } + + config = RLConfig.model_validate(payload) + assert config.trainer.model.cp == 16 + assert config.value is not None and config.value.model.cp == 8 + assert config.deployment.num_value_nodes == 1 + assert config.orchestrator.num_train_workers == 1 + assert config.orchestrator.value_rollout_transport.port == config.trainer.rollout_transport.port + 2 + + payload["deployment"]["num_value_nodes"] = 0 + with pytest.raises(ValidationError, match="num_value_nodes > 0"): + RLConfig.model_validate(payload) + + payload["deployment"]["num_value_nodes"] = 1 + payload["value"]["model"]["cp"] = 4 + with pytest.raises(ValidationError, match="one data-parallel worker"): + RLConfig.model_validate(payload) + + payload["value"]["model"]["cp"] = 8 + payload["trainer"]["model"]["cp"] = 8 + with pytest.raises(ValidationError, match="one policy data-parallel worker"): + RLConfig.model_validate(payload) + + def test_orchestrator_vlm_requires_renderer(): with pytest.raises(ValidationError, match="renderer"): OrchestratorConfig.model_validate( diff --git a/tests/unit/train/rl/test_loss.py b/tests/unit/train/rl/test_loss.py index de099aa7705..a185e554ce4 100644 --- a/tests/unit/train/rl/test_loss.py +++ b/tests/unit/train/rl/test_loss.py @@ -1,7 +1,13 @@ import pytest import torch -from prime_rl.configs.trainer import CustomLossConfig, IcePopLossConfig, IPOLossConfig +from prime_rl.configs.trainer import ( + CustomLossConfig, + IcePopLossConfig, + IPOLossConfig, + PPOClipLossConfig, + SAOLossConfig, +) from prime_rl.trainer.rl.loss import ( IcePopLoss, LossInputs, @@ -124,10 +130,42 @@ def test_icepop_loss_masks_extreme_ratio_without_nan(): result = IcePopLoss(IcePopLossConfig()).loss(inputs) assert torch.equal(result.loss, torch.zeros_like(result.loss)) + assert all(torch.isfinite(metric) for metric in result.metrics.values()) result.loss.backward() assert torch.equal(trainer_logprobs.grad, torch.zeros_like(trainer_logprobs.grad)) +def test_ppo_clipped_surrogate_gradients(): + ratios = torch.tensor([0.5, 0.5, 1.0, 2.0, 2.0, float("inf")], device="cuda") + logprobs = ratios.log().requires_grad_() + advantages = torch.tensor([1.0, -1.0, 1.0, 1.0, -1.0, 1.0], device="cuda") + mask = torch.tensor([True, True, True, True, True, False], device="cuda") + result = setup_rl_loss_fn(PPOClipLossConfig()).loss( + LossInputs(logprobs, torch.zeros_like(logprobs), None, advantages, mask) + ) + result.loss.backward() + torch.testing.assert_close(logprobs.grad, torch.tensor([-0.5, 0.0, -1.0, 0.0, 2.0, 0.0], device="cuda")) + torch.testing.assert_close(result.metrics["ppo/clip_fraction"], torch.tensor(0.4, device="cuda")) + assert all(torch.isfinite(metric) for metric in result.metrics.values()) + + +def test_sao_masks_both_sides_independent_of_advantage(): + ratios = torch.tensor([0.1, 0.7, 1.0, 6.0, 10.0], device="cuda") + logprobs = ratios.log().requires_grad_() + result = setup_rl_loss_fn(SAOLossConfig()).loss( + LossInputs( + logprobs, + torch.zeros_like(logprobs), + None, + torch.tensor([-1.0, 1.0, 1.0, -1.0, 1.0], device="cuda"), + torch.ones_like(logprobs, dtype=torch.bool), + ) + ) + result.loss.backward() + torch.testing.assert_close(logprobs.grad, torch.tensor([0.0, -0.7, -1.0, 6.0, 0.0], device="cuda")) + torch.testing.assert_close(result.metrics["is_masked"], torch.tensor(0.4, device="cuda")) + + def test_ce_component_matches_masked_nll(): trainer_logprobs = [torch.tensor([-0.1, -0.5, -0.2], dtype=torch.float32).cuda()] inference_logprobs = [torch.zeros(3, dtype=torch.float32).cuda()] From 7cb6aa33dab9f757bbbcf7aa43cb4f5dbcd72be3 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 08:01:53 +0000 Subject: [PATCH 2/9] fix(ppo): resolve remote value service URL during config validation --- .../prime-rl-configs/src/prime_rl/configs/orchestrator.py | 8 ++++++++ src/prime_rl/orchestrator/orchestrator.py | 5 ----- tests/unit/test_configs.py | 4 ++++ 3 files changed, 12 insertions(+), 5 deletions(-) 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 2a565a70965..4546aca6bbf 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -638,6 +638,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.""" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index ee79eb2302b..86b961ccdb1 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -206,11 +206,6 @@ async def setup(self) -> None: ) self.admin_plane = setup_admin_plane(config.model.client, config.model.name) - if config.value_service_url is not None: - for source in config.train.source: - if source.algo is not None and source.algo.type in ("ppo", "sao"): - source.algo.value_url = config.value_service_url - await monitors.setup( producer="orch", wandb=config.monitors.wandb, diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 9a8115c2602..8de1858b695 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -542,6 +542,10 @@ def test_multi_node_value_layout_keeps_policy_and_critic_parallelism_independent assert config.deployment.num_value_nodes == 1 assert config.orchestrator.num_train_workers == 1 assert config.orchestrator.value_rollout_transport.port == config.trainer.rollout_transport.port + 2 + remote_orchestrator = OrchestratorConfig.model_validate( + {**config.orchestrator.model_dump(), "value_service_url": "http://value-host:8123"} + ) + assert remote_orchestrator.train.source[0].algo.value_url == "http://value-host:8123" payload["deployment"]["num_value_nodes"] = 0 with pytest.raises(ValidationError, match="num_value_nodes > 0"): From 109d03c15515e505c2aaa7a3a26b50bae5bf6c52 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 22:07:14 +0000 Subject: [PATCH 3/9] Support batched PPO critic scoring and independent trainer layouts --- .../src/prime_rl/configs/orchestrator.py | 6 + .../src/prime_rl/configs/rl.py | 54 ++++- .../src/prime_rl/configs/value.py | 2 + src/prime_rl/entrypoints/rl.py | 16 ++ src/prime_rl/orchestrator/algo/base.py | 4 +- src/prime_rl/orchestrator/algo/ppo.py | 119 +++++---- src/prime_rl/orchestrator/orchestrator.py | 9 +- src/prime_rl/orchestrator/packing.py | 12 +- src/prime_rl/orchestrator/train_sink.py | 8 +- .../templates/multi_node_rl.sbatch.j2 | 40 ++- src/prime_rl/trainer/batch.py | 19 +- src/prime_rl/trainer/model.py | 17 +- src/prime_rl/trainer/rl/value.py | 146 +++++++++-- tests/unit/orchestrator/test_advantage.py | 59 +++++ tests/unit/orchestrator/test_batch.py | 20 ++ tests/unit/test_configs.py | 38 ++- tests/unit/train/rl/test_value_data_flow.py | 229 ++++++++++++++++++ .../train/rl/test_value_real_model_smoke.py | 84 +++++++ 18 files changed, 770 insertions(+), 112 deletions(-) create mode 100644 tests/unit/train/rl/test_value_data_flow.py create mode 100644 tests/unit/train/rl/test_value_real_model_smoke.py 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 4546aca6bbf..8053d9c1a0e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -571,6 +571,12 @@ class OrchestratorConfig(BaseConfig): """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``.""" 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 f65be56cf5c..7b4eb10e4a7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -214,6 +214,9 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): 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).""" @@ -226,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 @@ -355,14 +378,16 @@ def validate_deployment(self): 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.num_value_nodes == 0: - raise ValueError("A multi-node value trainer requires num_value_nodes > 0") + 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: - raise ValueError("num_value_nodes 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.") @@ -568,20 +593,26 @@ def auto_setup_value(self): 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.gpus_per_node - num_value_gpus = self.deployment.num_value_nodes * self.deployment.gpus_per_node + 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 on the Qwen3 MoE value model requires a rank divisible by 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") @@ -594,16 +625,15 @@ def auto_setup_value(self): 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 num_value_gpus != self.value.model.cp: - raise ValueError("The value trainer currently requires one data-parallel worker") - if num_train_gpus != self.trainer.model.cp: - raise ValueError("The value trainer currently requires one policy data-parallel worker") 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: @@ -814,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: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/value.py b/packages/prime-rl-configs/src/prime_rl/configs/value.py index 34de4efc859..b4c341da9c6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/value.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/value.py @@ -42,4 +42,6 @@ class ValueConfig(BaseConfig): def validate_pretraining(self): if self.pretrain_steps and self.pretrain_data is None: raise ValueError("value.pretrain_steps requires value.pretrain_data") + if self.model.cp > 1 and self.model.seq_len % (2 * self.model.cp): + raise ValueError("value.model.seq_len must be divisible by 2 * value.model.cp") return self diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index c0e3c308385..8b18b0da6fa 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -555,6 +555,14 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script launcher_log_dir=get_launcher_log_dir(config.run_dir), num_train_nodes=config.deployment.num_train_nodes, num_value_nodes=config.deployment.num_value_nodes, + value_trainer_nodes=config.deployment.value_trainer_nodes, + train_gpus_per_node=config.deployment.train_gpus_per_node, + value_gpus_per_node=config.deployment.value_gpus_per_node, + value_gpus_per_train_node=config.deployment.num_value_gpus_per_train_node, + value_gpu_offset=config.deployment.train_gpus_per_node + if config.deployment.num_value_gpus_per_train_node + else 0, + has_value=config.value is not None, num_infer_nodes=infer_deploy.num_nodes * config.deployment.num_infer_replicas, nodes_per_infer_replica=infer_deploy.num_nodes, num_infer_replicas=config.deployment.num_infer_replicas, @@ -603,6 +611,14 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script launcher_log_dir=get_launcher_log_dir(config.run_dir), num_train_nodes=config.deployment.num_train_nodes, num_value_nodes=config.deployment.num_value_nodes, + value_trainer_nodes=config.deployment.value_trainer_nodes, + train_gpus_per_node=config.deployment.train_gpus_per_node, + value_gpus_per_node=config.deployment.value_gpus_per_node, + value_gpus_per_train_node=config.deployment.num_value_gpus_per_train_node, + value_gpu_offset=config.deployment.train_gpus_per_node + if config.deployment.num_value_gpus_per_train_node + else 0, + has_value=config.value is not None, num_infer_nodes=config.deployment.total_infer_nodes, nodes_per_infer_replica=config.deployment.infer_nodes_per_replica, num_infer_replicas=config.deployment.num_infer_replicas, diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index b4c57492144..0b39c6a4f36 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -74,8 +74,8 @@ async def score_episode(self, episode: vf.Episode) -> None: async def score_group(self, episodes: list[vf.Episode]) -> None: """Assign group-relative annotations to a finalized cohort.""" - async def score_samples(self, trace: vf.Trace, samples: list[TrainingSample]) -> None: - """Assign token-level training data after branch samples are materialized.""" + async def score_samples(self, trace_samples: list[tuple[vf.Trace, list[TrainingSample]]]) -> None: + """Assign token-level training data after a group's branch samples are materialized.""" async def finalize_episode(self, episode: vf.Episode) -> None: """Run rollout-local scoring when the episode has trainable traces.""" diff --git a/src/prime_rl/orchestrator/algo/ppo.py b/src/prime_rl/orchestrator/algo/ppo.py index b680e20497e..c2c85abf1b9 100644 --- a/src/prime_rl/orchestrator/algo/ppo.py +++ b/src/prime_rl/orchestrator/algo/ppo.py @@ -40,50 +40,83 @@ async def setup(self) -> None: async def aclose(self) -> None: await self.value_client.aclose() - async def score_samples(self, trace: vf.Trace, samples: list[TrainingSample]) -> None: - branches = list(iter_trainable_branches(trace)) - if len(branches) != len(samples): - raise ValueError("PPO branch/sample count mismatch") - for (branch, _), sample in zip(branches, samples, strict=True): - scored_length = min(len(sample.token_ids), self.config.value_seq_len or len(sample.token_ids)) + async def score_samples_by_sample(self, trace: vf.Trace, samples: list[TrainingSample]) -> None: + for _, branch, sample, scored_length in self._scored_branches([(trace, samples)]): response = await self.value_client.post("/score", json={"token_ids": sample.token_ids[:scored_length]}) response.raise_for_status() score = response.json() - values = score["values"] - if len(values) != scored_length: - raise ValueError("Value service returned a misaligned token stream") - scored_mask = sample.mask[:scored_length] - n_actions = sum(scored_mask) - policy_lambda = self.config.policy_lambda - if self.config.length_adaptive_alpha is not None and n_actions: - policy_lambda = max(0.0, 1 - 1 / (self.config.length_adaptive_alpha * n_actions)) - advantages, targets = skip_observation_gae( - values, - scored_mask, - trace.reward if scored_length == len(sample.token_ids) else 0.0, - gamma=self.config.gamma, - policy_lambda=policy_lambda, - value_lambda=self.config.value_lambda, - bootstrap_value=score["bootstrap_value"] if scored_length < len(sample.token_ids) else 0.0, - ) - unscored = len(sample.token_ids) - scored_length - values += [0.0] * unscored - advantages += [0.0] * unscored - targets += [0.0] * unscored - sample.advantages = advantages - sample.old_values = values - sample.value_targets = targets - sample.value_mask = list(scored_mask) + [False] * unscored + self._apply_score(trace, branch, sample, score["values"], score["bootstrap_value"], scored_length) - offset = 0 - for node in branch.nodes: - end = offset + len(node.token_ids) - node_advantages = [ - advantage - for advantage, sampled in zip(advantages[offset:end], sample.mask[offset:end], strict=True) - if sampled - ] - if node_advantages: - node.advantages = node_advantages - offset = end - sample.mask = sample.value_mask.copy() + async def score_samples(self, trace_samples: list[tuple[vf.Trace, list[TrainingSample]]]) -> None: + scored = self._scored_branches(trace_samples) + if not scored: + return + response = await self.value_client.post( + "/score", + json={"batched_token_ids": [sample.token_ids[:length] for _, _, sample, length in scored]}, + ) + response.raise_for_status() + result = response.json() + values_batch = result["values"] + bootstraps = result["bootstrap_values"] + if len(values_batch) != len(scored) or len(bootstraps) != len(scored): + raise ValueError("Value service returned the wrong number of samples") + for (trace, branch, sample, length), values, bootstrap in zip(scored, values_batch, bootstraps, strict=True): + self._apply_score(trace, branch, sample, values, bootstrap, length) + + def _scored_branches( + self, trace_samples: list[tuple[vf.Trace, list[TrainingSample]]] + ) -> list[tuple[vf.Trace, vf.Branch, TrainingSample, int]]: + scored = [] + for trace, samples in trace_samples: + branches = list(iter_trainable_branches(trace)) + if len(branches) != len(samples): + raise ValueError("PPO branch/sample count mismatch") + for (branch, _), sample in zip(branches, samples, strict=True): + scored_length = min(len(sample.token_ids), self.config.value_seq_len or len(sample.token_ids)) + scored.append((trace, branch, sample, scored_length)) + return scored + + def _apply_score( + self, + trace: vf.Trace, + branch: vf.Branch, + sample: TrainingSample, + values: list[float], + bootstrap_value: float, + scored_length: int, + ) -> None: + if len(values) != scored_length: + raise ValueError("Value service returned a misaligned token stream") + scored_mask = sample.mask[:scored_length] + n_actions = sum(scored_mask) + policy_lambda = self.config.policy_lambda + if self.config.length_adaptive_alpha is not None and n_actions: + policy_lambda = max(0.0, 1 - 1 / (self.config.length_adaptive_alpha * n_actions)) + advantages, targets = skip_observation_gae( + values, + scored_mask, + trace.reward if scored_length == len(sample.token_ids) else 0.0, + gamma=self.config.gamma, + policy_lambda=policy_lambda, + value_lambda=self.config.value_lambda, + bootstrap_value=bootstrap_value if scored_length < len(sample.token_ids) else 0.0, + ) + unscored = len(sample.token_ids) - scored_length + sample.old_values = values + [0.0] * unscored + sample.advantages = advantages + [0.0] * unscored + sample.value_targets = targets + [0.0] * unscored + sample.value_mask = list(scored_mask) + [False] * unscored + + offset = 0 + for node in branch.nodes: + end = offset + len(node.token_ids) + node_advantages = [ + advantage + for advantage, sampled in zip(sample.advantages[offset:end], sample.mask[offset:end], strict=True) + if sampled + ] + if node_advantages: + node.advantages = node_advantages + offset = end + sample.mask = sample.value_mask.copy() diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 86b961ccdb1..a144bc2dafd 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -261,12 +261,15 @@ async def setup(self) -> None: # Transports are local setup — initialize them before the env and inference waits. self.packer = BatchPacker(config) + self.value_packer = BatchPacker(config, for_value=True) if config.value_rollout_transport is not None else None get_logger().info(f"Initializing micro batch sender ({config.rollout_transport})") self.sender = setup_batch_sender( config.output_dir, config.num_train_workers, self.progress.step, config.rollout_transport ) self.value_sender = ( - setup_batch_sender(config.output_dir, 1, self.progress.step, config.value_rollout_transport) + setup_batch_sender( + config.output_dir, config.value_num_train_workers, self.progress.step, config.value_rollout_transport + ) if config.value_rollout_transport is not None else None ) @@ -645,7 +648,9 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: micro_batch_grid = await asyncio.to_thread(self.packer.pack, batch.samples) pack_time = time.perf_counter() - pack_start_time if self.value_sender is not None: - await self.value_sender.send(micro_batch_grid) + assert self.value_packer is not None + value_grid = await asyncio.to_thread(self.value_packer.pack, batch.samples) + await self.value_sender.send(value_grid) assert self.value_client is not None value_deadline = time.monotonic() + 3600 while True: diff --git a/src/prime_rl/orchestrator/packing.py b/src/prime_rl/orchestrator/packing.py index c3716eb0701..b2ce239ea75 100644 --- a/src/prime_rl/orchestrator/packing.py +++ b/src/prime_rl/orchestrator/packing.py @@ -9,10 +9,13 @@ class BatchPacker: """Bin-packs a step's training samples into one micro-batch list per trainer DP rank.""" - def __init__(self, config: OrchestratorConfig): - self.seq_len = config.seq_len - self.num_train_workers = config.num_train_workers - self.pad_to_multiple_of = config.pad_to_multiple_of + def __init__(self, config: OrchestratorConfig, *, for_value: bool = False): + self.seq_len = config.value_seq_len if for_value else config.seq_len + self.num_train_workers = config.value_num_train_workers if for_value else config.num_train_workers + self.pad_to_multiple_of = config.value_pad_to_multiple_of if for_value else config.pad_to_multiple_of + self.for_value = for_value + if self.seq_len is None or self.num_train_workers is None or self.pad_to_multiple_of is None: + raise ValueError("Value packing requires critic sequence length, DP count, and CP padding") try: model_config = AutoConfig.from_pretrained( config.model.name, trust_remote_code=config.tokenizer.trust_remote_code @@ -32,4 +35,5 @@ def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]: num_train_workers=self.num_train_workers, bin_cost=self.bin_cost, pad_to_multiple_of=self.pad_to_multiple_of, + for_value=self.for_value, ) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 80c91c64806..c6b273b9092 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -277,11 +277,13 @@ async def process_group(self, group_id: str) -> None: ) return + trace_samples = [ + (trace, await asyncio.to_thread(trace_to_samples, trace, env_name=env_name)) for trace in survivors + ] + await env.algorithm.score_samples(trace_samples) samples_by_trace: dict[str, list[TrainingSample]] = {} temperature = env.sampling_args["temperature"] - for trace in survivors: - samples = await asyncio.to_thread(trace_to_samples, trace, env_name=env_name) - await env.algorithm.score_samples(trace, samples) + for trace, samples in trace_samples: for sample in samples: sample.temperatures = [temperature] * len(sample.token_ids) if env.requires_sampling_masks and sample.sampling_mask is None: diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 24dc1ed2c9f..f203b09c8a5 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -29,6 +29,11 @@ set -e export NUM_TRAIN_NODES={{ num_train_nodes }} export NUM_INFER_NODES={{ num_infer_nodes }} export NUM_VALUE_NODES={{ num_value_nodes }} +export NUM_VALUE_TRAINER_NODES={{ value_trainer_nodes }} +export TRAIN_GPUS_PER_NODE={{ train_gpus_per_node }} +export VALUE_GPUS_PER_NODE={{ value_gpus_per_node }} +export VALUE_GPUS_PER_TRAIN_NODE={{ value_gpus_per_train_node }} +export VALUE_GPU_OFFSET={{ value_gpu_offset }} export NODES_PER_INFER_REPLICA={{ nodes_per_infer_replica }} export NUM_INFER_REPLICAS={{ num_infer_replicas }} export GPUS_PER_NODE={{ gpus_per_node }} @@ -63,10 +68,10 @@ export LAUNCHER_DIR={{ launcher_dir }} export LOG_DIR={{ log_dir }} export PRL_ATTEMPT_CONFIG_DIR=$CONFIG_DIR export PRL_ATTEMPT_LOG_DIR=$LOG_DIR -mkdir -p $LOG_DIR/trainer $LOG_DIR/inference{% if num_value_nodes > 0 %} $LOG_DIR/value{% endif %} +mkdir -p $LOG_DIR/trainer $LOG_DIR/inference{% if has_value %} $LOG_DIR/value{% endif %} ln -sfn trainer/node_0.log $LOG_DIR/trainer.log ln -sfn inference/node_0.log $LOG_DIR/inference.log -{% if num_value_nodes > 0 -%} +{% if has_value -%} ln -sfn value/node_0.log $LOG_DIR/value.log {% endif -%} @@ -88,8 +93,12 @@ export WANDB_SHARED_MODE=1 # Networking HOSTNAMES=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) TRAIN_HOSTS=( ${HOSTNAMES[@]:$NUM_INFER_NODES:$NUM_TRAIN_NODES} ) -{% if num_value_nodes > 0 -%} +{% if has_value -%} +{% if value_gpus_per_train_node > 0 -%} +VALUE_HOSTS=( ${TRAIN_HOSTS[@]} ) +{% else -%} VALUE_HOSTS=( ${HOSTNAMES[@]:$((NUM_INFER_NODES + NUM_TRAIN_NODES)):$NUM_VALUE_NODES} ) +{% endif -%} export VALUE_ADDR="${VALUE_HOSTS[0]}" export VALUE_SERVICE_PORT={{ value_service_port }} {% endif -%} @@ -142,7 +151,7 @@ echo "ADMIN_URLS=${ADMIN_URLS}" echo "HOSTNAMES=${HOSTNAMES[*]}" echo "TRAIN_HOSTS=${TRAIN_HOSTS[*]}" -{% if num_value_nodes > 0 -%} +{% if has_value -%} echo "VALUE_HOSTS=${VALUE_HOSTS[*]}" {% endif -%} @@ -505,7 +514,7 @@ else {%- else %} TRAIN_NODE_RANK=$SLURM_PROCID {%- endif %} -{% if num_value_nodes > 0 %} +{% if has_value %} if [ "$TRAIN_NODE_RANK" -lt "$NUM_TRAIN_NODES" ]; then {% endif %} @@ -517,10 +526,10 @@ if [ "$TRAIN_NODE_RANK" -lt "$NUM_TRAIN_NODES" ]; then echo $HOSTNAMES_STR | tee $LOG_DIR/trainer/node_${TRAIN_NODE_RANK}.log ( - WANDB_SHARED_LABEL=trainer uv run torchrun \ + CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((TRAIN_GPUS_PER_NODE - 1))) WANDB_SHARED_LABEL=trainer uv run torchrun \ --role=trainer \ --nnodes=$NUM_TRAIN_NODES \ - --nproc-per-node=$GPUS_PER_NODE \ + --nproc-per-node=$TRAIN_GPUS_PER_NODE \ --node-rank=$TRAIN_NODE_RANK \ --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT \ --rdzv-id=job_$SLURM_JOB_ID \ @@ -539,21 +548,26 @@ if [ "$TRAIN_NODE_RANK" -lt "$NUM_TRAIN_NODES" ]; then fi exit "$COMPONENT_CODE" ) & -{% if num_value_nodes > 0 %}fi{% endif %} +{% if has_value %}fi{% endif %} {% if num_infer_nodes > 0 %} fi{% endif %} -{% if num_value_nodes > 0 %} +{% if has_value %} +{% if value_gpus_per_train_node > 0 %} +if [ "$SLURM_PROCID" -ge "$NUM_INFER_NODES" ] && [ "$SLURM_PROCID" -lt "$((NUM_INFER_NODES + NUM_TRAIN_NODES))" ]; then + VALUE_NODE_RANK=$((SLURM_PROCID - NUM_INFER_NODES)) +{% else %} if [ "$SLURM_PROCID" -ge "$((NUM_INFER_NODES + NUM_TRAIN_NODES))" ]; then VALUE_NODE_RANK=$((SLURM_PROCID - NUM_INFER_NODES - NUM_TRAIN_NODES)) +{% endif %} {%- for key, value in value_env_vars.items() %} export {{ key }}="{{ value }}" {%- endfor %} echo $HOSTNAMES_STR | tee $LOG_DIR/value/node_${VALUE_NODE_RANK}.log ( - WANDB_SHARED_LABEL=value uv run torchrun \ + CUDA_VISIBLE_DEVICES=$(seq -s, $VALUE_GPU_OFFSET $((VALUE_GPU_OFFSET + VALUE_GPUS_PER_NODE - 1))) WANDB_SHARED_LABEL=value uv run torchrun \ --role=value \ - --nnodes=$NUM_VALUE_NODES \ - --nproc-per-node=$GPUS_PER_NODE \ + --nnodes=$NUM_VALUE_TRAINER_NODES \ + --nproc-per-node=$VALUE_GPUS_PER_NODE \ --node-rank=$VALUE_NODE_RANK \ --rdzv-endpoint=$VALUE_ADDR:$((MASTER_PORT + 1)) \ --rdzv-id=value_$SLURM_JOB_ID \ @@ -612,7 +626,7 @@ if [ "$SLURM_PROCID" -eq "$ORCH_PROCID" ]; then {% if use_nccl_broadcast %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --weight_broadcast.host $MASTER_ADDR" {% elif use_nixl_broadcast %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --weight_broadcast.host $WEIGHT_BROADCAST_HOST" {% endif %}{% if use_zmq_transport %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --rollout_transport.host $ORCH_ADDR" - {% endif %}{% if num_value_nodes > 0 %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --value-rollout-transport.host $ORCH_ADDR --value-service-url http://$VALUE_ADDR:$VALUE_SERVICE_PORT" + {% endif %}{% if has_value %}ORCHESTRATOR_ARGS="$ORCHESTRATOR_ARGS --value-rollout-transport.host $ORCH_ADDR --value-service-url http://$VALUE_ADDR:$VALUE_SERVICE_PORT" {% endif %}WANDB_SHARED_LABEL=orchestrator uv run orchestrator $ORCHESTRATOR_ARGS \ 2>&1 | tee $LOG_DIR/orchestrator.log COMPONENT_CODE=$? diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 16b7a99d6a8..4e2feff8c0c 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -877,7 +877,7 @@ def _assert_token_arrays_aligned(micro_batch: MicroBatch) -> None: ) -def _make_dummy_batch(source: MicroBatch) -> MicroBatch: +def _make_dummy_batch(source: MicroBatch, *, for_value: bool = False) -> MicroBatch: """Create a zero-loss dummy batch from an existing batch, preserving its modality.""" dummy = copy.deepcopy(source) dummy.advantages = [0.0] * len(dummy.input_ids) @@ -887,9 +887,9 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch: dummy.rl_weights = None dummy.ce_weights = None dummy.ref_kl_weights = None - dummy.old_values = None - dummy.value_targets = None - dummy.value_mask = None + dummy.old_values = [0.0] * len(dummy.input_ids) if for_value else None + dummy.value_targets = [0.0] * len(dummy.input_ids) if for_value else None + dummy.value_mask = [False] * len(dummy.input_ids) if for_value else None # Fully loss-masked, so replaying sampling masks would be pure wasted work. dummy.sampling_mask = None # The copied identity would double-annotate the source's traces. @@ -898,11 +898,13 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch: return dummy -def _pad_group_for_distribution(group: list[MicroBatch], num_train_workers: int) -> list[MicroBatch]: +def _pad_group_for_distribution( + group: list[MicroBatch], num_train_workers: int, *, for_value: bool = False +) -> list[MicroBatch]: """Pad a group of micro batches so its length is divisible by num_train_workers.""" num_padding = -len(group) % num_train_workers if num_padding > 0 and len(group) > 0: - dummy = _make_dummy_batch(group[0]) + dummy = _make_dummy_batch(group[0], for_value=for_value) group.extend([dummy] * num_padding) return group @@ -913,6 +915,7 @@ def prepare_batch( num_train_workers: int, bin_cost: Callable[[Sequence[int]], int], pad_to_multiple_of: int = 1, + for_value: bool = False, ) -> list[list[MicroBatch]]: """ Prepare a batch of problems for each GPU. Each batch is a list of micro batches. @@ -933,8 +936,8 @@ def prepare_batch( text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)] # Pad each group independently so its count is divisible by num_train_workers - mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers) - text_batches = _pad_group_for_distribution(text_batches, num_train_workers) + mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers, for_value=for_value) + text_batches = _pad_group_for_distribution(text_batches, num_train_workers, for_value=for_value) # Alignment check after distribution padding so the dummy batches are covered too for micro_batch in (*mm_batches, *text_batches): diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 39ff3f5f44e..355cbf2da08 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -1060,9 +1060,7 @@ def setup_model( for param in model.lm_head.parameters(): param.requires_grad = False if freeze_attention: - for layer in model.model.layers: - for param in layer.self_attn.parameters(): - param.requires_grad = False + freeze_attention_modules(model) for param in model.value_head.parameters(): param.requires_grad = True @@ -1114,6 +1112,19 @@ def setup_model( return model +def freeze_attention_modules(model: nn.Module) -> None: + attention = [ + module + for name, module in model.named_modules() + if name.rsplit(".", 1)[-1] in {"self_attn", "linear_attn", "attn"} + ] + if not attention: + raise ValueError("value.freeze_attention requires named attention modules on the value backbone") + for module in attention: + for param in module.parameters(): + param.requires_grad = False + + def forward( model: nn.Module, input_ids: Int[Tensor, "batch seq"], diff --git a/src/prime_rl/trainer/rl/value.py b/src/prime_rl/trainer/rl/value.py index 7305e6c9aee..cf57ac06ded 100644 --- a/src/prime_rl/trainer/rl/value.py +++ b/src/prime_rl/trainer/rl/value.py @@ -48,20 +48,57 @@ @dataclass class ScoreRequest: - token_ids: list[int] + batched_token_ids: list[list[int]] + single: bool = False done: threading.Event = field(default_factory=threading.Event) - values: list[float] | None = None - bootstrap_value: float | None = None + values: list[list[float]] | None = None + bootstrap_values: list[float] | None = None error: BaseException | None = None +@dataclass +class ScoreBin: + indices: list[int] + token_ids: list[int] + lengths: list[int] + + +def pack_score_requests(batched_token_ids: list[list[int]], seq_len: int, dp: int, cp: int) -> list[list[ScoreBin]]: + """Pack whole scoring sequences, then give every DP replica the same number of bins.""" + if not batched_token_ids or any(not ids or len(ids) > seq_len for ids in batched_token_ids): + raise ValueError("Scoring requires nonempty sequences within value.model.seq_len") + divisor = 2 * cp if cp > 1 else 1 + if seq_len % divisor: + raise ValueError("value.model.seq_len must be divisible by the CP score divisor") + bins: list[ScoreBin] = [] + for index in sorted(range(len(batched_token_ids)), key=lambda i: len(batched_token_ids[i]), reverse=True): + tokens = batched_token_ids[index] + target = next((b for b in bins if len(b.token_ids) + len(tokens) <= seq_len), None) + if target is None: + target = ScoreBin([], [], []) + bins.append(target) + target.indices.append(index) + target.token_ids.extend(tokens) + target.lengths.append(len(tokens)) + ranks: list[list[ScoreBin]] = [[] for _ in range(dp)] + loads = [0] * dp + for score_bin in sorted(bins, key=lambda b: len(b.token_ids), reverse=True): + rank = min(range(dp), key=lambda i: (loads[i], len(ranks[i]))) + ranks[rank].append(score_bin) + loads[rank] += len(score_bin.token_ids) + steps = max(map(len, ranks)) + for rank_bins in ranks: + rank_bins.extend(ScoreBin([], [], []) for _ in range(steps - len(rank_bins))) + return ranks + + @dataclass class ServiceState: requests: queue.Queue[ScoreRequest] = field(default_factory=queue.Queue) completed_step: int = 0 -def _server_handler(state: ServiceState): +def _server_handler(state: ServiceState, max_seq_len: int | None = None): class Handler(BaseHTTPRequestHandler): def do_GET(self): if self.path == "/health": @@ -83,15 +120,23 @@ def do_POST(self): self.send_error(404) return body = self.rfile.read(int(self.headers["Content-Length"])) - token_ids = json.loads(body)["token_ids"] + data = json.loads(body) + single = "token_ids" in data + batched_token_ids = [data["token_ids"]] if single else data.get("batched_token_ids") if ( - not isinstance(token_ids, list) - or not token_ids - or not all(isinstance(token, int) for token in token_ids) + not isinstance(batched_token_ids, list) + or not batched_token_ids + or any( + not isinstance(ids, list) or not ids or any(not isinstance(token, int) for token in ids) + for ids in batched_token_ids + ) ): - self.send_error(400, "token_ids must be a nonempty list of integers") + self.send_error(400, "batched_token_ids must be a nonempty list of nonempty integer lists") + return + if max_seq_len is not None and any(len(ids) > max_seq_len for ids in batched_token_ids): + self.send_error(400, "scoring sequence exceeds value.model.seq_len") return - request = ScoreRequest(token_ids) + request = ScoreRequest(batched_token_ids, single=single) state.requests.put(request) if not request.done.wait(timeout=3600): self.send_error(504, "value scoring timed out") @@ -99,7 +144,12 @@ def do_POST(self): if request.error is not None: self.send_error(500, str(request.error)) return - payload = json.dumps({"values": request.values, "bootstrap_value": request.bootstrap_value}).encode() + if single: + payload = json.dumps( + {"values": request.values[0], "bootstrap_value": request.bootstrap_values[0]} + ).encode() + else: + payload = json.dumps({"values": request.values, "bootstrap_values": request.bootstrap_values}).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(payload))) @@ -160,6 +210,58 @@ def _score(model, token_ids: list[int], parallel_dims, cp_style: str, lora_enabl ) +def _score_batch( + model, batched_token_ids: list[list[int]], parallel_dims, seq_len: int, cp_style: str, lora_enabled: bool +) -> tuple[list[list[float]], list[float]] | None: + dp = parallel_dims.get_mesh("dp").size() + cp = parallel_dims.cp + rank = dist.get_rank() + dp_rank = rank // parallel_dims.non_data_parallel_size + bins_by_rank = pack_score_requests(batched_token_ids, seq_len, dp, cp) + local_results: list[tuple[int, list[float], float]] = [] + model.eval() + with torch.no_grad(): + for step in range(len(bins_by_rank[0])): + score_bin = bins_by_rank[dp_rank][step] + length = len(score_bin.token_ids) + max_length = max(len(bins[step].token_ids) for bins in bins_by_rank) + divisor = 2 * cp if cp > 1 else 1 + padded_length = ((max_length + divisor - 1) // divisor) * divisor + padding = padded_length - length + tokens = score_bin.token_ids + [0] * padding + positions = [position for n in score_bin.lengths for position in range(n)] + list(range(padding)) + segment_lengths = score_bin.lengths + ([padding] if padding else []) + input_ids = torch.tensor(tokens, device="cuda", dtype=torch.long).unsqueeze(0) + position_ids = torch.tensor(positions, device="cuda", dtype=torch.long).unsqueeze(0) + seq_lens = torch.tensor(segment_lengths, device="cuda", dtype=torch.long) + raw_values = _forward_values( + model, input_ids, position_ids, seq_lens, parallel_dims, cp_style, lora_enabled=lora_enabled + ) + if parallel_dims.cp_enabled: + raw_values = gather_for_cp_wo_grad(raw_values, cp, parallel_dims.get_mesh("cp").get_group()) + if rank % parallel_dims.non_data_parallel_size == 0: + offset = 0 + for index, n in zip(score_bin.indices, score_bin.lengths, strict=True): + segment = raw_values[:, offset : offset + n] + local_results.append( + (index, shift_tensor_right(segment)[0].float().cpu().tolist(), segment[0, -1].float().item()) + ) + offset += n + gathered = [None] * dist.get_world_size() if rank == 0 else None + dist.gather_object(local_results, gathered, dst=0) + if rank != 0: + return None + values: list[list[float] | None] = [None] * len(batched_token_ids) + bootstraps: list[float | None] = [None] * len(batched_token_ids) + for rank_results in gathered: + for index, token_values, bootstrap in rank_results: + values[index] = token_values + bootstraps[index] = bootstrap + if any(value is None for value in values) or any(value is None for value in bootstraps): + raise RuntimeError("Value scoring failed to return every sequence") + return values, bootstraps + + def _train_batch( model, optimizer, scheduler, gradient_manager, micro_batches, parallel_dims, config, head_only, updates=None ): @@ -322,14 +424,18 @@ def train(config: ValueConfig): if checkpoint_step is None: _pretrain(model, optimizer, scheduler, gradient_manager, parallel_dims, config) - dataloader = DataLoader(config.rollout_dir, progress.step, 1, config.rollout_transport) + dataloader = DataLoader( + config.rollout_dir, progress.step, parallel_dims.get_mesh("dp").size(), config.rollout_transport + ) dist.barrier() state = ServiceState(completed_step=progress.step - 1) server = None server_thread = None if world.is_master: - server = ThreadingHTTPServer((config.service_host, config.service_port), _server_handler(state)) + server = ThreadingHTTPServer( + (config.service_host, config.service_port), _server_handler(state, config.model.seq_len) + ) server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() logger.info(f"Value service ready on {config.service_host}:{config.service_port}") @@ -357,7 +463,7 @@ def train(config: ValueConfig): request = state.requests.get(timeout=0.1) except queue.Empty: continue - command = ("score", request.token_ids) + command = ("score", request.batched_token_ids) else: command = None objects = [command] @@ -384,12 +490,16 @@ def train(config: ValueConfig): continue if action == "score": try: - values, bootstrap_value = _score( - model, payload, parallel_dims, config.model.cp_style, config.model.lora is not None + result = _score_batch( + model, + payload, + parallel_dims, + config.model.seq_len, + config.model.cp_style, + config.model.lora is not None, ) if request is not None: - request.values = values - request.bootstrap_value = bootstrap_value + request.values, request.bootstrap_values = result except BaseException as error: if request is not None: request.error = error diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index a45fbba114f..48ea8572a26 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -1,5 +1,7 @@ import asyncio +import json +import httpx import pytest import verifiers.v1 as vf @@ -7,10 +9,12 @@ GRPOAlgoConfig, LinearLengthPenaltyConfig, MaxRLAlgoConfig, + PPOAlgoConfig, ) from prime_rl.orchestrator.algo.gae import skip_observation_gae from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm +from prime_rl.orchestrator.algo.ppo import PPOAlgorithm from prime_rl.orchestrator.algo.routing import assign_advantages from prime_rl.orchestrator.trajectories import trace_to_samples @@ -166,6 +170,61 @@ def _max_rl(group: list[vf.Episode]) -> list[float]: return [_scalar(episode) for episode in group] +def test_ppo_group_scoring_matches_per_sample_scoring(): + requests = [] + + def score(ids): + raw = [token * 0.1 for token in ids] + return [0.0] + raw[:-1], raw[-1] + + def respond(request): + body = json.loads(request.content) + requests.append(body) + if "batched_token_ids" in body: + scored = [score(ids) for ids in body["batched_token_ids"]] + return httpx.Response( + 200, json={"values": [v for v, _ in scored], "bootstrap_values": [b for _, b in scored]} + ) + values, bootstrap = score(body["token_ids"]) + return httpx.Response(200, json={"values": values, "bootstrap_value": bootstrap}) + + async def run(batched): + episodes = [ + _build_episode(1.0, sampled_lengths=[2, 2], obs_lengths=[1]), + _build_episode(0.0, sampled_lengths=[3]), + ] + trace_samples = [(e.traces[0], trace_to_samples(e.traces[0])) for e in episodes] + algo = PPOAlgorithm(PPOAlgoConfig(value_seq_len=4, length_adaptive_alpha=1.5), clients=None) + await algo.value_client.aclose() + algo.value_client = httpx.AsyncClient(base_url="http://value.test", transport=httpx.MockTransport(respond)) + if batched: + await algo.score_samples(trace_samples) + else: + for trace, samples in trace_samples: + await algo.score_samples_by_sample(trace, samples) + await algo.aclose() + return [ + ( + sample.advantages, + sample.old_values, + sample.value_targets, + sample.value_mask, + sample.mask, + [node.advantages for node in trace.nodes], + ) + for trace, samples in trace_samples + for sample in samples + ] + + by_sample = asyncio.run(run(False)) + assert len(requests) == 2 + requests.clear() + batched = asyncio.run(run(True)) + assert batched == by_sample + assert len(requests) == 1 + assert len(requests[0]["batched_token_ids"]) == 2 + + # -------------------------------------------------------------------------- # GRPO / MaxRL: group-relative credit, assigned in score_group. # -------------------------------------------------------------------------- diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 22127eedf22..cd3dec1d205 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -56,6 +56,26 @@ def make_sized_training_example(length: int, env_name: str = "test-env") -> Trai ) +def test_value_packing_uses_own_context_and_masked_dummies(): + samples = [make_sized_training_example(length) for length in (5, 4, 3)] + for sample in samples: + n = len(sample.token_ids) + sample.old_values = [0.1] * n + sample.value_targets = [0.2] * n + sample.value_mask = list(sample.mask) + policy_grid = prepare_batch(samples, seq_len=8, num_train_workers=1, bin_cost=sum) + value_grid = prepare_batch( + samples, seq_len=4, num_train_workers=2, bin_cost=sum, pad_to_multiple_of=2, for_value=True + ) + assert len(policy_grid) == 1 and len(value_grid) == 2 + assert all(len(worker) == len(value_grid[0]) for worker in value_grid) + assert all(len(batch.input_ids) <= 4 and len(batch.input_ids) % 2 == 0 for worker in value_grid for batch in worker) + assert sum(sum(batch.value_mask) for worker in value_grid for batch in worker) == 2 + dummies = [batch for worker in value_grid for batch in worker if batch.trace_ids is None] + assert dummies and all(batch.value_targets == [0.0] * len(batch.input_ids) for batch in dummies) + assert [len(sample.token_ids) for sample in samples] == [5, 4, 3] + + def _flatten_batches(batches_per_gpu): return [batch for worker_batches in batches_per_gpu for batch in worker_batches] diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 8de1858b695..5acf5dbfeb6 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -548,17 +548,47 @@ def test_multi_node_value_layout_keeps_policy_and_critic_parallelism_independent assert remote_orchestrator.train.source[0].algo.value_url == "http://value-host:8123" payload["deployment"]["num_value_nodes"] = 0 - with pytest.raises(ValidationError, match="num_value_nodes > 0"): + with pytest.raises(ValidationError, match="requires value nodes or GPUs"): RLConfig.model_validate(payload) + payload["deployment"]["num_value_gpus_per_train_node"] = 4 + payload["trainer"]["model"]["cp"] = 4 + payload["value"]["model"]["cp"] = 4 + config = RLConfig.model_validate(payload) + assert config.deployment.train_gpus_per_node == 4 + assert config.deployment.value_trainer_nodes == 2 + assert config.orchestrator.num_train_workers == 2 + assert config.orchestrator.value_num_train_workers == 2 + payload["deployment"].pop("num_value_gpus_per_train_node") + payload["trainer"]["model"]["cp"] = 16 + payload["value"]["model"]["cp"] = 8 + payload["deployment"]["num_value_nodes"] = 1 payload["value"]["model"]["cp"] = 4 - with pytest.raises(ValidationError, match="one data-parallel worker"): + config = RLConfig.model_validate(payload) + assert config.orchestrator.value_num_train_workers == 2 + assert config.orchestrator.value_pad_to_multiple_of == 4 + assert config.orchestrator.value_seq_len == config.value.model.seq_len + + payload["trainer"]["model"]["cp"] = 8 + config = RLConfig.model_validate(payload) + assert config.orchestrator.num_train_workers == 2 + assert config.orchestrator.value_num_train_workers == 2 + + payload["value"]["model"]["cp"] = 3 + payload["value"]["model"]["seq_len"] = 2046 + with pytest.raises(ValidationError, match="Value GPU count must be divisible"): + RLConfig.model_validate(payload) + + payload["value"]["model"]["cp"] = 4 + payload["value"]["model"].pop("seq_len") + payload["trainer"]["model"]["cp"] = 7 + with pytest.raises(ValidationError, match="Policy GPU count must be divisible"): RLConfig.model_validate(payload) - payload["value"]["model"]["cp"] = 8 payload["trainer"]["model"]["cp"] = 8 - with pytest.raises(ValidationError, match="one policy data-parallel worker"): + payload["trainer"]["model"]["seq_len"] = 2050 + with pytest.raises(ValidationError, match="value.model.seq_len must be divisible"): RLConfig.model_validate(payload) diff --git a/tests/unit/train/rl/test_value_data_flow.py b/tests/unit/train/rl/test_value_data_flow.py new file mode 100644 index 00000000000..96f0ad0d05b --- /dev/null +++ b/tests/unit/train/rl/test_value_data_flow.py @@ -0,0 +1,229 @@ +"""Critic score packing and distributed dummy-model data flow.""" + +import asyncio +import os +import socket +import threading +from http.server import ThreadingHTTPServer +from pathlib import Path +from types import SimpleNamespace + +import httpx +import torch +import torch.distributed as dist +import zmq + +from prime_rl.configs.shared import ZMQTransportConfig +from prime_rl.trainer.batch import prepare_batch +from prime_rl.trainer.model import freeze_attention_modules +from prime_rl.trainer.parallel_dims import ParallelDims +from prime_rl.trainer.rl.data import DataLoader +from prime_rl.trainer.rl.value import ServiceState, _score_batch, _server_handler, _train_batch, pack_score_requests +from prime_rl.transports.batch import setup_batch_receiver, setup_batch_sender +from prime_rl.transports.batch.types import TrainingSample + + +def test_score_bins_preserve_sequences_and_replica_steps(): + sequences = [[1, 2, 3], [4] * 8, [5, 6], [7] * 7, [8], [9] * 5] + for dp, cp in ((1, 1), (2, 2), (4, 4)): + ranks = pack_score_requests(sequences, seq_len=16, dp=dp, cp=cp) + assert len(ranks) == dp + assert len({len(bins) for bins in ranks}) == 1 + actual = {} + for bins in ranks: + for score_bin in bins: + assert len(score_bin.token_ids) <= 16 + offset = 0 + for index, length in zip(score_bin.indices, score_bin.lengths, strict=True): + actual[index] = score_bin.token_ids[offset : offset + length] + offset += length + assert actual == dict(enumerate(sequences)) + + +def test_score_http_accepts_batches_and_legacy_sequences(): + state = ServiceState() + server = ThreadingHTTPServer(("127.0.0.1", 0), _server_handler(state, max_seq_len=2)) + server_thread = threading.Thread(target=server.serve_forever) + server_thread.start() + + def serve_one(): + request = state.requests.get(timeout=5) + request.values = [[float(token) for token in ids] for ids in request.batched_token_ids] + request.bootstrap_values = [float(ids[-1]) for ids in request.batched_token_ids] + request.done.set() + + try: + with httpx.Client(base_url=f"http://127.0.0.1:{server.server_port}") as client: + for body, expected in ( + ( + {"batched_token_ids": [[1, 2], [3]]}, + {"values": [[1.0, 2.0], [3.0]], "bootstrap_values": [2.0, 3.0]}, + ), + ({"token_ids": [4, 5]}, {"values": [4.0, 5.0], "bootstrap_value": 5.0}), + ): + worker = threading.Thread(target=serve_one) + worker.start() + response = client.post("/score", json=body) + worker.join(timeout=5) + assert response.status_code == 200 and response.json() == expected + assert client.post("/score", json={"batched_token_ids": [[1], []]}).status_code == 400 + assert client.post("/score", json={"batched_token_ids": [[1, 2, 3]]}).status_code == 400 + finally: + server.shutdown() + server.server_close() + server_thread.join() + + +def test_freeze_attention_follows_modules_instead_of_model_type(): + class MixedBackbone(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = torch.nn.Linear(2, 2) + self.linear_attn = torch.nn.Linear(2, 2) + self.mlp = torch.nn.Linear(2, 2) + + model = MixedBackbone() + freeze_attention_modules(model) + assert not any(param.requires_grad for param in model.self_attn.parameters()) + assert not any(param.requires_grad for param in model.linear_attn.parameters()) + assert all(param.requires_grad for param in model.mlp.parameters()) + + +class DummyValueModel: + def __init__(self): + self.calls = 0 + + def eval(self): + return self + + def __call__(self, *, input_ids, position_ids, seq_lens, **kwargs): + self.calls += 1 + assert input_ids.shape == position_ids.shape + return {"values": input_ids.float() * 0.25 + position_ids.float() * 0.125} + + +class TrainableDummyValueModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(0.01, device="cuda")) + + def forward(self, *, input_ids, position_ids, seq_lens, **kwargs): + return {"values": input_ids.float() * self.weight + position_ids.float() * 0.0} + + +def _expected(ids): + raw = [token * 0.25 + index * 0.125 for index, token in enumerate(ids)] + return [0.0] + raw[:-1], raw[-1] + + +def _free_transport_port(): + with socket.socket() as first, socket.socket() as second: + first.bind(("127.0.0.1", 0)) + port = first.getsockname()[1] + second.bind(("127.0.0.1", port + 1)) + return port + + +def _check_training_traffic(critic_dp, world_size, dims): + rank = dist.get_rank() + policy_dp = 1 if critic_dp > 1 else world_size + samples = [] + for index, length in enumerate((9, 7, 5, 3, 2)): + samples.append( + TrainingSample( + token_ids=list(range(10 * index, 10 * index + length)), + mask=[False] * (length - 1) + [True], + logprobs=[0.0] * length, + temperatures=[1.0] * length, + env_name="test", + advantages=[0.0] * (length - 1) + [1.0], + old_values=[0.25] * length, + value_targets=[0.5] * length, + value_mask=[False] * (length - 1) + [True], + trace_id=str(index), + branch_index=0, + ) + ) + policy_grid = prepare_batch(samples, 16, policy_dp, sum, pad_to_multiple_of=world_size // policy_dp) + critic_grid = prepare_batch(samples, 8, critic_dp, sum, pad_to_multiple_of=world_size // critic_dp, for_value=True) + ports = [_free_transport_port(), _free_transport_port()] if rank == 0 else [0, 0] + dist.broadcast_object_list(ports, src=0, device=torch.device("cuda")) + policy_transport = ZMQTransportConfig(host="127.0.0.1", port=ports[0]) + critic_transport = ZMQTransportConfig(host="127.0.0.1", port=ports[1]) + senders = ( + [ + setup_batch_sender(Path("/tmp"), policy_dp, 0, policy_transport), + setup_batch_sender(Path("/tmp"), critic_dp, 0, critic_transport), + ] + if rank == 0 + else [] + ) + policy_rank = rank // (world_size // policy_dp) + critic_rank = rank // (world_size // critic_dp) + receivers = [ + setup_batch_receiver(Path("/tmp"), policy_rank, 0, policy_transport), + setup_batch_receiver(Path("/tmp"), critic_rank, 0, critic_transport), + ] + for receiver in receivers: + receiver.socket.setsockopt(zmq.RCVTIMEO, 20000) + dist.barrier() + if rank == 0: + + async def send_both(): + await asyncio.wait_for(senders[0].send(policy_grid), 20) + await asyncio.wait_for(senders[1].send(critic_grid), 20) + + asyncio.run(send_both()) + assert receivers[0].receive() == policy_grid[policy_rank] + received_value = receivers[1].receive() + assert received_value == critic_grid[critic_rank] + tensor_batches = [DataLoader._micro_batch_to_tensor(None, batch) for batch in received_value] + model = TrainableDummyValueModel() + optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) + config = SimpleNamespace( + model=SimpleNamespace(cp_style="ring", lora=None, full_offload=None), + optim=SimpleNamespace(max_norm=None), + updates_per_step=1, + ) + loss = _train_batch(model, optimizer, scheduler, None, tensor_batches, dims, config, head_only=False) + assert torch.isfinite(torch.tensor(loss)) + assert scheduler.last_epoch == 1 + dist.barrier() + for receiver in receivers: + receiver.close() + for sender in senders: + sender.close() + + +def _run_distributed(): + local_rank = int(os.environ["LOCAL_RANK"]) + cp = int(os.environ["TEST_VALUE_CP"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + world_size = dist.get_world_size() + dims = ParallelDims(dp_replicate=1, dp_shard=world_size // cp, cp=cp, pp=1, ep=1, world_size=world_size) + _check_training_traffic(world_size // cp, world_size, dims) + batches = [ + [[1, 2, 3], list(range(11, 19)), [21], list(range(31, 39)), [40, 41]], + [[2, 3, 4]], + [list(range(100 + i * 10, 100 + i * 10 + (i % 8) + 1)) for i in range(40)], + ] + for sequences in batches: + model = DummyValueModel() + result = _score_batch(model, sequences, dims, 16, "ring", False) + counts = [None] * world_size + dist.all_gather_object(counts, model.calls) + assert len(set(counts)) == 1 + if dist.get_rank() == 0: + assert result == ( + [_expected(ids)[0] for ids in sequences], + [_expected(ids)[1] for ids in sequences], + ) + else: + assert result is None + dist.destroy_process_group() + + +if __name__ == "__main__": + _run_distributed() diff --git a/tests/unit/train/rl/test_value_real_model_smoke.py b/tests/unit/train/rl/test_value_real_model_smoke.py new file mode 100644 index 00000000000..be21e4b9224 --- /dev/null +++ b/tests/unit/train/rl/test_value_real_model_smoke.py @@ -0,0 +1,84 @@ +"""Small Qwen3 MoE critic: packed scoring must match individual scoring on four GPUs.""" + +import os +from types import SimpleNamespace + +import torch +import torch.distributed as dist +from torch import nn + +from prime_rl.configs.trainer import ModelConfig +from prime_rl.trainer.batch import prepare_batch +from prime_rl.trainer.model import setup_fsdp +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM +from prime_rl.trainer.parallel_dims import ParallelDims +from prime_rl.trainer.rl.data import DataLoader +from prime_rl.trainer.rl.value import _score, _score_batch, _train_batch +from prime_rl.transports.batch.types import TrainingSample +from prime_rl.utils.cp import setup_context_parallel + + +def main(): + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group("nccl") + dims = ParallelDims(dp_replicate=1, dp_shard=2, cp=2, pp=1, ep=1, world_size=4) + torch.manual_seed(2026) + config = Qwen3MoeConfig( + vocab_size=256, + hidden_size=64, + head_dim=16, + intermediate_size=128, + moe_intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=2, + num_experts_per_tok=1, + max_position_embeddings=64, + mlp_only_layers=[], + ) + config._attn_implementation = "flash_attention_2" + with torch.device("cuda"): + model = Qwen3MoeForCausalLM._from_config(config) + model.value_head = nn.Linear(64, 1) + inject_prime_lm_head(model, chunk_size=None) + model_config = ModelConfig(name="Qwen/Qwen3-0.6B", seq_len=16, cp=2, ep=1, attn="flash_attention_2") + setup_fsdp(model, model_config, dims) + setup_context_parallel(model, model_config, dims) + + sequences = [list(range(start, start + length)) for start, length in ((1, 4), (20, 8), (40, 12), (60, 4), (80, 8))] + packed = _score_batch(model, sequences, dims, 16, "ring", False) + individual = [_score(model, ids, dims, "ring", False) for ids in sequences] + if dist.get_rank() == 0: + assert packed is not None + for index, (values, bootstrap) in enumerate(individual): + torch.testing.assert_close(torch.tensor(packed[0][index]), torch.tensor(values), atol=0.02, rtol=0.02) + torch.testing.assert_close(torch.tensor(packed[1][index]), torch.tensor(bootstrap), atol=0.02, rtol=0.02) + + samples = [ + TrainingSample( + token_ids=ids, + mask=[False] * (len(ids) - 1) + [True], + logprobs=[0.0] * len(ids), + temperatures=[1.0] * len(ids), + env_name="smoke", + advantages=[0.0] * (len(ids) - 1) + [1.0], + old_values=[0.0] * len(ids), + value_targets=[0.0] * (len(ids) - 1) + [1.0], + value_mask=[False] * (len(ids) - 1) + [True], + ) + for ids in sequences + ] + grid = prepare_batch(samples, 16, 2, sum, pad_to_multiple_of=2, for_value=True) + micro_batches = [DataLoader._micro_batch_to_tensor(None, batch) for batch in grid[dist.get_rank() // 2]] + optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) + value_config = SimpleNamespace(model=model_config, optim=SimpleNamespace(max_norm=None), updates_per_step=1) + loss = _train_batch(model, optimizer, scheduler, None, micro_batches, dims, value_config, head_only=False) + assert torch.isfinite(torch.tensor(loss)) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() From 8e5db9837526962a497e05d854fc1bf920e1a3c2 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 22:21:25 +0000 Subject: [PATCH 4/9] Avoid shared-node critic rendezvous port collision --- src/prime_rl/entrypoints/rl.py | 13 +++++++++++++ src/prime_rl/templates/multi_node_rl.sbatch.j2 | 5 ++++- tests/unit/test_configs.py | 5 +++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 8b18b0da6fa..d7cf5f7f036 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -54,6 +54,17 @@ ENVS_DIR = "envs" +def value_rendezvous_port(config: RLConfig) -> int: + assert config.value is not None + reserved = {29500, config.value.service_port} + if config.trainer.weight_broadcast.type == "nccl": + reserved.add(config.trainer.weight_broadcast.port) + for transport in (config.trainer.rollout_transport, config.value.rollout_transport): + if transport.type == "zmq": + reserved.update((transport.port, transport.port + 1)) + return next(port for port in range(29502, 29600) if port not in reserved) + + def env_servers(config: RLConfig) -> list[tuple[str, EnvConfig]]: """``(split, source)`` for every launcher-managed train/eval source. The launcher runs one env server per source; each binds an OS-assigned port and publishes it to @@ -585,6 +596,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script trainer_env_vars=trainer_env_vars, value_env_vars=value_env_vars, value_service_port=config.value.service_port if config.value is not None else 0, + value_rendezvous_port=value_rendezvous_port(config) if config.value is not None else 0, orchestrator_env_vars=orchestrator_env_vars, inference_env_vars=inference_env_vars, prefill_vllm_extra_json=vllm_overrides_fragment(infer_deploy.prefill_vllm_overrides), @@ -645,6 +657,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script trainer_env_vars=trainer_env_vars, value_env_vars=value_env_vars, value_service_port=config.value.service_port if config.value is not None else 0, + value_rendezvous_port=value_rendezvous_port(config) if config.value is not None else 0, orchestrator_env_vars=orchestrator_env_vars, inference_env_vars=inference_env_vars, train_env_names=train_env_names, diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index f203b09c8a5..726aff1cf1a 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -34,6 +34,9 @@ export TRAIN_GPUS_PER_NODE={{ train_gpus_per_node }} export VALUE_GPUS_PER_NODE={{ value_gpus_per_node }} export VALUE_GPUS_PER_TRAIN_NODE={{ value_gpus_per_train_node }} export VALUE_GPU_OFFSET={{ value_gpu_offset }} +{% if has_value -%} +export VALUE_MASTER_PORT={{ value_rendezvous_port }} +{% endif -%} export NODES_PER_INFER_REPLICA={{ nodes_per_infer_replica }} export NUM_INFER_REPLICAS={{ num_infer_replicas }} export GPUS_PER_NODE={{ gpus_per_node }} @@ -569,7 +572,7 @@ if [ "$SLURM_PROCID" -ge "$((NUM_INFER_NODES + NUM_TRAIN_NODES))" ]; then --nnodes=$NUM_VALUE_TRAINER_NODES \ --nproc-per-node=$VALUE_GPUS_PER_NODE \ --node-rank=$VALUE_NODE_RANK \ - --rdzv-endpoint=$VALUE_ADDR:$((MASTER_PORT + 1)) \ + --rdzv-endpoint=$VALUE_ADDR:$VALUE_MASTER_PORT \ --rdzv-id=value_$SLURM_JOB_ID \ --log-dir=$LOG_DIR/value/torchrun \ --tee=3 \ diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 5acf5dbfeb6..2ac4cb11a48 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -15,6 +15,7 @@ from prime_rl.configs.sft import SFTConfig from prime_rl.configs.trainer import ModelConfig as TrainerModelConfig from prime_rl.configs.trainer import TrainerConfig +from prime_rl.entrypoints.rl import value_rendezvous_port from prime_rl.utils.config import BaseConfig, cli, dump_resolved_config # All config config classes @@ -559,6 +560,10 @@ def test_multi_node_value_layout_keeps_policy_and_critic_parallelism_independent assert config.deployment.value_trainer_nodes == 2 assert config.orchestrator.num_train_workers == 2 assert config.orchestrator.value_num_train_workers == 2 + assert value_rendezvous_port(config) == 29502 + config.trainer.weight_broadcast.port = 29502 + config.value.service_port = 29503 + assert value_rendezvous_port(config) == 29504 payload["deployment"].pop("num_value_gpus_per_train_node") payload["trainer"]["model"]["cp"] = 16 payload["value"]["model"]["cp"] = 8 From dd06c0f10edf6d90d789f63429e8d993e5d11fa3 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 22:38:24 +0000 Subject: [PATCH 5/9] Discover distributed rendezvous ports on binding nodes --- src/prime_rl/entrypoints/rl.py | 20 ++++---- .../templates/multi_node_rl.sbatch.j2 | 42 +++++++++++++++-- src/prime_rl/utils/ports.py | 47 +++++++++++++++++++ tests/unit/test_configs.py | 21 +++++++-- .../train/rl/test_value_real_model_smoke.py | 19 +++++--- 5 files changed, 126 insertions(+), 23 deletions(-) create mode 100644 src/prime_rl/utils/ports.py diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index d7cf5f7f036..d4a69c02243 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -54,15 +54,19 @@ ENVS_DIR = "envs" -def value_rendezvous_port(config: RLConfig) -> int: - assert config.value is not None - reserved = {29500, config.value.service_port} - if config.trainer.weight_broadcast.type == "nccl": +def reserved_rendezvous_ports(config: RLConfig) -> list[int]: + reserved = set() + if hasattr(config.trainer.weight_broadcast, "port"): reserved.add(config.trainer.weight_broadcast.port) - for transport in (config.trainer.rollout_transport, config.value.rollout_transport): + if config.value is not None: + reserved.add(config.value.service_port) + transports = [config.trainer.rollout_transport] + if config.value is not None: + transports.append(config.value.rollout_transport) + for transport in transports: if transport.type == "zmq": reserved.update((transport.port, transport.port + 1)) - return next(port for port in range(29502, 29600) if port not in reserved) + return sorted(reserved) def env_servers(config: RLConfig) -> list[tuple[str, EnvConfig]]: @@ -596,7 +600,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script trainer_env_vars=trainer_env_vars, value_env_vars=value_env_vars, value_service_port=config.value.service_port if config.value is not None else 0, - value_rendezvous_port=value_rendezvous_port(config) if config.value is not None else 0, + reserved_rendezvous_ports=reserved_rendezvous_ports(config), orchestrator_env_vars=orchestrator_env_vars, inference_env_vars=inference_env_vars, prefill_vllm_extra_json=vllm_overrides_fragment(infer_deploy.prefill_vllm_overrides), @@ -657,7 +661,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, log_dir: Path, script trainer_env_vars=trainer_env_vars, value_env_vars=value_env_vars, value_service_port=config.value.service_port if config.value is not None else 0, - value_rendezvous_port=value_rendezvous_port(config) if config.value is not None else 0, + reserved_rendezvous_ports=reserved_rendezvous_ports(config), orchestrator_env_vars=orchestrator_env_vars, inference_env_vars=inference_env_vars, train_env_names=train_env_names, diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 726aff1cf1a..41cb2fb2062 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -34,9 +34,6 @@ export TRAIN_GPUS_PER_NODE={{ train_gpus_per_node }} export VALUE_GPUS_PER_NODE={{ value_gpus_per_node }} export VALUE_GPUS_PER_TRAIN_NODE={{ value_gpus_per_train_node }} export VALUE_GPU_OFFSET={{ value_gpu_offset }} -{% if has_value -%} -export VALUE_MASTER_PORT={{ value_rendezvous_port }} -{% endif -%} export NODES_PER_INFER_REPLICA={{ nodes_per_infer_replica }} export NUM_INFER_REPLICAS={{ num_infer_replicas }} export GPUS_PER_NODE={{ gpus_per_node }} @@ -159,9 +156,7 @@ echo "VALUE_HOSTS=${VALUE_HOSTS[*]}" {% endif -%} export MASTER_ADDR="${HOSTNAMES[$NUM_INFER_NODES]}" -export MASTER_PORT=29500 echo "MASTER_ADDR=${MASTER_ADDR}" -echo "MASTER_PORT=${MASTER_PORT}" {% if use_nixl_broadcast -%} {% if launch_modelexpress -%} export WEIGHT_BROADCAST_HOST="$MASTER_ADDR" @@ -244,6 +239,43 @@ srun bash -s <<'CLEANUP_SH' echo "[node-cleanup] $(hostname) procs=$procs gpu_mem=${gpu}MiB" CLEANUP_SH +# Discover rendezvous ports on the hosts that will bind them. A shared policy/value +# node gets both ports in one probe, so its two groups cannot choose the same port. +discover_rendezvous_ports() { + local host="$1" count="$2" attempt output + for attempt in 1 2 3 4 5; do + if output=$(srun --nodes=1 --ntasks=1 --nodelist="$host" uv run python -m prime_rl.utils.ports "$count" {{ reserved_rendezvous_ports | join(' ') }}); then + if [[ "$output" =~ ^[0-9]+(\ [0-9]+)*$ ]]; then + set -- $output + if [ "$#" -eq "$count" ]; then + echo "$output" + return 0 + fi + fi + fi + echo "Port discovery attempt $attempt failed on $host; retrying" >&2 + sleep "$attempt" + done + echo "Could not discover $count rendezvous port(s) on $host" >&2 + return 1 +} + +{% if has_value and value_gpus_per_train_node > 0 -%} +RENDEZVOUS_PORTS=$(discover_rendezvous_ports "$MASTER_ADDR" 2) || exit 1 +read -r MASTER_PORT VALUE_MASTER_PORT <<< "$RENDEZVOUS_PORTS" +{% else -%} +MASTER_PORT=$(discover_rendezvous_ports "$MASTER_ADDR" 1) || exit 1 +{% if has_value -%} +VALUE_MASTER_PORT=$(discover_rendezvous_ports "$VALUE_ADDR" 1) || exit 1 +{% endif -%} +{% endif -%} +export MASTER_PORT +echo "MASTER_PORT=$MASTER_PORT" +{% if has_value -%} +export VALUE_MASTER_PORT +echo "VALUE_MASTER_PORT=$VALUE_MASTER_PORT" +{% endif -%} + # Run RL # --kill-on-bad-exit=1 propagates a non-zero exit on any node to all other # tasks so a crashed component terminates the SLURM job instead of leaving the diff --git a/src/prime_rl/utils/ports.py b/src/prime_rl/utils/ports.py new file mode 100644 index 00000000000..54bc563f028 --- /dev/null +++ b/src/prime_rl/utils/ports.py @@ -0,0 +1,47 @@ +"""Find distinct free TCP ports on the node that will run a distributed service.""" + +import argparse +import errno +import random +import socket +from collections.abc import Iterable + + +def find_available_ports( + count: int, excluded: Iterable[int] = (), *, start: int = 29500, stop: int = 30000 +) -> list[int]: + if count < 1 or stop - start < count: + raise ValueError("Port range cannot satisfy the requested count") + reserved = set(excluded) + sockets: list[socket.socket] = [] + try: + for port in random.SystemRandom().sample(range(start, stop), stop - start): + if port in reserved: + continue + candidate = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + candidate.bind(("0.0.0.0", port)) + except OSError as error: + candidate.close() + if error.errno in (errno.EADDRINUSE, errno.EACCES): + continue + raise + sockets.append(candidate) + if len(sockets) == count: + return [sock.getsockname()[1] for sock in sockets] + finally: + for sock in sockets: + sock.close() + raise RuntimeError(f"Could not find {count} free TCP ports in [{start}, {stop})") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("count", type=int) + parser.add_argument("excluded", type=int, nargs="*") + args = parser.parse_args() + print(*find_available_ports(args.count, args.excluded)) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 2ac4cb11a48..9c0fe22b926 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -1,4 +1,5 @@ import os +import socket from pathlib import Path from typing import Annotated, Literal @@ -15,8 +16,9 @@ from prime_rl.configs.sft import SFTConfig from prime_rl.configs.trainer import ModelConfig as TrainerModelConfig from prime_rl.configs.trainer import TrainerConfig -from prime_rl.entrypoints.rl import value_rendezvous_port +from prime_rl.entrypoints.rl import reserved_rendezvous_ports from prime_rl.utils.config import BaseConfig, cli, dump_resolved_config +from prime_rl.utils.ports import find_available_ports # All config config classes CONFIG_CLASSES = [ @@ -30,6 +32,18 @@ ] +def test_rendezvous_ports_are_distinct_and_skip_bound_ports(): + with socket.socket() as occupied: + occupied.bind(("0.0.0.0", 0)) + port = occupied.getsockname()[1] + with pytest.raises(RuntimeError, match="Could not find"): + find_available_ports(1, start=port, stop=port + 1) + + ports = find_available_ports(2, excluded=[29501], start=29500, stop=29505) + assert len(set(ports)) == 2 + assert 29501 not in ports + + def get_config_files() -> list[Path]: """Any TOML file inside `configs/`, `examples/` or `k8s/`.""" config_files = list(Path("configs").rglob("*.toml")) @@ -560,10 +574,11 @@ def test_multi_node_value_layout_keeps_policy_and_critic_parallelism_independent assert config.deployment.value_trainer_nodes == 2 assert config.orchestrator.num_train_workers == 2 assert config.orchestrator.value_num_train_workers == 2 - assert value_rendezvous_port(config) == 29502 + assert config.trainer.weight_broadcast.port in reserved_rendezvous_ports(config) + assert config.value.service_port in reserved_rendezvous_ports(config) config.trainer.weight_broadcast.port = 29502 config.value.service_port = 29503 - assert value_rendezvous_port(config) == 29504 + assert {29502, 29503}.issubset(reserved_rendezvous_ports(config)) payload["deployment"].pop("num_value_gpus_per_train_node") payload["trainer"]["model"]["cp"] = 16 payload["value"]["model"]["cp"] = 8 diff --git a/tests/unit/train/rl/test_value_real_model_smoke.py b/tests/unit/train/rl/test_value_real_model_smoke.py index be21e4b9224..6b7175f79fe 100644 --- a/tests/unit/train/rl/test_value_real_model_smoke.py +++ b/tests/unit/train/rl/test_value_real_model_smoke.py @@ -22,7 +22,10 @@ def main(): torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) dist.init_process_group("nccl") - dims = ParallelDims(dp_replicate=1, dp_shard=2, cp=2, pp=1, ep=1, world_size=4) + cp = int(os.environ.get("TEST_VALUE_CP", "2")) + cp_style = os.environ.get("TEST_VALUE_CP_STYLE", "ring") + dp = 4 // cp + dims = ParallelDims(dp_replicate=1, dp_shard=dp, cp=cp, pp=1, ep=1, world_size=4) torch.manual_seed(2026) config = Qwen3MoeConfig( vocab_size=256, @@ -32,7 +35,7 @@ def main(): moe_intermediate_size=64, num_hidden_layers=1, num_attention_heads=4, - num_key_value_heads=2, + num_key_value_heads=4, num_experts=2, num_experts_per_tok=1, max_position_embeddings=64, @@ -43,13 +46,15 @@ def main(): model = Qwen3MoeForCausalLM._from_config(config) model.value_head = nn.Linear(64, 1) inject_prime_lm_head(model, chunk_size=None) - model_config = ModelConfig(name="Qwen/Qwen3-0.6B", seq_len=16, cp=2, ep=1, attn="flash_attention_2") + model_config = ModelConfig( + name="Qwen/Qwen3-0.6B", seq_len=16, cp=cp, cp_style=cp_style, ep=1, attn="flash_attention_2" + ) setup_fsdp(model, model_config, dims) setup_context_parallel(model, model_config, dims) sequences = [list(range(start, start + length)) for start, length in ((1, 4), (20, 8), (40, 12), (60, 4), (80, 8))] - packed = _score_batch(model, sequences, dims, 16, "ring", False) - individual = [_score(model, ids, dims, "ring", False) for ids in sequences] + packed = _score_batch(model, sequences, dims, 16, cp_style, False) + individual = [_score(model, ids, dims, cp_style, False) for ids in sequences] if dist.get_rank() == 0: assert packed is not None for index, (values, bootstrap) in enumerate(individual): @@ -70,8 +75,8 @@ def main(): ) for ids in sequences ] - grid = prepare_batch(samples, 16, 2, sum, pad_to_multiple_of=2, for_value=True) - micro_batches = [DataLoader._micro_batch_to_tensor(None, batch) for batch in grid[dist.get_rank() // 2]] + grid = prepare_batch(samples, 16, dp, sum, pad_to_multiple_of=cp, for_value=True) + micro_batches = [DataLoader._micro_batch_to_tensor(None, batch) for batch in grid[dist.get_rank() // cp]] optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) value_config = SimpleNamespace(model=model_config, optim=SimpleNamespace(max_norm=None), updates_per_step=1) From a995e3959ff3212a89810fca3d0187d144e2f0a3 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 23:07:35 +0000 Subject: [PATCH 6/9] Wait for inference workers before router discovery --- src/prime_rl/templates/_launch_router.sh.j2 | 25 ++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/templates/_launch_router.sh.j2 b/src/prime_rl/templates/_launch_router.sh.j2 index a32c2600d5c..2d4c2448ef9 100644 --- a/src/prime_rl/templates/_launch_router.sh.j2 +++ b/src/prime_rl/templates/_launch_router.sh.j2 @@ -81,6 +81,29 @@ LLMD_EOF else cmd+=(--worker-urls $worker_args) fi - "${cmd[@]}" >> "$log" 2>&1 & + if [ "$mode" = regular ]; then + ( + read -ra workers <<< "$worker_args" + deadline=$((SECONDS + 4200)) + while true; do + ready=1 + for url in "${workers[@]}"; do + if ! curl -fsS --max-time 5 -o /dev/null "$url/v1/models"; then + ready=0 + break + fi + done + if [ "$ready" -eq 1 ]; then break; fi + if [ "$SECONDS" -ge "$deadline" ]; then + echo "Timed out waiting for all inference workers before starting router" >&2 + exit 1 + fi + sleep 5 + done + exec "${cmd[@]}" + ) >> "$log" 2>&1 & + else + "${cmd[@]}" >> "$log" 2>&1 & + fi {%- endif %} } From ccd5b1032d58e1591a394f195241d09ee48f75c8 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 23:21:40 +0000 Subject: [PATCH 7/9] Allow slower health probes during long inference --- src/prime_rl/templates/_launch_router.sh.j2 | 1 + 1 file changed, 1 insertion(+) diff --git a/src/prime_rl/templates/_launch_router.sh.j2 b/src/prime_rl/templates/_launch_router.sh.j2 index 2d4c2448ef9..5d86885d72a 100644 --- a/src/prime_rl/templates/_launch_router.sh.j2 +++ b/src/prime_rl/templates/_launch_router.sh.j2 @@ -75,6 +75,7 @@ LLMD_EOF --request-id-headers x-session-id --request-timeout-secs {{ router.request_timeout_secs }} --prometheus-port "$((port + 21000))" + --health-check-timeout-secs 30 --worker-startup-timeout-secs 4200 --log-level debug) if [ "$mode" = pd ]; then cmd+=(--vllm-pd-disaggregation $worker_args) From f0c4d0b1c80ec247a21d5bd2739470653bea0ec6 Mon Sep 17 00:00:00 2001 From: Phi Date: Sun, 27 Sep 2026 23:27:47 +0000 Subject: [PATCH 8/9] Probe rendezvous ports on target host address --- .../templates/multi_node_rl.sbatch.j2 | 2 +- src/prime_rl/utils/ports.py | 31 ++++++++++++++----- tests/unit/test_configs.py | 6 ++-- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 41cb2fb2062..639f1bd1c07 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -244,7 +244,7 @@ CLEANUP_SH discover_rendezvous_ports() { local host="$1" count="$2" attempt output for attempt in 1 2 3 4 5; do - if output=$(srun --nodes=1 --ntasks=1 --nodelist="$host" uv run python -m prime_rl.utils.ports "$count" {{ reserved_rendezvous_ports | join(' ') }}); then + if output=$(srun --nodes=1 --ntasks=1 --nodelist="$host" uv run python -m prime_rl.utils.ports "$host" "$count" {{ reserved_rendezvous_ports | join(' ') }}); then if [[ "$output" =~ ^[0-9]+(\ [0-9]+)*$ ]]; then set -- $output if [ "$#" -eq "$count" ]; then diff --git a/src/prime_rl/utils/ports.py b/src/prime_rl/utils/ports.py index 54bc563f028..a7eacdd6e36 100644 --- a/src/prime_rl/utils/ports.py +++ b/src/prime_rl/utils/ports.py @@ -8,27 +8,41 @@ def find_available_ports( - count: int, excluded: Iterable[int] = (), *, start: int = 29500, stop: int = 30000 + host: str, count: int, excluded: Iterable[int] = (), *, start: int = 29500, stop: int = 30000 ) -> list[int]: if count < 1 or stop - start < count: raise ValueError("Port range cannot satisfy the requested count") + address = socket.gethostbyname(host) + if address == "0.0.0.0": + raise ValueError("Port discovery requires a specific host address") + addresses = (address,) if address == "127.0.0.1" else (address, "127.0.0.1") reserved = set(excluded) sockets: list[socket.socket] = [] + ports: list[int] = [] try: for port in random.SystemRandom().sample(range(start, stop), stop - start): if port in reserved: continue - candidate = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + candidates: list[socket.socket] = [] try: - candidate.bind(("0.0.0.0", port)) + for bind_address in addresses: + candidate = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + candidate.bind((bind_address, port)) + except OSError: + candidate.close() + raise + candidates.append(candidate) except OSError as error: - candidate.close() + for candidate in candidates: + candidate.close() if error.errno in (errno.EADDRINUSE, errno.EACCES): continue raise - sockets.append(candidate) - if len(sockets) == count: - return [sock.getsockname()[1] for sock in sockets] + sockets.extend(candidates) + ports.append(port) + if len(ports) == count: + return ports finally: for sock in sockets: sock.close() @@ -37,10 +51,11 @@ def find_available_ports( def main() -> None: parser = argparse.ArgumentParser() + parser.add_argument("host") parser.add_argument("count", type=int) parser.add_argument("excluded", type=int, nargs="*") args = parser.parse_args() - print(*find_available_ports(args.count, args.excluded)) + print(*find_available_ports(args.host, args.count, args.excluded)) if __name__ == "__main__": diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 9c0fe22b926..b2e2ebf14fc 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -34,12 +34,12 @@ def test_rendezvous_ports_are_distinct_and_skip_bound_ports(): with socket.socket() as occupied: - occupied.bind(("0.0.0.0", 0)) + occupied.bind(("127.0.0.1", 0)) port = occupied.getsockname()[1] with pytest.raises(RuntimeError, match="Could not find"): - find_available_ports(1, start=port, stop=port + 1) + find_available_ports("127.0.0.1", 1, start=port, stop=port + 1) - ports = find_available_ports(2, excluded=[29501], start=29500, stop=29505) + ports = find_available_ports("127.0.0.1", 2, excluded=[29501], start=29500, stop=29505) assert len(set(ports)) == 2 assert 29501 not in ports From be2fc8850e294a9c50d1ea7fa3487382146852e8 Mon Sep 17 00:00:00 2001 From: Phi Date: Tue, 29 Sep 2026 07:29:51 +0000 Subject: [PATCH 9/9] fix slowness bugs in ppo --- .../src/prime_rl/configs/shared.py | 4 ++ src/prime_rl/orchestrator/orchestrator.py | 10 +++- src/prime_rl/orchestrator/train_sink.py | 5 ++ .../templates/multi_node_rl.sbatch.j2 | 5 +- src/prime_rl/trainer/optim/offload.py | 2 + src/prime_rl/trainer/rl/train.py | 2 +- src/prime_rl/trainer/rl/value.py | 59 +++++++++++++++++-- src/prime_rl/trainer/utils.py | 18 ++++-- tests/unit/train/rl/test_value_data_flow.py | 2 +- .../train/rl/test_value_real_model_smoke.py | 2 +- 10 files changed, 94 insertions(+), 15 deletions(-) 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 40c7aa2162a..99516a4c7b7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/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, diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index a144bc2dafd..2b733241856 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -649,6 +649,7 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: pack_time = time.perf_counter() - pack_start_time if self.value_sender is not None: assert self.value_packer is not None + value_wait_start = time.perf_counter() value_grid = await asyncio.to_thread(self.value_packer.pack, batch.samples) await self.value_sender.send(value_grid) assert self.value_client is not None @@ -656,11 +657,13 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: while True: response = await self.value_client.get("/status") response.raise_for_status() - if response.json()["completed_step"] >= self.progress.step: + value_status = response.json() + if value_status["completed_step"] >= self.progress.step: break if time.monotonic() >= value_deadline: raise TimeoutError(f"Value trainer did not complete step {self.progress.step}") await asyncio.sleep(0.1) + value_wait_time = time.perf_counter() - value_wait_start await self.sender.send(micro_batch_grid) self.progress.step += 1 self.update_dispatch_gate() @@ -716,6 +719,11 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None: "time/wait_for_policy": self.wait_for_policy_time, "step": step, } + if self.value_sender is not None: + metrics["time/value_score"] = self.train_sink.value_score_seconds + metrics["time/value_update_wait"] = value_wait_time + metrics |= {f"critic/{name}": value for name, value in value_status["metrics"].items()} + self.train_sink.value_score_seconds = 0.0 # Staleness of the shipped cohort, decomposed into its in-flight and # in-queue shares; ``dropped`` counts queued traces the sink voided # since the last ship. diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index c6b273b9092..7fb76f15a0b 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import time from collections import defaultdict from collections.abc import Callable, Iterable @@ -106,6 +107,7 @@ def __init__( self._swept_step = 0 self.zero_output_units = 0 self.reported_zero_output_windows = 0 + self.value_score_seconds = 0.0 def group_size_for(self, env_name: str) -> int: return self.train_envs.get(env_name).config.group_size @@ -280,7 +282,10 @@ async def process_group(self, group_id: str) -> None: trace_samples = [ (trace, await asyncio.to_thread(trace_to_samples, trace, env_name=env_name)) for trace in survivors ] + score_start = time.perf_counter() await env.algorithm.score_samples(trace_samples) + if self.config.value_service_url is not None: + self.value_score_seconds += time.perf_counter() - score_start samples_by_trace: dict[str, list[TrainingSample]] = {} temperature = env.sampling_args["temperature"] for trace, samples in trace_samples: diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 639f1bd1c07..76235a26a01 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -3,6 +3,9 @@ #SBATCH --job-name={{ job_name }} #SBATCH --nodes={{ num_train_nodes + num_infer_nodes + num_value_nodes }} #SBATCH --ntasks-per-node=1 +{%- if cpus_per_task %} +#SBATCH --cpus-per-task={{ cpus_per_task }} +{%- endif %} #SBATCH --gres=gpu:{{ gpus_per_node }} #SBATCH --partition={{ partition }} {%- if account %} @@ -286,7 +289,7 @@ echo "VALUE_MASTER_PORT=$VALUE_MASTER_PORT" # quoted delimiter passes the script through verbatim, so apostrophes and quotes # (including in the shared _launch_*.sh.j2 includes) need no escaping. Nothing in # the body reads stdin, which would otherwise eat the rest of the script. -srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' +srun --cpu-bind=none --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' # pipefail ensures `cmd | tee log` propagates `cmd` failures, and the # wait/cleanup block at the end of this script surfaces background # process crashes (router, orchestrator) as task failures. diff --git a/src/prime_rl/trainer/optim/offload.py b/src/prime_rl/trainer/optim/offload.py index e944b95afdd..20da87a6244 100644 --- a/src/prime_rl/trainer/optim/offload.py +++ b/src/prime_rl/trainer/optim/offload.py @@ -1079,6 +1079,7 @@ def __init__( self.optimizer = optimizer self.offload_config = offload_config self._initialized = False + self.last_step_timings: dict[str, float] = {} self._master_weights = master_weights self._chunks = self._build_chunks() # Reuse the transfer streams across steps: fresh streams each step land @@ -1335,6 +1336,7 @@ def step(self, closure=None): if isinstance(self._gradient_manager, BoundedGradientOffloadManager): timings = self._gradient_manager.consume_timings() timings["drain"] = time.perf_counter() - drain_start + self.last_step_timings = timings get_logger().debug( "Offload pipeline: " + " ".join(f"{key}={value:.3f}s" for key, value in sorted(timings.items())) ) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index bb2495a23b1..5ec62313031 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -499,7 +499,7 @@ def train(config: TrainerConfig): with maybe_record_function("backward"): begin_backward(gradient_manager, final_backward=micro_step == len(micro_batches) - 1) loss.backward() - finish_backward(gradient_manager, wait_for_copies=config.model.full_offload is not None) + finish_backward(gradient_manager) # Add relevant tensors to tensor dict for logging purposes entropy = out["entropy"][loss_mask].detach().to("cpu") diff --git a/src/prime_rl/trainer/rl/value.py b/src/prime_rl/trainer/rl/value.py index cf57ac06ded..9aaa474de43 100644 --- a/src/prime_rl/trainer/rl/value.py +++ b/src/prime_rl/trainer/rl/value.py @@ -23,7 +23,7 @@ from prime_rl.trainer.lora import get_lora_state from prime_rl.trainer.model import get_full_offload_dtype_policy, setup_model from prime_rl.trainer.models.layers.lora import set_lora_num_tokens -from prime_rl.trainer.optim import setup_optimizer +from prime_rl.trainer.optim import FullCPUOffloadOptimizer, setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.rl.data import DataLoader, TensorMicroBatch from prime_rl.trainer.rl.loss import shift_tensor_right @@ -50,6 +50,7 @@ class ScoreRequest: batched_token_ids: list[list[int]] single: bool = False + enqueued_at: float = field(default_factory=time.perf_counter) done: threading.Event = field(default_factory=threading.Event) values: list[list[float]] | None = None bootstrap_values: list[float] | None = None @@ -96,6 +97,7 @@ def pack_score_requests(batched_token_ids: list[list[int]], seq_len: int, dp: in class ServiceState: requests: queue.Queue[ScoreRequest] = field(default_factory=queue.Queue) completed_step: int = 0 + completed_metrics: dict[str, float] = field(default_factory=dict) def _server_handler(state: ServiceState, max_seq_len: int | None = None): @@ -108,7 +110,7 @@ def do_GET(self): if self.path != "/status": self.send_error(404) return - payload = json.dumps({"completed_step": state.completed_step}).encode() + payload = json.dumps({"completed_step": state.completed_step, "metrics": state.completed_metrics}).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(payload))) @@ -278,7 +280,12 @@ def _train_batch( denominator = total_count.clamp_min(1) last_loss = 0.0 + forward_backward_seconds = 0.0 + optimizer_seconds = 0.0 + offload_materialize_seconds = 0.0 + offload_optimizer_kernel_seconds = 0.0 for _ in range(updates if updates is not None else config.updates_per_step): + update_start = time.perf_counter() update_loss = 0.0 prepare_gradient_offload(gradient_manager, parallel_dims.fsdp_gradient_divide_factor, overlap_optimizer=True) for micro_step, batch in enumerate(micro_batches): @@ -303,17 +310,28 @@ def _train_batch( loss = (((values.float() - targets.float()) ** 2) * mask).sum() / denominator begin_backward(gradient_manager, final_backward=micro_step == len(micro_batches) - 1) loss.backward() - finish_backward(gradient_manager, wait_for_copies=config.model.full_offload is not None) + finish_backward(gradient_manager) update_loss += loss.detach().item() + forward_backward_seconds += time.perf_counter() - update_start + optimizer_start = time.perf_counter() if gradient_manager is None: scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor) if config.optim.max_norm is not None: clip_grad_norm_(gradient_manager, model, config.optim.max_norm, parallel_dims.ep_enabled) optimizer.step() + if isinstance(optimizer, FullCPUOffloadOptimizer): + offload_materialize_seconds += optimizer.last_step_timings.get("materialize", 0.0) + offload_optimizer_kernel_seconds += optimizer.last_step_timings.get("optimizer_kernel", 0.0) optimizer.zero_grad() scheduler.step() + optimizer_seconds += time.perf_counter() - optimizer_start last_loss = update_loss - return last_loss + return last_loss, { + "forward_backward_seconds": forward_backward_seconds, + "optimizer_seconds": optimizer_seconds, + "offload_materialize_seconds": offload_materialize_seconds, + "offload_optimizer_kernel_seconds": offload_optimizer_kernel_seconds, + } def _pretrain(model, optimizer, scheduler, gradient_manager, parallel_dims, config): @@ -444,6 +462,10 @@ def train(config: ValueConfig): final_step = None last_batch = None last_sync_step = progress.step - 1 + score_requests = 0 + score_tokens = 0 + score_queue_seconds = 0.0 + score_compute_seconds = 0.0 while True: request = None if world.is_master: @@ -489,6 +511,11 @@ def train(config: ValueConfig): logger.info(f"Loaded policy backbone at step {payload} and retuned critic LoRA") continue if action == "score": + if request is not None: + score_requests += 1 + score_tokens += sum(map(len, request.batched_token_ids)) + score_queue_seconds += time.perf_counter() - request.enqueued_at + score_start = time.perf_counter() try: result = _score_batch( model, @@ -506,12 +533,13 @@ def train(config: ValueConfig): raise finally: if request is not None: + score_compute_seconds += time.perf_counter() - score_start request.done.set() continue micro_batches: list[TensorMicroBatch] = dataloader.get_batch() t0 = time.perf_counter() - loss = _train_batch( + loss, train_metrics = _train_batch( model, optimizer, scheduler, @@ -523,7 +551,26 @@ def train(config: ValueConfig): ) last_batch = micro_batches if world.is_master: - logger.info(f"Value step {progress.step} | loss={loss:.5f} | time={time.perf_counter() - t0:.1f}s") + train_seconds = time.perf_counter() - t0 + state.completed_metrics = { + "train_seconds": train_seconds, + "score_queue_seconds": score_queue_seconds, + "score_compute_seconds": score_compute_seconds, + **train_metrics, + } + logger.info( + f"Value step {progress.step} | loss={loss:.5f} | time={train_seconds:.1f}s " + f"forward_backward={train_metrics['forward_backward_seconds']:.1f}s " + f"optimizer={train_metrics['optimizer_seconds']:.1f}s " + f"offload_materialize={train_metrics['offload_materialize_seconds']:.1f}s " + f"offload_kernel={train_metrics['offload_optimizer_kernel_seconds']:.1f}s | " + f"score_requests={score_requests} score_tokens={score_tokens} " + f"score_queue={score_queue_seconds:.1f}s score_compute={score_compute_seconds:.1f}s" + ) + score_requests = 0 + score_tokens = 0 + score_queue_seconds = 0.0 + score_compute_seconds = 0.0 is_last_step = config.max_steps is not None and progress.step >= config.max_steps if ( config.ckpt is not None diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index a9cf91c2abb..019f1935ac2 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -129,9 +129,16 @@ def bind_process_to_gpu_numa_node() -> None: logger = get_logger() device_id = torch.cuda.current_device() + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + visible_device = visible_devices.split(",")[device_id].strip() if visible_devices else str(device_id) pynvml.nvmlInit() try: - bus_id = pynvml.nvmlDeviceGetPciInfo(pynvml.nvmlDeviceGetHandleByIndex(device_id)).busId + handle = ( + pynvml.nvmlDeviceGetHandleByIndex(int(visible_device)) + if visible_device.isdecimal() + else pynvml.nvmlDeviceGetHandleByUUID(visible_device) + ) + bus_id = pynvml.nvmlDeviceGetPciInfo(handle).busId finally: pynvml.nvmlShutdown() if isinstance(bus_id, bytes): @@ -149,13 +156,16 @@ def bind_process_to_gpu_numa_node() -> None: cpus.update(range(int(start), int(end) + 1)) else: cpus.add(int(part)) - os.sched_setaffinity(0, cpus) - logger.info(f"Bound rank with GPU {device_id} to NUMA node {numa_node} ({len(cpus)} CPUs)") + allowed_cpus = cpus & os.sched_getaffinity(0) + if not allowed_cpus: + raise RuntimeError(f"No allocated CPUs on GPU {visible_device}'s NUMA node {numa_node}") + os.sched_setaffinity(0, allowed_cpus) + logger.info(f"Bound rank with GPU {visible_device} to NUMA node {numa_node} ({len(allowed_cpus)} CPUs)") def configure_cpu_optimizer_threads() -> None: available = os.sched_getaffinity(0) - fair_share = (os.cpu_count() or len(available)) // get_world().local_world_size + fair_share = len(available) // get_world().local_world_size threads = max(1, min(len(available), fair_share)) torch.set_num_threads(threads) get_logger().info( diff --git a/tests/unit/train/rl/test_value_data_flow.py b/tests/unit/train/rl/test_value_data_flow.py index 96f0ad0d05b..b8f87d2deaa 100644 --- a/tests/unit/train/rl/test_value_data_flow.py +++ b/tests/unit/train/rl/test_value_data_flow.py @@ -186,7 +186,7 @@ async def send_both(): optim=SimpleNamespace(max_norm=None), updates_per_step=1, ) - loss = _train_batch(model, optimizer, scheduler, None, tensor_batches, dims, config, head_only=False) + loss, _ = _train_batch(model, optimizer, scheduler, None, tensor_batches, dims, config, head_only=False) assert torch.isfinite(torch.tensor(loss)) assert scheduler.last_epoch == 1 dist.barrier() diff --git a/tests/unit/train/rl/test_value_real_model_smoke.py b/tests/unit/train/rl/test_value_real_model_smoke.py index 6b7175f79fe..723ef5d0933 100644 --- a/tests/unit/train/rl/test_value_real_model_smoke.py +++ b/tests/unit/train/rl/test_value_real_model_smoke.py @@ -80,7 +80,7 @@ def main(): optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda _: 1.0) value_config = SimpleNamespace(model=model_config, optim=SimpleNamespace(max_norm=None), updates_per_step=1) - loss = _train_batch(model, optimizer, scheduler, None, micro_batches, dims, value_config, head_only=False) + loss, _ = _train_batch(model, optimizer, scheduler, None, micro_batches, dims, value_config, head_only=False) assert torch.isfinite(torch.tensor(loss)) dist.destroy_process_group()