fix: forward beta to GRPOConfig (GRPO ran with no KL term) and stop growing the vocab every run

[kayyar-roblox]

Fix GRPO KL regularization and per-checkpoint vocab growth

Found while running a GRPO/RLVR smoke test on Apple Silicon. Two independent bugs, each verified by execution, plus a reproducible ablation showing the first one silently destroys held-out accuracy.

Bug 1 — GRPO ran with no KL penalty, and no knob could enable one

grpo_trainer.py:setup_training_args forwarded num_generations and max_completion_length but never beta. TRL 1.2.0's GRPOConfig defaults to beta=0.0, so every GRPO run trained with the KL term switched off and training.beta was silently inert. gspo_trainer.py already forwarded it; GRPO did not.

With nothing anchoring the policy to its reference, the shaping rewards are free to dominate. On GSM8K the policy learns that an empty <think> block plus a bare integer collects the full format reward (0.5) and digit reward (0.5):

[BASE, 216 tok]                          [GRPO beta=0.0, 22 tok]
1. Donuts: 3 × $68 = $204                <think>
2. Mini cupcakes: 2 × $80 = $160         1400
3. Mini cheesecakes: 6 × $55 = $330      </think>
Total = $204 + $160 + $330 = $694  ✓     <answer>
                                         1400
                                         </answer>  ✗

Training-time telemetry, same 32 steps and same data, only beta differs — the kl channel is exactly zero before the fix:

               beta=0.0 (before)      |     beta=0.04 (after)
 steps      len  entropy       kl     |    len  entropy       kl
  1-8     125.6    0.899   0.0000     |  138.3    1.034   0.0224
  9-16     92.6    0.750   0.0000     |  118.4    0.774   0.0488
 17-24     72.4    0.859   0.0000     |  100.6    0.603   0.0469
 25-32     86.4    0.809   0.0000     |  124.9    0.644   0.0440

epsilon / epsilon_high / steps_per_generation are deliberately not forwarded for GRPO: their TrainingConfig defaults (3e-4 / 4e-4 / 4) are the GSPO paper's values, and applying them to GRPO would silently tighten clipping from TRL's 0.2 to 3e-4. They stay GSPO-only until GRPO gets its own defaults.

Bug 2 — every adapter checkpoint carried the full embedding matrix

core/trainer_base.py:setup_tokenizer_with_model called add_special_tokens({"pad_token": "[PAD]"}) unconditionally, which made the preceding if tokenizer.pad_token is None check dead code. Qwen2.5 already pads with <|endoftext|>, so this appended a redundant token (len(tokenizer) 151665 → 151666). The unconditional resize_token_embeddings(len(tokenizer)) then changed the embedding row count — for Qwen2.5-0.5B a shrink of the published matrix from 151936 to 151666 rows — so PEFT flagged the embedding as modified and serialized it into every checkpoint:

adapter total: 276,184,576 params
  135,892,736  base_model.model.model.embed_tokens.weight  [151666, 896]
  135,892,736  base_model.model.lm_head.weight             [151666, 896]

98.4% embeddings, for 4.4 M of actual LoRA. 561 MB per checkpoint on a 0.5B model; ~2 GB on an 8B one, multiplied by save_steps and by S3 checkpoint upload where enabled.

