304 lines
8.4 KiB
Python
304 lines
8.4 KiB
Python
"""Alignment tests for lesson 24 GRPO/CISPO helpers.
|
|
|
|
Run after implementing:
|
|
|
|
python course/impl/tests/test_grpo.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[3]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from course.impl.core.grpo import cispo_policy_loss # noqa: E402
|
|
from course.impl.core.grpo import grpo_policy_loss # noqa: E402
|
|
from course.impl.core.grpo import group_relative_advantages # noqa: E402
|
|
from course.impl.core.grpo import token_kl_penalty # noqa: E402
|
|
|
|
|
|
def assert_close(name: str, actual: torch.Tensor, expected: torch.Tensor, tol: float = 1e-7) -> None:
|
|
diff = (actual - expected).abs().max().item()
|
|
print(f"{name}_max_abs_diff={diff:.12f}")
|
|
assert diff < tol, f"{name} diff {diff} >= {tol}"
|
|
|
|
|
|
def expected_group_relative_advantages(
|
|
rewards: torch.Tensor,
|
|
num_generations: int,
|
|
eps: float = 1e-4,
|
|
) -> torch.Tensor:
|
|
grouped_rewards = rewards.view(-1, num_generations)
|
|
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(num_generations)
|
|
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(num_generations)
|
|
return (rewards - mean_r) / (std_r + eps)
|
|
|
|
|
|
def expected_token_kl_penalty(policy_logps: torch.Tensor, ref_logps: torch.Tensor) -> torch.Tensor:
|
|
kl_div = ref_logps - policy_logps
|
|
return torch.exp(kl_div) - kl_div - 1.0
|
|
|
|
|
|
def reduce_like_source(per_token_loss: torch.Tensor, completion_mask: torch.Tensor) -> torch.Tensor:
|
|
per_response = (per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)
|
|
return per_response.mean()
|
|
|
|
|
|
def expected_grpo_policy_loss(
|
|
policy_logps: torch.Tensor,
|
|
old_logps: torch.Tensor,
|
|
ref_logps: torch.Tensor,
|
|
advantages: torch.Tensor,
|
|
completion_mask: torch.Tensor,
|
|
epsilon: float,
|
|
beta: float,
|
|
) -> torch.Tensor:
|
|
per_token_kl = expected_token_kl_penalty(policy_logps, ref_logps)
|
|
ratio = torch.exp(policy_logps - old_logps)
|
|
clipped_ratio = torch.clamp(ratio, 1.0 - epsilon, 1.0 + epsilon)
|
|
per_token_loss1 = ratio * advantages.unsqueeze(1)
|
|
per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
|
|
per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - beta * per_token_kl)
|
|
return reduce_like_source(per_token_loss, completion_mask)
|
|
|
|
|
|
def expected_cispo_policy_loss(
|
|
policy_logps: torch.Tensor,
|
|
old_logps: torch.Tensor,
|
|
ref_logps: torch.Tensor,
|
|
advantages: torch.Tensor,
|
|
completion_mask: torch.Tensor,
|
|
epsilon_high: float,
|
|
beta: float,
|
|
) -> torch.Tensor:
|
|
per_token_kl = expected_token_kl_penalty(policy_logps, ref_logps)
|
|
ratio = torch.exp(policy_logps - old_logps)
|
|
clamped_ratio = torch.clamp(ratio, max=epsilon_high).detach()
|
|
per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * policy_logps - beta * per_token_kl)
|
|
return reduce_like_source(per_token_loss, completion_mask)
|
|
|
|
|
|
def test_group_relative_advantages_alignment() -> None:
|
|
rewards = torch.tensor([1.0, 3.0, 2.0, 2.0, -1.0, 1.0])
|
|
expected = expected_group_relative_advantages(rewards, num_generations=2)
|
|
actual = group_relative_advantages(rewards, num_generations=2)
|
|
assert_close("group_relative_advantages", actual, expected)
|
|
print("group_relative_advantages=passed")
|
|
|
|
|
|
def test_token_kl_penalty_alignment() -> None:
|
|
policy_logps = torch.tensor(
|
|
[
|
|
[-1.0, -0.7, -2.2],
|
|
[-0.5, -1.3, -0.9],
|
|
]
|
|
)
|
|
ref_logps = torch.tensor(
|
|
[
|
|
[-1.1, -0.9, -2.0],
|
|
[-0.4, -1.0, -0.8],
|
|
]
|
|
)
|
|
expected = expected_token_kl_penalty(policy_logps, ref_logps)
|
|
actual = token_kl_penalty(policy_logps, ref_logps)
|
|
assert_close("token_kl_penalty", actual, expected)
|
|
print("token_kl_penalty=passed")
|
|
|
|
|
|
def test_grpo_policy_loss_alignment() -> None:
|
|
policy_logps = torch.tensor(
|
|
[
|
|
[-1.00, -0.70, -2.20, -0.30],
|
|
[-0.50, -1.30, -0.90, -2.00],
|
|
[-1.50, -0.20, -1.10, -0.80],
|
|
[-0.90, -1.80, -0.40, -1.20],
|
|
]
|
|
)
|
|
old_logps = torch.tensor(
|
|
[
|
|
[-1.20, -0.80, -1.90, -0.40],
|
|
[-0.70, -1.00, -0.90, -1.70],
|
|
[-1.60, -0.50, -1.00, -1.10],
|
|
[-1.20, -1.60, -0.60, -1.00],
|
|
]
|
|
)
|
|
ref_logps = torch.tensor(
|
|
[
|
|
[-1.10, -0.90, -2.00, -0.30],
|
|
[-0.40, -1.00, -0.80, -2.10],
|
|
[-1.30, -0.40, -1.30, -0.70],
|
|
[-1.00, -1.70, -0.70, -1.40],
|
|
]
|
|
)
|
|
advantages = torch.tensor([1.2, -0.8, 0.5, -1.1])
|
|
completion_mask = torch.tensor(
|
|
[
|
|
[1.0, 1.0, 1.0, 1.0],
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[1.0, 1.0, 0.0, 0.0],
|
|
[1.0, 1.0, 1.0, 0.0],
|
|
]
|
|
)
|
|
epsilon = 0.2
|
|
beta = 0.03
|
|
|
|
expected = expected_grpo_policy_loss(
|
|
policy_logps,
|
|
old_logps,
|
|
ref_logps,
|
|
advantages,
|
|
completion_mask,
|
|
epsilon,
|
|
beta,
|
|
)
|
|
actual = grpo_policy_loss(
|
|
policy_logps,
|
|
old_logps,
|
|
ref_logps,
|
|
advantages,
|
|
completion_mask,
|
|
epsilon,
|
|
beta,
|
|
)
|
|
assert_close("grpo_policy_loss", actual, expected)
|
|
|
|
|
|
def test_cispo_policy_loss_alignment() -> None:
|
|
policy_logps = torch.tensor(
|
|
[
|
|
[-1.00, -0.70, -2.20, -0.30],
|
|
[-0.50, -1.30, -0.90, -2.00],
|
|
[-1.50, -0.20, -1.10, -0.80],
|
|
[-0.90, -1.80, -0.40, -1.20],
|
|
]
|
|
)
|
|
old_logps = torch.tensor(
|
|
[
|
|
[-1.20, -0.80, -1.90, -0.40],
|
|
[-0.70, -1.00, -0.90, -1.70],
|
|
[-1.60, -0.50, -1.00, -1.10],
|
|
[-1.20, -1.60, -0.60, -1.00],
|
|
]
|
|
)
|
|
ref_logps = torch.tensor(
|
|
[
|
|
[-1.10, -0.90, -2.00, -0.30],
|
|
[-0.40, -1.00, -0.80, -2.10],
|
|
[-1.30, -0.40, -1.30, -0.70],
|
|
[-1.00, -1.70, -0.70, -1.40],
|
|
]
|
|
)
|
|
advantages = torch.tensor([1.2, -0.8, 0.5, -1.1])
|
|
completion_mask = torch.tensor(
|
|
[
|
|
[1.0, 1.0, 1.0, 1.0],
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[1.0, 1.0, 0.0, 0.0],
|
|
[1.0, 1.0, 1.0, 0.0],
|
|
]
|
|
)
|
|
epsilon_high = 1.3
|
|
beta = 0.03
|
|
|
|
expected = expected_cispo_policy_loss(
|
|
policy_logps,
|
|
old_logps,
|
|
ref_logps,
|
|
advantages,
|
|
completion_mask,
|
|
epsilon_high,
|
|
beta,
|
|
)
|
|
actual = cispo_policy_loss(
|
|
policy_logps,
|
|
old_logps,
|
|
ref_logps,
|
|
advantages,
|
|
completion_mask,
|
|
epsilon_high,
|
|
beta,
|
|
)
|
|
assert_close("cispo_policy_loss", actual, expected)
|
|
|
|
|
|
def test_cispo_clamped_ratio_detached_gradient() -> None:
|
|
policy_logps = torch.tensor(
|
|
[
|
|
[-1.00, -0.70],
|
|
[-0.50, -1.30],
|
|
],
|
|
requires_grad=True,
|
|
)
|
|
expected_policy_logps = policy_logps.detach().clone().requires_grad_(True)
|
|
old_logps = torch.tensor(
|
|
[
|
|
[-1.20, -0.80],
|
|
[-0.70, -1.00],
|
|
]
|
|
)
|
|
ref_logps = torch.tensor(
|
|
[
|
|
[-1.10, -0.90],
|
|
[-0.40, -1.00],
|
|
]
|
|
)
|
|
advantages = torch.tensor([1.2, -0.8])
|
|
completion_mask = torch.ones_like(old_logps)
|
|
epsilon_high = 1.3
|
|
beta = 0.0
|
|
|
|
actual = cispo_policy_loss(
|
|
policy_logps,
|
|
old_logps,
|
|
ref_logps,
|
|
advantages,
|
|
completion_mask,
|
|
epsilon_high,
|
|
beta,
|
|
)
|
|
actual.backward()
|
|
|
|
expected = expected_cispo_policy_loss(
|
|
expected_policy_logps,
|
|
old_logps,
|
|
ref_logps,
|
|
advantages,
|
|
completion_mask,
|
|
epsilon_high,
|
|
beta,
|
|
)
|
|
expected.backward()
|
|
|
|
assert_close("cispo_detached_gradient", policy_logps.grad, expected_policy_logps.grad)
|
|
|
|
|
|
def main() -> None:
|
|
tests = [
|
|
test_group_relative_advantages_alignment,
|
|
test_token_kl_penalty_alignment,
|
|
test_grpo_policy_loss_alignment,
|
|
test_cispo_policy_loss_alignment,
|
|
test_cispo_clamped_ratio_detached_gradient,
|
|
]
|
|
failures = []
|
|
for test in tests:
|
|
try:
|
|
test()
|
|
except Exception as exc: # noqa: BLE001 - report all lesson test failures.
|
|
failures.append((test.__name__, exc))
|
|
print(f"{test.__name__}=failed: {type(exc).__name__}: {exc}")
|
|
|
|
if failures:
|
|
names = ", ".join(name for name, _ in failures)
|
|
raise SystemExit(f"grpo_core=failed: {names}")
|
|
print("grpo_core=passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|