fix: tighten runtime and GRPO metrics

.gitignore

@@ -1,5 +1,7 @@
 __pycache__/
 *.pyc
+.pytest_cache/
+.codex_pytest_tmp*/
 .DS_Store
 out
 website/

requirements.txt

@@ -1,5 +1,6 @@
 datasets==3.6.0
 datasketch==1.6.4
+fastapi==0.115.6
 Flask==3.0.3
 Flask_Cors==4.0.0
 jieba==0.42.1
@@ -23,6 +24,7 @@ jinja2==3.1.2
 jsonlines==4.0.0
 trl==0.13.0
 ujson==5.1.0
+uvicorn==0.34.0
 wandb==0.18.3
 streamlit==1.50.0
 einops==0.8.1

trainer/train_grpo.py

@@ -168,7 +168,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
             current_aux_loss = aux_loss.item()
             avg_reward_val = rewards.mean().item()
             avg_len_val = completion_mask.sum(dim=1).float().mean().item()
-            kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
+            kl_ref_val = (per_token_kl * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
             advantages_mean_val = advantages.mean().item()
             advantages_std_val = advantages.std().item()
             current_lr = optimizer.param_groups[0]['lr']