| |
| """Sharded checkpoint tools: inspect, verify, export to HF safetensors, reshard. |
| |
| python scripts/ckpt_convert.py inspect RUN/checkpoints/micro_N |
| python scripts/ckpt_convert.py verify RUN/checkpoints/micro_N |
| python scripts/ckpt_convert.py export RUN/checkpoints/micro_N OUT_DIR |
| python scripts/ckpt_convert.py reshard RUN/checkpoints/micro_N NEW_DIR --world-size 2 |
| |
| CPU only, no GPU, no process group, no heavy-job lock: `export` is a streaming |
| byte copy (peak RSS ~ a few MiB plus the Python runtime, independent of model |
| size) and `reshard` re-cuts the weights the same way. The ONE part that holds |
| real memory is the optimizer re-cut inside `reshard` (one NEW rank's optimizer |
| state, e.g. 12 B/param/new_ws for fp32-master AdamW); it prints that estimate |
| and refuses to start when the box cannot afford it (CLAUDE.md: 45 GB |
| MemAvailable floor for anything heavy, no swap). |
| |
| WHY THIS EXISTS. `train.checkpoint_format = "sharded"` writes per-rank files |
| that eval, the rollout weight sync and external tools cannot read. `export` |
| turns one into exactly the directory the gathered writer would have produced |
| (byte-identical; tested by file hash) without ever holding the model in RAM -- |
| the gathered path's rank-0 gather measured +28 GB RSS on OLMoE (docs/ |
| MEASUREMENTS.md, Track M). `reshard` is the ONLY way to resume at a different |
| world size: resume refuses it, because sharded weights are rank-local row |
| slices and silently loading them under another partitioning trains plausibly |
| on scrambled weights. |
| |
| Exit status: 0 ok, 1 refused/corrupt (message on stderr), 2 usage. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
|
|
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| from align.train import sharded_ckpt as sc |
|
|
|
|
| def _mem_available_gib() -> float: |
| with open("/proc/meminfo") as f: |
| for line in f: |
| if line.startswith("MemAvailable:"): |
| return int(line.split()[1]) / 2**20 |
| return 0.0 |
|
|
|
|
| def _optimizer_bytes(src: str, ws: int) -> int: |
| return sum(os.path.getsize(os.path.join(src, f"optimizer_rank{r}.pt")) for r in range(ws) if os.path.exists(os.path.join(src, f"optimizer_rank{r}.pt"))) |
|
|
|
|
| def cmd_inspect(a) -> int: |
| m = sc.read_manifest(a.src) |
| total = sum(f["bytes"] for f in m["files"].values()) |
| print( |
| json.dumps( |
| { |
| "world_size": m["world_size"], |
| "tensors": len(m["tensors"]), |
| "weight_bytes": total, |
| "export_family": (m.get("export") or {}).get("family"), |
| "resharded_from": m.get("resharded_from"), |
| "problem": sc.sharded_problem(a.src, lambda p: None), |
| }, |
| indent=1, |
| ) |
| ) |
| return 0 |
|
|
|
|
| def cmd_verify(a) -> int: |
| m = sc.read_manifest(a.src) |
| bad = [] |
| for r in range(m["world_size"]): |
| p = sc.verify_rank_file(a.src, r) |
| print(f"rank {r}: {'OK' if p is None else p}") |
| if p: |
| bad.append(p) |
| return 1 if bad else 0 |
|
|
|
|
| def cmd_export(a) -> int: |
| from align.utils.guard import assert_paths_off_root |
|
|
| assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))]) |
| res = sc.export_hf(a.src, a.dst, verify=not a.no_verify) |
| print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()})) |
| return 0 |
|
|
|
|
| def cmd_reshard(a) -> int: |
| from align.utils.guard import assert_paths_off_root |
|
|
| assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))]) |
| m = sc.read_manifest(a.src) |
| if not a.drop_optimizer: |
| est = _optimizer_bytes(a.src, m["world_size"]) / max(1, a.world_size) / 2**30 |
| avail = _mem_available_gib() |
| print(f"optimizer re-cut peak ~ {est:.1f} GiB; MemAvailable {avail:.0f} GiB", file=sys.stderr) |
| if not a.allow_large and avail < est + 45: |
| print( |
| f"refusing: need MemAvailable >= {est + 45:.0f} GiB (estimate + the 45 GiB box floor) and have " |
| f"{avail:.0f}. Wait for the box, or pass --allow-large if you accept the risk, or --drop-optimizer " |
| "(fresh moments; NOT lossless).", |
| file=sys.stderr, |
| ) |
| return 1 |
| res = sc.reshard(a.src, a.dst, a.world_size, drop_optimizer=a.drop_optimizer) |
| print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()})) |
| return 0 |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| for name in ("inspect", "verify"): |
| p = sub.add_parser(name) |
| p.add_argument("src") |
| p = sub.add_parser("export") |
| p.add_argument("src") |
| p.add_argument("dst") |
| p.add_argument("--no-verify", action="store_true", help="skip the per-file sha256 pass (not recommended)") |
| p = sub.add_parser("reshard") |
| p.add_argument("src") |
| p.add_argument("dst") |
| p.add_argument("--world-size", type=int, required=True) |
| p.add_argument("--drop-optimizer", action="store_true", help="omit optimizer files: resume starts FRESH moments (NOT lossless)") |
| p.add_argument("--allow-large", action="store_true") |
| a = ap.parse_args(argv) |
| try: |
| return {"inspect": cmd_inspect, "verify": cmd_verify, "export": cmd_export, "reshard": cmd_reshard}[a.cmd](a) |
| except (sc.ShardedCheckpointError, FileExistsError, FileNotFoundError, KeyError) as e: |
| print(f"{type(e).__name__}: {e}", file=sys.stderr) |
| return 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
