ViperEkura
350e4a1849
- 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
470 lines
17 KiB
Python
470 lines
17 KiB
Python
"""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"]
|