The fix reuses eos as pad when a tokenizer has none (so Llama 3 / GPT-2 don't grow either) and only ever grows the embedding, never shrinks it — a naive != guard would truncate Qwen's published 151936 rows down to 151665.

	before	after
adapter_model.safetensors	561 MB	17.6 MB
vocab growth, Qwen2.5	151665 → 151666	none
vocab growth, GPT-2	50257 → 50258	none
embedding resize	151936 → 151666 (shrink)	not triggered

Ablation

Held-out GSM8K test[:100], greedy, paired (identical problems every arm), 32-step runs over train[:64], Qwen2.5-0.5B-Instruct + LoRA. Reproduce with scripts/ab_eval_grpo.py.

Two scorings are reported because the repo's extract_answer requires <answer> tags, which the base model never emits — scoring it strictly reports 0% no matter how good its arithmetic is:

• strict — the repo's <answer>-tag extractor, i.e. what GRPO actually optimizes
• loose — last number in the completion, i.e. format-blind true accuracy

  base   strict=  0.0% (  0/100)   loose= 44.0% ( 44/100)   xml=0.000  mean_tok= 209.9
  nokl   strict= 16.0% ( 16/100)   loose= 17.0% ( 17/100)   xml=0.500  mean_tok=  68.3
  kl     strict= 29.0% ( 29/100)   loose= 37.0% ( 37/100)   xml=0.424  mean_tok= 134.9

McNemar exact tests on the paired outcomes:

comparison	metric	fixed / broken	p	verdict
nokl vs base	loose	5 / 32	<0.0001	significant regression
kl vs base	loose	11 / 18	0.265	no measurable regression
kl vs nokl	loose	23 / 3	0.00009	significant improvement
kl vs nokl	strict	17 / 4	0.0072	significant improvement

What this shows: forwarding beta converts a statistically significant −27 point accuracy regression into one that is statistically indistinguishable from the untrained base (−7 points, p=0.27), while significantly improving both the optimized metric (+13 points strict) and true accuracy (+20 points loose) over the current trainer.

What this does not show: the fixed run does not beat the base model on format-blind math (37% vs 44%). 32 steps on 64 prompts is a smoke test, not a training budget — the claim here is that the KL term stops GRPO from actively destroying the policy, not that this recipe improves reasoning.

Also included

recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml — a GRPO recipe sized for a single Apple Silicon box (0.5B model, train[:64], 32 steps ≈ 5 min, 1.4 GB unified memory), with measured throughput annotated. Useful as a CI-able smoke test for the RLVR path; the stock GRPO recipe pulls OpenMathInstruct-2's 1M-row split as a second dataset, which is not a smoke test.

scripts/ab_eval_grpo.py — the paired evaluator used above. Scores with the repo's own template and reward functions, and auto-detects whether an adapter carries a legacy resized embedding so pre- and post-fix checkpoints can be compared in one run.

Reproduction

uv venv --python 3.12 .venv-mps
VIRTUAL_ENV=.venv-mps uv pip install torch torchvision==0.28.0
VIRTUAL_ENV=.venv-mps uv pip install -e .        # no [cuda] extra on macOS

# beta=0.0 arm (pre-fix behavior) vs beta=0.04 arm
python trainers/train.py --recipe recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml \
  --num-gpus 1 training.beta=0.0 training.output_dir=/tmp/nokl
python trainers/train.py --recipe recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml \
  --num-gpus 1 training.output_dir=/tmp/kl

python scripts/ab_eval_grpo.py --arm nokl=/tmp/nokl --arm kl=/tmp/kl --n 100

Environment: M3 Max / 64 GB, macOS, torch 2.13.0, transformers 5.8.0, trl 1.2.0, peft 0.19.0.

Note on scope of testing: the kl arm above was trained on the exact code path in this PR for Qwen2.5 (which has a pad token, so the eos fallback branch is not reached). The eos fallback and the resize guard were verified separately against Qwen2.5 and GPT-2.

Not addressed here

Found during the same audit, left out to keep this PR reviewable:

1. DPOConfig never receives beta either — dpo_trainer.py has zero references to it, so DPO always trains at TRL's default 0.1 and training.beta is silently ignored. Latent (no shipped DPO recipe sets it) but it is the DPO hyperparameter.
2. The documented data.datasets[0].name= override syntax does not work. set_nested_value (utils/recipe_overrides.py:35) splits only on ., producing a bogus datasets[0] key and a TypeError. Documented in 7 places across trainers/README.md and recipes/training/README.md.
3. data.max_prompt_length never truncates prompts for GRPO/GSPO — TRL 1.2.0's GRPOConfig has no such field. It only acts as the subtrahend for max_completion_length, so the recipe comment "Maximum length for prompts" is misleading.
4. SFT / CPT / SFT_VLM silently drop save_only_model and prediction_loss_only.
5. Reward-function edge cases: exact_match is raw string equality, so 8/1000 GSM8K gold answers containing commas can never match; count_xml's trailing penalty is unbounded (5000 trailing chars → −4.5, swamping the +2.0 correctness reward); extract_answer returns the entire completion when no <answer> tag is present; digit_reward rejects negatives and decimals.
6. custom_reward_func ships in the live reward list returning a constant 0.0.

Happy to split any of these into follow-ups.