Spaces:
Running
Running
| """Bounded OpenEnv RL pilot: frozen baseline -> GRPO LoRA -> identical evaluation. | |
| TRAIN_GPU=H100:2 modal run modal_app.py::job --max-turns 2 --reward roofline \ | |
| --cmd 'timeout 5400 python small_run.py --out /runs/proof-20261019' | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import subprocess | |
| import sys | |
| import time | |
| from pathlib import Path | |
| from baseline_eval import split_task_ids, summarize | |
| TRAIN = ["level1/19_ReLU", "level1/20_LeakyReLU", "level1/21_Sigmoid", "level1/22_Tanh"] | |
| TRANSFER = ["level1/25_Swish", "level1/28_HardSigmoid", "level1/32_HardTanh"] | |
| HELDOUT = ["level1/23_Softmax", "level1/24_LogSoftmax", "level1/88_MinGPTNewGelu"] | |
| def compare(before, after, groups): | |
| """Compare the same tasks; never infer a training gain from within-episode refinement.""" | |
| if sorted(e["task_id"] for e in before) != sorted(e["task_id"] for e in after): | |
| raise ValueError("before/after task sets differ") | |
| if any(e.get("error") for e in before + after): | |
| raise ValueError("evaluation contains infrastructure errors") | |
| result = {} | |
| for name, ids in groups.items(): | |
| result[name] = {} | |
| for label, episodes in [("before", before), ("after", after)]: | |
| chosen = [e for e in episodes if e["task_id"] in ids] | |
| if sorted(e["task_id"] for e in chosen) != sorted(ids): | |
| raise ValueError(f"missing or repeated task in {name}/{label}") | |
| summary = summarize(chosen) | |
| # Including failures prevents a policy solving fewer easy tasks from appearing faster. | |
| summary["mean_best_speedup_including_failures"] = sum( | |
| max((t["speedup"] for t in e["turns"] if t["correct"]), default=0) | |
| for e in chosen | |
| ) / len(chosen) | |
| result[name][label] = summary | |
| return result | |
| def main(): | |
| p = argparse.ArgumentParser() | |
| p.add_argument("--out", required=True) | |
| p.add_argument("--model", default="Qwen/Qwen2.5-Coder-7B-Instruct") | |
| p.add_argument("--steps", type=int, default=16) | |
| a = p.parse_args() | |
| out = Path(a.out) | |
| out.mkdir(parents=True, exist_ok=False) | |
| groups = {"training_tasks": TRAIN, "reserved_transfer": TRANSFER, "official_heldout": HELDOUT} | |
| assert set(TRAIN + TRANSFER) <= set(split_task_ids("train", [1])) | |
| assert set(HELDOUT) <= set(split_task_ids("eval", [1])) | |
| assert not set(TRAIN) & set(TRANSFER + HELDOUT) | |
| manifest = {"model": a.model, "steps": a.steps, "groups": groups, | |
| "eval_temperature": 0.0, "max_new_tokens": 2048, "seed": 42, | |
| "created_at_unix": time.time(), "torch_baseline": "eager"} | |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2)) | |
| task_file = out / "train_tasks.json" | |
| task_file.write_text(json.dumps(TRAIN)) | |
| def run(args, name): | |
| print(f"PHASE {name}: {args}", flush=True) | |
| start = time.monotonic() | |
| with (out / f"{name}.log").open("w") as log: | |
| proc = subprocess.Popen([sys.executable, "-u", *args], stdout=subprocess.PIPE, | |
| stderr=subprocess.STDOUT, text=True, bufsize=1) | |
| for line in proc.stdout: | |
| print(line, end="", flush=True) | |
| log.write(line) | |
| log.flush() | |
| if proc.wait(): | |
| raise RuntimeError(f"{name} failed; see {log.name}") | |
| manifest.setdefault("phase_seconds", {})[name] = time.monotonic() - start | |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2)) | |
| def evaluate(model, label): | |
| run(["baseline_eval.py", "--env-url", os.environ["KERNEL_ENV_URL"], "--model", model, | |
| "--backend", "vllm", "--temperature", "0", "--max-new-tokens", "2048", | |
| "--max-model-len", "8192", "--task-ids", *TRAIN, *TRANSFER, *HELDOUT, | |
| "--out", str(out / f"{label}.jsonl")], label) | |
| evaluate(a.model, "before") | |
| run(["grpo_kernel.py", "--model_name_or_path", a.model, "--dtype", "bfloat16", "--bf16", "true", | |
| "--env_url", os.environ["KERNEL_ENV_URL"], "--task_file", str(task_file), | |
| "--use_peft", "--lora_r", "16", "--lora_alpha", "32", "--lora_target_modules", "all-linear", | |
| "--use_vllm", "--vllm_mode", "colocate", "--vllm_gpu_memory_utilization", "0.25", | |
| "--vllm_max_model_length", "8192", "--num_generations", "4", | |
| "--per_device_train_batch_size", "4", "--gradient_accumulation_steps", "1", | |
| "--max_completion_length", "2048", "--max_steps", str(a.steps), "--learning_rate", "5e-5", | |
| "--gradient_checkpointing", "true", "--max_grad_norm", "0.05", "--loss_type", "dr_grpo", | |
| "--return_mode", "discounted", "--gamma", "0.4", "--temperature", "0.8", | |
| "--seed", "42", "--logging_steps", "1", "--report_to", "none", "--save_steps", "8", | |
| "--save_total_limit", "2", "--output_dir", str(out / "adapter")], "train") | |
| run(["merge_lora.py", "--base", a.model, "--adapter", str(out / "adapter"), | |
| "--out", str(out / "merged")], "merge") | |
| evaluate(str(out / "merged"), "after") | |
| episodes = [[json.loads(l) for l in (out / f"{label}.jsonl").read_text().splitlines()] | |
| for label in ("before", "after")] | |
| result = compare(*episodes, groups) | |
| (out / "comparison.json").write_text(json.dumps(result, indent=2)) | |
| print("COMPARISON", json.dumps(result, indent=2), flush=True) | |
| if __name__ == "__main__": | |
| main() | |