SWD / examples /load_checkpoint.py
XuhanH's picture
Publish SWD factor checkpoints
43bcf0c verified
Raw
History Blame Contribute Delete
564 Bytes
"""Minimal example for applying an SWD checkpoint to a base model."""
from pathlib import Path
from transformers import AutoModelForCausalLM
from swd_loader import apply_swd_checkpoint
repo_root = Path(__file__).resolve().parents[1]
checkpoint = repo_root / "checkpoints/gpt2-small/layer8-cproj/s0p5-tokens16384"
model = AutoModelForCausalLM.from_pretrained("gpt2")
# Use mode="folded" for a conventional dense module after loading.
applied = apply_swd_checkpoint(model, checkpoint, mode="factorized")
print(applied)
print(model.transformer.h[8].mlp.c_proj)