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 new file mode 100644 index 0000000..3de4184 --- /dev/null +++ b/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 + # , 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 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 diff --git a/scripts/ab_eval_grpo.py b/scripts/ab_eval_grpo.py new file mode 100644 index 0000000..773f16e --- /dev/null +++ b/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 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("--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 -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() 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):