If you are training a model with RLVR/GRPO and you keep hitting the same roadblocks—unstable rewards, memory spikes, or confusing loss numbers—here is a practical, step‑by‑step fix list that works for most setups.
1. Verify your terminator list. The code builds TERMINATORS from EOS, PAD and any generation‑config EOS IDs. If you add special tokens later, update this set before the rollout; otherwise the mask will cut off generations too early or too late, hurting reward signals.
2. Keep the prompt length fixed per batch. The snippet pads all prompts to the longest length (P) and repeats them G times for sampling. If you change the batch size or number of samples per prompt, recompute P and the padding tensors; mismatched shapes cause silent errors in the attention mask.
3. Use micro‑batching consistently. Both the old‑log‑prob and new‑log‑prob calculations split the sequence into chunks of size grpo_micro_bs. Make sure the same slice indices are used for old_lp, ref_lp, adv and mask; otherwise the ratio will be computed on mismatched tokens and the KL term will explode.
4. Scale the loss by the number of micro‑batches. The script divides the final loss by n_chunks before backward(). If you change the gradient accumulation strategy, adjust this divisor so that the effective batch size stays what you expect.
5. Clear GPU cache only when needed. The empty_cache() call after each optimizer step is useful on limited VRAM, but it adds overhead. If you observe steady memory usage, you can remove it to speed up training.
6. Checkpoint merging. After training, merge_and_unload() writes a full model to the folder pointed by OUT. Verify that the tokenizer is saved in the same directory; otherwise downstream generation will fail due to missing vocab files.
7. Log the right metrics. The print statement shows pg, kl and clipfrac per epoch. If any of these values diverge (e.g., kl > 0.1 or clipfrac > 0.3), lower the learning rate, increase the KL beta, or tighten the clip range.
By auditing each of these points you will eliminate most sources of instability and get smoother, reproducible RLVR/GRPO runs.
#AI #ML #LLM #RL #DeepLearning #Productivity