diff --git a/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml b/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml new file mode 100644 index 000000000..84d86d663 --- /dev/null +++ b/examples/ar/qwen3_ppo_4b_base_dapo_sglang.yaml @@ -0,0 +1,156 @@ +# @package _global_ +# PPO + GAE on Qwen3-4B-Base / DAPO-Math with colocated SGLang rollout. +# +# Prepare data: +# python -m unirl.utils.prepare_dapo_math --out-dir data/dapo_math +# +# Run: +# DATA_PATH=data/dapo_math/train.jsonl \ +# EVAL_DATA_PATH=data/dapo_math/aime_eval.jsonl \ +# python -m unirl.train_ar --config-name=ar/qwen3_ppo_4b_base_dapo_sglang + +num_devices: 32 +batch_size: 64 +num_rollouts: 800 +weight_sync_interval: 1 +eval_interval: 10 + +# PPO prepares per-token GAE on each train worker after the micro-batched +# pre-update critic replay. The driver must not replace it with GRPO scores. +advantage_mode: gae +balance_shards: true + +logging: + report_to_wandb: true + project_name: unirl-ppo + run_name: ppo_qwen3-4b-base_dapo_sglang + entity: ${oc.env:WANDB_ENTITY,null} + tags: [ppo, gae, qwen3, 4b-base, dapo, sglang, v2] + +bundle: + _target_: unirl.models.qwen3.bundle.Qwen3Bundle.from_config + config: + _target_: unirl.models.qwen3.config.Qwen3PipelineConfig + pretrained_model_ckpt_path: ${oc.env:QWEN3_PATH,Qwen/Qwen3-4B-Base} + model_precision: fp32 + use_gradient_checkpointing: true + use_value_head: true + attn_implementation: flex_attention + +pipeline: + _target_: unirl.models.qwen3.pipeline.Qwen3Pipeline.from_bundle + enable_thinking: true + autocast_precision: bf16 + logprob_precision: fp32 + +backend: + _target_: unirl.train.backend.fsdp.FSDPBackend + block_class_names: ["Qwen3DecoderLayer"] + trainable_attr: transformer + fsdp_cfg: + _target_: unirl.train.configs.FSDPConfig + param_dtype: bf16 + cpu_offload: false + mixed_precision: true + fsdp_mode: full + reshard_after_forward: true + activation_checkpointing: true + use_torch_compile: false + forward_prefetch: false + defer_grad_sync: false + optimizer_cfg: + _target_: unirl.train.backend.base.OptimizerConfig + learning_rate: 1.0e-6 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-8 + weight_decay: 0.01 + scheduler_cfg: + _target_: unirl.train.backend.base.LrSchedulerConfig + type: constant + warmup_steps: 0 + total_steps: 1000 + +rollout: + _target_: unirl.rollout.engine.sglang.engine.SGLangRolloutEngine + config: + _target_: unirl.rollout.engine.sglang.config.SGLangEngineConfig + backend: native + pretrained_model_ckpt_path: ${oc.env:QWEN3_PATH,Qwen/Qwen3-4B-Base} + tp_size: 1 + max_new_tokens: 8192 + temperature: 1.0 + top_p: 1.0 + concurrency: 16 + samples_pre_expanded: true + chat_template_kwargs: + enable_thinking: true + engine_kwargs: + rl_on_policy_target: fsdp + mem_fraction_static: 0.3 + skip_server_warmup: true + attention_backend: triton + disable_cuda_graph: false + cuda_graph_max_bs: 16 + enable_lora: false + +reward: + _target_: unirl.reward.service.RewardService + truncated_reward: keep + backend: + _target_: unirl.reward.local.mathverify.MathVerifyRewardScorer + base_device: cpu + config: + _target_: unirl.reward.local.mathverify.MathVerifySpec + +algorithm: + _target_: unirl.algorithms.ppo.PPO + stage_attr: ar + clip_range: 0.2 + clip_range_high: null + clip_schedule: constant + cliprange_value: 0.2 + vf_coef: 0.5 + gae_gamma: 1.0 + gae_lambda: 0.95 + loss_agg_mode: seq-mean-token-mean + horizon: 8192 + sampling_temperature: 1.0 + conditions_cls: + _target_: hydra.utils.get_class + path: unirl.models.qwen3.conditions.Qwen3ARConditions + +sync: + _target_: unirl.distributed.weight_sync.full.tensor.TensorWeightSync + lora_merged: false + bucket_size_mb: 64 + flush_cache: true + # The rollout model has no critic; never send train-only value-head tensors. + name_remap: {"value_head.*": null} + +stack: + _target_: unirl.train.stack.TrainStack + micro_batch_size: 1 + max_grad_norm: 1.0 + num_updates_per_batch: 4 + micro_planner: + _target_: unirl.train.stack.TokenBudgetPlanner + token_budget: 10240 + +data_source: + _target_: unirl.data.data_source.MultimodalRLDataSource + args: + run: + data_path: ${oc.env:DATA_PATH} + eval_data_path: ${oc.env:EVAL_DATA_PATH,${oc.env:DATA_PATH}} + seed: 42 + algorithm: + prompts_per_rollout: 64 + +sampling: + _target_: unirl.types.sampling.ARSamplingParams + samples_per_prompt: 8 + temperature: 1.0 + top_p: 1.0 + top_k: 0 + max_new_tokens: 8192 diff --git a/unirl/algorithms/__init__.py b/unirl/algorithms/__init__.py index 8604d0b50..da5886314 100644 --- a/unirl/algorithms/__init__.py +++ b/unirl/algorithms/__init__.py @@ -15,6 +15,7 @@ from .flowgrpo import FlowGRPO, FlowGRPOConfig from .grpo import GRPO, GRPOConfig from .gspo import GSPO, GSPOConfig +from .ppo import PPO, PPOConfig from .sft import SFT, FlowMatchSFT __all__ = [ @@ -24,6 +25,8 @@ "GRPOConfig", "GSPO", "GSPOConfig", + "PPO", + "PPOConfig", "CPPO", "CPPOConfig", "DPPO", diff --git a/unirl/algorithms/base.py b/unirl/algorithms/base.py index e329743d1..8db2286c2 100644 --- a/unirl/algorithms/base.py +++ b/unirl/algorithms/base.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from unirl.types.conditions import Condition + from unirl.types.sample import Part from unirl.types.segments.base import Segment @@ -440,6 +441,15 @@ def prepare_segment( """ return None + def prepare_part(self, part: "Part") -> "Part": + """Optional post-anchor hook over the complete arranged worker shard. + + Runs after per-micro anchor fields have been reassembled and before any + optimizer update. PPO uses it to derive GAE from frozen critic values; + other algorithms keep the part unchanged. + """ + return part + @abstractmethod def compute_loss_and_backward( self, diff --git a/unirl/algorithms/ppo.py b/unirl/algorithms/ppo.py new file mode 100644 index 000000000..c393a177e --- /dev/null +++ b/unirl/algorithms/ppo.py @@ -0,0 +1,310 @@ +"""PPO with per-token GAE and a clipped value objective for AR training.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Mapping, Optional, Type + +import torch + +from unirl.models.types.replay_result import ReplayResult +from unirl.types.conditions import Condition +from unirl.types.sample import Part +from unirl.types.segments.text import TextSegment + +from .base import ( + AlgorithmStepResult, + BaseAlgorithmConfig, + StageAlgorithm, + _grpo_clip_loss, + _resolve_clip_range_from_schedule, + rollout_replay_logp_absdiff, + typed_conditions, +) + +_LOSS_AGG_MODES = frozenset({"token-mean", "seq-mean-token-sum-norm", "seq-mean-token-mean"}) + + +@dataclass +class PPOConfig(BaseAlgorithmConfig): + stage_attr: str = "ar" + conditions_cls: str = "" + clip_range: float = 0.2 + clip_range_high: Optional[float] = None + clip_schedule: str = "constant" + cliprange_value: float = 0.2 + vf_coef: float = 0.5 + gae_gamma: float = 1.0 + gae_lambda: float = 0.95 + loss_agg_mode: str = "token-mean" + horizon: int = 8192 + + +def _ppo_clipped_value_loss( + *, + values: torch.Tensor, + old_values: torch.Tensor, + returns: torch.Tensor, + clip_range: float, +) -> torch.Tensor: + """Elementwise ``0.5 * max((V-R)^2, (V_clipped-R)^2)``.""" + values_f = values.float() + old_f = old_values.detach().float() + returns_f = returns.detach().float() + clipped = old_f + (values_f - old_f).clamp(-clip_range, clip_range) + return 0.5 * torch.maximum((values_f - returns_f).square(), (clipped - returns_f).square()) + + +def _aggregate_token_loss( + loss_per_token: torch.Tensor, + *, + active: torch.Tensor, + segment: TextSegment, + loss_agg_mode: str, + horizon: int, +) -> torch.Tensor: + """Reduce only trainable tokens while preserving the configured sequence weighting.""" + if loss_agg_mode == "token-mean": + return loss_per_token[active].mean() + if segment.lengths is None: + raise ValueError(f"PPO: loss_agg_mode={loss_agg_mode!r} requires packed segment lengths") + + loss_chunks = torch.split(loss_per_token, segment.lengths.tolist()) + mask_chunks = torch.split(active, segment.lengths.tolist()) + sequence_losses = [] + for losses, mask in zip(loss_chunks, mask_chunks): + selected = losses[mask] + if loss_agg_mode == "seq-mean-token-sum-norm": + sequence_losses.append(selected.sum() / float(horizon)) + else: + sequence_losses.append(selected.mean() if selected.numel() else losses.new_zeros(())) + return torch.stack(sequence_losses).mean() + + +class PPO(StageAlgorithm): + """PPO over an AR ``TextSegment`` with a train-side value head.""" + + supports_multi_update = True + anchor_fields = ("values",) + + def __init__( + self, + *, + stage: Any = None, + pipeline: Any = None, + stage_attr: str = "ar", + clip_range: float = 0.2, + clip_schedule: str = "constant", + clip_range_high: Optional[float] = None, + cliprange_value: float = 0.2, + vf_coef: float = 0.5, + gae_gamma: float = 1.0, + gae_lambda: float = 0.95, + loss_agg_mode: str = "token-mean", + horizon: int = 8192, + conditions_cls: Optional[Type[Any]] = None, + sampling_temperature: Optional[float] = None, + ) -> None: + super().__init__() + if stage is None and pipeline is None: + raise ValueError("PPO: either `stage` or `pipeline` must be provided") + if stage is None: + stage = getattr(pipeline, stage_attr) + if float(clip_range) < 0.0 or float(cliprange_value) < 0.0: + raise ValueError("PPO: policy and value clip ranges must be non-negative") + if float(vf_coef) < 0.0: + raise ValueError("PPO: vf_coef must be non-negative") + if not (0.0 <= float(gae_gamma) <= 1.0 and 0.0 <= float(gae_lambda) <= 1.0): + raise ValueError("PPO: gae_gamma and gae_lambda must be in [0, 1]") + if str(loss_agg_mode) not in _LOSS_AGG_MODES: + raise ValueError(f"PPO: unsupported loss_agg_mode={loss_agg_mode!r}") + if int(horizon) <= 0: + raise ValueError("PPO: horizon must be positive") + + self.stage = stage + self.clip_range = float(clip_range) + self.clip_range_high = None if clip_range_high is None else float(clip_range_high) + self.clip_schedule = str(clip_schedule) + self.cliprange_value = float(cliprange_value) + self.vf_coef = float(vf_coef) + self.gae_gamma = float(gae_gamma) + self.gae_lambda = float(gae_lambda) + self.loss_agg_mode = str(loss_agg_mode) + self.loss_weighting = "token" if self.loss_agg_mode == "token-mean" else "sample" + self.horizon = int(horizon) + self.conditions_cls = conditions_cls + if sampling_temperature is None: + from unirl.types.sampling import ARSamplingParams + + sampling_temperature = ARSamplingParams.__dataclass_fields__["temperature"].default + self.sampling_temperature = float(sampling_temperature) + + def recomputes_anchor(self) -> bool: + """Critic anchors must use the exact micro geometry of the train forward.""" + return True + + def prepare_segment( + self, + *, + conditions: Mapping[str, Condition], + segment: TextSegment, + ) -> None: + """Freeze pre-update critic values for one planned replay micro.""" + if segment.tokens is None or segment.log_probs is None: + raise ValueError("PPO.prepare_segment: segment requires tokens and rollout log_probs") + typed_conds = typed_conditions(conditions, self.conditions_cls) + with torch.no_grad(): + replay = self.stage.replay( + typed_conds, + segment=segment, + temperature=self.sampling_temperature, + return_values=True, + ) + values = _replay_values(replay) + if values.shape != segment.tokens.shape: + raise ValueError( + f"PPO.prepare_segment: values shape {tuple(values.shape)} " + f"!= packed tokens shape {tuple(segment.tokens.shape)}" + ) + segment.values = values.detach() + + def prepare_part(self, part: Part) -> Part: + """Compute GAE after all per-micro critic anchors have been reassembled.""" + return part.compute_gae_advantages(gamma=self.gae_gamma, gae_lambda=self.gae_lambda) + + def compute_loss_and_backward( + self, + *, + conditions: Mapping[str, Condition], + segment: TextSegment, + advantages: torch.Tensor, + training_progress: float, + loss_scale: float, + ) -> AlgorithmStepResult: + del advantages # PPO consumes packed GAE, not the compatibility summary on Part. + if segment.tokens is None or segment.lengths is None or segment.log_probs is None: + return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False) + if int(segment.tokens.numel()) == 0: + return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False) + if segment.token_advantages is None or segment.returns is None or segment.values is None: + raise ValueError( + "PPO.compute_loss_and_backward: token_advantages, returns, and frozen values " + "must be prepared before training" + ) + + typed_conds = typed_conditions(conditions, self.conditions_cls) + replay = self.stage.replay( + typed_conds, + segment=segment, + temperature=self.sampling_temperature, + return_values=True, + ) + if not isinstance(replay, ReplayResult): + raise TypeError("PPO: replay with return_values=True must return ReplayResult") + new_logp = replay.log_probs + new_values = _replay_values(replay) + + old_logp = segment.log_probs.to(dtype=new_logp.dtype, device=new_logp.device) + old_values = segment.values.to(dtype=new_values.dtype, device=new_values.device) + returns = segment.returns.to(dtype=new_values.dtype, device=new_values.device) + token_advantages = segment.token_advantages.to(dtype=new_logp.dtype, device=new_logp.device).detach() + expected = new_logp.shape + named_tensors = { + "new_values": new_values, + "old_logp": old_logp, + "old_values": old_values, + "returns": returns, + "token_advantages": token_advantages, + } + mismatched = {name: tuple(tensor.shape) for name, tensor in named_tensors.items() if tensor.shape != expected} + if mismatched: + raise ValueError( + f"PPO.compute_loss_and_backward: packed tensor shape mismatch; expected {expected}, got {mismatched}" + ) + + active = torch.ones(expected, dtype=torch.bool, device=new_logp.device) + if segment.loss_mask is not None: + active = segment.loss_mask.to(device=new_logp.device, dtype=torch.bool) + if active.shape != expected: + raise ValueError( + f"PPO.compute_loss_and_backward: loss_mask shape {tuple(active.shape)} != {tuple(expected)}" + ) + active_count = int(active.sum().item()) + if active_count == 0: + return AlgorithmStepResult(loss=0.0, metrics={}, num_steps_or_tokens=0, has_backward=False) + + clip_range = _resolve_clip_range_from_schedule(self.clip_range, self.clip_schedule, training_progress) + clip_high = ( + None + if self.clip_range_high is None + else _resolve_clip_range_from_schedule(self.clip_range_high, self.clip_schedule, training_progress) + ) + policy_active, ratio_metrics = _grpo_clip_loss( + new_logp=new_logp[active], + old_logp=old_logp[active], + advantages=token_advantages[active], + clip_range=clip_range, + clip_range_high=clip_high, + ) + value_active = _ppo_clipped_value_loss( + values=new_values[active], + old_values=old_values[active], + returns=returns[active], + clip_range=self.cliprange_value, + ) + + policy_per_token = torch.zeros_like(new_logp).masked_scatter(active, policy_active) + value_per_token = torch.zeros_like(new_values, dtype=torch.float32).masked_scatter(active, value_active) + policy_loss = _aggregate_token_loss( + policy_per_token, + active=active, + segment=segment, + loss_agg_mode=self.loss_agg_mode, + horizon=self.horizon, + ) + value_loss = _aggregate_token_loss( + value_per_token, + active=active, + segment=segment, + loss_agg_mode=self.loss_agg_mode, + horizon=self.horizon, + ) + loss = policy_loss + self.vf_coef * value_loss + (loss * loss_scale).backward() + + active_returns = returns[active].float() + active_values = new_values[active].float() + return_var = active_returns.var(unbiased=False) + explained_variance = ( + 1.0 - (active_returns - active_values).var(unbiased=False) / return_var + if float(return_var) > 0.0 + else return_var.new_zeros(()) + ) + value_clip_fraction = ((active_values - old_values[active].float()).abs() > self.cliprange_value).float().mean() + metrics: Dict[str, Any] = { + "policy_loss": float(policy_loss.detach()), + "value_loss": float(value_loss.detach()), + "value_mean": float(active_values.detach().mean()), + "return_mean": float(active_returns.detach().mean()), + "explained_variance": float(explained_variance.detach()), + "value_clip_fraction": float(value_clip_fraction.detach()), + "clip_range": float(clip_range), + "cliprange_value": self.cliprange_value, + **rollout_replay_logp_absdiff(new_logp[active], old_logp[active]), + **{name: float(value) for name, value in ratio_metrics.items()}, + } + return AlgorithmStepResult( + loss=float(loss.detach()), + metrics=metrics, + num_steps_or_tokens=active_count, + has_backward=True, + ) + + +def _replay_values(replay: Any) -> torch.Tensor: + if not isinstance(replay, ReplayResult) or replay.values is None: + raise TypeError("PPO: replay with return_values=True must return ReplayResult.values") + return replay.values + + +__all__ = ["PPO", "PPOConfig"] diff --git a/unirl/train/stack/base.py b/unirl/train/stack/base.py index b4160fec7..2f4816a2c 100644 --- a/unirl/train/stack/base.py +++ b/unirl/train/stack/base.py @@ -505,6 +505,7 @@ def train_track( profiler = self._train_step_profiler() if profile_scope() == "train" else None with profiler.record("train_track") if profiler is not None else nullcontext(): self.prepare_segment(part, plans=plans) + part = self.algorithm.prepare_part(part) self.fsdp_backend.model.train() result = self._run_updates(part, plans=plans, training_progress=float(training_progress)) if profiler is not None: diff --git a/unirl/train_ar.py b/unirl/train_ar.py index 063418718..3919b6b08 100755 --- a/unirl/train_ar.py +++ b/unirl/train_ar.py @@ -51,6 +51,7 @@ def teardown() -> None: logging_cfg=cfg.get("logging"), adv_normalization_scope=cfg.get("adv_normalization_scope", "group"), normalize_adv_by_std=cfg.get("normalize_adv_by_std", True), + advantage_mode=cfg.get("advantage_mode", "grpo"), balance_shards=cfg.get("balance_shards", False), eval_interval=cfg.get("eval_interval", 0), eval_num_prompts=cfg.get("eval_num_prompts", -1), diff --git a/unirl/train_async_ar.py b/unirl/train_async_ar.py index 511e5e089..37c359ab9 100755 --- a/unirl/train_async_ar.py +++ b/unirl/train_async_ar.py @@ -47,6 +47,7 @@ def main(cfg: DictConfig) -> None: logging_cfg=cfg.get("logging"), adv_normalization_scope=cfg.get("adv_normalization_scope", "group"), normalize_adv_by_std=bool(cfg.get("normalize_adv_by_std", True)), + advantage_mode=cfg.get("advantage_mode", "grpo"), balance_shards=bool(cfg.get("balance_shards", False)), eval_interval=int(cfg.get("eval_interval", 0)), eval_num_prompts=int(cfg.get("eval_num_prompts", -1)), diff --git a/unirl/trainer/ar.py b/unirl/trainer/ar.py index 2371dfada..9af9eeed4 100644 --- a/unirl/trainer/ar.py +++ b/unirl/trainer/ar.py @@ -60,6 +60,7 @@ def __init__( logging_cfg: Optional[DictConfig] = None, adv_normalization_scope: str = "group", normalize_adv_by_std: bool = True, + advantage_mode: str = "grpo", balance_shards: bool = False, eval_interval: int = 0, eval_num_prompts: int = -1, @@ -83,6 +84,9 @@ def __init__( # group std. False = mean-center only (reward - group_mean), NO std division — # removes the difficulty bias that over-amplifies low-std (hard) prompts. self.normalize_adv_by_std = normalize_adv_by_std + self.advantage_mode = str(advantage_mode).strip().lower() + if self.advantage_mode not in ("grpo", "gae"): + raise ValueError(f"ARTrainer: advantage_mode must be 'grpo' or 'gae', got {advantage_mode!r}") # verl trainer.balance_batch parity: driver-side reorder of the rollout # batch so each DP shard receives a similar total-token workload. FSDP # collectives sync all ranks every micro, so a step runs at the SLOWEST @@ -456,7 +460,11 @@ def train_step( if isinstance(part.component_rewards, dict): part.component_rewards = {name: hydrate(value) for name, value in part.component_rewards.items()} mean_reward = float(part.rewards.to(torch.float32).mean().item()) - part = part.compute_advantages(normalize=self.normalize_adv_by_std, scope=self.adv_normalization_scope) + if self.advantage_mode == "grpo": + part = part.compute_advantages( + normalize=self.normalize_adv_by_std, + scope=self.adv_normalization_scope, + ) sample = sample.with_parts([*sample.parts[:-1], part]) self._dump_rollout_samples(sample, rollout_id) diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index a64e517d2..936e43859 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -96,6 +96,7 @@ def __init__( logging_cfg: Optional[DictConfig] = None, adv_normalization_scope: str = "group", normalize_adv_by_std: bool = True, + advantage_mode: str = "grpo", balance_shards: bool = False, eval_interval: int = 0, eval_num_prompts: int = -1, @@ -122,6 +123,9 @@ def __init__( self.batch_size = batch_size self.adv_normalization_scope = adv_normalization_scope self.normalize_adv_by_std = normalize_adv_by_std + self.advantage_mode = str(advantage_mode).strip().lower() + if self.advantage_mode not in ("grpo", "gae"): + raise ValueError(f"AsyncARTrainer: advantage_mode must be 'grpo' or 'gae', got {advantage_mode!r}") self.balance_shards = bool(balance_shards) self.eval_interval = int(eval_interval) _num = int(eval_num_prompts) @@ -293,7 +297,11 @@ def _advantage_and_train( if part.rewards is not None: part.rewards = hydrate(part.rewards) mean_reward = float(part.rewards.to(torch.float32).mean().item()) - part = part.compute_advantages(normalize=self.normalize_adv_by_std, scope=self.adv_normalization_scope) + if self.advantage_mode == "grpo": + part = part.compute_advantages( + normalize=self.normalize_adv_by_std, + scope=self.adv_normalization_scope, + ) sample = sample.with_parts([*sample.parts[:-1], part]) train_part = part if self.balance_shards: