File size: 1,602 Bytes
2ae6691 | 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 | """Convert a sharded training checkpoint into one plain model file for eval.
Training saves model + optimizer split across all GPUs. Eval only needs the
model weights, in a single file that one GPU can load.
"""
import argparse
import torch
import torch.distributed.checkpoint as dcp
from ar.model.config import ARConfig
from ar.model.model import ARModel
def main(ckpt_dir, out_path):
# build on meta then allocate on CPU: skips a slow 4B-param random init
with torch.device("meta"):
model = ARModel(ARConfig())
model.to_empty(device="cpu")
sd = model.state_dict()
# see which weights the checkpoint actually holds (its keys look like "model.<name>")
meta = dcp.FileSystemReader(ckpt_dir).read_metadata()
saved = {k[len("model."):] for k in meta.state_dict_metadata if k.startswith("model.")}
wanted = {k: v for k, v in sd.items() if k in saved}
missing = [k for k in sd if k not in saved]
print(f"loading {len(wanted)}/{len(sd)} tensors; not in checkpoint: {missing}")
# load ONLY the model part, filling the tensors in place (optimizer is skipped)
dcp.load(state_dict={"model": wanted}, checkpoint_id=ckpt_dir)
# bf16 halves the file (~8 GB); eval runs in bf16 anyway
out = {k: v.to(torch.bfloat16) for k, v in model.state_dict().items()}
torch.save(out, out_path)
print(f"saved {len(out)} tensors -> {out_path}")
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--ckpt", required=True)
p.add_argument("--out", required=True)
args = p.parse_args()
main(args.ckpt, args.out) |