| |
| """挑蒸馏用的双向 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() |
|
|
