chat_rl.py — 強化學習實作解說

scripts/chat_rl.py 實作了 nanochat 的強化學習階段，以 GSM8K 數學題為訓練任務。

演算法：簡化版 GRPO / REINFORCE

Karpathy 在註釋中說明這是 GRPO 的簡化版，實際上更接近 REINFORCE：

1. 無 trust region：刪除 KL 正則化項，沒有 reference model
2. On-policy 無 PPO ratio+clip：因為每次更新都使用當前 policy 產生的資料
3. DAPO 風格的 normalization：token-level，非 sequence-level
4. Advantage：僅減去 mean（r - μ），不使用 z-score (r - μ)/σ

流程

1. 載入 SFT 階段的模型
2. 對每個 GSM8K 題目產生多個回答（--num-samples=16）
3. 計算每個回答的 reward（答案是否正確）
4. 計算 log probabilities，以 advantage 加權計算 policy gradient 損失
5. 更新模型參數

損失函數

logp = -model(inputs, targets, loss_reduction='none')  # log probability
pg_obj = (logp * advantages.unsqueeze(-1)).sum()        # policy gradient objective
loss = -pg_obj / (num_valid * num_passes * examples_per_rank)  # minimize neg PG

監控指標

• reward：平均 reward（越高越好）
• pass@k：top-k 回答的正確率
• sequence_length：生成序列平均長度