File size: 2,074 Bytes
187a491 | 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 | """Extract model weights from a Lightning checkpoint into a standalone model.pt."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
# Ensure project root is on sys.path so `src` package is importable
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import torch
from src.model.config import BoccaccioConfig
from src.model.transformer import BoccaccioForCausalLM
def main() -> None:
parser = argparse.ArgumentParser(description="Extract model.pt from Lightning checkpoint")
parser.add_argument("--ckpt", required=True, help="Path to Lightning .ckpt file")
parser.add_argument("--output", required=True, help="Path to save model.pt")
parser.add_argument("--variant", default="model", help="Model variant in configs/model.yaml")
args = parser.parse_args()
config = BoccaccioConfig.from_yaml("configs/model.yaml", variant=args.variant)
ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
state_dict = ckpt.get("state_dict", ckpt)
# Remap prefixes from Lightning checkpoint to BoccaccioForCausalLM
# Possible prefixes: "model._orig_mod.X" (torch.compile) or "model.X" (standard)
cleaned = {}
for key, value in state_dict.items():
new_key = key
if new_key.startswith("model._orig_mod."):
new_key = new_key[len("model._orig_mod."):]
elif new_key.startswith("model."):
new_key = new_key[len("model."):]
elif new_key.startswith("_orig_mod."):
new_key = new_key[len("_orig_mod."):]
cleaned[new_key] = value
model = BoccaccioForCausalLM(config)
missing, unexpected = model.load_state_dict(cleaned, strict=False)
if missing:
print(f" WARNING: {len(missing)} missing keys: {missing[:3]}")
if unexpected:
print(f" WARNING: {len(unexpected)} unexpected keys: {unexpected[:3]}")
torch.save(model.state_dict(), args.output)
print(f"model.pt saved to {args.output} (step {ckpt.get('global_step', '?')})")
if __name__ == "__main__":
main()
|