From 350e4a18493a4222f84b93abfb88d52bdd25a1b1 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 5 Sep 2026 01:59:50 +0800 Subject: [PATCH] feat: add online ppo with value-model critic and gae advantages - register online_ppo train type backed by PPOStrategy: token-level clipped surrogate over GAE advantages plus masked value regression against rollout-pinned returns, with explained-variance metrics - fold the reference-KL penalty (k3 estimator) into per-token rewards before GAE and pin advantages/returns on RolloutResult so replayed gradient steps optimize fixed targets - add self-contained ValueModel critic with a zero-initialized value head and backbone warm-started from policy weights; AutoRegressiveLM stays untouched and trunk parity is pinned by tests - step the critic's own optimizer outside the policy-version lock with the same max_grad_norm clipping as the policy - persist critic state as value_model.pt/value_optimizer.pt checkpoint extras; resume restores it, fails loudly when missing, and the train.sh completeness check requires the extras for online_ppo configs - extract shared rollout sequence/logprob helpers from GRPO (behavior unchanged) and add ppo_gamma/ppo_gae_lambda/ppo_vf_coef CLI options --- astrai/config/train_config.py | 12 +- astrai/model/__init__.py | 2 + astrai/model/value.py | 55 +++ astrai/trainer/rollout.py | 7 + astrai/trainer/strategy.py | 450 ++++++++++++++++++++--- astrai/trainer/train_callback.py | 6 + astrai/trainer/train_context.py | 65 +++- docs/developer/docker-training.md | 4 + docs/guides/params.md | 16 +- docs/guides/training.md | 12 + scripts/docker/lib/train-common.sh | 15 +- scripts/docker/train-entrypoint.sh | 1 + scripts/tools/train.py | 38 +- scripts/train.sh | 1 + tests/trainer/test_online_e2e.py | 38 +- tests/trainer/test_ppo_strategy.py | 469 ++++++++++++++++++++++++ tests/trainer/test_ppo_train_context.py | 271 ++++++++++++++ 17 files changed, 1390 insertions(+), 72 deletions(-) create mode 100644 astrai/model/value.py create mode 100644 tests/trainer/test_ppo_strategy.py create mode 100644 tests/trainer/test_ppo_train_context.py diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index b27bd88..c53f30b 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -11,7 +11,9 @@ from torch.utils.data import Dataset from astrai.config.base import BaseConfig from astrai.model.components.lora import LoRAConfig -TRAIN_TYPES = frozenset({"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"}) +TRAIN_TYPES = frozenset( + {"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo", "online_ppo"} +) PARALLEL_MODES = frozenset({"none", "ddp", "fsdp"}) BACKENDS = frozenset({"nccl", "gloo"}) START_METHODS = frozenset({"spawn", "fork", "forkserver"}) @@ -70,6 +72,8 @@ class TrainConfig(BaseConfig): rollout_top_p (float): Top-p (nucleus) filtering for online rollout. Defaults to 0.9. rollout_max_tokens (int): Maximum generated tokens per response in rollout. Defaults to 1024. reward_model_fn (Optional[Callable]): Factory for reward model, required for online RL strategies. Defaults to None. + critic_model_fn (Optional[Callable]): Factory for the value (critic) model, required for online_ppo. Defaults to None. + critic_optimizer_fn (Optional[Callable]): Factory for the critic optimizer; None reuses optimizer_fn. Defaults to None. executor_kwargs (Dict[str, Any]): Extra kwargs passed to ExecutorFactory.create(). Defaults to {}. strategy_kwargs (Dict[str, Any]): Extra strategy arguments. Defaults to {}. """ @@ -125,6 +129,8 @@ class TrainConfig(BaseConfig): rollout_top_p: float = 0.9 rollout_max_tokens: int = 1024 reward_model_fn: Optional[Callable] = None + critic_model_fn: Optional[Callable] = None + critic_optimizer_fn: Optional[Callable] = None executor_kwargs: Dict[str, Any] = field(default_factory=dict) strategy_kwargs: Dict[str, Any] = field(default_factory=dict) @@ -227,6 +233,10 @@ class TrainConfig(BaseConfig): f"reward_model_fn is required for online RL strategy " f"{self.strategy!r}" ) + if self.strategy == "online_ppo" and self.critic_model_fn is None: + raise ValueError( + "critic_model_fn is required for online RL strategy 'online_ppo'" + ) if self.nprocs > 1: raise ValueError( f"online RL strategy {self.strategy!r} requires single-process " diff --git a/astrai/model/__init__.py b/astrai/model/__init__.py index 22b2283..87e0383 100644 --- a/astrai/model/__init__.py +++ b/astrai/model/__init__.py @@ -13,6 +13,7 @@ from astrai.model.components.mlp import MLP, DeepSeekMoE from astrai.model.components.norm import RMSNorm from astrai.model.encoder import EmbeddingEncoder from astrai.model.transformer import AutoRegressiveLM +from astrai.model.value import ValueModel __all__ = [ # Modules @@ -26,6 +27,7 @@ __all__ = [ "AutoRegressiveLM", "EmbeddingEncoder", "AutoModel", + "ValueModel", # LoRA "LoRAConfig", "inject_lora", diff --git a/astrai/model/value.py b/astrai/model/value.py new file mode 100644 index 0000000..de12261 --- /dev/null +++ b/astrai/model/value.py @@ -0,0 +1,55 @@ +"""Value (critic) model for actor-critic RL training.""" + +from typing import Dict, Optional + +import torch.nn as nn +from torch import Tensor + +from astrai.config.model_config import AutoRegressiveLMConfig +from astrai.model.automodel import ModelFactory +from astrai.model.components.linear import Linear +from astrai.model.transformer import AutoRegressiveLM, process_attention_mask + + +@ModelFactory.register("value_model") +class ValueModel(AutoRegressiveLM): + """Critic scoring each state with a scalar instead of vocab logits. + + Inherits the ``AutoRegressiveLM`` components so a policy checkpoint can + warm-start the critic backbone (``load_state_dict(..., strict=False)``); + only ``value_head`` keeps its fresh initialization. The inherited + ``lm_head`` parameters stay dormant — the forward below never projects + through them — so checkpoints round-trip with stable keys. The trunk + pass mirrors ``AutoRegressiveLM.forward`` for training-style input; + ``tests/trainer/test_ppo_strategy.py`` pins the two to identical + hidden states. + """ + + def __init__(self, config: AutoRegressiveLMConfig): + super().__init__(config) + self.value_head = Linear(config.hidden_size, 1, bias=True) + # Zero head so training starts from V(s) == 0 and the first GAE + # advantages are driven purely by rewards. + nn.init.zeros_(self.value_head.weight) + nn.init.zeros_(self.value_head.bias) + + def forward( + self, + input_ids: Tensor, + input_mask: Optional[Tensor] = None, + position_ids: Optional[Tensor] = None, + ) -> Dict[str, Tensor]: + if input_ids.ndim != 2: + raise ValueError("critic input_ids must be [batch, seq_len]") + x = self.embed_tokens(input_ids) + rotary_emb = self.rotary_embedding(x, position_ids) + attn_mask = process_attention_mask(input_mask) + use_sdpa_causal_mask = attn_mask is None + + for layer in self.layers: + x = layer(x, rotary_emb, attn_mask, None, use_sdpa_causal_mask, None)[ + "hidden_states" + ] + hidden_states = self.norm(x) + values = self.value_head(hidden_states).squeeze(-1) + return {"values": values} diff --git a/astrai/trainer/rollout.py b/astrai/trainer/rollout.py index 6686c46..3459cce 100644 --- a/astrai/trainer/rollout.py +++ b/astrai/trainer/rollout.py @@ -68,9 +68,16 @@ class RolloutResult(RawRollout): Fields: rewards: Reward per response, shape ``[B, G]``. + advantages: Optional GAE advantages ``[B, G, R_max]`` pinned at + rollout time by actor-critic strategies (PPO). ``None`` until + a strategy computes them. + returns: Optional GAE value targets ``[B, G, R_max]`` matching + ``advantages``. """ rewards: Tensor + advantages: Optional[Tensor] = None + returns: Optional[Tensor] = None class BaseRewardModel(ABC): diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 13274d9..69619a5 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -1,7 +1,7 @@ """Training strategy implementations with factory pattern.""" from abc import ABC -from typing import Any, Callable, Dict, List, Optional, TypedDict, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, TypedDict, Union import torch import torch.nn as nn @@ -85,6 +85,166 @@ def get_logprobs( } +def rollout_sequences( + prompts: Tensor, + prompt_mask: Tensor, + responses: Tensor, + response_masks: Tensor, +) -> Tuple[Tensor, Tensor]: + """Concatenate grouped prompts with responses for sequence scoring. + + Expands ``prompts`` [B, P] across the group dimension of ``responses`` + [B, G, R] and builds the combined key-padding + causal attention mask. + + Returns: + ``(full_sequences, attn_mask)`` each shaped [B*G, P + R]; the + attention mask is 4-D boolean. + """ + group_size = responses.size(1) + responses_flat = responses.view(-1, responses.size(-1)) + masks_flat = response_masks.view(-1, responses.size(-1)).bool() + prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1) + prompt_mask_expanded = ( + prompt_mask.unsqueeze(1).expand(-1, group_size, -1).flatten(0, 1).bool() + ) + + full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1) + # Build full attention mask: key-padding + causal + key_pad = torch.cat([prompt_mask_expanded, masks_flat], dim=-1)[:, None, None, :] + S = key_pad.shape[-1] + causal = torch.tril( + torch.ones(S, S, dtype=torch.bool, device=full_sequences.device) + )[None, None, :, :] + attn_mask = key_pad & causal + return full_sequences, attn_mask + + +def rollout_token_logprobs( + model: nn.Module, + prompts: Tensor, + prompt_mask: Tensor, + responses: Tensor, + response_masks: Tensor, +) -> LogprobsOutput: + """Per-response-token log probabilities for a grouped rollout batch. + + Prompt tokens are masked out (0) so logprobs are computed only for + response tokens. ``get_logprobs`` shifts the mask by one position, so + the first response token's logprob (predicted from the last prompt + token) is correctly included. + + Returns: + ``logprobs`` reshaped to [B, G, R]: position j is the log-probability + of response token j under ``model``. + """ + batch_size, group_size, response_len = responses.shape + prompt_len = prompts.size(1) + full_sequences, attn_mask = rollout_sequences( + prompts, prompt_mask, responses, response_masks + ) + masks_flat = response_masks.view(-1, response_len) + full_masks = torch.cat( + [ + torch.zeros( + batch_size * group_size, + prompt_len, + dtype=torch.bool, + device=full_sequences.device, + ), + masks_flat, + ], + dim=-1, + ) + + # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len). + # Response token logprobs occupy the last ``response_len`` positions. + output = get_logprobs(model, full_sequences, attn_mask, full_masks, "none") + output["logprobs"] = output["logprobs"][:, prompt_len - 1 :].view( + batch_size, group_size, response_len + ) + return output + + +def rollout_token_values( + model: nn.Module, + prompts: Tensor, + prompt_mask: Tensor, + responses: Tensor, + response_masks: Tensor, +) -> Tensor: + """Critic values [B, G, R] aligned with response token positions. + + Position j holds V(s_j) — the value of the state right before response + token j is emitted — matching the logprob alignment of + :func:`rollout_token_logprobs`. + """ + prompt_len = prompts.size(1) + full_sequences, attn_mask = rollout_sequences( + prompts, prompt_mask, responses, response_masks + ) + output = model(full_sequences, input_mask=attn_mask) + values = output["values"].float()[ + :, prompt_len - 1 : prompt_len - 1 + responses.size(-1) + ] + return values.view(responses.shape) + + +def compute_gae( + rewards: Tensor, + values: Tensor, + mask: Tensor, + gamma: float, + gae_lambda: float, +) -> Tuple[Tensor, Tensor]: + """Generalized advantage estimation over padded response tokens. + + Args: + rewards: [B, G, R] per-token rewards; the terminal reward must sit + at each response's last valid position, padded positions 0. + values: [B, G, R] rollout-time critic values V(s_t) (see + :func:`rollout_token_values`). + mask: [B, G, R] valid-token mask; padded positions are excluded + and cannot leak into valid advantages. + gamma: Discount factor. + gae_lambda: GAE bias/variance trade-off. + + Returns: + ``(advantages, returns)`` shaped [B, G, R]. The episode ends at the + last valid token (no bootstrap value beyond truncation). + """ + response_len = rewards.size(-1) + flat_rewards = rewards.reshape(-1, response_len) + flat_mask = mask.reshape(-1, response_len).to(values.dtype) + # Padded values must be zero or the backward scan would leak them. + flat_values = values.reshape(-1, response_len) * flat_mask + + advantages = torch.zeros_like(flat_rewards) + gae = torch.zeros_like(flat_values[:, 0]) + for t in range(response_len - 1, -1, -1): + next_values = ( + flat_values[:, t + 1] + if t + 1 < response_len + else torch.zeros_like(flat_values[:, t]) + ) + delta = flat_rewards[:, t] + gamma * next_values - flat_values[:, t] + gae = flat_mask[:, t] * (delta + gamma * gae_lambda * gae) + advantages[:, t] = gae + returns = advantages + flat_values + return advantages.view_as(rewards), returns.view_as(rewards) + + +def _validate_behavior_logprobs(behavior_logprobs: Tensor, responses: Tensor) -> None: + """Reject behaviour-policy logprobs that do not match the responses.""" + if behavior_logprobs.shape != responses.shape: + raise ValueError( + "logprobs_old shape must match responses: " + f"got {tuple(behavior_logprobs.shape)}, " + f"expected {tuple(responses.shape)}" + ) + if not torch.isfinite(behavior_logprobs).all(): + raise ValueError("logprobs_old must contain only finite values") + + def make_doc_boundary_mask(position_ids: Tensor) -> Tensor: S = position_ids.size(1) device = position_ids.device @@ -630,82 +790,36 @@ class GRPOStrategy(BaseStrategy): masks = batch["masks"] rewards = batch["rewards"] - batch_size, group_size, response_len = responses.shape behavior_logprobs = batch.get("logprobs_old") if behavior_logprobs is not None: - if behavior_logprobs.shape != responses.shape: - raise ValueError( - "logprobs_old shape must match responses: " - f"got {tuple(behavior_logprobs.shape)}, " - f"expected {tuple(responses.shape)}" - ) - if not torch.isfinite(behavior_logprobs).all(): - raise ValueError("logprobs_old must contain only finite values") + _validate_behavior_logprobs(behavior_logprobs, responses) behavior_logprobs = behavior_logprobs.detach().float() elif self.old_model is None: raise ValueError( "GRPO batches must provide logprobs_old when no old_model is configured" ) - responses_flat = responses.view(-1, response_len) - masks_flat = masks.view(-1, response_len) - prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1) prompt_mask = batch.get("prompt_mask") if prompt_mask is None: prompt_mask = prompts.ne(0) - prompt_mask_expanded = ( - prompt_mask.unsqueeze(1).expand(-1, group_size, -1).flatten(0, 1) - ) - prompt_len = prompt_expanded.size(1) - full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1) - # Prompt tokens are masked out (0) so logprobs are computed only for - # response tokens. get_logprobs shifts the mask by one position, so - # the first response token's logprob (predicted from the last prompt - # token) is correctly included. - full_masks = torch.cat( - [torch.zeros_like(prompt_expanded, dtype=torch.bool), masks_flat], dim=-1 - ) - - # Build full attention mask: key-padding + causal - key_pad = torch.cat([prompt_mask_expanded, masks_flat.bool()], dim=-1)[ - :, None, None, : - ] - S = key_pad.shape[-1] - causal = torch.tril( - torch.ones(S, S, dtype=torch.bool, device=full_sequences.device) - )[None, None, :, :] - attn_mask = key_pad & causal - - # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len). - # Response token logprobs occupy the last ``response_len`` positions - # (the first response token is predicted from the last prompt token). - policy_output = get_logprobs( - self.model, full_sequences, attn_mask, full_masks, "none" + policy_output = rollout_token_logprobs( + self.model, prompts, prompt_mask, responses, masks ) token_log_probs_policy = policy_output["logprobs"] aux_loss = policy_output["aux_loss"] - token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :] with torch.no_grad(): if behavior_logprobs is None: - old_output = get_logprobs( - self.old_model, full_sequences, attn_mask, full_masks, "none" - ) - token_log_probs_old = old_output["logprobs"] - token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :] + token_log_probs_old = rollout_token_logprobs( + self.old_model, prompts, prompt_mask, responses, masks + )["logprobs"] else: token_log_probs_old = behavior_logprobs - ref_output = get_logprobs( - self.ref_model, full_sequences, attn_mask, full_masks, "none" - ) - token_log_probs_ref = ref_output["logprobs"] - token_log_probs_ref = token_log_probs_ref[:, prompt_len - 1 :] + token_log_probs_ref = rollout_token_logprobs( + self.ref_model, prompts, prompt_mask, responses, masks + )["logprobs"] - # Reshape to [B, G, response_len] - token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1) - token_log_probs_old = token_log_probs_old.view(batch_size, group_size, -1) - token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1) - token_masks = masks_flat.view(batch_size, group_size, -1).float() + token_masks = masks.float() # Group-normalized advantages from scalar per-response rewards. eps = 1e-8 @@ -754,8 +868,228 @@ class GRPOStrategy(BaseStrategy): } +@StrategyFactory.register("online_ppo") +class PPOStrategy(BaseStrategy): + """Proximal Policy Optimization with a learned critic (actor-critic). + + Uses the same token-level clipped surrogate as GRPO, but advantages + come from GAE(λ) over a :class:`~astrai.model.value.ValueModel` critic + instead of group-normalized rewards. + + Roles: + + * **Policy** ``self.model`` — the actor being trained. + * **Behaviour policy** — per-token ``logprobs_old`` captured by the + rollout sampler; always required (no offline old-model fallback). + * **Critic** ``self.critic`` — trained jointly by masked MSE regression + against the GAE returns. It owns a separate optimizer, stepped by + :meth:`optimizer_step` *outside* the policy-version lock: the critic + never serves generation, so its weights are not part of a policy + version publication. + * **Reference model** ``self.ref_model`` — optional frozen copy of the + initial policy. When set (and ``kl_coef > 0``), a per-token KL + penalty (k3 estimator, ``logπ_old − logπ_ref``) is folded into the + rewards before GAE, following InstructGPT-style reward shaping. + + Advantages and returns are computed once per rollout — with rollout-time + critic values — and pinned on the :class:`RolloutResult`, so every + replayed gradient step optimizes the same fixed targets, mirroring + classic PPO's multiple epochs over one batch. + """ + + def __init__( + self, + model: nn.Module, + device: str, + critic: nn.Module, + critic_optimizer: Optional[Optimizer] = None, + ref_model: Optional[nn.Module] = None, + clip_eps: float = 0.2, + kl_coef: float = 0.01, + gamma: float = 1.0, + gae_lambda: float = 0.95, + vf_coef: float = 0.5, + max_grad_norm: Optional[float] = None, + **kwargs, + ): + super().__init__(model, device, **kwargs) + self.critic = critic + self.critic_optimizer = critic_optimizer + self.ref_model = ref_model + self.clip_eps = clip_eps + self.kl_coef = kl_coef + self.gamma = gamma + self.gae_lambda = gae_lambda + self.vf_coef = vf_coef + self.max_grad_norm = max_grad_norm + + def optimizer_step(self, optimizer: Optimizer): + """Step the policy under the version lock, then the critic. + + The critic step runs after the policy's atomic version publication: + a concurrent rollout may already observe the new policy version, + but only the (unused-for-generation) critic lags by one step. + """ + result = super().optimizer_step(optimizer) + if self.critic_optimizer is not None: + if self.max_grad_norm is not None: + torch.nn.utils.clip_grad_norm_( + self.critic.parameters(), self.max_grad_norm + ) + self.critic_optimizer.step() + self.critic_optimizer.zero_grad() + return result + + @torch.no_grad() + def _compute_advantages( + self, + prompts: Tensor, + prompt_mask: Tensor, + responses: Tensor, + response_masks: Tensor, + rewards: Tensor, + behavior_logprobs: Tensor, + ) -> Tuple[Tensor, Tensor]: + """GAE advantages/returns pinned to rollout-time critic values.""" + device = self.device + prompts = prompts.to(device) + prompt_mask = prompt_mask.to(device) + responses = responses.to(device) + response_masks = response_masks.to(device) + rewards = rewards.to(device) + behavior_logprobs = behavior_logprobs.to(device) + + values = rollout_token_values( + self.critic, prompts, prompt_mask, responses, response_masks + ) + + # Terminal reward lands on each response's last valid token. + token_rewards = torch.zeros_like(values) + lengths = response_masks.long().sum(dim=-1) + terminal = (lengths - 1).clamp(min=0) + token_rewards.view(-1, values.size(-1)).scatter_( + 1, terminal.reshape(-1, 1), rewards.reshape(-1, 1).to(values.dtype) + ) + + if self.ref_model is not None and self.kl_coef > 0: + ref_logprobs = rollout_token_logprobs( + self.ref_model, prompts, prompt_mask, responses, response_masks + )["logprobs"] + # k3 per-token KL estimator folded into the reward. + kl_penalty = behavior_logprobs.float() - ref_logprobs + token_rewards = ( + token_rewards + - self.kl_coef * kl_penalty * response_masks.to(values.dtype) + ) + + return compute_gae( + token_rewards, values, response_masks, self.gamma, self.gae_lambda + ) + + def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]: + _validate_behavior_logprobs(result.logprobs_old, result.responses) + if result.advantages is None or result.returns is None: + result.advantages, result.returns = self._compute_advantages( + result.prompts, + result.prompt_mask, + result.responses, + result.response_mask, + result.rewards, + result.logprobs_old, + ) + return { + "prompts": result.prompts, + "prompt_mask": result.prompt_mask, + "responses": result.responses, + "masks": result.response_mask, + "rewards": result.rewards, + "logprobs_old": result.logprobs_old, + "advantages": result.advantages, + "returns": result.returns, + } + + def supports_online(self) -> bool: + return True + + def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: + batch = move_to_device(batch, self.device) + prompts = batch["prompts"] + responses = batch["responses"] + masks = batch["masks"] + + behavior_logprobs = batch.get("logprobs_old") + if behavior_logprobs is None: + raise ValueError( + "PPO batches must provide logprobs_old captured at rollout" + ) + _validate_behavior_logprobs(behavior_logprobs, responses) + behavior_logprobs = behavior_logprobs.detach().float() + + prompt_mask = batch.get("prompt_mask") + if prompt_mask is None: + prompt_mask = prompts.ne(0) + + advantages = batch.get("advantages") + returns = batch.get("returns") + if advantages is None or returns is None: + advantages, returns = self._compute_advantages( + prompts, + prompt_mask, + responses, + masks, + batch["rewards"], + behavior_logprobs, + ) + + policy_output = rollout_token_logprobs( + self.model, prompts, prompt_mask, responses, masks + ) + token_log_probs_policy = policy_output["logprobs"] + + # Critic forward with gradients: value regression against the + # rollout-pinned GAE returns. + values = rollout_token_values( + self.critic, prompts, prompt_mask, responses, masks + ) + token_masks = masks.float() + token_count = token_masks.sum().clamp(min=1.0) + + # Token-level ratio (π_θ / π_old) and PPO clipping. + log_ratio = token_log_probs_policy - behavior_logprobs + ratio = torch.exp(log_ratio) + + surr1 = ratio * advantages + surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages + per_token_policy_loss = -torch.min(surr1, surr2) + policy_loss = (per_token_policy_loss * token_masks).sum() / token_count + + value_loss = ((values - returns) ** 2 * token_masks).sum() / token_count + task_loss = policy_loss + self.vf_coef * value_loss + + def masked_variance(x: Tensor) -> Tensor: + mean = (x * token_masks).sum() / token_count + return ((x - mean) ** 2 * token_masks).sum() / token_count + + with torch.no_grad(): + explained_variance = 1.0 - masked_variance( + values - returns + ) / masked_variance(returns).clamp(min=1e-8) + + return self._loss_output( + task_loss, + { + "policy_loss": policy_loss, + "value_loss": value_loss, + "explained_variance": explained_variance, + }, + policy_output["aux_loss"], + policy_output.get("router_stats"), + ) + + # Factory aliases: online variants use the same strategy class; the # ``RolloutRunner`` is injected by ``TrainContextBuilder`` to enable -# online mode, so no separate subclass is needed. +# online mode, so no separate subclass is needed. ``PPOStrategy`` is +# registered under its sole name ``online_ppo`` — it has no offline mode. StrategyFactory.register("online_grpo")(GRPOStrategy) StrategyFactory.register("online_dpo")(DPOStrategy) diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index a18fdcb..7f582e4 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -231,6 +231,12 @@ class CheckpointCallback(TrainCallback): obj = getattr(context, name, None) if obj: extra[name] = obj.state_dict() + critic = getattr(context.strategy, "critic", None) + if critic is not None: + extra["value_model"] = critic.state_dict() + critic_optimizer = getattr(context.strategy, "critic_optimizer", None) + if critic_optimizer is not None: + extra["value_optimizer"] = critic_optimizer.state_dict() return extra diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 37ff5ca..7770df2 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -16,6 +16,7 @@ from astrai.model.components.lora import inject_lora from astrai.parallel.executor import ( BaseExecutor, ExecutorFactory, + broadcast_state_dict, create_ref_model, strip_compile_prefix, ) @@ -297,7 +298,7 @@ class TrainContextBuilder: cfg = self.config kwargs = dict(cfg.strategy_kwargs) kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef) - if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo"): + if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo", "online_ppo"): kwargs["ref_model"] = create_ref_model( cfg.model_fn, executor=executor, @@ -313,6 +314,11 @@ class TrainContextBuilder: ) elif cfg.strategy == "online_grpo": kwargs["old_model"] = None + if cfg.strategy == "online_ppo": + critic, critic_optimizer = self._create_critic(context, executor) + kwargs["critic"] = critic + kwargs["critic_optimizer"] = critic_optimizer + kwargs.setdefault("max_grad_norm", cfg.max_grad_norm) context.strategy = StrategyFactory.create( cfg.strategy, model=context.model, @@ -322,6 +328,63 @@ class TrainContextBuilder: ) return kwargs + def _create_critic( + self, context: TrainContext, executor: BaseExecutor + ) -> tuple[nn.Module, OptimizerProtocol]: + """Build the PPO critic and its optimizer, restoring persisted state. + + The critic backbone warm-starts from the policy weights (standard + actor-critic initialization; the fresh value head is the only + randomly-initialized part). On resume, ``value_model`` / + ``value_optimizer`` checkpoint extras override the warm start — + and their absence is fatal rather than a silent fresh critic. + """ + cfg = self.config + device = get_current_device() + checkpoint = context.checkpoint + if checkpoint is not None: + missing = [ + name + for name in ("value_model", "value_optimizer") + if name not in checkpoint.extra + ] + if missing: + raise ValueError( + "online_ppo resume requires critic state in the " + f"checkpoint; missing extras: {', '.join(missing)}" + ) + + state_dict = executor.unwrap_model(context.model) + if executor.use_distributed: + state_dict = broadcast_state_dict(state_dict) + critic = cfg.critic_model_fn() + if state_dict is not None: + state_dict = strip_compile_prefix(state_dict) + result = critic.load_state_dict(state_dict, strict=False) + if result.unexpected_keys: + raise ValueError( + "critic model received unexpected keys from the policy " + f"state dict: {result.unexpected_keys[:3]}" + ) + unexpected_missing = [ + key for key in result.missing_keys if not key.startswith("value_head.") + ] + if unexpected_missing: + raise ValueError( + "critic backbone is missing policy parameters: " + f"{unexpected_missing[:3]}" + ) + if checkpoint is not None: + critic.load_state_dict(checkpoint.extra["value_model"]) + critic = critic.to(device) + critic.train() + + optimizer_factory = cfg.critic_optimizer_fn or cfg.optimizer_fn + critic_optimizer = optimizer_factory(critic) + if checkpoint is not None: + critic_optimizer.load_state_dict(checkpoint.extra["value_optimizer"]) + return critic, critic_optimizer + def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> None: cfg = self.config if not cfg.strategy.startswith("online_"): diff --git a/docs/developer/docker-training.md b/docs/developer/docker-training.md index 991ee53..9a3a8d5 100644 --- a/docs/developer/docker-training.md +++ b/docs/developer/docker-training.md @@ -186,6 +186,10 @@ scheduler.pt manifest.json ``` +`online_ppo` jobs additionally require `value_model.pt` and +`value_optimizer.pt` (the critic state); the completeness check derives this +from the training config's `train_type`. + New checkpoints write `manifest.json` after every payload file, sync the complete staging directory, and then atomically rename that directory into place. Legacy checkpoints without a manifest remain resumable when the original required files diff --git a/docs/guides/params.md b/docs/guides/params.md index 7dd8ce8..6d9fb9c 100644 --- a/docs/guides/params.md +++ b/docs/guides/params.md @@ -14,7 +14,7 @@ | Parameter | Description | Default | |-----------|-------------|---------| | `--config`, `-c` | YAML config file; explicit CLI options override YAML values | None | -| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`) | required | +| `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`, `online_grpo`, `online_dpo`, `online_ppo`) | required | | `--data_root_path` | Dataset root directory | required | | `--param_path` | Model parameters or checkpoint path | required | | `--resume` | Resume training from `--param_path` | False | @@ -139,16 +139,22 @@ with `--optimizer=muon_adamw`. |-----------|-------------|---------|---------| | `--dpo_beta` | DPO beta value | 0.1 | `dpo`, `online_dpo` | | `--label_smoothing` | Label smoothing for cross-entropy loss | 0.0 | `seq`, `sft` | -| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo` | -| `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo`, `online_grpo` | -| `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo`, `online_grpo` | +| `--group_size` | GRPO/rollout group size | 4 | `grpo`, `online_grpo`, `online_dpo`, `online_ppo` | +| `--grpo_clip_eps` | Clipping epsilon for the PPO-style surrogate loss | 0.2 | `grpo`, `online_grpo`, `online_ppo` | +| `--grpo_kl_coef` | KL penalty coefficient | 0.01 | `grpo`, `online_grpo`, `online_ppo` | +| `--ppo_gamma` | PPO reward discount factor | 1.0 | `online_ppo` | +| `--ppo_gae_lambda` | PPO GAE bias/variance trade-off | 0.95 | `online_ppo` | +| `--ppo_vf_coef` | PPO value-loss coefficient | 0.5 | `online_ppo` | | `--neftune_alpha` | NEFTune noise alpha (0=disabled, typical: 5.0) | 0.0 | `sft` | ### Online Rollout `online_grpo` and `online_dpo` are factory aliases for the existing `grpo` and `dpo` strategy classes; online behavior is enabled by rollout components rather -than separate strategy subclasses. These options apply to the online aliases. +than separate strategy subclasses. `online_ppo` is a dedicated actor-critic +strategy: a `ValueModel` critic supplies GAE advantages, and its state persists +as `value_model.pt`/`value_optimizer.pt` checkpoint extras (required for +resume). These options apply to the online strategies. Online strategies require a `BaseRewardModel` factory in `TrainConfig`; `train.py` does not currently provide a command-line option for configuring one. diff --git a/docs/guides/training.md b/docs/guides/training.md index bba80aa..fa3b016 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -170,6 +170,18 @@ them with a `BaseRewardModel`. It refreshes cached rollouts every behaviour log-probabilities into the loss, so it does not allocate or synchronize a separate old-policy model. +`online_ppo` is actor-critic PPO on the same rollout pipeline. A `ValueModel` +critic (backbone warm-started from the policy, zero-initialized value head) +scores the rollout states; advantages come from GAE(`--ppo_gamma`, +`--ppo_gae_lambda`) with the terminal reward on each response's last token and +the reference-KL penalty (k3 estimator, `--grpo_kl_coef`) folded into per-token +rewards. Advantages and returns are computed once per rollout and pinned on the +`RolloutResult`, so replayed steps optimize fixed targets. The critic has its +own optimizer, stepped outside the policy-version lock, and persists as +`value_model.pt`/`value_optimizer.pt` checkpoint extras — resume without them +fails loudly, and `scripts/train.sh` treats a PPO checkpoint as incomplete when +they are missing. + Every successful optimizer step mutates the shared model and advances its monotonic `policy_version` under the same generation lock. The scheduler invalidates reusable KV prefixes before accepting the new version, so an async diff --git a/scripts/docker/lib/train-common.sh b/scripts/docker/lib/train-common.sh index 8d6f56a..c1dfab1 100755 --- a/scripts/docker/lib/train-common.sh +++ b/scripts/docker/lib/train-common.sh @@ -22,13 +22,26 @@ validate_job_name() { die "Invalid TRAIN_JOB_NAME '$1'; use letters, numbers, dot, underscore, or dash" } +checkpoint_extra_files() { + # Additional files a complete checkpoint must contain for the strategy + # configured in the given training YAML. PPO persists critic state as + # checkpoint extras (value_model.pt / value_optimizer.pt); a resume + # without them must not look complete. + local config="$1" + + [[ -n "${config}" && -f "${config}" ]] || return 0 + if grep -Eq '^[[:space:]]*train_type:[[:space:]]*["'\'']?online_ppo' "${config}"; then + printf 'value_model.pt value_optimizer.pt' + fi +} + checkpoint_is_complete() { local checkpoint="$1" local file [[ -d "${checkpoint}" ]] || return 1 - for file in meta.json config.json model.safetensors optimizer.pt scheduler.pt; do + for file in meta.json config.json model.safetensors optimizer.pt scheduler.pt ${CHECKPOINT_EXTRA_FILES:-}; do [[ -s "${checkpoint}/${file}" ]] || return 1 done diff --git a/scripts/docker/train-entrypoint.sh b/scripts/docker/train-entrypoint.sh index 8d65a08..0e9754e 100755 --- a/scripts/docker/train-entrypoint.sh +++ b/scripts/docker/train-entrypoint.sh @@ -34,6 +34,7 @@ fi if [[ -n "${TRAIN_CONFIG}" ]]; then [[ -f "${TRAIN_CONFIG}" ]] || die "Training config not found: ${TRAIN_CONFIG}" fi +export CHECKPOINT_EXTRA_FILES="$(checkpoint_extra_files "${TRAIN_CONFIG}")" [[ -r /data ]] || die "Training data directory is not readable: /data" mkdir -p "${CHECKPOINT_DIR}" diff --git a/scripts/tools/train.py b/scripts/tools/train.py index b0d6759..80f7fee 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -21,7 +21,7 @@ from astrai.config.train_config import ( TRAIN_TYPES, ) from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn -from astrai.model import AutoRegressiveLM +from astrai.model import AutoRegressiveLM, ValueModel from astrai.model.components.decoder_block import DecoderBlock from astrai.optim import OptimizerFactory from astrai.trainer import SchedulerFactory, Trainer @@ -212,6 +212,27 @@ _SPECS = [ default=0.01, help="GRPO KL penalty coefficient.", ), + OptSpec( + "ppo_gamma", + "Algorithm", + type=float, + default=1.0, + help="PPO reward discount factor.", + ), + OptSpec( + "ppo_gae_lambda", + "Algorithm", + type=float, + default=0.95, + help="PPO GAE bias/variance trade-off.", + ), + OptSpec( + "ppo_vf_coef", + "Algorithm", + type=float, + default=0.5, + help="PPO value-loss coefficient.", + ), OptSpec( "moe_aux_loss_coef", "Algorithm", @@ -408,6 +429,10 @@ def create_model(config): return AutoRegressiveLM(config).to(dtype=torch.bfloat16) +def create_value_model(config): + return ValueModel(config).to(dtype=torch.bfloat16) + + def create_optimizer( model, optimizer_name: str = "muon_adamw", **kwargs ) -> optim.Optimizer: @@ -502,6 +527,9 @@ def train( "clip_eps": kwargs.pop("grpo_clip_eps"), "kl_coef": kwargs.pop("grpo_kl_coef"), "group_size": kwargs.pop("group_size"), + "gamma": kwargs.pop("ppo_gamma"), + "gae_lambda": kwargs.pop("ppo_gae_lambda"), + "vf_coef": kwargs.pop("ppo_vf_coef"), } rollout_interval = kwargs.pop("rollout_interval", 512) @@ -511,6 +539,11 @@ def train( rollout_top_p = kwargs.pop("rollout_top_p", 0.9) rollout_max_tokens = kwargs.pop("rollout_max_tokens", 1024) reward_model_fn: Callable[[], BaseRewardModel] | None = None + critic_model_fn = None + if train_type == "online_ppo": + # The optimizer defaults to the policy's; critic_optimizer_fn can + # override it in the TrainConfig. + critic_model_fn = partial(create_value_model, config) executor_kwargs = {} if parallel_mode == "ddp": @@ -622,7 +655,7 @@ def train( collate_fn = dpo_collate_fn elif train_type == "grpo": collate_fn = grpo_collate_fn - elif train_type in ("online_grpo", "online_dpo"): + elif train_type in ("online_grpo", "online_dpo", "online_ppo"): collate_fn = None train_config = TrainConfig( @@ -668,6 +701,7 @@ def train( rollout_top_p=rollout_top_p, rollout_max_tokens=rollout_max_tokens, reward_model_fn=reward_model_fn, + critic_model_fn=critic_model_fn, moe_aux_loss_coef=kwargs.pop("moe_aux_loss_coef", 0.01), ) diff --git a/scripts/train.sh b/scripts/train.sh index 8b78511..89a60ea 100755 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -55,6 +55,7 @@ load_config() { die "Failed to load runtime configuration" eval "${exports}" validate_job_name "${TRAIN_JOB_NAME}" + export CHECKPOINT_EXTRA_FILES="$(checkpoint_extra_files "${CONFIG_FILE}")" } compose() { diff --git a/tests/trainer/test_online_e2e.py b/tests/trainer/test_online_e2e.py index e88e650..fc739c1 100644 --- a/tests/trainer/test_online_e2e.py +++ b/tests/trainer/test_online_e2e.py @@ -10,6 +10,7 @@ from torch.utils.data import Dataset import astrai.trainer.train_context as train_context from astrai.config import TrainConfig from astrai.model.transformer import AutoRegressiveLM +from astrai.model.value import ValueModel from astrai.serialization import Checkpoint from astrai.trainer.rollout import BaseRewardModel from astrai.trainer.schedule import SchedulerFactory @@ -67,6 +68,10 @@ def _model_fn(model_config): return AutoRegressiveLM(model_config).to(dtype=torch.float32) +def _value_model_fn(model_config): + return ValueModel(model_config).to(dtype=torch.float32) + + def _optimizer_fn(m): return torch.optim.AdamW(m.parameters(), lr=1e-4) @@ -81,16 +86,32 @@ _ONLINE_STRATEGIES = [ pytest.param( "online_grpo", {"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2}, + None, id="grpo", ), - pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, id="dpo"), + pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, None, id="dpo"), + pytest.param( + "online_ppo", + { + "clip_eps": 0.2, + "kl_coef": 0.01, + "group_size": 2, + "gamma": 1.0, + "gae_lambda": 0.95, + "vf_coef": 0.5, + }, + True, + id="ppo", + ), ] @pytest.mark.integration -@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES) +@pytest.mark.parametrize( + ("strategy", "strategy_kwargs", "with_critic"), _ONLINE_STRATEGIES +) def test_online_rollout_end_to_end( - base_test_env, strategy, strategy_kwargs, monkeypatch + base_test_env, strategy, strategy_kwargs, with_critic, monkeypatch ): """Run one epoch of online RL rollout with KV-cache-backed generation.""" created_reference_models = [] @@ -110,7 +131,7 @@ def test_online_rollout_end_to_end( tokenizer.set_chat_template(CHAT_TEMPLATE) tokenizer.save_pretrained(test_dir) - train_config = TrainConfig( + config_kwargs = dict( strategy=strategy, model_fn=partial(_model_fn, model_config), dataset=InstructionDataset(), @@ -135,6 +156,10 @@ def test_online_rollout_end_to_end( reward_model_fn=LengthRewardModel, collate_fn=instruction_collate_fn, ) + if with_critic: + config_kwargs["critic_model_fn"] = partial(_value_model_fn, model_config) + config_kwargs["critic_optimizer_fn"] = _optimizer_fn + train_config = TrainConfig(**config_kwargs) trainer = Trainer(train_config) trainer.train(param_path=test_dir) @@ -144,6 +169,11 @@ def test_online_rollout_end_to_end( checkpoint = Checkpoint.load(checkpoint_dir) assert checkpoint.meta["policy_version"] == 2 assert len(created_reference_models) == 1 + if with_critic: + assert "value_model" in checkpoint.extra + assert "value_optimizer" in checkpoint.extra + else: + assert "value_model" not in checkpoint.extra def _minimal_online_config(**overrides): diff --git a/tests/trainer/test_ppo_strategy.py b/tests/trainer/test_ppo_strategy.py new file mode 100644 index 0000000..42c5486 --- /dev/null +++ b/tests/trainer/test_ppo_strategy.py @@ -0,0 +1,469 @@ +"""Unit tests for PPO: GAE numerics, the ValueModel critic, and PPOStrategy.""" + +import pytest +import torch + +import astrai.trainer.strategy as strategy_module +from astrai.model.transformer import AutoRegressiveLM +from astrai.model.value import ValueModel +from astrai.trainer.rollout import RolloutResult +from astrai.trainer.strategy import ( + PPOStrategy, + StrategyFactory, + compute_gae, +) +from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config + + +def _make_batch( + batch_size=2, group_size=4, prompt_len=8, response_len=12, device="cpu" +): + """Construct a PPO batch with deterministic shapes. + + Returns dict with prompts [B, P], responses [B, G, R], masks [B, G, R], + rewards [B, G], logprobs_old [B, G, R]. + """ + return { + "prompts": torch.randint(0, 200, (batch_size, prompt_len), device=device), + "responses": torch.randint( + 0, 200, (batch_size, group_size, response_len), device=device + ), + "masks": torch.ones(batch_size, group_size, response_len, device=device), + "rewards": torch.randn(batch_size, group_size, device=device), + "logprobs_old": torch.zeros( + batch_size, group_size, response_len, device=device + ), + } + + +def _make_value_model(policy_model, device): + """Build a ValueModel whose backbone warm-starts from the policy.""" + critic = ValueModel(policy_model.config).to(device=device) + result = critic.load_state_dict(policy_model.state_dict(), strict=False) + assert not result.unexpected_keys + assert all(key.startswith("value_head.") for key in result.missing_keys) + return critic + + +@pytest.fixture +def ppo_strategy(device): + model, _ = make_model(device) + critic = _make_value_model(model, device) + strategy = PPOStrategy( + model=model, + device=device, + critic=critic, + critic_optimizer=torch.optim.AdamW(critic.parameters(), lr=1e-3), + ref_model=make_frozen(model, device), + clip_eps=0.2, + kl_coef=0.01, + gamma=1.0, + gae_lambda=0.95, + vf_coef=0.5, + executor=FakeExecutor(), + ) + return strategy, device + + +# ============== compute_gae ============== + + +def test_gae_monte_carlo_when_values_zero(device): + """γ=1, λ=1, V=0: advantage and return equal the terminal reward at + every valid position (Monte-Carlo return).""" + B, G, R = 2, 3, 4 + rewards = torch.zeros(B, G, R, device=device) + rewards[..., -1] = 1.0 + values = torch.zeros(B, G, R, device=device) + mask = torch.ones(B, G, R, dtype=torch.bool, device=device) + + advantages, returns = compute_gae(rewards, values, mask, gamma=1.0, gae_lambda=1.0) + + assert torch.allclose(advantages, torch.full_like(rewards, 1.0)) + assert torch.allclose(returns, torch.full_like(rewards, 1.0)) + + +def test_gae_lambda_zero_is_one_step_td(device): + """λ=0: advantage degenerates to the TD residual δ_t.""" + torch.manual_seed(0) + rewards = torch.zeros(1, 1, 3, device=device) + rewards[0, 0, -1] = 2.0 + values = torch.tensor([[[0.5, 1.0, -0.5]]], device=device) + mask = torch.ones(1, 1, 3, dtype=torch.bool, device=device) + + advantages, returns = compute_gae(rewards, values, mask, gamma=0.9, gae_lambda=0.0) + + # δ_2 = r + 0 - V_2 = 2.5; δ_1 = 0 + 0.9·V_2 - V_1 = -1.45; + # δ_0 = 0 + 0.9·V_1 - V_0 = 0.4 + expected = torch.tensor([[[0.4, -1.45, 2.5]]], device=device) + assert torch.allclose(advantages, expected, atol=1e-6) + assert torch.allclose(returns, advantages + values, atol=1e-6) + + +def test_gae_hand_computed_discounted_case(device): + """γ=0.9, λ=0.8 against a hand-rolled backward accumulation.""" + rewards = torch.tensor([[[0.0, 0.0, 1.0]]], device=device) + values = torch.tensor([[[0.1, 0.2, 0.3]]], device=device) + mask = torch.ones(1, 1, 3, dtype=torch.bool, device=device) + gamma, lam = 0.9, 0.8 + + advantages, returns = compute_gae(rewards, values, mask, gamma, lam) + + delta2 = 1.0 + 0.0 - 0.3 + gae2 = delta2 + delta1 = 0.0 + gamma * 0.3 - 0.2 + gae1 = delta1 + gamma * lam * gae2 + delta0 = 0.0 + gamma * 0.2 - 0.1 + gae0 = delta0 + gamma * lam * gae1 + expected = torch.tensor([[[gae0, gae1, gae2]]], device=device) + assert torch.allclose(advantages, expected, atol=1e-6) + assert torch.allclose(returns, expected + values, atol=1e-6) + + +def test_gae_padding_does_not_leak(device): + """Garbage values at padded positions must not change valid outputs.""" + torch.manual_seed(1) + B, G, R = 2, 2, 5 + rewards = torch.zeros(B, G, R, device=device) + rewards[0, 0, 2] = 1.0 # terminal at position 2 of a length-3 response + values = torch.randn(B, G, R, device=device) + mask = torch.ones(B, G, R, dtype=torch.bool, device=device) + mask[0, 0, 3:] = False + mask[1, :, 2:] = False + rewards[0, 0, 3:] = 100.0 # reward garbage in padding must be ignored + + advantages, returns = compute_gae(rewards, values, mask, gamma=0.9, gae_lambda=0.9) + + assert torch.allclose(advantages[0, 0, 3:], torch.zeros_like(advantages[0, 0, 3:])) + assert torch.allclose(returns[0, 0, 3:], torch.zeros_like(returns[0, 0, 3:])) + # The terminal reward at position 2 still drives a finite advantage. + assert advantages[0, 0, 2] != 0.0 + + +def test_gae_empty_response_is_all_zero(device): + """A fully padded response yields zero advantages and returns.""" + rewards = torch.zeros(1, 1, 3, device=device) + values = torch.randn(1, 1, 3, device=device) + mask = torch.zeros(1, 1, 3, dtype=torch.bool, device=device) + + advantages, returns = compute_gae(rewards, values, mask, 1.0, 0.95) + + assert torch.count_nonzero(advantages) == 0 + assert torch.count_nonzero(returns) == 0 + + +# ============== ValueModel ============== + + +def test_value_model_trunk_matches_policy_hidden_states(device): + """ValueModel's trunk reproduces AutoRegressiveLM's hidden states. + + Pins the duplicated trunk pass in ``ValueModel.forward`` to the policy + forward: a ones-initialized value head must return the row-wise sum of + the policy's ``hidden_states``. + """ + model, _ = make_model(device) + critic = _make_value_model(model, device) + with torch.no_grad(): + critic.value_head.weight.fill_(1.0) + critic.value_head.bias.zero_() + + torch.manual_seed(2) + input_ids = torch.randint(0, 200, (2, 10), device=device) + input_mask = torch.ones(2, 10, dtype=torch.bool, device=device) + input_mask[1, :3] = False + + with torch.no_grad(): + policy_hidden = model(input_ids, input_mask=input_mask)["hidden_states"] + values = critic(input_ids, input_mask=input_mask)["values"] + + assert values.shape == (2, 10) + assert torch.allclose(values, policy_hidden.sum(dim=-1), atol=1e-5) + + +def test_value_model_zero_head_outputs_zero(device): + model, _ = make_model(device) + critic = _make_value_model(model, device) + input_ids = torch.randint(0, 200, (2, 8), device=device) + with torch.no_grad(): + values = critic(input_ids)["values"] + assert torch.count_nonzero(values) == 0 + + +def test_value_model_rejects_packed_inference_input(device): + critic = ValueModel(make_rollout_config()).to(device=device) + with pytest.raises(ValueError, match="critic input_ids"): + critic(torch.randint(0, 200, (16,), device=device)) + + +# ============== PPOStrategy ============== + + +def test_factory_registers_online_ppo(): + assert StrategyFactory.is_registered("online_ppo") + assert StrategyFactory.get_component_class("online_ppo") is PPOStrategy + + +def test_ppo_supports_online(ppo_strategy): + strategy, _ = ppo_strategy + assert strategy.supports_online() is True + + +def test_ppo_loss_is_finite_and_differentiable(ppo_strategy): + strategy, device = ppo_strategy + batch = _make_batch(device=device) + loss = strategy.compute_loss(batch) + assert loss.dim() == 0 + assert torch.isfinite(loss).item() + loss.backward() + assert any( + p.grad is not None and p.grad.abs().sum().item() > 0 + for p in strategy.model.parameters() + ) + assert any( + p.grad is not None and p.grad.abs().sum().item() > 0 + for p in strategy.critic.parameters() + ) + + +def test_ppo_requires_behavior_logprobs(ppo_strategy): + strategy, device = ppo_strategy + batch = _make_batch(device=device) + del batch["logprobs_old"] + with pytest.raises(ValueError, match="logprobs_old"): + strategy.compute_loss(batch) + + +@pytest.mark.parametrize("invalid", ["shape", "nonfinite"]) +def test_ppo_rejects_invalid_behavior_logprobs(ppo_strategy, invalid): + strategy, device = ppo_strategy + batch = _make_batch(device=device) + if invalid == "shape": + batch["logprobs_old"] = torch.zeros(1, device=device) + match = "shape must match responses" + else: + batch["logprobs_old"] = torch.zeros_like(batch["responses"], dtype=torch.float) + batch["logprobs_old"][0, 0, 0] = float("nan") + match = "only finite values" + with pytest.raises(ValueError, match=match): + strategy.compute_loss(batch) + + +def test_ppo_ref_model_not_updated_by_backward(ppo_strategy): + strategy, device = ppo_strategy + loss = strategy.compute_loss(_make_batch(device=device)) + loss.backward() + for p in strategy.ref_model.parameters(): + assert p.grad is None + + +def test_ppo_zero_advantage_and_zero_critic_gives_zero_loss(ppo_strategy): + """A zero-head critic, zero advantages, and zero returns → zero loss.""" + strategy, device = ppo_strategy + batch = _make_batch(device=device) + batch["advantages"] = torch.zeros_like(batch["responses"], dtype=torch.float) + batch["returns"] = torch.zeros_like(batch["responses"], dtype=torch.float) + loss = strategy.compute_loss(batch) + assert loss.item() == pytest.approx(0.0, abs=1e-6) + + +def test_ppo_all_masked_response_tokens_zero_loss(ppo_strategy): + strategy, device = ppo_strategy + batch = _make_batch(device=device) + batch["masks"] = torch.zeros_like(batch["masks"]) + loss = strategy.compute_loss(batch) + assert loss.item() == pytest.approx(0.0, abs=1e-6) + + +def test_ppo_uses_supplied_advantages_without_recomputation(ppo_strategy): + """Explicit advantages/returns must short-circuit GAE computation.""" + strategy, device = ppo_strategy + batch = _make_batch(device=device) + + def _fail(*args, **kwargs): + raise AssertionError("advantages were supplied; GAE must not run") + + strategy._compute_advantages = _fail + batch["advantages"] = torch.ones_like(batch["responses"], dtype=torch.float) + batch["returns"] = torch.zeros_like(batch["responses"], dtype=torch.float) + loss = strategy.compute_loss(batch) + assert torch.isfinite(loss).item() + + +def test_ppo_optimizer_step_updates_policy_and_critic(ppo_strategy): + """optimizer_step steps the policy optimizer and then the critic's.""" + strategy, device = ppo_strategy + batch = _make_batch(device=device) + loss = strategy.compute_loss(batch) + loss.backward() + + policy_optimizer = torch.optim.SGD(strategy.model.parameters(), lr=0.1) + policy_before = next(strategy.model.parameters()).detach().clone() + critic_before = next(strategy.critic.parameters()).detach().clone() + + strategy.optimizer_step(policy_optimizer) + + assert not torch.equal(next(strategy.model.parameters()), policy_before) + assert not torch.equal(next(strategy.critic.parameters()), critic_before) + # Critic gradients are cleared after its step. + assert all(p.grad is None for p in strategy.critic.parameters()) + + +def test_ppo_optimizer_step_clips_critic_gradients(ppo_strategy): + strategy, device = ppo_strategy + batch = _make_batch(device=device) + loss = strategy.compute_loss(batch) + loss.backward() + strategy.critic_optimizer = torch.optim.SGD(strategy.critic.parameters(), lr=1.0) + strategy.max_grad_norm = 1e-8 + before = next(strategy.critic.parameters()).detach().clone() + + strategy.optimizer_step(torch.optim.SGD(strategy.model.parameters(), lr=0.0)) + + # Clipped-to-zero critic gradients under SGD (no momentum) leave the + # parameters unchanged. + assert torch.equal(next(strategy.critic.parameters()), before) + + +# ============== prepare_from_rollout / GAE integration ============== + + +def _make_rollout_result(B=2, G=2, P=6, R=5, device="cpu"): + return RolloutResult( + prompts=torch.randint(3, 200, (B, P), device=device), + prompt_mask=torch.ones(B, P, dtype=torch.bool, device=device), + responses=torch.randint(3, 200, (B, G, R), device=device), + response_mask=torch.ones(B, G, R, dtype=torch.bool, device=device), + rewards=torch.randn(B, G, device=device), + logprobs_old=torch.zeros(B, G, R, device=device), + ) + + +def test_prepare_from_rollout_computes_and_pins_gae(ppo_strategy, monkeypatch): + """prepare attaches GAE tensors to the result once, then reuses them.""" + strategy, device = ppo_strategy + result = _make_rollout_result(device=device) + + batch = strategy.prepare_from_rollout(result) + assert result.advantages is not None and result.returns is not None + assert batch["advantages"] is result.advantages + assert batch["returns"] is result.returns + assert batch["advantages"].shape == result.responses.shape + + calls = [] + original = strategy._compute_advantages + monkeypatch.setattr( + strategy, + "_compute_advantages", + lambda *a, **k: calls.append(1) or original(*a, **k), + ) + strategy.prepare_from_rollout(result) + assert not calls, "pinned advantages must not be recomputed on replay" + + +def test_prepare_from_rollout_respects_response_padding(ppo_strategy): + """Padded response positions get zero advantages and returns.""" + strategy, device = ppo_strategy + result = _make_rollout_result(device=device) + result.response_mask[0, 0, 3:] = False + + batch = strategy.prepare_from_rollout(result) + + assert torch.count_nonzero(batch["advantages"][0, 0, 3:]) == 0 + assert torch.count_nonzero(batch["returns"][0, 0, 3:]) == 0 + assert torch.count_nonzero(batch["advantages"][0, 0, :3]) > 0 + + +def test_compute_advantages_matches_hand_computed_gae(ppo_strategy, monkeypatch): + """_compute_advantages applies terminal rewards and GAE faithfully.""" + strategy, device = ppo_strategy + result = _make_rollout_result(B=1, G=1, P=4, R=3, device=device) + result.rewards = torch.tensor([[2.0]], device=device) + result.logprobs_old = torch.zeros(1, 1, 3, device=device) + # ref_model == policy at init → zero KL reward shaping only if the + # policy and ref agree on logprobs; keep ref out of the picture here. + strategy.ref_model = None + + fixed_values = torch.tensor([[[0.1, 0.2, 0.3]]], device=device) + monkeypatch.setattr( + strategy_module, + "rollout_token_values", + lambda *args, **kwargs: fixed_values.clone(), + ) + + advantages, returns = strategy._compute_advantages( + result.prompts, + result.prompt_mask, + result.responses, + result.response_mask, + result.rewards, + result.logprobs_old, + ) + + rewards = torch.tensor([[[0.0, 0.0, 2.0]]], device=device) + expected_adv, expected_ret = compute_gae( + rewards, fixed_values, result.response_mask, 1.0, 0.95 + ) + assert torch.allclose(advantages, expected_adv, atol=1e-6) + assert torch.allclose(returns, expected_ret, atol=1e-6) + + +def test_compute_advantages_folds_kl_penalty_into_rewards(ppo_strategy, monkeypatch): + """With a ref model, each valid token's reward loses kl_coef·k3.""" + strategy, device = ppo_strategy + result = _make_rollout_result(B=1, G=1, P=4, R=2, device=device) + result.rewards = torch.tensor([[1.0]], device=device) + # behaviour policy disagrees with ref by +1 logprob on every token + result.logprobs_old = torch.ones(1, 1, 2, device=device) + + fixed_values = torch.zeros(1, 1, 2, device=device) + fixed_ref_logprobs = torch.zeros(1, 1, 2, device=device) + monkeypatch.setattr( + strategy_module, + "rollout_token_values", + lambda *args, **kwargs: fixed_values.clone(), + ) + monkeypatch.setattr( + strategy_module, + "rollout_token_logprobs", + lambda *args, **kwargs: {"logprobs": fixed_ref_logprobs.clone()}, + ) + + advantages, _ = strategy._compute_advantages( + result.prompts, + result.prompt_mask, + result.responses, + result.response_mask, + result.rewards, + result.logprobs_old, + ) + + # per-token reward = -kl_coef·(1 - 0) = -0.01; terminal adds 1.0 + expected_rewards = torch.tensor([[-0.01, 0.99]], device=device) + expected_adv, _ = compute_gae( + expected_rewards.unsqueeze(0), fixed_values, result.response_mask, 1.0, 0.95 + ) + assert torch.allclose(advantages, expected_adv, atol=1e-6) + + +def test_online_call_returns_finite_loss(ppo_strategy): + strategy, device = ppo_strategy + + class _RecordingRunner: + policy_version = 0 + + def __call__(self, batch): + return _make_rollout_result(device=device), True + + def step(self): + pass + + def apply_weight_update(self, policy_version, update): + return update() + + strategy.set_rollout_runner(_RecordingRunner()) + out = strategy({"instruction": ["x"]}) + assert torch.isfinite(out["loss"]).item() + assert "policy_loss" in out["metrics"] + assert "value_loss" in out["metrics"] + assert "explained_variance" in out["metrics"] diff --git a/tests/trainer/test_ppo_train_context.py b/tests/trainer/test_ppo_train_context.py new file mode 100644 index 0000000..9c9ef2a --- /dev/null +++ b/tests/trainer/test_ppo_train_context.py @@ -0,0 +1,271 @@ +"""Wiring tests for online PPO: config validation, critic assembly, and +checkpoint round-trip of critic state.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch + +from astrai.config import TrainConfig +from astrai.model.transformer import AutoRegressiveLM +from astrai.model.value import ValueModel +from astrai.serialization import Checkpoint +from astrai.trainer.rollout import BaseRewardModel +from astrai.trainer.schedule import SchedulerFactory +from astrai.trainer.train_callback import CheckpointCallback +from astrai.trainer.train_context import TrainContext, TrainContextBuilder +from astrai.trainer.trainer import Trainer +from tests.helpers import ( + FakeExecutor, + build_test_tokenizer, + make_model, + make_rollout_config, +) + + +class _StubRewardModel(BaseRewardModel): + def score(self, prompts, responses): + return torch.zeros(len(prompts), len(responses[0]) if prompts else 0) + + +class _StubDataset(torch.utils.data.Dataset): + def __len__(self): + return 2 + + def __getitem__(self, idx): + return {"instruction": "hello", "input": ""} + + +def _stub_collate(batch): + return { + "instruction": [b["instruction"] for b in batch], + "input": [b.get("input", "") for b in batch], + } + + +def _ppo_config(device, **overrides): + defaults = dict( + strategy="online_ppo", + model_fn=lambda: AutoRegressiveLM(make_rollout_config()), + dataset=_StubDataset(), + optimizer_fn=lambda m: torch.optim.SGD(m.parameters(), lr=0.0), + scheduler_fn=lambda o: SchedulerFactory.create( + "cosine", o, warmup_steps=1, lr_decay_steps=4, min_rate=0.05 + ), + reward_model_fn=_StubRewardModel, + critic_model_fn=lambda: ValueModel(make_rollout_config()), + collate_fn=_stub_collate, + device_type=device, + nprocs=1, + parallel_mode="none", + strategy_kwargs={"clip_eps": 0.2, "group_size": 2}, + rollout_interval=1, + rollout_max_policy_lag=0, + rollout_max_tokens=4, + rollout_temperature=1.0, + ) + defaults.update(overrides) + return TrainConfig(**defaults) + + +def test_online_ppo_config_requires_critic_model_fn(device): + with pytest.raises(ValueError, match="critic_model_fn is required"): + _ppo_config(device, critic_model_fn=None) + + +def test_online_ppo_config_accepts_critic(device): + config = _ppo_config(device) + assert config.strategy == "online_ppo" + + +def test_create_critic_warm_starts_backbone_from_policy(device, monkeypatch): + monkeypatch.setenv("LOCAL_DEVICE", device) + model, config = make_model(device) + cfg = _ppo_config(device) + builder = TrainContextBuilder(cfg) + context = TrainContext(model=model) + + critic, _ = builder._create_critic(context, FakeExecutor()) + + policy_sd = model.state_dict() + critic_sd = critic.state_dict() + for key in policy_sd: + assert torch.equal(critic_sd[key], policy_sd[key]) + assert torch.count_nonzero(critic_sd["value_head.weight"]) == 0 + assert torch.count_nonzero(critic_sd["value_head.bias"]) == 0 + + +def test_create_critic_restores_checkpoint_extras(device, monkeypatch): + monkeypatch.setenv("LOCAL_DEVICE", device) + model, config = make_model(device) + cfg = _ppo_config(device) + builder = TrainContextBuilder(cfg) + + saved_critic = ValueModel(config).to(device) + with torch.no_grad(): + saved_critic.value_head.weight.fill_(1.0) + saved_optimizer = torch.optim.SGD(saved_critic.parameters(), lr=0.1) + checkpoint = Checkpoint( + state_dict=model.state_dict(), + config=config.to_dict(), + extra={ + "optimizer": {}, + "scheduler": {}, + "value_model": saved_critic.state_dict(), + "value_optimizer": saved_optimizer.state_dict(), + }, + ) + context = TrainContext(model=model, checkpoint=checkpoint) + + critic, critic_optimizer = builder._create_critic(context, FakeExecutor()) + + assert torch.equal( + critic.state_dict()["value_head.weight"], + saved_critic.state_dict()["value_head.weight"], + ) + assert ( + critic_optimizer.state_dict()["param_groups"] + == saved_optimizer.state_dict()["param_groups"] + ) + + +def test_create_critic_resume_without_extras_fails_loudly(device, monkeypatch): + monkeypatch.setenv("LOCAL_DEVICE", device) + model, _ = make_model(device) + cfg = _ppo_config(device) + builder = TrainContextBuilder(cfg) + checkpoint = Checkpoint( + state_dict=model.state_dict(), + extra={"optimizer": {}, "scheduler": {}}, + ) + context = TrainContext(model=model, checkpoint=checkpoint) + + with pytest.raises( + ValueError, match="missing extras: value_model, value_optimizer" + ): + builder._create_critic(context, FakeExecutor()) + + +def test_builder_resumes_critic_from_checkpoint(device, temp_dir, monkeypatch): + """A full TrainContextBuilder resume restores the persisted critic.""" + monkeypatch.setenv("LOCAL_DEVICE", device) + model, config = make_model(device) + saved_critic = ValueModel(config).to(device) + with torch.no_grad(): + saved_critic.value_head.weight.fill_(2.0) + saved_optimizer = torch.optim.SGD(saved_critic.parameters(), lr=0.1) + policy_optimizer = torch.optim.SGD(model.parameters(), lr=0.0) + policy_scheduler = SchedulerFactory.create( + "cosine", policy_optimizer, warmup_steps=1, lr_decay_steps=4, min_rate=0.05 + ) + checkpoint = Checkpoint( + state_dict=model.state_dict(), + epoch=0, + consumed_samples=2, + config=config.to_dict(), + extra={ + "optimizer": policy_optimizer.state_dict(), + "scheduler": policy_scheduler.state_dict(), + "value_model": saved_critic.state_dict(), + "value_optimizer": saved_optimizer.state_dict(), + }, + meta={"policy_version": 3}, + ) + checkpoint.save(temp_dir) + build_test_tokenizer(vocab_size=200).save_pretrained(temp_dir) + + cfg = _ppo_config( + device, + model_fn=lambda: AutoRegressiveLM(config), + critic_model_fn=lambda: ValueModel(config), + ckpt_dir=os.path.join(temp_dir, "ckpt"), + ) + context = TrainContextBuilder(cfg).with_param_path(temp_dir, resume=True).build() + + assert isinstance(context.strategy.critic, ValueModel) + assert torch.equal( + context.strategy.critic.state_dict()["value_head.weight"], + saved_critic.state_dict()["value_head.weight"], + ) + assert context.strategy.policy_version == 3 + + +def test_save_extra_persists_critic_state(device): + model, _ = make_model(device) + critic = ValueModel(make_rollout_config()).to(device) + from astrai.trainer.strategy import PPOStrategy + + strategy = PPOStrategy( + model=model, + device=device, + critic=critic, + critic_optimizer=torch.optim.SGD(critic.parameters(), lr=0.0), + executor=FakeExecutor(), + ) + context = TrainContext(strategy=strategy) + + extra = CheckpointCallback.save_extra(context) + + assert set(extra) == {"value_model", "value_optimizer"} + saved = extra["value_model"] + live = critic.state_dict() + assert set(saved) == set(live) + for key in saved: + assert torch.equal(saved[key], live[key]) + + +def test_save_extra_without_critic_has_no_value_entries(device): + model, _ = make_model(device) + from astrai.trainer.strategy import GRPOStrategy + from tests.helpers import make_frozen + + strategy = GRPOStrategy( + model=model, + device=device, + old_model=None, + ref_model=make_frozen(model, device), + executor=FakeExecutor(), + ) + context = TrainContext(strategy=strategy) + + extra = CheckpointCallback.save_extra(context) + + assert "value_model" not in extra + assert "value_optimizer" not in extra + + +def test_trainer_default_callbacks_do_not_break_ppo(device, temp_dir): + """The Trainer's default callback set constructs fine for online_ppo.""" + cfg = _ppo_config(device, ckpt_dir=os.path.join(temp_dir, "ckpt")) + trainer = Trainer(cfg) + assert trainer.callbacks + + +def test_sh_checkpoint_extra_files_detects_online_ppo(temp_dir): + """The shell completeness helper derives PPO's extra required files.""" + lib = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "docker" + / "lib" + / "train-common.sh" + ) + ppo_yaml = Path(temp_dir) / "ppo.yaml" + ppo_yaml.write_text("train_type: online_ppo\n") + grpo_yaml = Path(temp_dir) / "grpo.yaml" + grpo_yaml.write_text('train_type: "online_grpo"\n') + quoted_yaml = Path(temp_dir) / "quoted.yaml" + quoted_yaml.write_text(' train_type: "online_ppo"\n') + + def extra_files(yaml_path): + script = f'source "{lib}"; checkpoint_extra_files "{yaml_path}"' + result = subprocess.run( + ["bash", "-c", script], capture_output=True, text=True, check=True + ) + return result.stdout.strip() + + assert extra_files(ppo_yaml) == "value_model.pt value_optimizer.pt" + assert extra_files(quoted_yaml) == "value_model.pt value_optimizer.pt" + assert extra_files(grpo_yaml) == ""