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/models/qwen3/ar.py b/unirl/models/qwen3/ar.py index 1e072ac76..793e06827 100644 --- a/unirl/models/qwen3/ar.py +++ b/unirl/models/qwen3/ar.py @@ -23,7 +23,7 @@ from dataclasses import dataclass from dataclasses import field as dc_field from types import MethodType -from typing import Any, List, Optional, Tuple +from typing import Any, List, Optional, Tuple, Union import torch import torch.distributed as dist @@ -31,6 +31,7 @@ from torch.utils.checkpoint import checkpoint from unirl.models.types.ar import ARSamplingParams, ARStage, ARStep, left_pad_prompt +from unirl.models.types.replay_result import ReplayResult from unirl.types.segments import TextSegment from unirl.utils.dtypes import parse_torch_dtype @@ -87,6 +88,7 @@ def _replay_aware_forward( temperature: float = 1.0, autocast_dtype: Optional[torch.dtype] = None, packed_predict_index: Optional[torch.Tensor] = None, + return_values: bool = False, **kw: Any, ) -> Any: """Dual-mode ``forward`` installed on the Qwen3 CausalLM instance. @@ -108,6 +110,8 @@ def _replay_aware_forward( return f(self, **kw) raise RuntimeError("_replay_aware_forward: no class-level forward found in the MRO") + _require_value_head_for_replay(self, return_values) + # cuDNN's fused SDPA backward (ScaledDotProductCudnnAttentionBackward0) returns # NaN grads on some bf16 sequences while the forward stays finite (confirmed via # torch.autograd.detect_anomaly): it floods every parameter grad and forces the @@ -129,6 +133,7 @@ def _replay_aware_forward( # [B, chunk, vocab] FP32 transient stays ~1.2 GiB, and each chunk is # gradient-checkpointed (recomputed in backward rather than held). T = float(temperature) if float(temperature) > 0.0 else 1.0 + value_head = getattr(self, "value_head", None) if return_values else None if packed_predict_index is not None: # Packed varlen replay: ``hidden`` is one packed row [1, L_total, H] @@ -155,8 +160,16 @@ def _flat_logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: else: flat_parts.append(_flat_logp_chunk(h, tok)) if not flat_parts: - return hidden.new_zeros((0,), dtype=torch.float32) - return torch.cat(flat_parts, dim=0) + empty = hidden.new_zeros((0,), dtype=torch.float32) + if value_head is None: + return empty + return ReplayResult(log_probs=empty, values=empty) + log_probs = torch.cat(flat_parts, dim=0) + if value_head is None: + return log_probs + value_parts = [value_head(h_pred[s : s + flat_chunk]) for s in range(0, int(h_pred.size(0)), flat_chunk)] + values = torch.cat(value_parts, dim=0) if value_parts else log_probs.new_zeros(0) + return ReplayResult(log_probs=log_probs, values=values) T_max = int(response_tokens.size(1)) resp_hidden = hidden[:, prompt_len - 1 : prompt_len - 1 + T_max, :] @@ -176,8 +189,51 @@ def _logp_chunk(h: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: else: parts.append(_logp_chunk(h, tok)) if not parts: - return resp_hidden.new_zeros((bsz, 0), dtype=torch.float32) # T_max == 0 - return torch.cat(parts, dim=1) + empty = resp_hidden.new_zeros((bsz, 0), dtype=torch.float32) + if value_head is None: + return empty + return ReplayResult(log_probs=empty, values=empty) + log_probs = torch.cat(parts, dim=1) + if value_head is None: + return log_probs + value_parts = [value_head(resp_hidden[:, s : s + chunk, :]) for s in range(0, T_max, chunk)] + values = torch.cat(value_parts, dim=1) if value_parts else log_probs.new_zeros((bsz, 0)) + return ReplayResult(log_probs=log_probs, values=values) + + +def _require_value_head_for_replay(model: Any, return_values: bool) -> None: + if return_values and getattr(model, "value_head", None) is None: + raise ValueError( + "Qwen3 replay: return_values=True requires a value head (set use_value_head=True in the pipeline config)" + ) + + +def _finalize_replay_output( + out: Union[torch.Tensor, ReplayResult], + *, + segment: TextSegment, + return_values: bool, + logprob_dtype: torch.dtype, +) -> Union[torch.Tensor, ReplayResult]: + """Cast log-probs and flatten padded critic values to segment order.""" + if not isinstance(out, ReplayResult): + if return_values: + raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values") + return out.to(dtype=logprob_dtype) + + log_probs = out.log_probs.to(dtype=logprob_dtype) + if not return_values: + return log_probs + if out.values is None: + raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values") + if log_probs.ndim == 1: + return ReplayResult(log_probs=log_probs, values=out.values.float()) + if segment.lengths is None: + raise ValueError("Qwen3ARStage.replay: segment requires lengths to flatten critic values") + + flat_values = [out.values[b, : int(length)] for b, length in enumerate(segment.lengths.tolist()) if int(length) > 0] + values = torch.cat(flat_values, dim=0) if flat_values else out.values.new_zeros(0) + return ReplayResult(log_probs=log_probs, values=values.float()) # Attention backends with a sparse packed kernel (skip cross-sequence blocks): @@ -421,22 +477,32 @@ def replay( *, segment: TextSegment, temperature: float = 1.0, - ) -> torch.Tensor: + return_values: bool = False, + ) -> Union[torch.Tensor, ReplayResult]: """Per-token log-prob replay over a stored rollout segment. Branch: prefer :meth:`packed_replay` (packed-varlen, zero padding, B > 1) and fall back to :meth:`padding_replay` (the dense ``[B, P_max + T_max]`` - padded path) when packing does not apply. Returns packed varlen - ``[total_tokens]`` aligned with ``segment.log_probs``; caller controls - grad / ``.train()`` scope. ``temperature`` divides logits before - ``log_softmax`` to match SGLang's sampler (``1.0`` is a no-op). + padded path) when packing does not apply. Returns packed varlen log-probs, + or a :class:`ReplayResult` with aligned critic values when requested. """ + _require_value_head_for_replay(self.model.transformer, return_values) attn_impl = getattr(getattr(self.model.transformer, "config", None), "_attn_implementation", None) if _packed_replay_supported(attn_impl): - packed = self.packed_replay(conditions, segment=segment, temperature=temperature) + packed = self.packed_replay( + conditions, + segment=segment, + temperature=temperature, + return_values=return_values, + ) if packed is not None: return packed - return self.padding_replay(conditions, segment=segment, temperature=temperature) + return self.padding_replay( + conditions, + segment=segment, + temperature=temperature, + return_values=return_values, + ) def packed_replay( self, @@ -444,7 +510,8 @@ def packed_replay( *, segment: TextSegment, temperature: float = 1.0, - ) -> Optional[torch.Tensor]: + return_values: bool = False, + ) -> Optional[Union[torch.Tensor, ReplayResult]]: """Packed-varlen replay (B > 1): zero padding anywhere. Concatenate every sample's REAL prompt tokens + its flat response tokens @@ -520,9 +587,15 @@ def packed_replay( packed_predict_index=predict_index, prompt_len=0, temperature=temperature, + return_values=return_values, autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None), ) - return per_token_flat.to(dtype=self.logprob_dtype) + return _finalize_replay_output( + per_token_flat, + segment=segment, + return_values=return_values, + logprob_dtype=self.logprob_dtype, + ) def padding_replay( self, @@ -530,7 +603,8 @@ def padding_replay( *, segment: TextSegment, temperature: float = 1.0, - ) -> torch.Tensor: + return_values: bool = False, + ) -> Union[torch.Tensor, ReplayResult]: """Dense ``[B, P_max + T_max]`` padded replay — the default / fallback path. One teacher-forced forward over padded ``prompt + response``; gather @@ -629,28 +703,44 @@ def padding_replay( # root-wrapped or plain) and never materializes [B, L, vocab] logits. # The cuda-vs-cpu autocast decision lives here; dtype validity and the # autocast scope live in the patched forward. - per_token = self.model.transformer( + out = self.model.transformer( input_ids=full_ids, attention_mask=full_mask, position_ids=position_ids, response_tokens=response_tokens, prompt_len=prompt_len, temperature=temperature, + return_values=return_values, autocast_dtype=(self.autocast_dtype if device.type == "cuda" else None), - ) # [B, T_max] FP32 - - if T_max == 0: - return torch.zeros(0, dtype=self.logprob_dtype, device=device) + ) # [B, T_max] FP32 tensor or ReplayResult - flat: List[torch.Tensor] = [] - for b in range(batch_size): - n = lengths[b] - if n == 0: - continue - flat.append(per_token[b, :n]) - if not flat: - return torch.zeros(0, dtype=self.logprob_dtype, device=device) - return torch.cat(flat, dim=0).to(dtype=self.logprob_dtype) + finalized = _finalize_replay_output( + out, + segment=segment, + return_values=return_values, + logprob_dtype=self.logprob_dtype, + ) + if isinstance(finalized, ReplayResult): + per_token = finalized.log_probs + values = finalized.values + else: + per_token = finalized + values = None + + flat_log_probs: List[torch.Tensor] = [] + for b, n in enumerate(lengths): + if n > 0: + flat_log_probs.append(per_token[b, :n]) + log_probs = ( + torch.cat(flat_log_probs, dim=0) + if flat_log_probs + else torch.zeros(0, dtype=self.logprob_dtype, device=device) + ) + if not return_values: + return log_probs + if values is None: + raise ValueError("Qwen3ARStage.replay: return_values=True but critic returned no values") + return ReplayResult(log_probs=log_probs, values=values) def _resolve_stop_ids( self, diff --git a/unirl/models/qwen3/bundle.py b/unirl/models/qwen3/bundle.py index 3c13308a8..d44201f0d 100644 --- a/unirl/models/qwen3/bundle.py +++ b/unirl/models/qwen3/bundle.py @@ -28,6 +28,7 @@ from unirl.models.types.bundle import Bundle from unirl.models.types.meta_init import build_meta_init_transformer +from unirl.models.types.value_head import ValueHead from unirl.utils.dtypes import parse_torch_dtype from .config import Qwen3PipelineConfig @@ -35,6 +36,24 @@ logger = logging.getLogger(__name__) +def _stamp_value_head_reset(transformer: nn.Module) -> None: + """Zero checkpoint-absent value-head params after meta materialization.""" + from unirl.train.deferred import _stamp + + def _reset(model: nn.Module) -> None: + reset: list[str] = [] + with torch.no_grad(): + for name, param in model.named_parameters(): + if name.startswith("value_head."): + param.zero_() + reset.append(name) + if not reset: + raise RuntimeError("Qwen3 meta-init: value_head parameters disappeared before post-load reset") + logger.info("Qwen3 meta-init: zero-initialized checkpoint-absent value head: %s", reset) + + _stamp(transformer, _reset) + + class Qwen3Bundle(Bundle): """Qwen3 bundle: causal-LM transformer + matching tokenizer.""" @@ -115,6 +134,16 @@ def from_config(cls, config: Qwen3PipelineConfig) -> "Qwen3Bundle": if tokenizer.pad_token is None and tokenizer.eos_token is not None: tokenizer.pad_token = tokenizer.eos_token + if config.use_value_head: + hidden_size = int(getattr(transformer.config, "hidden_size")) + transformer_device = next(transformer.parameters()).device + transformer.value_head = ValueHead(hidden_size, device=transformer_device) + if config.meta_init_transformer: + # The base checkpoint has no value_head.* tensors. ``to_empty`` + # materializes the meta head as uninitialized storage, so reset + # it only after the sharded checkpoint load has completed. + _stamp_value_head_reset(transformer) + bundle = cls( transformer=transformer, tokenizer=tokenizer, diff --git a/unirl/models/qwen3/config.py b/unirl/models/qwen3/config.py index 37f63b5ef..2fa5d3504 100644 --- a/unirl/models/qwen3/config.py +++ b/unirl/models/qwen3/config.py @@ -69,6 +69,9 @@ class Qwen3PipelineConfig: use_lora: bool = False lora_target_modules: Optional[List[str]] = None + # Attach a scalar critic to the causal LM for PPO / GAE training. + use_value_head: bool = False + system_instruction: Optional[str] = None # Chat-template thinking switch; MUST agree with the rollout engine's # chat_template_kwargs.enable_thinking or train/rollout prompts diverge. diff --git a/unirl/models/types/replay_result.py b/unirl/models/types/replay_result.py index 3bf7a7222..381d513c5 100644 --- a/unirl/models/types/replay_result.py +++ b/unirl/models/types/replay_result.py @@ -8,10 +8,8 @@ Diffusion stages populate ``log_probs`` and ``prev_sample_means`` (the mean of the SDE Gaussian — μ_θ — used as the second moment in the KL -penalty). AR stages currently return a plain ``Tensor`` (signature -divergence with diffusion is intentional for now); when AR replay grows -``logits``-based KL support, the ``logits`` field on this result will be -the canonical home. +penalty). AR stages return a plain ``Tensor`` for policy-only replay, or +a :class:`ReplayResult` when optional critic ``values`` are requested. """ from __future__ import annotations @@ -42,5 +40,9 @@ class ReplayResult: or entropy penalty support; not needed for Binary KL (which uses only per-token log-probs). Currently not populated.""" + values: Optional[torch.Tensor] = None + """Per-token critic predictions ``V_t``. Packed ``[total_tokens]`` for AR. + ``None`` when replay did not request a value head.""" + __all__ = ["ReplayResult"] diff --git a/unirl/models/types/value_head.py b/unirl/models/types/value_head.py new file mode 100644 index 000000000..1193a8017 --- /dev/null +++ b/unirl/models/types/value_head.py @@ -0,0 +1,35 @@ +"""Scalar value head for PPO-style critic training on AR hidden states.""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ValueHead(nn.Module): + """Linear critic ``V(h)`` on the hidden state that predicts each action.""" + + def __init__(self, hidden_size: int, *, device: Any = None) -> None: + super().__init__() + self.proj = nn.Linear(hidden_size, 1, bias=True, dtype=torch.float32, device=device) + self.reset_parameters() + + def reset_parameters(self) -> None: + """Start from ``V(s)=0``; deterministic and safe after sharded meta-init.""" + nn.init.zeros_(self.proj.weight) + if self.proj.bias is not None: + nn.init.zeros_(self.proj.bias) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + """Map ``[..., H]`` hidden states to FP32 scalar values ``[...]``.""" + # FSDP mixed precision may expose gathered parameters in its compute + # dtype even when the sharded masters are FP32. + weight = self.proj.weight.float() + bias = self.proj.bias.float() if self.proj.bias is not None else None + return F.linear(hidden.float(), weight, bias).squeeze(-1) + + +__all__ = ["ValueHead"] 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: diff --git a/unirl/types/advantages.py b/unirl/types/advantages.py index f49f96b6e..d4c3d4b35 100644 --- a/unirl/types/advantages.py +++ b/unirl/types/advantages.py @@ -71,6 +71,34 @@ def compute_gae_advantages( return advantages, returns +def scatter_terminal_rewards( + rewards_per_sample: torch.Tensor, + *, + cu_seqlens: torch.Tensor, +) -> torch.Tensor: + """Scatter each trajectory reward onto its final packed response token.""" + if rewards_per_sample.ndim != 1: + raise ValueError(f"scatter_terminal_rewards: expected 1D rewards, got shape {tuple(rewards_per_sample.shape)}") + if cu_seqlens.ndim != 1 or cu_seqlens.numel() == 0: + raise ValueError("scatter_terminal_rewards: cu_seqlens must be a non-empty 1D tensor") + batch_size = int(cu_seqlens.numel()) - 1 + if int(rewards_per_sample.numel()) != batch_size: + raise ValueError( + f"scatter_terminal_rewards: rewards batch ({int(rewards_per_sample.numel())}) " + f"!= packed batch ({batch_size})" + ) + + cu = [int(offset) for offset in cu_seqlens.tolist()] + if cu[0] != 0 or any(end < start for start, end in zip(cu, cu[1:])): + raise ValueError(f"scatter_terminal_rewards: invalid cumulative offsets {cu}") + + token_rewards = rewards_per_sample.new_zeros(cu[-1]) + for reward, start, end in zip(rewards_per_sample, cu, cu[1:]): + if end > start: + token_rewards[end - 1] = reward + return token_rewards + + def _gae_1d( rewards: torch.Tensor, values: torch.Tensor, diff --git a/unirl/types/sample.py b/unirl/types/sample.py index fede101f6..797207d22 100644 --- a/unirl/types/sample.py +++ b/unirl/types/sample.py @@ -31,12 +31,14 @@ shared_field, ) from unirl.distributed.tensor.ref import hydrate +from unirl.types.advantages import compute_gae_advantages as _compute_gae +from unirl.types.advantages import scatter_terminal_rewards from unirl.types.conditions import Condition from unirl.types.media_preview import MediaPreview from unirl.types.primitives import Audios, Images, Texts, Videos, primitive_modality_key from unirl.types.sample_id import ancestor_id, child_id, parent_id from unirl.types.sampling import BaseSamplingParams -from unirl.types.segments import Segment +from unirl.types.segments import Segment, TextSegment from unirl.utils.shard_balance import lpt_shard_permutation, shard_token_spread logger = logging.getLogger(__name__) @@ -418,6 +420,89 @@ def compute_advantages( adv = reshaped - mean return _part_with_field(self, "advantages", adv.flatten()) + def compute_gae_advantages( + self, + *, + gamma: float = 1.0, + gae_lambda: float = 0.95, + ) -> "Part": + """Attach packed token advantages and returns for an AR PPO update. + + ``loss_mask`` controls which tokens later contribute to the policy and + value losses. It deliberately does not control the GAE recursion: + masked actions are still states in the same trajectory, so terminal + reward must propagate across them. + """ + if self.rewards is None: + raise ValueError("Part.compute_gae_advantages: part has no rewards") + if not isinstance(self.segment, TextSegment): + raise ValueError("Part.compute_gae_advantages: requires a TextSegment") + segment = self.segment + if segment.values is None: + raise ValueError("Part.compute_gae_advantages: segment.values is None") + if segment.cu_seqlens is None or segment.lengths is None: + raise ValueError( + "Part.compute_gae_advantages: segment requires framework-managed " + "cu_seqlens (construct via TextSegment.pack)" + ) + + values = hydrate(segment.values).to(torch.float32) + cu_seqlens = segment.cu_seqlens.to(device=values.device) + total_tokens = int(cu_seqlens[-1].item()) + if values.ndim != 1 or int(values.numel()) != total_tokens: + raise ValueError( + "Part.compute_gae_advantages: values must be a packed 1D tensor " + f"with {total_tokens} elements, got shape {tuple(values.shape)}" + ) + + rewards = hydrate(self.rewards).to(device=values.device, dtype=torch.float32) + token_rewards = scatter_terminal_rewards(rewards, cu_seqlens=cu_seqlens) + token_advantages = values.new_zeros(values.shape) + token_returns = values.new_zeros(values.shape) + cu = [int(offset) for offset in cu_seqlens.tolist()] + for start, end in zip(cu, cu[1:]): + if end <= start: + continue + advantages, returns = _compute_gae( + token_rewards[start:end], + values[start:end], + gamma=gamma, + gae_lambda=gae_lambda, + ) + token_advantages[start:end] = advantages + token_returns[start:end] = returns + + segment_fields = {f.name: getattr(segment, f.name) for f in dc_fields(segment)} + segment_fields["token_advantages"] = token_advantages + segment_fields["returns"] = token_returns + updated_segment = segment._rebuild(segment_fields) + + # Keep the existing per-sample advantage metric meaningful. The loss + # mask applies here only as a reduction mask, never as a done signal. + loss_mask = None + if segment.loss_mask is not None: + loss_mask = hydrate(segment.loss_mask).to(device=values.device, dtype=torch.bool) + if loss_mask.shape != values.shape: + raise ValueError( + "Part.compute_gae_advantages: loss_mask shape " + f"{tuple(loss_mask.shape)} != values shape {tuple(values.shape)}" + ) + sample_advantages: List[torch.Tensor] = [] + for start, end in zip(cu, cu[1:]): + if end <= start: + sample_advantages.append(values.new_zeros(())) + continue + selected = token_advantages[start:end] + if loss_mask is not None: + selected = selected[loss_mask[start:end]] + sample_advantages.append(selected.mean() if selected.numel() else values.new_zeros(())) + mean_advantages = ( + torch.stack(sample_advantages) if sample_advantages else values.new_zeros((0,), dtype=torch.float32) + ) + + updated = _part_with_field(self, "segment", updated_segment) + return _part_with_field(updated, "advantages", mean_advantages) + def _part_with_field(part: Part, field_name: str, value: Any) -> Part: """Copy of ``part`` with one field replaced.""" diff --git a/unirl/types/segments/text.py b/unirl/types/segments/text.py index 7a1570f23..8850f902d 100644 --- a/unirl/types/segments/text.py +++ b/unirl/types/segments/text.py @@ -43,6 +43,10 @@ class TextSegment(Segment): # Original engine emission retained when an algorithm replaces ``log_probs`` # with a train-side replay anchor. rollout_log_probs: Optional[torch.Tensor] = packed_field(default=None) + # Optional PPO critic state, aligned one-to-one with packed response tokens. + values: Optional[torch.Tensor] = packed_field(default=None) + returns: Optional[torch.Tensor] = packed_field(default=None) + token_advantages: Optional[torch.Tensor] = packed_field(default=None) def as_condition_with(self, encoder: Callable[..., Any]) -> Condition: """Re-embed packed tokens via the supplied encoder into a TextEmbedCondition.