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 0f57abcced2..a45a60a08dc 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -631,6 +631,17 @@ class OrchestratorConfig(BaseConfig): rollout_transport: TransportConfig = ZMQTransportConfig() """Transport used to ship rollouts from orchestrator to trainer.""" + value_rollout_transport: TransportConfig | None = None + """Transport used to ship PPO batches to the separate value trainer.""" + value_service_url: str | None = None + """HTTP endpoint for PPO value scoring and update status.""" + value_seq_len: int | None = Field(None, ge=1) + """Critic sequence length for independent value batch packing.""" + value_num_train_workers: int | None = Field(None, ge=1) + """Critic DP rank count; one ZMQ stream is sent per rank.""" + value_pad_to_multiple_of: int | None = Field(None, ge=1) + """Critic CP degree for padding value batches.""" + output_dir: Path = Field(default_factory=default_output_dir) """Directory to write outputs to — checkpoints, weights, rollouts, and logs are written as subdirectories. Shared with the trainer; should be a persistent directory with enough disk space and unique per experiment running on a single node. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" @@ -697,6 +708,14 @@ def inherit_env_algorithms(self): env_cfg.algo = self.algo.model_copy(deep=True) return self + @model_validator(mode="after") + def sync_value_service_url(self): + if self.value_service_url is not None: + for source in self.train.source: + if source.algo is not None and source.algo.type in ("ppo", "sao"): + source.algo.value_url = self.value_service_url + return self + @model_validator(mode="after") def validate_env_algorithms(self): """Let each algorithm reject environments it cannot score correctly.""" 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..9af6c7fa7a6 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, + PPOLossConfig, + SAOLossConfig, + TokenizerConfig, + TrainerConfig, ) from prime_rl.configs.trainer import ( - NIXLWeightBroadcastConfig as TrainerNIXLWeightBroadcastConfig, + NCCLWeightBroadcastConfig as TrainerNCCLWeightBroadcastConfig, ) from prime_rl.configs.trainer import ( - TokenizerConfig, - TrainerConfig, + NIXLWeightBroadcastConfig as TrainerNIXLWeightBroadcastConfig, ) +from prime_rl.configs.value import ValueConfig from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource from prime_rl.utils.validation import ( propagate_shared_fields, @@ -186,12 +190,16 @@ class SingleNodeDeploymentConfig(BaseDeploymentConfig): num_infer_gpus: int = 1 """GPUs allocated to inference.""" + num_value_gpus: int = Field(0, ge=0) + """GPUs allocated to a separate PPO value trainer.""" + @model_validator(mode="after") def validate_gpu_count(self): - total = self.num_train_gpus + self.num_infer_gpus + total = self.num_train_gpus + self.num_infer_gpus + self.num_value_gpus if total > self.gpus_per_node: raise ValueError( - f"Total GPU count ({total} = {self.num_train_gpus} train + {self.num_infer_gpus} infer)" + f"Total GPU count ({total} = {self.num_train_gpus} train + {self.num_infer_gpus} infer + " + f"{self.num_value_gpus} value)" f" exceeds gpus_per_node ({self.gpus_per_node})." ) return self @@ -203,6 +211,12 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): num_train_nodes: int """Training nodes.""" + num_value_nodes: int = Field(0, ge=0) + """Dedicated value-trainer nodes for PPO/SAO.""" + + num_value_gpus_per_train_node: int = Field(0, ge=0) + """GPUs reserved for the critic on each training node instead of dedicated value nodes.""" + num_infer_nodes: int | None = Field(None, ge=0) """Inference nodes per replica. If unset, inferred from ``inference.deployment``. Set to 0 to skip inference and orchestrator (requires fake data).""" @@ -215,6 +229,26 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): orchestrator_on_inference: bool = False """Run the orchestrator on the last inference node instead of trainer rank 0 (frees host RAM on the trainer node).""" + @model_validator(mode="after") + def validate_value_gpu_layout(self): + if self.num_value_gpus_per_train_node >= self.gpus_per_node: + raise ValueError("num_value_gpus_per_train_node must leave at least one policy GPU per training node") + if self.num_value_nodes and self.num_value_gpus_per_train_node: + raise ValueError("Use dedicated value nodes or shared training-node GPUs, not both") + return self + + @property + def train_gpus_per_node(self) -> int: + return self.gpus_per_node - self.num_value_gpus_per_train_node + + @property + def value_gpus_per_node(self) -> int: + return self.num_value_gpus_per_train_node or self.gpus_per_node + + @property + def value_trainer_nodes(self) -> int: + return self.num_train_nodes if self.num_value_gpus_per_train_node else self.num_value_nodes + @property def infer_nodes_per_replica(self) -> int: return self.num_infer_nodes or 0 @@ -232,6 +266,9 @@ def total_infer_nodes(self) -> int: class RLConfig(BaseConfig): trainer: TrainerConfig + value: ValueConfig | None = None + """Separate PPO/SAO value trainer.""" + orchestrator: OrchestratorConfig inference: InferenceConfig | None = None @@ -337,6 +374,20 @@ def auto_setup_infer_nodes(self): @model_validator(mode="after") def validate_deployment(self): + if self.value is not None: + if self.deployment.type == "single_node" and self.deployment.num_value_gpus == 0: + raise ValueError("A value trainer requires num_value_gpus > 0") + if self.deployment.type == "multi_node": + if self.deployment.value_trainer_nodes == 0: + raise ValueError("A multi-node value trainer requires value nodes or GPUs on training nodes") + if self.deployment.infer_nodes_per_replica == 0: + raise ValueError("A multi-node value trainer requires inference nodes for the orchestrator") + elif self.deployment.type == "single_node" and self.deployment.num_value_gpus: + raise ValueError("num_value_gpus requires a [value] configuration") + elif self.deployment.type == "multi_node" and ( + self.deployment.num_value_nodes or self.deployment.num_value_gpus_per_train_node + ): + raise ValueError("Value GPUs require a [value] configuration") if self.deployment.type == "multi_node": if self.slurm is None: raise ValueError("Must use SLURM for multi-node deployment.") @@ -393,6 +444,9 @@ def auto_setup_run_dir(self): "output_dir / run.name — set those instead." ) sub.output_dir = run_dir + if self.value is not None: + self.value.output_dir = run_dir / "value" + self.value.rollout_dir = run_dir return self @model_validator(mode="after") @@ -402,6 +456,8 @@ def auto_setup_resume(self): return self self.trainer.resume = self.resume.model_copy() self.orchestrator.resume = self.resume.model_copy() + if self.value is not None: + self.value.resume = self.resume.model_copy() return self @model_validator(mode="after") @@ -516,6 +572,84 @@ def auto_setup_rollout_transport(self): self.rollout_transport = self.trainer.rollout_transport return self + @model_validator(mode="after") + def auto_setup_value(self): + if self.value is None: + if any( + source.algo is not None and source.algo.type in ("ppo", "sao") + for source in self.orchestrator.train.source + ): + raise ValueError("PPO/SAO requires a [value] configuration") + return self + algo_types = {source.algo.type for source in self.orchestrator.train.source if source.algo is not None} + if len(algo_types) != 1 or not algo_types <= {"ppo", "sao"}: + raise ValueError("A value trainer requires one PPO or SAO algorithm across all training sources") + algorithm_type = next(iter(algo_types)) + if algorithm_type == "sao" and "freeze_attention" not in self.value.model_fields_set: + self.value.freeze_attention = True + if "loss" not in self.trainer.model_fields_set: + self.trainer.loss = PPOLossConfig() if algorithm_type == "ppo" else SAOLossConfig() + elif self.trainer.loss.type != algorithm_type: + raise ValueError("trainer.loss.type must match the PPO or SAO algorithm") + model_overrides = self.value.model.model_dump(exclude_unset=True) + self.value.model = ModelConfig.model_validate({**self.trainer.model.model_dump(), **model_overrides}) + if self.value.model.cp > 1 and self.value.model.seq_len % (2 * self.value.model.cp): + raise ValueError("value.model.seq_len must be divisible by 2 * value.model.cp") + if self.value.model.name != self.trainer.model.name: + raise ValueError("value.model.name must match trainer.model.name") + if self.deployment.type == "single_node": + num_train_gpus = self.deployment.num_train_gpus + num_value_gpus = self.deployment.num_value_gpus + else: + num_train_gpus = self.deployment.num_train_nodes * self.deployment.train_gpus_per_node + num_value_gpus = self.deployment.value_trainer_nodes * self.deployment.value_gpus_per_node + if num_train_gpus % self.trainer.model.cp: + raise ValueError("Policy GPU count must be divisible by trainer.model.cp") + if num_value_gpus % self.value.model.cp: + raise ValueError("Value GPU count must be divisible by value.model.cp") + if ( + self.value.model.lora is not None + and "experts" in self.value.model.lora.target_modules + and self.value.model.lora.rank % 8 + ): + raise ValueError("Expert LoRA requires a rank divisible by 8") + if self.value.policy_sync_interval is not None: + if self.value.model.lora is None: + raise ValueError("value.policy_sync_interval requires value.model.lora") + if self.trainer.model.lora is not None: + raise ValueError("Policy backbone sync requires a full-weight policy trainer") + if self.value.model.cp != self.trainer.model.cp or self.value.model.ep != self.trainer.model.ep: + raise ValueError("Policy and value models must use the same CP and EP for backbone sync") + if num_train_gpus != num_value_gpus: + raise ValueError("Policy and value models must use the same GPU count for backbone sync") + self.value.policy_sync_dir = self.run_dir / "policy_sync" + self.trainer.ppo_policy_sync_dir = self.value.policy_sync_dir + self.trainer.ppo_policy_sync_interval = self.value.policy_sync_interval + if self.trainer.rollout_transport.type != "zmq": + raise ValueError("The separate value trainer currently requires ZMQ rollout transport") + self.value.rollout_transport = self.trainer.rollout_transport.model_copy(deep=True) + self.value.rollout_transport.port += 2 + self.orchestrator.value_rollout_transport = self.value.rollout_transport.model_copy(deep=True) + self.orchestrator.value_service_url = f"http://{self.value.service_host}:{self.value.service_port}" + self.orchestrator.value_seq_len = self.value.model.seq_len + self.orchestrator.value_num_train_workers = num_value_gpus // self.value.model.cp + self.orchestrator.value_pad_to_multiple_of = self.value.model.cp + self.value.max_steps = self.trainer.max_steps + self.value.ckpt = self.value.ckpt or self.trainer.ckpt + if self.value.resume is not None and self.value.resume.dir is not None: + source_run_dir = self.value.resume.dir.parent.parent + self.value.resume.dir = source_run_dir / "value" / "checkpoints" / self.value.resume.dir.name + self.trainer.ppo_head_warmup_steps = max(0, self.value.head_warmup_steps - self.value.pretrain_steps) + for source in self.orchestrator.train.source: + algo = source.algo + if algo is not None and "value_url" not in algo.model_fields_set: + algo.value_url = f"http://{self.value.service_host}:{self.value.service_port}" + if algo is not None and algo.value_seq_len is None: + algo.value_seq_len = min(self.value.model.seq_len, self.orchestrator.seq_len) + if algo is not None and algo.value_seq_len > self.value.model.seq_len: + raise ValueError("algo.value_seq_len cannot exceed value.model.seq_len") + return self + @model_validator(mode="after") def validate_eplb(self): if self.inference is not None and self.inference.vllm.enable_eplb: @@ -710,7 +844,7 @@ def auto_setup_deployment(self): elif self.deployment.type == "multi_node": # multi-node self.orchestrator.num_train_workers = ( - self.deployment.num_train_nodes * self.deployment.gpus_per_node // self.trainer.model.cp + self.deployment.num_train_nodes * self.deployment.train_gpus_per_node // self.trainer.model.cp ) if self.deployment.nodes_per_fsdp_group is not None: 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/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 70100bee117..a53fd55ed9c 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: @@ -638,6 +641,12 @@ class CISPOLossConfig(BaseConfig): """Temperature for the advantage term.""" +class SAOLossConfig(IcePopLossConfig): + type: Literal["sao"] = "sao" # type: ignore[assignment] + ratio_low: float = Field(0.7, gt=0) + ratio_high: float = Field(6.0, gt=0) + + class CustomLossConfig(BaseConfig): type: Literal["custom"] = "custom" @@ -649,7 +658,8 @@ class CustomLossConfig(BaseConfig): LossConfig: TypeAlias = Annotated[ - IPOLossConfig | IcePopLossConfig | PPOLossConfig | CISPOLossConfig | CustomLossConfig, Field(discriminator="type") + IPOLossConfig | IcePopLossConfig | PPOLossConfig | CISPOLossConfig | SAOLossConfig | CustomLossConfig, + Field(discriminator="type"), ] @@ -718,6 +728,11 @@ class TrainerConfig(BaseConfig): loss: LossConfig = IPOLossConfig() """Loss config for the rl loss component (see ``setup_rl_loss_fn``). The ce / ref_kl components are fixed and do not read this.""" + ppo_head_warmup_steps: int = Field(0, ge=0) + """Live critic head warmup steps before policy-backbone snapshots begin.""" + ppo_policy_sync_interval: int | None = Field(None, ge=1) + ppo_policy_sync_dir: Path | None = None + optim: OptimizerConfig = AdamWConfig() scheduler: SchedulerConfig = ConstantSchedulerConfig() 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..b4c341da9c6 --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/value.py @@ -0,0 +1,47 @@ +"""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") + 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 2c9832ddedb..d4a69c02243 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -41,17 +41,34 @@ 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" ENVS_DIR = "envs" +def reserved_rendezvous_ports(config: RLConfig) -> list[int]: + reserved = set() + if hasattr(config.trainer.weight_broadcast, "port"): + reserved.add(config.trainer.weight_broadcast.port) + 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 sorted(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 @@ -81,6 +98,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 +112,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 +170,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 +185,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 +267,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 +390,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 +514,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 +569,15 @@ 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, + 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, @@ -521,6 +598,9 @@ 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, + 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), @@ -546,6 +626,15 @@ 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, + 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, @@ -570,6 +659,9 @@ 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, + 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/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..0b39c6a4f36 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_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.""" 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..c2c85abf1b9 --- /dev/null +++ b/src/prime_rl/orchestrator/algo/ppo.py @@ -0,0 +1,122 @@ +"""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_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() + self._apply_score(trace, branch, sample, score["values"], score["bootstrap_value"], scored_length) + + 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 9e2cdb98dae..cfa13956532 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 @@ -260,10 +261,23 @@ 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, config.value_num_train_workers, 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 +647,23 @@ 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: + 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 + value_deadline = time.monotonic() + 3600 + while True: + response = await self.value_client.get("/status") + response.raise_for_status() + 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() @@ -688,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. @@ -1013,6 +1049,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() @@ -1040,6 +1080,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/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 50ec66ce30b..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 @@ -33,6 +34,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 @@ -104,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 @@ -275,10 +279,16 @@ 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 + ] + 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 in survivors: - samples = await asyncio.to_thread(trace_to_samples, trace, env_name=env_name) + 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/_launch_router.sh.j2 b/src/prime_rl/templates/_launch_router.sh.j2 index a32c2600d5c..5d86885d72a 100644 --- a/src/prime_rl/templates/_launch_router.sh.j2 +++ b/src/prime_rl/templates/_launch_router.sh.j2 @@ -75,12 +75,36 @@ 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) 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 %} } diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index b1c225cd995..76235a26a01 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -1,8 +1,11 @@ #!/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 +{%- if cpus_per_task %} +#SBATCH --cpus-per-task={{ cpus_per_task }} +{%- endif %} #SBATCH --gres=gpu:{{ gpus_per_node }} #SBATCH --partition={{ partition }} {%- if account %} @@ -28,6 +31,12 @@ 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 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 }} @@ -62,9 +71,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 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 has_value -%} +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 +95,16 @@ 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 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 -%} export HOSTNAMES_STR="${HOSTNAMES[*]}" {% if num_infer_nodes > 0 -%} @@ -133,11 +154,12 @@ echo "ADMIN_URLS=${ADMIN_URLS}" echo "HOSTNAMES=${HOSTNAMES[*]}" echo "TRAIN_HOSTS=${TRAIN_HOSTS[*]}" +{% if has_value -%} +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" @@ -220,6 +242,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 "$host" "$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 @@ -230,7 +289,7 @@ CLEANUP_SH # 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. @@ -493,6 +552,9 @@ else {%- else %} TRAIN_NODE_RANK=$SLURM_PROCID {%- endif %} +{% if has_value %} +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() %} @@ -502,10 +564,10 @@ else 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 \ @@ -524,7 +586,46 @@ else fi exit "$COMPONENT_CODE" ) & +{% if has_value %}fi{% endif %} {% if num_infer_nodes > 0 %} fi{% endif %} +{% 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 + + ( + 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_TRAINER_NODES \ + --nproc-per-node=$VALUE_GPUS_PER_NODE \ + --node-rank=$VALUE_NODE_RANK \ + --rdzv-endpoint=$VALUE_ADDR:$VALUE_MASTER_PORT \ + --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 +664,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 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 925e65a507d..ad45f588a8a 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -360,6 +360,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_refs = training_example.mm_refs @@ -397,6 +402,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: @@ -421,6 +430,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)}" @@ -457,6 +469,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], @@ -522,6 +537,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} @@ -537,6 +553,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_refs: MMRefs | None = None streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL} @@ -557,6 +576,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: @@ -614,6 +637,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, @@ -716,6 +742,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 @@ -760,6 +790,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) @@ -791,7 +824,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) @@ -801,6 +834,9 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch: dummy.rl_weights = None dummy.ce_weights = None dummy.ref_kl_weights = 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. @@ -809,11 +845,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 @@ -824,6 +862,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. @@ -844,8 +883,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 b75a6fc4ee3..db186cac752 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 @@ -700,29 +702,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() @@ -737,6 +747,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( @@ -970,7 +982,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(): @@ -1014,6 +1031,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) @@ -1035,6 +1057,14 @@ 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: + freeze_attention_modules(model) + for param in model.value_head.parameters(): + param.requires_grad = True + if frozen_vision_encoder is not None: freeze_vision_encoder( model, @@ -1076,9 +1106,26 @@ 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 +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/models/layers/lm_head.py b/src/prime_rl/trainer/models/layers/lm_head.py index 46820fd674d..bc0223ae2a5 100644 --- a/src/prime_rl/trainer/models/layers/lm_head.py +++ b/src/prime_rl/trainer/models/layers/lm_head.py @@ -22,6 +22,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: @@ -35,6 +36,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")), ) @@ -420,6 +422,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 @@ -433,6 +437,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..20da87a6244 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__( @@ -1080,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 @@ -1113,7 +1113,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, @@ -1336,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/data.py b/src/prime_rl/trainer/rl/data.py index e2b360dc617..fcd3dc9849d 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -57,6 +57,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: @@ -137,6 +140,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: @@ -169,6 +175,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, } @@ -252,6 +261,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 b1f59b836bb..dd083be6633 100644 --- a/src/prime_rl/trainer/rl/loss.py +++ b/src/prime_rl/trainer/rl/loss.py @@ -13,6 +13,7 @@ IPOLossConfig, LossConfig, PPOLossConfig, + SAOLossConfig, ) from prime_rl.trainer.models.layers.lm_head import sampling_replay_mask from prime_rl.utils.utils import import_object @@ -204,7 +205,9 @@ def __init__(self, config: IcePopLossConfig): def loss(self, inputs: LossInputs) -> LossOutputs: loss_config = self.config - log_importance_ratio = inputs.trainer_logprobs[inputs.loss_mask] - inputs.inference_logprobs[inputs.loss_mask] + log_importance_ratio = ( + inputs.trainer_logprobs[inputs.loss_mask].float() - inputs.inference_logprobs[inputs.loss_mask].float() + ) advantages = inputs.advantages[inputs.loss_mask] weights = inputs.loss_weights[inputs.loss_mask] if inputs.loss_weights is not None else None @@ -220,7 +223,6 @@ def loss(self, inputs: LossInputs) -> LossOutputs: per_token_loss = per_token_loss * weights[keep_mask] mismatch_kl = _mismatch_kl_from_log_ratio(log_importance_ratio) - metrics = { "masked_mismatch_kl": _safe_mean(mismatch_kl, is_masked), "unmasked_mismatch_kl": _safe_mean(mismatch_kl, keep_mask), @@ -237,8 +239,8 @@ def __init__(self, config: PPOLossConfig): def loss(self, inputs: LossInputs) -> LossOutputs: config = self.config - logprobs = inputs.trainer_logprobs[inputs.loss_mask] - log_ratio = logprobs - inputs.inference_logprobs[inputs.loss_mask] + logprobs = inputs.trainer_logprobs[inputs.loss_mask].float() + log_ratio = logprobs - inputs.inference_logprobs[inputs.loss_mask].float() advantages = config.adv_tau * inputs.advantages[inputs.loss_mask] ratio = _capped_importance_ratio(log_ratio, config.max_importance_ratio) @@ -250,10 +252,14 @@ def loss(self, inputs: LossInputs) -> LossOutputs: clipped = ((advantages > 0) & (log_ratio.detach() > log_ratio.new_tensor(config.ratio_high).log())) | ( (advantages < 0) & (log_ratio.detach() < log_ratio.new_tensor(config.ratio_low).log()) ) + clip_fraction = clipped.sum() / max(clipped.numel(), 1) + approx_kl = _mismatch_kl_from_log_ratio(log_ratio.detach()) metrics = { - "is_clipped": clipped.sum() / max(clipped.numel(), 1), + "is_clipped": clip_fraction, "ratio_capped": (log_ratio.detach() > log_ratio.new_tensor(config.max_importance_ratio).log()).sum() / max(log_ratio.numel(), 1), + "ppo/clip_fraction": clip_fraction, + "ppo/approx_kl": approx_kl.sum() / max(approx_kl.numel(), 1), } return LossOutputs(loss=per_token_loss.sum(), metrics=metrics) @@ -371,6 +377,8 @@ def setup_rl_loss_fn(loss_config: LossConfig) -> Loss: return CustomLoss(loss_config) case IPOLossConfig(): return IPOLoss(loss_config) + case SAOLossConfig(): + return IcePopLoss(loss_config) case IcePopLossConfig(): return IcePopLoss(loss_config) case PPOLossConfig(): 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 990004e2fae..5674aeb3468 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, @@ -612,6 +613,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..9aaa474de43 --- /dev/null +++ b/src/prime_rl/trainer/rl/value.py @@ -0,0 +1,607 @@ +"""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 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 +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: + 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 + 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 + completed_metrics: dict[str, float] = field(default_factory=dict) + + +def _server_handler(state: ServiceState, max_seq_len: int | None = None): + 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, "metrics": state.completed_metrics}).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"])) + 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(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, "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(batched_token_ids, single=single) + 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 + 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))) + 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 _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 +): + 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 + 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): + 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) + 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, { + "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): + 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, 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, 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}") + + try: + 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: + 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.batched_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": + 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, + payload, + parallel_dims, + config.model.seq_len, + config.model.cp_style, + config.model.lora is not None, + ) + if request is not None: + request.values, request.bootstrap_values = result + except BaseException as error: + if request is not None: + request.error = error + 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_metrics = _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: + 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 + 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/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/src/prime_rl/transports/batch/types.py b/src/prime_rl/transports/batch/types.py index 4fe5225dd6c..58afe23c6cb 100644 --- a/src/prime_rl/transports/batch/types.py +++ b/src/prime_rl/transports/batch/types.py @@ -79,6 +79,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): @@ -115,3 +121,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/src/prime_rl/utils/ports.py b/src/prime_rl/utils/ports.py new file mode 100644 index 00000000000..a7eacdd6e36 --- /dev/null +++ b/src/prime_rl/utils/ports.py @@ -0,0 +1,62 @@ +"""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( + 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 + candidates: list[socket.socket] = [] + try: + 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: + for candidate in candidates: + candidate.close() + if error.errno in (errno.EADDRINUSE, errno.EACCES): + continue + raise + sockets.extend(candidates) + ports.append(port) + if len(ports) == count: + return ports + 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("host") + parser.add_argument("count", type=int) + parser.add_argument("excluded", type=int, nargs="*") + args = parser.parse_args() + print(*find_available_ports(args.host, args.count, args.excluded)) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index f8aa3a5c75f..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,9 +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 @@ -165,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. # -------------------------------------------------------------------------- @@ -181,6 +241,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 4622b1146f6..ae35e236168 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -51,6 +51,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] @@ -310,6 +330,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 81e0b65e364..0d9a2d2d36b 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,7 +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 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 = [ @@ -29,6 +32,18 @@ ] +def test_rendezvous_ports_are_distinct_and_skip_bound_ports(): + with socket.socket() as occupied: + occupied.bind(("127.0.0.1", 0)) + port = occupied.getsockname()[1] + with pytest.raises(RuntimeError, match="Could not find"): + find_available_ports("127.0.0.1", 1, start=port, stop=port + 1) + + 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 + + def get_config_files() -> list[Path]: """Any TOML file inside `configs/`, `examples/` or `k8s/`.""" config_files = list(Path("configs").rglob("*.toml")) @@ -522,6 +537,87 @@ 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 + 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="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 + 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 {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 + + payload["deployment"]["num_value_nodes"] = 1 + payload["value"]["model"]["cp"] = 4 + 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["trainer"]["model"]["cp"] = 8 + payload["trainer"]["model"]["seq_len"] = 2050 + with pytest.raises(ValidationError, match="value.model.seq_len must be divisible"): + 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 8cb2c572f5b..57217b3b584 100644 --- a/tests/unit/train/rl/test_loss.py +++ b/tests/unit/train/rl/test_loss.py @@ -1,7 +1,14 @@ import pytest import torch -from prime_rl.configs.trainer import CISPOLossConfig, CustomLossConfig, IcePopLossConfig, IPOLossConfig, PPOLossConfig +from prime_rl.configs.trainer import ( + CISPOLossConfig, + CustomLossConfig, + IcePopLossConfig, + IPOLossConfig, + PPOLossConfig, + SAOLossConfig, +) from prime_rl.trainer.rl.loss import ( IcePopLoss, LossInputs, @@ -126,11 +133,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(value) for value in result.metrics.values()) + 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(PPOLossConfig()).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")) + + @pytest.mark.parametrize("config", [IPOLossConfig(kl_tau=0.01), IcePopLossConfig()]) def test_ipo_icepop_match_original_on_finite_ratios(config): torch.manual_seed(23) 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..b8f87d2deaa --- /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..723ef5d0933 --- /dev/null +++ b/tests/unit/train/rl/test_value_real_model_smoke.py @@ -0,0 +1,89 @@ +"""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") + 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, + 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=4, + 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=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, 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): + 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, 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) + 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()