v5d: save best-epoch checkpoint during SFT

training/train_sft.py

@@ -382,6 +382,8 @@ def main() -> None:
 
     t_start = time.time()
     global_step = 0
+    best_eval_score = -1.0
+    best_epoch_dir: Optional[Path] = None
     for epoch in range(cfg.num_epochs):
         random.shuffle(examples)
         total_loss = 0.0
@@ -399,7 +401,16 @@ def main() -> None:
                 log({"step": global_step, "train/epoch": epoch + 1, "train/example": i, "train/loss": total_loss / i})
         log({"step": global_step, "train/epoch": epoch + 1, "train/loss": total_loss / max(len(examples), 1)})
         if cfg.eval_every_epoch:
-            log(evaluate(epoch + 1))
+            eval_result = evaluate(epoch + 1)
+            log(eval_result)
+            score = eval_result.get("eval/avg_grader_score", 0.0)
+            if score > best_eval_score:
+                best_eval_score = score
+                best_epoch_dir = artifact_dir / f"best_epoch_{epoch+1}"
+                best_epoch_dir.mkdir(parents=True, exist_ok=True)
+                model.save_pretrained(str(best_epoch_dir))
+                tokenizer.save_pretrained(str(best_epoch_dir))
+                print(f"[checkpoint] new best epoch {epoch+1}: score={score:.4f}", flush=True)
 
     summary = {
         "run_finished_at": datetime.now(timezone.utc).isoformat(),
@@ -430,6 +441,18 @@ def main() -> None:
         )
         print(f"[push] done -> https://huggingface.co/{repo_id}", flush=True)
 
+        if best_epoch_dir and best_epoch_dir.exists():
+            best_repo = f"{repo_id}-best"
+            create_repo(repo_id=best_repo, repo_type="model", exist_ok=True, private=False, token=token)
+            HfApi(token=token).upload_folder(
+                folder_path=str(best_epoch_dir),
+                repo_id=best_repo,
+                repo_type="model",
+                token=token,
+                commit_message=f"Best epoch checkpoint (score={best_eval_score:.4f})",
+            )
+            print(f"[push] best checkpoint -> https://huggingface.co/{best_repo}", flush=True)
+
 
 def _parse_args() -> SftConfig:
     p = argparse.ArgumentParser()