[feat] GenRL: fix PPO loop cadence and diagnostics - #1403
Conversation
Extracted from hao-ai-lab#1391. Addresses review feedback on prompt group validation, accumulated microbatch zero_grad, per-timestep probe args, and post-update KL logging. GenRL-Stack: 4/6
|
There was a problem hiding this comment.
Code Review
This pull request introduces several enhancements to the GenRL training pipeline, including configuration validation, periodic evaluation during training, logging of trainable parameters, and support for PPO microbatch accumulation. However, the reviewer identified critical issues that must be addressed: an ImportError caused by importing the undefined function clear_reward_models, a potential TypeError when calling eval_once with unsupported arguments (max_batches and seed), and a potential shape mismatch crash when slicing self._train_neg_embeds during microbatch accumulation.
Co-authored-by: Davids048 <jundasu@ucsd.edu>
Extracted from #1391.
GenRL-Stack: 4/6
Purpose
Fix GenRL PPO/GRPO training-loop behavior and diagnostics without mixing in reward-model or LoRA changes.
This also addresses review feedback from Gemini/David on sampler validation, accumulated microbatch gradients, and post-update KL diagnostics.
Fixes #
Changes
zero_grad()whenoptimizer_step_per_timestep=false.0.Test Plan
Test Results
Test output