| |
| """r5_sweep.py — 第五轮总控: 方向①消融 + ②思考块适配 + ③答案级 RL, 训练↔评估交错 |
| |
| 阶段 (每阶段幂等: head+dump 均存在则跳过): |
| A1 分桶加权 easy:8 (head-only, round4 目标) |
| A2 全局权重 10 (head-only, round4 目标) |
| A3 sharp 目标 + map (head-only) |
| A4 MLP-256 头 (head-only, sharp + map) |
| B 逐步CE+LoRA (方向②, head+lora 续训 out_zh, lr 5e-5, step-ce 1.0) |
| C 答案级 RL (方向③, REINFORCE, 96 样本, λ=0.3) |
| R4 复评基线 (out_r4 head 同日复测, 校准对比口径) |
| """ |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
|
|
| os.chdir("/content/pondernet_r4") |
| PY = sys.executable |
| MODEL = "models/MiniCPM5-2B" |
| TRAIN = "data_zh/boot_train.jsonl" |
| EVAL = "data_zh/boot_eval.jsonl" |
|
|
| HEAD_COMMON = ("--train-mode head --init-adapter out_zh --step-supervise " |
| "--difficulty-prior easy:0.9,medium:0.5,hard:0.05 " |
| "--epochs 2 --batch-size 2 --accum 8 --max-steps 6 " |
| "--max-len 768 --grad-ckpt --dtype float16 --log-every 10") |
| TRAIN_COMMON = f"{PY} train_ponder_head.py --model {MODEL} --data {TRAIN}" |
|
|
| STAGES = [ |
| ("A1_map8", f"{TRAIN_COMMON} {HEAD_COMMON} --sup-weight-map easy:8,medium:2,hard:1 --output out_r5_A1"), |
| ("A2_w10", f"{TRAIN_COMMON} {HEAD_COMMON} --sup-weight 10.0 --output out_r5_A2"), |
| ("A3_sharp", f"{TRAIN_COMMON} {HEAD_COMMON} --sup-target-preset sharp --sup-weight-map easy:4,medium:2,hard:1 --output out_r5_A3"), |
| ("A4_mlp", f"{TRAIN_COMMON} {HEAD_COMMON} --sup-target-preset sharp --sup-weight-map easy:4,medium:2,hard:1 --head-mlp 256 --output out_r5_A4"), |
| ("B_stepce", f"{TRAIN_COMMON} --train-mode head+lora --init-adapter out_zh --step-supervise " |
| f"--step-ce-weight 1.0 --sup-weight-map easy:4,medium:2,hard:1 --lr-lora 5e-5 " |
| f"--epochs 1 --batch-size 2 --accum 8 --max-steps 6 --max-len 768 " |
| f"--grad-ckpt --dtype float16 --log-every 5 --output out_r5_B"), |
| ("C_rl", f"{PY} train_ponder_rl.py --model {MODEL} --adapter out_zh " |
| f"--head out_r4/ponder_head.safetensors --data {TRAIN} " |
| f"--n-samples 96 --batch-size 4 --iters 24 --step-penalty 0.3 " |
| f"--ent-coef 0.01 --lr 5e-4 --max-steps 6 --max-len 768 " |
| f"--grad-ckpt --output out_r5_C"), |
| ] |
| EVALS = { |
| "A1_map8": "--head out_r5_A1/ponder_head.safetensors --adapter out_zh", |
| "A2_w10": "--head out_r5_A2/ponder_head.safetensors --adapter out_zh", |
| "A3_sharp": "--head out_r5_A3/ponder_head.safetensors --adapter out_zh", |
| "A4_mlp": "--head out_r5_A4/ponder_head.safetensors --adapter out_zh", |
| "B_stepce": "--head out_r5_B/ponder_head.safetensors --adapter out_r5_B", |
| "C_rl": "--head out_r5_C/ponder_head.safetensors --adapter out_zh", |
| "R4_baseline": "--head out_r4/ponder_head.safetensors --adapter out_zh", |
| } |
|
|
|
|
| def gpu_wait(min_free_mb=9000, tries=12): |
| """共享 T4 守卫: 等到空闲显存够用再开下一阶段。""" |
| for _ in range(tries): |
| r = subprocess.run( |
| "nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits", |
| shell=True, capture_output=True, text=True) |
| try: |
| free = int(r.stdout.strip().splitlines()[0]) |
| except Exception: |
| free = min_free_mb |
| if free >= min_free_mb: |
| return True |
| print(f"[gpu_wait] 空闲 {free}MB < {min_free_mb}MB, 等 120s", flush=True) |
| time.sleep(120) |
| return False |
|
|
|
|
| def sh(cmd, timeout=7200): |
| print(f"\n$ {cmd}\n", flush=True) |
| r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) |
| return r.returncode, (r.stdout or "") + (r.stderr or "") |
|
|
|
|
| def log_lines(tag, text, keys, limit=18): |
| keep = [l for l in text.splitlines() if any(k in l for k in keys)] |
| print(f"----- [{tag}] 关键日志 -----", flush=True) |
| for l in keep[-limit:]: |
| print(l, flush=True) |
|
|
|
|
| def main(): |
| summary = {} |
| t0 = time.time() |
| for name, cmd in STAGES: |
| head = f"out_r5_{name.split('_')[0]}" |
| head_path = os.path.join(head, "ponder_head.safetensors") |
| dump_path = os.path.join(head, "eval_dump.jsonl") |
| print(f"\n########## STAGE {name} ({time.time()-t0:.0f}s) ##########", flush=True) |
| gpu_wait() |
| |
| if not os.path.exists(head_path): |
| try: |
| rc, out = sh(cmd, timeout=10800) |
| log_lines(name, out, ("步数硬监督", "逐步 CE", "可训练参数", "step ", |
| "iter ", "完成:", "已保存", "Error", "Traceback", |
| "最近10迭代")) |
| summary[name] = f"train rc={rc}" |
| if rc != 0: |
| summary[name] += " (TRAIN_FAIL, 跳过该阶段 eval)" |
| continue |
| except Exception as e: |
| summary[name] = f"train EXC={e} (跳过 eval)" |
| continue |
| else: |
| print(f"[skip-train] {head_path} 已存在", flush=True) |
| summary[name] = "train skipped (exists)" |
| |
| if not os.path.exists(dump_path): |
| ecfg = EVALS[name] |
| ecmd = (f"{PY} eval_ponder.py --model {MODEL} --data {EVAL} {ecfg} " |
| f"--max-steps 6 --max-len 768 --batch-size 4 --tag r5_{name} " |
| f"--dump {dump_path}") |
| try: |
| rc, out = sh(ecmd, timeout=3600) |
| log_lines(f"eval:{name}", out, ("全局 CE", "难度", "easy", "medium", "hard", |
| "逐样本", "Error", "Traceback"), limit=12) |
| summary[name] += f" | eval rc={rc}" |
| except Exception as e: |
| summary[name] += f" | eval EXC={e}" |
| else: |
| summary[name] += " | eval skipped (exists)" |
|
|
| |
| dump_path = "out_r4/eval_dump_r5day.jsonl" |
| if not os.path.exists(dump_path): |
| print(f"\n########## EVAL R4_baseline ({time.time()-t0:.0f}s) ##########", flush=True) |
| gpu_wait() |
| ecmd = (f"{PY} eval_ponder.py --model {MODEL} --data {EVAL} {EVALS['R4_baseline']} " |
| f"--max-steps 6 --max-len 768 --batch-size 4 --tag r4_baseline " |
| f"--dump {dump_path}") |
| try: |
| rc, out = sh(ecmd, timeout=3600) |
| log_lines("eval:R4", out, ("全局 CE", "难度", "easy", "medium", "hard", "逐样本"), 10) |
| summary["R4_baseline"] = f"eval rc={rc}" |
| except Exception as e: |
| summary["R4_baseline"] = f"eval EXC={e}" |
|
|
| with open("r5_summary.json", "w", encoding="utf-8") as f: |
| json.dump({"wall_s": round(time.time() - t0), "stages": summary}, |
| f, ensure_ascii=False, indent=1) |
| print(f"\n########## SWEEP DONE ({time.time()-t0:.0f}s) ##########", flush=True) |
| for k, v in summary.items(): |
| print(f" {k}: {v}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
