AR / ar /eval /export_ckpt.py

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 18 days ago

1.6 kB

"""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)

	"""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)