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
This commit is contained in:
@@ -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 "
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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}
|
||||
@@ -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):
|
||||
|
||||
+392
-58
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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_"):
|
||||
|
||||
@@ -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
|
||||
|
||||
+11
-5
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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}"
|
||||
|
||||
+36
-2
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"]
|
||||
@@ -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) == ""
|
||||
Reference in New Issue
Block a user