From edc3dc1b31f4221ec4723c7f47412d64a5d9e167 Mon Sep 17 00:00:00 2001 From: Kartik Ayyar Date: Sat, 22 Aug 2026 13:17:33 -0700 Subject: [PATCH 1/4] Add MPS-sized GRPO smoke recipe for Apple Silicon --- .../training/grpo/qwen2_5_0_5B_mps_smoke.yaml | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml diff --git a/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml b/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml new file mode 100644 index 0000000..bf53c8c --- /dev/null +++ b/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml @@ -0,0 +1,82 @@ +# GRPO / RLVR smoke test sized for a single Apple Silicon Mac (MPS). +# +# Binding constraint on a Mac is NOT memory (64 GB unified is plenty here) — +# it is rollout generation throughput. GRPO generates num_generations +# completions per prompt per step through plain HF `generate` (no vLLM on +# MPS), so wall-clock scales as: +# +# max_steps x (per_device_train_batch_size x gradient_accumulation_steps) +# x max_completion_length +# +# where max_completion_length = data.max_length - data.max_prompt_length. +# +# 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_length: 768 # 768 - 512 = 256-token completions (vs 1024 in the stock recipe) + 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 + 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 + max_steps: 10 # hard cap so the smoke run is bounded; raise to 100-200 to see reward move + 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 From c2219482c178d3c8a27d192c003247b52438fe39 Mon Sep 17 00:00:00 2001 From: Kartik Ayyar Date: Sat, 22 Aug 2026 13:26:25 -0700 Subject: [PATCH 2/4] Annotate MPS GRPO smoke recipe with measured throughput; raise completion cap to 512 --- .../training/grpo/qwen2_5_0_5B_mps_smoke.yaml | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml b/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml index bf53c8c..a5da365 100644 --- a/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml +++ b/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml @@ -1,14 +1,14 @@ # GRPO / RLVR smoke test sized for a single Apple Silicon Mac (MPS). # -# Binding constraint on a Mac is NOT memory (64 GB unified is plenty here) — -# it is rollout generation throughput. GRPO generates num_generations -# completions per prompt per step through plain HF `generate` (no vLLM on -# MPS), so wall-clock scales as: +# 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. # -# max_steps x (per_device_train_batch_size x gradient_accumulation_steps) -# x max_completion_length -# -# where max_completion_length = data.max_length - data.max_prompt_length. +# 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 @@ -45,7 +45,12 @@ data: prompt_column: "question" answer_column: "answer" - max_length: 768 # 768 - 512 = 256-token completions (vs 1024 in the stock recipe) + # 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 + # , 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 @@ -58,7 +63,10 @@ training: gradient_accumulation_steps: 2 # global batch 8 -> 2 unique prompts per optimizer step learning_rate: 1.0e-4 num_train_epochs: 1 - max_steps: 10 # hard cap so the smoke run is bounded; raise to 100-200 to see reward move + # ~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 From f0b78f4c8b883e4cfd734c6a44204d12d9863956 Mon Sep 17 00:00:00 2001 From: Kartik Ayyar Date: Sat, 22 Aug 2026 14:24:31 -0700 Subject: [PATCH 3/4] Add paired base-vs-GRPO GSM8K eval with strict and format-blind scoring --- scripts/ab_eval_grpo.py | 174 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 174 insertions(+) create mode 100644 scripts/ab_eval_grpo.py diff --git a/scripts/ab_eval_grpo.py b/scripts/ab_eval_grpo.py new file mode 100644 index 0000000..e608db3 --- /dev/null +++ b/scripts/ab_eval_grpo.py @@ -0,0 +1,174 @@ +"""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 --adapter /tmp/grpo_trained --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): + """Replicate core/trainer_base.py:setup_tokenizer_with_model exactly, so the + adapter's resized embedding matches and both arms are prompted identically.""" + tok = AutoTokenizer.from_pretrained(base_model) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + tok.padding_side = "left" + tok.add_special_tokens({"pad_token": "[PAD]"}) + return tok + + +def load_model(base_model, tok, adapter=None, device="mps"): + model = AutoModelForCausalLM.from_pretrained(base_model, dtype=torch.float16) + model.resize_token_embeddings(len(tok)) + if adapter: + 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 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 -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("--adapter", required=True) + 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") + results = {} + for name, adapter in [("base", None), ("grpo", args.adapter)]: + 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 -tag extractor (what GRPO optimized)") + print(" loose = last-number-in-text (format-blind true accuracy)\n") + sb, lb = summarize("base", results["base"]) + sg, lg = summarize("grpo", results["grpo"]) + + for lbl, key, ab, ag in [("strict", "strict_correct", sb, sg), + ("loose ", "loose_correct", lb, lg)]: + gained = sum(1 for b, g in zip(results["base"], results["grpo"]) + if g[key] and not b[key]) + lost = sum(1 for b, g in zip(results["base"], results["grpo"]) + if b[key] and not g[key]) + print(f" {lbl}: {ag - ab:+6.1%} ({gained} fixed, {lost} broken, " + f"net {gained - lost:+d}/{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() From 57670de9e0e7fb53bf3853a9f803089670c634b6 Mon Sep 17 00:00:00 2001 From: Kartik Ayyar Date: Sat, 22 Aug 2026 14:48:50 -0700 Subject: [PATCH 4/4] fix: forward beta to GRPOConfig and stop growing the vocab on every run Two independent bugs found while running a GRPO smoke test on Apple Silicon. 1. GRPO never forwarded `beta` to GRPOConfig, so training.beta was inert and every GRPO run used TRL's default beta=0.0 -- no KL term at all. With no anchor to the reference policy, GSM8K training collapses to emitting bare digits inside empty tags, which earns full format+digit reward while abandoning chain-of-thought. Measured: held-out accuracy 44% -> 17%. GSPO already forwarded beta; GRPO did not. epsilon/epsilon_high/steps_per_generation are deliberately left un-forwarded for GRPO: their TrainingConfig defaults are GSPO paper values and would silently tighten GRPO clipping from TRL's 0.2 to 3e-4. 2. setup_tokenizer_with_model called add_special_tokens({"pad_token": "[PAD]"}) unconditionally, making the preceding `if pad_token is None` check dead code. For models that already pad (Qwen2.5 pads with <|endoftext|>) this appended a redundant token, and the unconditional resize_token_embeddings then changed the embedding row count (Qwen2.5-0.5B: 151936 -> 151666, a shrink of the published matrix). PEFT therefore serialized the full embed_tokens + lm_head into every adapter: 561 MB of which 98.4% was embeddings, for 4.4 MB of LoRA. Now pad falls back to eos when absent (no vocab growth for Llama 3 / GPT-2 either), and the embedding is only ever grown, never shrunk. Adapter size for the smoke recipe: 561 MB -> 17.6 MB. --- core/trainer_base.py | 37 ++++++++--- .../training/grpo/qwen2_5_0_5B_mps_smoke.yaml | 4 ++ scripts/ab_eval_grpo.py | 61 +++++++++++++------ trainers/grpo_trainer.py | 14 +++++ 4 files changed, 89 insertions(+), 27 deletions(-) diff --git a/core/trainer_base.py b/core/trainer_base.py index 34845b8..2278ad2 100644 --- a/core/trainer_base.py +++ b/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 diff --git a/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml b/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml index a5da365..3de4184 100644 --- a/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml +++ b/recipes/training/grpo/qwen2_5_0_5B_mps_smoke.yaml @@ -59,6 +59,10 @@ training: 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 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 diff --git a/scripts/ab_eval_grpo.py b/scripts/ab_eval_grpo.py index e608db3..773f16e 100644 --- a/scripts/ab_eval_grpo.py +++ b/scripts/ab_eval_grpo.py @@ -9,7 +9,7 @@ comparison is deterministic. Usage: - python scripts/ab_eval_grpo.py --adapter /tmp/grpo_trained --n 100 + python scripts/ab_eval_grpo.py --arm nokl=/tmp/grpo_trained --arm kl=/tmp/grpo_fixed --n 100 """ import argparse import json @@ -31,20 +31,39 @@ def build_tokenizer(base_model): - """Replicate core/trainer_base.py:setup_tokenizer_with_model exactly, so the - adapter's resized embedding matches and both arms are prompted identically.""" + """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.pad_token = tok.eos_token + tok.add_special_tokens({"pad_token": "[PAD]"}) tok.padding_side = "left" - tok.add_special_tokens({"pad_token": "[PAD]"}) 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) - model.resize_token_embeddings(len(tok)) 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() @@ -118,7 +137,8 @@ def summarize(name, rows): def main(): p = argparse.ArgumentParser() p.add_argument("--base", default="Qwen/Qwen2.5-0.5B-Instruct") - p.add_argument("--adapter", required=True) + 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) @@ -139,8 +159,9 @@ def main(): 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 [("base", None), ("grpo", args.adapter)]: + for name, adapter in arms: print(f" generating: {name}") t0 = time.time() model = load_model(args.base, tok, adapter, args.device) @@ -153,17 +174,19 @@ def main(): print("\nRESULTS strict = repo -tag extractor (what GRPO optimized)") print(" loose = last-number-in-text (format-blind true accuracy)\n") - sb, lb = summarize("base", results["base"]) - sg, lg = summarize("grpo", results["grpo"]) - - for lbl, key, ab, ag in [("strict", "strict_correct", sb, sg), - ("loose ", "loose_correct", lb, lg)]: - gained = sum(1 for b, g in zip(results["base"], results["grpo"]) - if g[key] and not b[key]) - lost = sum(1 for b, g in zip(results["base"], results["grpo"]) - if b[key] and not g[key]) - print(f" {lbl}: {ag - ab:+6.1%} ({gained} fixed, {lost} broken, " - f"net {gained - lost:+d}/{len(results['base'])})") + 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) diff --git a/trainers/grpo_trainer.py b/trainers/grpo_trainer.py index 60287ed..cc31752 100644 --- a/trainers/grpo_trainer.py +++ b/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 + # 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):