# -*- coding: utf-8 -*- """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)" # ---- R4 同日复评基线 ---- 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()