Spaces:
Running
Running
| """GRPO on kernel_env: TRL GRPOTrainer + multi-turn rollout_func over the KernelEnv client. | |
| accelerate launch grpo_kernel.py --model_name_or_path Qwen/Qwen2.5-Coder-7B-Instruct \ | |
| --env_url http://EVAL_BOX:8000 --num_generations 8 --use_vllm --output_dir runs/grpo | |
| All GRPOConfig / ModelConfig fields are CLI flags (TrlParser). Dataset rows are KernelBench task | |
| ids; rollout_func resets one env session per row (TRL already repeats each row `num_generations` | |
| times = one GRPO group per task), then alternates model turn -> env.step -> feedback turn until the | |
| env ends the episode. The whole episode is one completion: model tokens train (env_mask=1), feedback | |
| tokens don't (env_mask=0). Reward = discounted sum of per-turn env rewards, or best turn. | |
| Held-out eval = TRL's own eval loop over eval-split task ids (`--eval_strategy steps --eval_steps N`). | |
| """ | |
| import json | |
| import logging | |
| import os | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dataclasses import dataclass, field | |
| from datasets import Dataset | |
| from transformers.trainer_utils import get_last_checkpoint | |
| from trl import GRPOConfig, GRPOTrainer, ModelConfig, TrlParser, get_peft_config | |
| from baseline_eval import extract_code, feedback_message, split_task_ids, summarize | |
| from kernel_env import KernelAction, KernelEnv | |
| log = logging.getLogger("grpo_kernel") | |
| class KernelArgs: | |
| env_url: str = field(default="http://localhost:8000", metadata={"help": "kernel_env server (on eval GPUs)"}) | |
| levels: list[int] = field(default_factory=lambda: [1], metadata={"help": "KernelBench levels to use"}) | |
| kernelbench_dir: str = field(default=os.environ.get("KERNELBENCH_DIR", ""), metadata={"help": "local copy, for task ids"}) | |
| return_mode: str = field(default="discounted", metadata={"help": "discounted | best"}) | |
| enable_thinking: bool = field(default=False, metadata={"help": "Qwen3-style <think> in the chat template (baseline: off)"}) | |
| gamma: float = field(default=0.9, metadata={"help": "discount over turns (return_mode=discounted)"}) | |
| env_timeout_s: float = field(default=900.0, metadata={"help": "per env message; covers GPU-pool queueing"}) | |
| max_eval_tasks: int = field(default=0, metadata={"help": "cap held-out tasks per eval (0 = all)"}) | |
| task_file: str = field(default="", metadata={"help": "JSON list of training task ids (curriculum); default: whole train split"}) | |
| def episode_return(rewards: list, mode: str, gamma: float) -> float: | |
| if mode == "best": | |
| return max(rewards, default=0.0) | |
| return sum(gamma**t * r for t, r in enumerate(rewards)) | |
| def text_tokenizer(processing_class): | |
| """VLM checkpoints (e.g. Qwen3.5) load a Processor: its chat template returns nested lists and it has no | |
| eos_token_id. Rollouts are text-only, so use the wrapped tokenizer.""" | |
| return getattr(processing_class, "tokenizer", processing_class) | |
| def feedback_ids(tok, text: str, enable_thinking: bool = False) -> list: | |
| """Tokens after an assistant turn's EOS: the feedback user turn + next assistant header. | |
| Cut after the assistant turn's EOS (second-to-last EOS; a default system prompt may add one in front) | |
| rather than diffing two renders: templates such as | |
| Qwen3.5's re-render earlier assistant turns differently (they drop the empty think block), so the | |
| conversation is not prefix-stable. The rollout stitches the model's own tokens, so only this suffix matters. | |
| """ | |
| msgs = [{"role": "user", "content": "x"}, {"role": "assistant", "content": "y"}, {"role": "user", "content": text}] | |
| full = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True, return_dict=False, | |
| enable_thinking=enable_thinking) # templates without the switch ignore it | |
| eos = [i for i, t in enumerate(full) if t == tok.eos_token_id] | |
| return full[eos[-2] + 1:] | |
| def make_rollout_func(args: KernelArgs, suspicious_path: str): | |
| def rollout_func(prompts: list, trainer) -> dict: | |
| tok = text_tokenizer(trainer.processing_class) | |
| mode = "train" if trainer.model.training else "eval" | |
| split = "train" if mode == "train" else "eval" | |
| envs = [KernelEnv(base_url=args.env_url, message_timeout_s=args.env_timeout_s).sync() for _ in prompts] | |
| pool = ThreadPoolExecutor(max_workers=len(envs)) | |
| try: | |
| obs = list(pool.map(lambda e, t: e.reset(task_id=t, split=split).observation, envs, prompts)) | |
| prompt_ids = [tok.apply_chat_template([{"role": "user", "content": o.problem}], add_generation_prompt=True, | |
| tokenize=True, return_dict=False, enable_thinking=args.enable_thinking) | |
| for o in obs] | |
| n = len(prompts) | |
| comp, mask, lps = [[] for _ in range(n)], [[] for _ in range(n)], [[] for _ in range(n)] | |
| eps = [{"task_id": t, "turns": [], "codes": [], "completions": [], "step": trainer.state.global_step, | |
| "mode": mode} for t in prompts] | |
| have_lp, live = True, list(range(n)) | |
| while live: | |
| new_ids, new_lps = trainer._generate_single_turn([prompt_ids[i] + comp[i] for i in live], None, {}, 1) | |
| have_lp &= new_lps is not None # HF generate gives none; vLLM does | |
| texts = tok.batch_decode(new_ids, skip_special_tokens=True) | |
| codes = [extract_code(t) for t in texts] | |
| results = list(pool.map(lambda i, c: envs[i].step(KernelAction(code=c)), live, codes)) | |
| still = [] | |
| for j, (i, r) in enumerate(zip(live, results)): | |
| ids, o = list(new_ids[j]), r.observation | |
| comp[i] += ids | |
| mask[i] += [1] * len(ids) | |
| lps[i] += list(new_lps[j]) if have_lp else [] | |
| truncated = not ids or ids[-1] != tok.eos_token_id | |
| eps[i]["codes"].append(codes[j]) | |
| eps[i]["completions"].append(texts[j]) | |
| eps[i]["turns"].append({"stage": o.stage, "correct": bool(o.correct), | |
| "speedup": o.speedup or 0.0, "reward": r.reward or 0.0, | |
| "roofline_frac": o.roofline_frac, "ref_roofline_frac": o.ref_roofline_frac, | |
| "t_ref_ms": o.t_ref_ms, "t_new_ms": o.t_new_ms, "t_min_ms": o.t_min_ms, | |
| "bound": o.bound, "feedback": o.feedback, "truncated": truncated, | |
| "completion_tokens": len(ids)}) | |
| if r.done: | |
| continue | |
| # close a truncated turn ourselves so the template stays well-formed | |
| extra = ([] if not truncated else [tok.eos_token_id]) + feedback_ids(tok, feedback_message(o, truncated), args.enable_thinking) | |
| comp[i] += extra | |
| mask[i] += [0] * len(extra) | |
| lps[i] += [0.0] * len(extra) if have_lp else [] | |
| still.append(i) | |
| live = still | |
| for ep, env in zip(eps, envs): | |
| st = env.state() | |
| ep.update(episode_id=st.episode_id, history=st.history) | |
| # every episode, per rank, for report.py and for reading what the policy is learning | |
| rank = getattr(getattr(trainer, "accelerator", None), "process_index", 0) | |
| with open(os.path.join(os.path.dirname(suspicious_path), f"rollouts-rank{rank}.jsonl"), "a") as f: | |
| for ep in eps: | |
| f.write(json.dumps(ep) + "\n") | |
| finally: | |
| for e in envs: | |
| e.close() | |
| pool.shutdown() | |
| returns = [episode_return([t["reward"] for t in ep["turns"]], args.return_mode, args.gamma) for ep in eps] | |
| log_episodes(trainer, mode, eps, suspicious_path) | |
| return {"prompt_ids": prompt_ids, "completion_ids": comp, "logprobs": lps if have_lp else None, | |
| "env_mask": mask, "episode_return": returns} | |
| return rollout_func | |
| def log_episodes(trainer, mode: str, eps: list, suspicious_path: str): | |
| s = summarize(eps) | |
| for k in ("compile_pct", "correct_pct", "correct_turn1_pct", "faster_than_ref_pct", "mean_speedup", | |
| "median_roofline_frac", "suspicious"): | |
| if s[k] is not None: | |
| trainer._metrics[mode][f"kernel/{k}"].append(float(s[k])) | |
| trainer._metrics[mode]["kernel/turns"].append(sum(len(ep["turns"]) for ep in eps) / max(len(eps), 1)) | |
| # env flags speedup > 50x as possible reward hack: keep the code for manual review | |
| for ep in eps: | |
| for h in ep["history"]: | |
| if h.get("suspicious") or h.get("roofline_violation"): | |
| log.warning("suspicious %.1fx speedup on %s turn %d", h["speedup"], ep["task_id"], h["turn"]) | |
| with open(suspicious_path, "a") as f: | |
| f.write(json.dumps({"step": trainer.state.global_step, "mode": mode, "task_id": ep["task_id"], | |
| **h, "code": ep["codes"][h["turn"] - 1]}) + "\n") | |
| def env_return(prompts, completions, episode_return, **kwargs): | |
| return episode_return | |
| def main(): | |
| args, cfg, model_args = TrlParser((KernelArgs, GRPOConfig, ModelConfig)).parse_args_and_config() | |
| cfg.model_init_kwargs = {**(cfg.model_init_kwargs or {}), "dtype": model_args.dtype or "auto"} | |
| train_ids = split_task_ids("train", args.levels, args.kernelbench_dir) | |
| eval_ids = split_task_ids("eval", args.levels, args.kernelbench_dir) | |
| if args.task_file: | |
| with open(args.task_file) as f: | |
| chosen = json.load(f) | |
| leaked = set(chosen) - set(train_ids) # never train on held-out (or unknown) tasks | |
| if leaked: | |
| raise ValueError(f"task_file has tasks outside the train split: {sorted(leaked)[:5]}") | |
| train_ids = chosen | |
| eval_ids = eval_ids[: args.max_eval_tasks] if args.max_eval_tasks else eval_ids | |
| log.info("%d train / %d eval tasks", len(train_ids), len(eval_ids)) | |
| os.makedirs(cfg.output_dir, exist_ok=True) | |
| trainer = GRPOTrainer( | |
| model=model_args.model_name_or_path, | |
| reward_funcs=env_return, | |
| args=cfg, | |
| train_dataset=Dataset.from_dict({"prompt": train_ids}), | |
| eval_dataset=Dataset.from_dict({"prompt": eval_ids}) if cfg.eval_strategy != "no" else None, | |
| peft_config=get_peft_config(model_args), | |
| rollout_func=make_rollout_func(args, os.path.join(cfg.output_dir, "suspicious.jsonl")), | |
| ) | |
| # resume after a crash or a Modal retry: newest checkpoint-* in output_dir, if any | |
| last = get_last_checkpoint(cfg.output_dir) | |
| if last: | |
| log.info("resuming from %s", last) | |
| trainer.train(resume_from_checkpoint=last) | |
| trainer.save_model(cfg.output_dir) | |
| if __name__ == "__main__": | |
| logging.basicConfig(level=logging.INFO) | |
| main() | |