wehub-resource-sync
32 lines
1.7 KiB
Python
32 lines
1.7 KiB
Python
"""Live-progress GRPO-ascends-reward proof (prints every few iters). See verify_rl_optimizes.py."""
|
|
import torch
|
|
from src.models.transformer import Transformer
|
|
from src.post_training.grpo import group_advantages, grpo_loss
|
|
from src.post_training.rollout import generate_with_logprobs, compute_logprobs
|
|
from src.post_training.utils import make_frozen_copy, set_seed
|
|
|
|
VOCAB, TARGET, PL, GEN, G = 64, 7, 4, 8, 8
|
|
set_seed(0)
|
|
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
policy = Transformer(n_head=4, n_embed=64, context_length=PL + GEN + 2, vocab_size=VOCAB, N_BLOCKS=2).to(dev)
|
|
ref = make_frozen_copy(policy, device=dev)
|
|
opt = torch.optim.Adam(policy.parameters(), lr=1e-2)
|
|
prompt = torch.zeros(G, PL, dtype=torch.long, device=dev)
|
|
hist = []
|
|
for it in range(40):
|
|
rb = generate_with_logprobs(policy, prompt, GEN, temperature=1.0)
|
|
seqs, rmask = rb.sequences, rb.response_mask
|
|
rew = torch.tensor([(seqs[i, PL:] == TARGET).float().mean().item() for i in range(G)], device=dev)
|
|
adv = group_advantages(rew, G)
|
|
old, _ = compute_logprobs(policy, seqs, rmask, requires_grad=False)
|
|
rf, _ = compute_logprobs(ref, seqs, rmask, requires_grad=False)
|
|
new, _ = compute_logprobs(policy, seqs, rmask, requires_grad=True)
|
|
loss, _ = grpo_loss(new, old, rf, adv, rmask[:, 1:], clip=0.2, kl_coef=0.0)
|
|
opt.zero_grad(); loss.backward(); opt.step()
|
|
hist.append(rew.mean().item())
|
|
if it % 5 == 0 or it == 39:
|
|
print(f"iter {it:2d} | mean reward (P[target token]) = {rew.mean().item():.3f}", flush=True)
|
|
start, end = sum(hist[:5]) / 5, sum(hist[-5:]) / 5
|
|
verdict = "PASS (RL ascends reward)" if end > start + 0.2 else "NO IMPROVEMENT"
|
|
print(f"\nRESULT: reward {start:.3f} -> {end:.3f} => {verdict}", flush=True)
|