File size: 564 Bytes
43bcf0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | """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)
|