File size: 8,280 Bytes
246cc9a | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 | #!/usr/bin/env python
"""Inspect / recover GRPO training checkpoints from the Hugging Face Hub.
Use this to answer "did Colab/Kaggle save anything before it died?" and to
pull the latest checkpoint locally for manual inspection or DGX resume.
Examples
--------
List what's on the Hub:
python scripts/check_hub_checkpoints.py --hub-model-id noanya/zombiee
Show training progress (step, loss, learning rate from trainer_state.json):
python scripts/check_hub_checkpoints.py --hub-model-id noanya/zombiee --info
Download the latest checkpoint to ./recovered/:
python scripts/check_hub_checkpoints.py --hub-model-id noanya/zombiee \\
--download ./recovered
Then resume training from it:
python -m training.train \\
--resume-from-checkpoint ./recovered \\
--push-to-hub --hub-model-id noanya/zombiee \\
--max-steps 4000 --output-dir ./lora_v1
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from datetime import datetime, timezone
def parse_args():
p = argparse.ArgumentParser(
description="List / download GRPO training checkpoints from HF Hub.",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__,
)
p.add_argument(
"--hub-model-id", default=os.environ.get("HUB_MODEL_ID", "noanya/zombiee"),
help="HF Hub repo id, e.g. 'noanya/zombiee' (default: $HUB_MODEL_ID or noanya/zombiee).",
)
p.add_argument(
"--info", action="store_true",
help="Read trainer_state.json from the latest checkpoint and print training progress.",
)
p.add_argument(
"--download", metavar="DIR", default=None,
help="Download the latest checkpoint to this directory.",
)
p.add_argument(
"--checkpoint", metavar="N", type=int, default=None,
help="Operate on checkpoint-N specifically instead of the latest.",
)
p.add_argument(
"--token", default=os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HF_TOKEN"),
help="HF token (default: $HUGGINGFACE_TOKEN / $HF_TOKEN). Required for private repos.",
)
return p.parse_args()
def list_checkpoints(api, repo_id, token):
"""Return (sorted list of checkpoint step numbers, list of root files)."""
try:
files = api.list_repo_files(repo_id, token=token)
except Exception as e:
print(f"ERROR: could not list {repo_id}: {e}", file=sys.stderr)
print(
" If the repo is private, set HUGGINGFACE_TOKEN. If it doesn't exist yet,\n"
" no training run has pushed to it.",
file=sys.stderr,
)
sys.exit(1)
steps = set()
root_files = []
for f in files:
if f.startswith("checkpoint-"):
try:
steps.add(int(f.split("/", 1)[0].split("-", 1)[1]))
except ValueError:
pass
elif "/" not in f:
root_files.append(f)
return sorted(steps), sorted(root_files)
def fetch_trainer_state(api, repo_id, step, token, work_dir):
"""Download trainer_state.json from checkpoint-step and return parsed dict."""
from huggingface_hub import hf_hub_download
path = hf_hub_download(
repo_id=repo_id,
filename=f"checkpoint-{step}/trainer_state.json",
local_dir=work_dir,
token=token,
)
with open(path) as f:
return json.load(f)
def fmt_age(iso_or_dt):
"""Render 'X minutes/hours/days ago' from an HF datetime."""
if isinstance(iso_or_dt, str):
try:
dt = datetime.fromisoformat(iso_or_dt.replace("Z", "+00:00"))
except ValueError:
return iso_or_dt
else:
dt = iso_or_dt
if dt.tzinfo is None:
dt = dt.replace(tzinfo=timezone.utc)
delta = datetime.now(timezone.utc) - dt
s = int(delta.total_seconds())
if s < 60:
return f"{s}s ago"
if s < 3600:
return f"{s // 60}m ago"
if s < 86400:
return f"{s // 3600}h {(s % 3600) // 60}m ago"
return f"{s // 86400}d {(s % 86400) // 3600}h ago"
def cmd_list(api, repo_id, token):
steps, root_files = list_checkpoints(api, repo_id, token)
print(f"Repo: https://huggingface.co/{repo_id}")
try:
info = api.repo_info(repo_id, token=token)
print(f"Last commit: {info.sha[:8]} ({fmt_age(info.lastModified)})")
except Exception:
pass
print()
if not steps:
print("No checkpoint-* directories found.")
if root_files:
print(f"Root files present: {', '.join(root_files)}")
print("(Looks like only a final-model push, no intermediate checkpoints.)")
else:
print("Repo is empty — training never reached the first save.")
return
print(f"Found {len(steps)} checkpoint(s): {', '.join(f'checkpoint-{s}' for s in steps)}")
print(f"Latest: checkpoint-{steps[-1]}")
if root_files:
print(f"Root files: {', '.join(root_files)}")
def cmd_info(api, repo_id, token, step):
steps, _ = list_checkpoints(api, repo_id, token)
if not steps:
print("No checkpoints to inspect.", file=sys.stderr)
sys.exit(1)
target = step if step is not None else steps[-1]
if target not in steps:
print(f"checkpoint-{target} not on hub. Available: {steps}", file=sys.stderr)
sys.exit(1)
print(f"Inspecting checkpoint-{target}...")
state = fetch_trainer_state(api, repo_id, target, token, "/tmp/_hub_inspect")
print()
print(f" global_step : {state.get('global_step')}")
print(f" epoch : {state.get('epoch'):.4f}" if state.get("epoch") is not None else " epoch : ?")
print(f" max_steps : {state.get('max_steps')}")
print(f" best_metric : {state.get('best_metric')}")
print(f" total_flos : {state.get('total_flos')}")
log_history = state.get("log_history", [])
if log_history:
print(f" log entries : {len(log_history)}")
last = log_history[-1]
print()
print(" Most recent log entry:")
for k in ("loss", "learning_rate", "grad_norm", "reward", "kl", "step"):
if k in last:
v = last[k]
if isinstance(v, float):
print(f" {k:18}: {v:.6f}")
else:
print(f" {k:18}: {v}")
pct = (target / state["max_steps"] * 100) if state.get("max_steps") else None
if pct is not None:
print()
print(f"Progress: {target} / {state['max_steps']} steps ({pct:.1f}% done)")
def cmd_download(api, repo_id, token, target_dir, step):
from huggingface_hub import snapshot_download
steps, _ = list_checkpoints(api, repo_id, token)
if not steps:
print("Nothing to download — no checkpoints on hub.", file=sys.stderr)
sys.exit(1)
chosen = step if step is not None else steps[-1]
if chosen not in steps:
print(f"checkpoint-{chosen} not on hub. Available: {steps}", file=sys.stderr)
sys.exit(1)
os.makedirs(target_dir, exist_ok=True)
print(f"Downloading checkpoint-{chosen} from {repo_id} -> {target_dir}/")
local = snapshot_download(
repo_id=repo_id,
allow_patterns=[f"checkpoint-{chosen}/*"],
local_dir=target_dir,
token=token,
)
final = os.path.join(local, f"checkpoint-{chosen}")
print()
print(f"Done. Local path: {final}")
print()
print("To resume training from this checkpoint:")
print(f" python -m training.train \\")
print(f" --resume-from-checkpoint {final} \\")
print(f" --push-to-hub --hub-model-id {repo_id} \\")
print(f" --output-dir ./lora_v1")
def main():
args = parse_args()
try:
from huggingface_hub import HfApi
except ImportError:
print("pip install huggingface_hub", file=sys.stderr)
sys.exit(1)
api = HfApi()
if args.download:
cmd_download(api, args.hub_model_id, args.token, args.download, args.checkpoint)
elif args.info:
cmd_info(api, args.hub_model_id, args.token, args.checkpoint)
else:
cmd_list(api, args.hub_model_id, args.token)
if __name__ == "__main__":
main()
|