File size: 5,666 Bytes
bcc8d35 | 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 | #!/usr/bin/env python
"""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", "") # never touch a GPU from a converter
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from align.train import sharded_ckpt as sc # noqa: E402
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())
|