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

core/trainer_base.py

@@ -488,16 +488,37 @@ def setup_tokenizer_with_model(self, model, model_name: Optional[str] = None):
             model_name = self.config.model.base_model_name
 
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-        # Set pad token if not present
+
+        # Establish a pad token WITHOUT growing the vocabulary. The previous code
+        # called add_special_tokens({"pad_token": "[PAD]"}) unconditionally, which
+        # made the `is None` check below dead and appended a brand-new token even
+        # to models that already pad (Qwen2.5 pads with <|endoftext|>). That grows
+        # len(tokenizer), forces the embedding resize below, and makes PEFT
+        # serialize the full embed_tokens + lm_head into every adapter checkpoint
+        # -- measured 561 MB instead of 18 MB for a 0.5B model, ~2 GB for an 8B.
+        #
+        # Reusing eos as pad keeps the vocab untouched for the common case of a
+        # tokenizer with no pad token (Llama 3, GPT-2); a fresh [PAD] is only
+        # added when there is no eos to borrow.
         if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
+            if tokenizer.eos_token is not None:
+                tokenizer.pad_token = tokenizer.eos_token
+            else:
+                tokenizer.add_special_tokens({"pad_token": "[PAD]"})
         tokenizer.padding_side = "left"
-        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-
-        # Resize embeddings
-        model.resize_token_embeddings(len(tokenizer))
-
+
+        # Only GROW the embedding matrix, never shrink it. Models are commonly
+        # published with vocab_size padded past len(tokenizer) for kernel
+        # alignment (Qwen2.5-0.5B: 151936 vs 151665), so an unconditional
+        # resize to len(tokenizer) would truncate real rows.
+        current_rows = model.get_input_embeddings().weight.shape[0]
+        if len(tokenizer) > current_rows:
+            self.logger.info(
+                f"Resizing embeddings {current_rows} -> {len(tokenizer)} "
+                f"(tokenizer has more tokens than the model's embedding matrix)"
+            )
+            model.resize_token_embeddings(len(tokenizer))
+
         return tokenizer
 
     @staticmethod

recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml

@@ -0,0 +1,94 @@
+# GRPO / RLVR smoke test sized for a single Apple Silicon Mac (MPS).
+#
+# MEASURED on an M3 Max / 64 GB (torch 2.13.0, transformers 5.8.0, trl 1.2.0):
+#   10 steps = 99 s wall clock (~10 s/step), 1.38 GB MPS allocated.
+#   Step time tracks ACTUAL completion length, not the cap: it fell 16 s -> 5 s
+#   over 10 steps as the policy learned to emit short well-formed answers.
+#   1 epoch over 64 rows = 32 optimizer steps (2 unique prompts/step) ~= 5.5 min.
+#
+# Memory is nowhere near the limit — the constraint is rollout generation
+# throughput, since GRPO samples num_generations completions per prompt through
+# plain HF `generate` (no vLLM on MPS). Scale by steps, not by dataset size.
+#
+# Run:
+#   python trainers/train.py --recipe recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml --num-gpus 1
+
+model:
+  base_model_name: "Qwen/Qwen2.5-0.5B-Instruct"   # ~1 GB download, open weights, has a chat template
+  torch_dtype: "float16"                           # MPS prefers fp16; fall back to float32 if you see NaN losses
+  low_cpu_mem_usage: true
+  load_in_8bit: false                              # bitsandbytes is CUDA-only; the launcher force-disables these on MPS
+  load_in_4bit: false
+  use_flash_attention: false                       # flash-attn is CUDA-only
+
+  use_lora: true
+  lora_r: 8
+  lora_alpha: 16
+  lora_dropout: 0.05
+  lora_target_modules:
+    - q_proj
+    - v_proj
+    - k_proj
+    - o_proj
+    - gate_proj
+    - up_proj
+    - down_proj
+  lora_bias: "none"
+
+data:
+  datasets:
+    # HF slice syntax passes straight through to load_dataset()
+    # (utils/dataset_utils.py:106), so this only downloads/keeps 64 rows.
+    - name: "openai/gsm8k"
+      subset: "main"
+      split: "train[:64]"
+      prompt_column: "question"
+      answer_column: "answer"
+
+  # max_completion_length = max_length - max_prompt_length = 512.
+  # Measured: raising this cap from 256 -> 512 cost nothing (105 s vs 99 s for
+  # 10 steps) because the policy generates ~50-120 tokens anyway. At 256 the cap
+  # clipped 12-37% of completions early on, and a clipped completion never emits
+  # </answer>, so it silently forfeits the format reward. Keep the headroom.
+  max_length: 1024
+  max_prompt_length: 512                           # GSM8K template prompt is ~200 tokens incl. system + 1-shot
+  remove_unused_columns: false                     # must stay false: the reward fn needs the `answer` column
+
+training:
+  algorithm: "grpo"
+  output_dir: "models/qwen2_5-0_5B-grpo-mps-smoke"
+
+  num_generations: 4                               # group size; keep global batch divisible by this
+  # KL coefficient toward the reference policy. Non-zero is essential here: at
+  # beta=0.0 this exact run drives held-out GSM8K accuracy DOWN from 44% to 17%
+  # by learning to emit bare digits in empty <think> tags. See scripts/ab_eval_grpo.py.
+  beta: 0.04
+  per_device_train_batch_size: 4
+  gradient_accumulation_steps: 2                   # global batch 8 -> 2 unique prompts per optimizer step
+  learning_rate: 1.0e-4
+  num_train_epochs: 1
+  # ~10 s/step measured. 10 steps (~100 s) is already enough to watch the XML
+  # format reward climb 0.03 -> ~0.50 (its ceiling). For a longer ablation:
+  # 128 steps ~= 22 min, and bump the slice above to train[:256] to stay 1 epoch.
+  max_steps: 10
+  warmup_steps: 2                                  # stock recipe's 50 would exceed max_steps entirely
+
+  logging_steps: 1
+  save_steps: 1000                                 # effectively "final save only" for a 10-step run
+  eval_steps: 1000
+
+  bf16: false                                      # bf16 is auto-switched to fp16 on MPS (utils/device_utils.py)
+  fp16: true                                       # auto-disabled if torch < 2.8.0 (GradScaler+MPS requirement)
+  gradient_checkpointing: false                    # pure slowdown here; you have the memory
+
+  dataloader_num_workers: 0
+  dataloader_pin_memory: false
+  dataloader_drop_last: true
+  save_only_model: true
+  prediction_loss_only: true
+
+wandb:
+  enabled: false
+
+s3:
+  enabled: false

scripts/ab_eval_grpo.py

