File size: 1,608 Bytes
51ac7cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | #!/usr/bin/env python3
"""挑蒸馏用的双向 SFT 检查点：val_log.csv 里 cond 的 mse（各 t 取平均）最低的已存 step；与最优差 < TOL 的里取最晚的
（val 噪声内多训的更好）。只看 step ≥ --min_step（v8 的 SFT 教师是 step-6000）。
usage: pick_sft.py <sft_out> [--min_step 6000] [--tol 0.005] → stdout 打印 ckpt 绝对路径，<sft_out>/sft_pick.json 记理由"""
import argparse
import csv
import collections
import json
import os
import sys
def main():
ap = argparse.ArgumentParser()
ap.add_argument("out")
ap.add_argument("--min_step", type=int, default=6000)
ap.add_argument("--tol", type=float, default=0.005)
a = ap.parse_args()
rows = list(csv.DictReader(open(os.path.join(a.out, "val_log.csv"))))
by = collections.defaultdict(list)
for r in rows:
by[int(r["step"])].append(float(r["mse"]))
cand = {s: sum(v) / len(v) for s, v in by.items()
if s >= a.min_step and os.path.exists(os.path.join(a.out, f"step-{s}.safetensors"))}
if not cand:
sys.exit(f"pick_sft: {a.out} 里没有 step ≥ {a.min_step} 且有 val 记录的检查点")
best = min(cand.values())
pick = max(s for s, v in cand.items() if v <= best * (1 + a.tol))
path = os.path.abspath(os.path.join(a.out, f"step-{pick}.safetensors"))
json.dump({"pick": pick, "path": path, "val_mse_mean": cand[pick], "best": best, "tol": a.tol,
"candidates": dict(sorted(cand.items()))}, open(os.path.join(a.out, "sft_pick.json"), "w"), indent=1)
print(path)
if __name__ == "__main__":
main()
|