fix: tighten runtime and GRPO metrics

pull/808/head
RileyChen23 2026-07-14 18:41:40 +08:00
parent 6a98206aef
commit 8cbbda005c
3 changed files with 7 additions and 3 deletions

2
.gitignore vendored
View File

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

View File

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

View File

@ -168,7 +168,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
current_aux_loss = aux_loss.item() current_aux_loss = aux_loss.item()
avg_reward_val = rewards.mean().item() avg_reward_val = rewards.mean().item()
avg_len_val = completion_mask.sum(dim=1).float().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_mean_val = advantages.mean().item()
advantages_std_val = advantages.std().item() advantages_std_val = advantages.std().item()
current_lr = optimizer.param_groups[0]['lr'] current_lr = optimizer.param_groups[0]['lr']