@@ -0,0 +1,197 @@
+"""Paired base-vs-GRPO evaluation on held-out GSM8K.
+
+Scores both models with the repo's OWN template and reward functions
+(trainers/templates/gsm8k_template.py, trainers/rewards/*), so the numbers
+measure exactly what GRPO optimized rather than a separate eval regex.
+
+Both arms see identical prompts and identical tokenizer treatment; the only
+difference is whether the LoRA adapter is attached. Greedy decoding, so the
+comparison is deterministic.
+
+Usage:
+  python scripts/ab_eval_grpo.py --arm nokl=/tmp/grpo_trained --arm kl=/tmp/grpo_fixed --n 100
+"""
+import argparse
+import json
+import os
+import re
+import sys
+import time
+
+import torch
+from datasets import load_dataset
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from trainers.rewards.accuracy_rewards import extract_answer
+from trainers.rewards.format_rewards import count_xml
+from trainers.templates.gsm8k_template import GSM8KTemplate
+
+
+def build_tokenizer(base_model):
+    """Mirror the FIXED core/trainer_base.py:setup_tokenizer_with_model: add a pad
+    token only when the tokenizer lacks one, so the vocab is never grown."""
+    tok = AutoTokenizer.from_pretrained(base_model)
+    if tok.pad_token is None:
+        tok.add_special_tokens({"pad_token": "[PAD]"})
+    tok.padding_side = "left"
+    return tok
+
+
+def adapter_embed_rows(adapter):
+    """Embedding rows baked into an adapter, or None if it carries none.
+
+    Pre-fix checkpoints contain a full resized embed_tokens; post-fix ones don't.
+    Detecting it lets one script score both layouts with no manual flag.
+    """
+    path = os.path.join(adapter, "adapter_model.safetensors")
+    if not os.path.exists(path):
+        return None
+    from safetensors import safe_open
+    with safe_open(path, "pt") as f:
+        for k in f.keys():
+            if k.endswith("embed_tokens.weight"):
+                return f.get_slice(k).get_shape()[0]
+    return None
+
+
+def load_model(base_model, tok, adapter=None, device="mps"):
+    model = AutoModelForCausalLM.from_pretrained(base_model, dtype=torch.float16)
+    if adapter:
+        rows = adapter_embed_rows(adapter)
+        if rows and rows != model.get_input_embeddings().weight.shape[0]:
+            print(f"    (legacy adapter carries a resized embedding: {rows} rows)")
+            model.resize_token_embeddings(rows)
+        model = PeftModel.from_pretrained(model, adapter)
+        model = model.merge_and_unload()
+    return model.to(device).eval()
+
+
+@torch.no_grad()
+def generate(model, tok, prompts, max_new_tokens, batch_size, device):
+    out = []
+    for i in range(0, len(prompts), batch_size):
+        batch = prompts[i : i + batch_size]
+        enc = tok(batch, return_tensors="pt", padding=True, truncation=True,
+                  max_length=1024).to(device)
+        gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,
+                             pad_token_id=tok.pad_token_id)
+        for j in range(len(batch)):
+            out.append(tok.decode(gen[j][enc["input_ids"].shape[1]:],
+                                  skip_special_tokens=True))
+        print(f"    {min(i + batch_size, len(prompts))}/{len(prompts)}", flush=True)
+    return out
+
+
+def norm(s):
+    """Normalize a numeric answer so 18,000 == 18000 and 7.0 == 7."""
+    s = s.replace(",", "").replace("$", "").strip().rstrip(".")
+    try:
+        f = float(s)
+        return str(int(f)) if f == int(f) else str(f)
+    except ValueError:
+        return s
+
+
+def last_number(text):
+    """Format-agnostic GSM8K extraction: the last number in the completion.
+    Needed because the repo's extract_answer requires <answer> tags, which the
+    BASE model never emits -- scoring it with the strict extractor reports 0%
+    regardless of whether the arithmetic was right."""
+    m = re.findall(r"-?\d[\d,]*\.?\d*", text)
+    return norm(m[-1]) if m else ""
+
+
+def score(completions, answers, tok):
+    """Score each completion two ways:
+      strict_correct  - the repo's <answer>-tag extractor (what GRPO optimized)
+      loose_correct   - last-number-in-text (true math accuracy, format-blind)
+    """
+    rows = []
+    for comp, ans in zip(completions, answers):
+        pred = extract_answer(comp)
+        rows.append({
+            "exact": 2.0 if pred == ans else 0.0,
+            "xml": count_xml(comp),
+            "digit": 0.5 if pred.isdigit() else 0.0,
+            "ntok": len(tok(comp)["input_ids"]),
+            "strict_correct": pred == ans,
+            "loose_correct": last_number(comp) == norm(ans),
+        })
+    return rows
+
+
+def summarize(name, rows):
+    n = len(rows)
+    mean = lambda k: sum(r[k] for r in rows) / n
+    strict = sum(r["strict_correct"] for r in rows) / n
+    loose = sum(r["loose_correct"] for r in rows) / n
+    print(f"  {name:<6} strict={strict:6.1%} ({sum(r['strict_correct'] for r in rows):>3}/{n})   "
+          f"loose={loose:6.1%} ({sum(r['loose_correct'] for r in rows):>3}/{n})   "
+          f"xml={mean('xml'):.3f}  digit={mean('digit'):.3f}  mean_tok={mean('ntok'):6.1f}")
+    return strict, loose
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--base", default="Qwen/Qwen2.5-0.5B-Instruct")
+    p.add_argument("--arm", action="append", default=[], metavar="NAME=PATH",
+                   help="repeatable; a 'base' arm (no adapter) is always run first")
+    p.add_argument("--split", default="test")
+    p.add_argument("--n", type=int, default=100)
+    p.add_argument("--max-new-tokens", type=int, default=512)
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--device", default="mps")
+    p.add_argument("--dump", default="/tmp/ab_eval_rows.json")
+    args = p.parse_args()
+    dump_path = args.dump
+
+    ds = load_dataset("openai/gsm8k", "main", split=f"{args.split}[:{args.n}]")
+    tok = build_tokenizer(args.base)
+
+    prompts, answers = [], []
+    for ex in ds:
+        f = GSM8KTemplate.format_for_training(ex, "question", "answer")
+        prompts.append(tok.apply_chat_template(f["prompt"], tokenize=False,
+                                               add_generation_prompt=True))
+        answers.append(f["answer"])
+
+    print(f"Held-out GSM8K {args.split}[:{args.n}], greedy, max_new_tokens={args.max_new_tokens}\n")
+    arms = [("base", None)] + [tuple(a.split("=", 1)) for a in args.arm]
+    results = {}
+    for name, adapter in arms:
+        print(f"  generating: {name}")
+        t0 = time.time()
+        model = load_model(args.base, tok, adapter, args.device)
+        comps = generate(model, tok, prompts, args.max_new_tokens, args.batch_size, args.device)
+        results[name] = score(comps, answers, tok)
+        json.dump(comps, open(f"/tmp/ab_completions_{name}.json", "w"), indent=1)
+        del model
+        torch.mps.empty_cache()
+        print(f"    done in {time.time() - t0:.0f}s")
+
+    print("\nRESULTS   strict = repo <answer>-tag extractor (what GRPO optimized)")
+    print("          loose  = last-number-in-text (format-blind true accuracy)\n")
+    accs = {name: summarize(name, results[name]) for name, _ in arms}
+
+    print("\n  vs base (paired, identical problems):")
+    for name, _ in arms[1:]:
+        for i, (lbl, key) in enumerate([("strict", "strict_correct"),
+                                        ("loose ", "loose_correct")]):
+            gained = sum(1 for b, g in zip(results["base"], results[name])
+                         if g[key] and not b[key])
+            lost = sum(1 for b, g in zip(results["base"], results[name])
+                       if b[key] and not g[key])
+            print(f"    {name:<6} {lbl}: {accs[name][i] - accs['base'][i]:+6.1%}   "
+                  f"({gained} fixed, {lost} broken, net {gained - lost:+d}"
+                  f"/{len(results['base'])})")
+
+    with open(dump_path, "w") as f:
+        json.dump({k: [dict(r) for r in v] for k, v in results.items()}, f, indent=1)
+    print(f"\n  per-example scores -> {dump_path}")
+
+
+if __name__ == "__main__":
+    main()

trainers/grpo_trainer.py

@@ -175,6 +175,20 @@ def setup_training_args(self) -> GRPOConfig:
             # GRPO specific parameters
             num_generations=self.config.training.num_generations,
             max_completion_length=self.config.data.max_length - self.config.data.max_prompt_length,
+            # KL regularization toward the reference policy. Previously not
+            # forwarded at all, so GRPO always ran at TRL's default beta=0.0
+            # (no KL term) and training.beta was silently inert. With no anchor,
+            # the policy free-runs toward whatever the shaping rewards favor --
+            # on GSM8K it collapses to emitting bare digits inside empty <think>
+            # tags, which scores full format+digit reward while abandoning
+            # chain-of-thought.
+            #
+            # epsilon / epsilon_high / steps_per_generation are deliberately NOT
+            # forwarded here: their defaults in TrainingConfig (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.
+            beta=float(self.config.training.beta),
         )
 
     def setup_trainer(self):