ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add online ppo with value-model critic and gae advantages

Browse Source 
- 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:
ViperEkura
2026-09-05 01:59:50 +08:00
parent 816b96a58a
commit 350e4a1849
17 changed files with 1390 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/trainer/test_online_e2e.py
+34 -4
View File
@@ -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):
tests/trainer/test_ppo_strategy.py
+469
View File
@@ -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"]
tests/trainer/test_ppo_train_context.py
+271
View File
@@ -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) == ""