Buckets:
| """GRPO on the reasoner with TRL (vLLM colocated on the same GPU). | |
| Reward per chain, read from its answer letter (one token) after `</think>`: | |
| 2 * target[letter] - 1 for a letter: +1 correct, -1 wrong on exact and hard labels; on soft-label items the share of | |
| annotators who chose that answer, rescaled (a 60/40 split gives +0.2 / -0.2), so a split item cannot push hard. | |
| -1.5 for no letter or a chain cut by the rollout cap (loops); minus 0.25 per 4,096 tokens past 4,096. | |
| GRPO turns the reward into a group-relative advantage that every token of that chain receives. Advantages are | |
| scaled by the batch's spread, not each group's (scale_rewards="batch", as Dr. GRPO argues), so a soft item's small | |
| differences stay small instead of being blown up to the size of a clear right/wrong group. | |
| Prompts: training items where the 4B itself is sometimes right and sometimes wrong (1-3 of its own 4 chains correct; | |
| teacher records are left out), so groups carry signal. The rollout cap applies to RL rollouts only; evaluation stays | |
| uncapped. | |
| Budget (measured on those items: chain median 5,280 tokens, mean 6,985, 21% past 8,192): each generation round makes | |
| 128 chains (16 prompts x 8) and feeds 4 optimizer steps of 32 chains, so the slow tail of a round is paid once per 4 | |
| steps; about 1.6 min/step, 120 steps in ~3.2 h. One chain per forward (full-vocabulary logits of an 8K chain are | |
| ~4 GB in bf16). --model must be a text-only checkpoint (architecture Qwen3_5ForCausalLM, e.g. a merge_lora.py output): | |
| TRL trains it and pushes its weights into vLLM by name, and the hub repo's multimodal layout would not match. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import random | |
| from common import read_jsonl, render, parse_letter, shuffled_order, split_chain | |
| def build_prompts(pool: dict, gen: list[dict], tok, rng: random.Random, n: int | None, allowed: set | None = None, | |
| fresh: dict | None = None) -> list[dict]: | |
| """allowed: ids chosen by rl_select.py (label-noise rules and the V2 readout check); None = the mixed-item rule only. | |
| fresh: never-trained questions (rl_fresh.py); used when rl_select.py chose them, with no chain check.""" | |
| rows = [] | |
| def add(it): | |
| order = shuffled_order(it, rng) | |
| rows.append({"prompt": render(tok, it, think=True, order=order), "gold": order.index(it["gold"]), "k": len(it["options"]), | |
| "target": [float(it["target"][j]) for j in order], "item_id": it["id"]}) | |
| for rec in gen: | |
| if allowed is not None and rec["id"] not in allowed: | |
| continue | |
| it = pool.get(rec["id"]) | |
| if it is None or rec.get("source") == "teacher": | |
| continue | |
| correct = sum(1 for s in rec["samples"] if s["closed"] and s["answer"] is not None and rec["order"][s["answer"]] == it["gold"]) | |
| if not 1 <= correct <= 3: | |
| continue | |
| add(it) | |
| for it in (fresh or {}).values(): | |
| if allowed is not None and it["id"] in allowed: | |
| add(it) | |
| rng.shuffle(rows) | |
| return rows[:n] if n else rows | |
| def make_reward(tok, cap_tokens: int): | |
| def reward(completions, target, k, **_): | |
| out = [] | |
| for text, t, kk in zip(completions, target, k): | |
| text = text if isinstance(text, str) else text[0]["content"] | |
| n_tok = len(tok(text, add_special_tokens=False)["input_ids"]) | |
| chain, answer = split_chain(text) | |
| letter = parse_letter(answer, kk) if chain is not None else None | |
| if letter is None or n_tok >= cap_tokens: | |
| r = -1.5 | |
| else: | |
| r = 2.0 * t[letter] - 1.0 | |
| r -= 0.25 * max(0, n_tok - 4096) / 4096 | |
| out.append(float(r)) | |
| return out | |
| return reward | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", required=True) | |
| ap.add_argument("--pool", default="data/train_pool.jsonl") | |
| ap.add_argument("--gen", default="data/gen.jsonl") | |
| ap.add_argument("--select", default="data/rl_prompts.jsonl", help="rl_select.py output; 'none' = mixed-item rule only") | |
| ap.add_argument("--fresh-pool", default="data/rl_fresh_pool.jsonl", help="rl_fresh.py output; '' = in-domain only") | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--steps", type=int, default=120) | |
| ap.add_argument("--prompts", type=int, default=None, help="number of distinct prompts to draw from") | |
| ap.add_argument("--num-gen", type=int, default=8) | |
| ap.add_argument("--chains-per-step", type=int, default=32, help="completions per optimizer step") | |
| ap.add_argument("--steps-per-gen", type=int, default=4, help="optimizer steps fed by one generation round") | |
| ap.add_argument("--micro", type=int, default=1, help="completions per forward/backward") | |
| ap.add_argument("--max-completion", type=int, default=8192) | |
| ap.add_argument("--lr", type=float, default=1e-5) | |
| ap.add_argument("--beta", type=float, default=0.0) | |
| ap.add_argument("--vllm-mem", type=float, default=0.5) | |
| ap.add_argument("--is-mode", default="token_truncate", | |
| help="vLLM importance-sampling mode; TRL's default sequence_mask multiplies per-token ratios over a whole " | |
| "chain, so at ~4K tokens tiny engine differences zero out long chains (the ones the cap penalty is for)") | |
| ap.add_argument("--seed", type=int, default=17) | |
| ap.add_argument("--resume", default=None, help="checkpoint dir to continue from (optimizer, step and data order restored)") | |
| ap.add_argument("--init-adapter", default=None, help="start from this LoRA adapter (a checkpoint dir) with a fresh optimizer, " | |
| "e.g. to continue at a different learning rate") | |
| ap.add_argument("--exclude-seen", default=None, help="completions dir of an earlier run: its prompts are left out") | |
| args = ap.parse_args() | |
| from datasets import Dataset | |
| from peft import LoraConfig | |
| from transformers import AutoTokenizer | |
| from trl import GRPOConfig, GRPOTrainer | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| rng = random.Random(args.seed) | |
| pool = {x["id"]: x for x in read_jsonl(args.pool)} | |
| allowed = None if args.select == "none" else {r["id"] for r in read_jsonl(args.select)} | |
| fresh = {x["id"]: x for x in read_jsonl(args.fresh_pool)} if args.fresh_pool and os.path.exists(args.fresh_pool) else {} | |
| rows = build_prompts(pool, read_jsonl(args.gen), tok, rng, args.prompts, allowed, fresh) | |
| if args.exclude_seen: | |
| import glob | |
| import pandas as pd | |
| seen = set() | |
| for f in glob.glob(os.path.join(args.exclude_seen, "*.parquet")): | |
| seen |= set(pd.read_parquet(f, columns=["prompt"])["prompt"]) | |
| before = len(rows) | |
| rows = [r for r in rows if r["prompt"] not in seen] | |
| print(json.dumps({"excluded_seen_prompts": before - len(rows), "prompts_left": len(rows)}), flush=True) | |
| print(json.dumps({"prompts": len(rows)}), flush=True) | |
| from transformers import AutoConfig | |
| arch = AutoConfig.from_pretrained(args.model).architectures | |
| assert arch == ["Qwen3_5ForCausalLM"], f"--model must be a text-only checkpoint, got {arch}" | |
| ga = args.chains_per_step // args.micro | |
| print(json.dumps({"chains_per_round": args.micro * ga * args.steps_per_gen, | |
| "prompts_per_round": args.micro * ga * args.steps_per_gen // args.num_gen}), flush=True) | |
| cfg = GRPOConfig( | |
| output_dir=args.out, max_steps=args.steps, learning_rate=args.lr, lr_scheduler_type="constant_with_warmup", warmup_steps=5, | |
| per_device_train_batch_size=args.micro, gradient_accumulation_steps=ga, steps_per_generation=ga * args.steps_per_gen, | |
| num_generations=args.num_gen, max_completion_length=args.max_completion, temperature=0.6, top_p=0.95, top_k=20, | |
| use_vllm=True, vllm_mode="colocate", vllm_gpu_memory_utilization=args.vllm_mem, vllm_max_model_length=32768, | |
| vllm_importance_sampling_correction=True, vllm_importance_sampling_mode=args.is_mode, beta=args.beta, loss_type="dapo", mask_truncated_completions=False, | |
| scale_rewards="batch", | |
| bf16=True, gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, | |
| logging_steps=1, save_steps=max(10, args.steps // 6), save_total_limit=3, log_completions=True, num_completions_to_print=2, | |
| report_to="none", seed=args.seed) | |
| peft_cfg = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, bias="none", task_type="CAUSAL_LM", | |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", | |
| "in_proj_qkv", "in_proj_z", "out_proj"]) | |
| model = args.model | |
| if args.init_adapter: | |
| import torch | |
| from peft import PeftModel | |
| from transformers import AutoModelForCausalLM | |
| base = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16) | |
| model, peft_cfg = PeftModel.from_pretrained(base, args.init_adapter, is_trainable=True), None | |
| trainer = GRPOTrainer(model=model, reward_funcs=[make_reward(tok, args.max_completion)], args=cfg, | |
| train_dataset=Dataset.from_list(rows), processing_class=tok, peft_config=peft_cfg) | |
| trainer.train(resume_from_checkpoint=args.resume) | |
| trainer.save_model(os.path.join(args.out, "final")) | |
| with open(os.path.join(args.out, "log_history.json"), "w") as f: | |
| json.dump(trainer.state.log_history, f, indent=1) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 9.55 kB
- Xet hash:
- 17ce88cb6d2fdc4a10cbff3becaadc7dd51c3268d9438df59e8d2e25d30c259e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.