231 lines
8.3 KiB
Python
231 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
||
"""Migrate checkpoint from separate Q/K/V projections to fused QKV.
|
||
|
||
Usage:
|
||
python3 scripts/migrate_qkv_checkpoint.py <checkpoint_dir>
|
||
|
||
Migrates both model.pt AND optimizer.pt:
|
||
- model.pt: q_proj/k_proj/v_proj weights → qkv_proj weight
|
||
- optimizer.pt: exp_avg/exp_avg_sq states fused, param indices re-mapped
|
||
|
||
The concatenation order is [Q ; K ; V] along the output (dim-0) axis,
|
||
which matches the split in MultiHeadAttention.forward:
|
||
q, k, v = qkv.split([_q_dim, _kv_dim, _kv_dim], dim=-1)
|
||
|
||
Optimizer layout (group 0 = weight_decay, per layer × 28):
|
||
[i*6+0] q_proj.weight [3072, 3072]
|
||
[i*6+1] k_proj.weight [1024, 3072]
|
||
[i*6+2] v_proj.weight [1024, 3072]
|
||
[i*6+3] out_proj.weight [3072, 3072]
|
||
[i*6+4] fc1_weight [16384, 3072]
|
||
[i*6+5] fc2_weight [3072, 8192]
|
||
After fusion: indices 0,1,2 → single qkv_proj → 4 params per layer.
|
||
"""
|
||
import sys
|
||
import torch
|
||
from pathlib import Path
|
||
|
||
N_LAYERS = 28
|
||
OLD_PARAMS_PER_LAYER = 6 # q, k, v, out, fc1, fc2
|
||
NEW_PARAMS_PER_LAYER = 4 # qkv, out, fc1, fc2
|
||
|
||
|
||
def migrate_model(state: dict) -> dict:
|
||
"""Fuse Q/K/V projection weights into QKV in model state dict."""
|
||
new_state: dict = {}
|
||
layers_done: set = set()
|
||
|
||
for key, val in state.items():
|
||
if ".q_proj." not in key and ".k_proj." not in key and ".v_proj." not in key:
|
||
new_state[key] = val
|
||
continue
|
||
|
||
if ".q_proj." not in key:
|
||
continue
|
||
|
||
prefix = key.rsplit(".", 2)[0]
|
||
suffix = key.rsplit(".", 1)[-1]
|
||
|
||
tag = (prefix, suffix)
|
||
if tag in layers_done:
|
||
continue
|
||
layers_done.add(tag)
|
||
|
||
q_key = f"{prefix}.q_proj.{suffix}"
|
||
k_key = f"{prefix}.k_proj.{suffix}"
|
||
v_key = f"{prefix}.v_proj.{suffix}"
|
||
|
||
missing = [k for k in (q_key, k_key, v_key) if k not in state]
|
||
if missing:
|
||
raise KeyError(f"Expected keys not found in checkpoint: {missing}")
|
||
|
||
q_w, k_w, v_w = state[q_key], state[k_key], state[v_key]
|
||
fused = torch.cat([q_w, k_w, v_w], dim=0)
|
||
fused_key = f"{prefix}.qkv_proj.{suffix}"
|
||
new_state[fused_key] = fused
|
||
print(f" Fused {fused_key}: {list(fused.shape)}"
|
||
f" (q={list(q_w.shape)}, k={list(k_w.shape)}, v={list(v_w.shape)})")
|
||
|
||
leaked = [k for k in new_state if ".q_proj." in k or ".k_proj." in k or ".v_proj." in k]
|
||
if leaked:
|
||
raise RuntimeError(f"BUG: old projection keys still present: {leaked}")
|
||
|
||
return new_state
|
||
|
||
|
||
def migrate_optimizer(opt_state: dict) -> dict:
|
||
"""Fuse optimizer states for Q/K/V → QKV and re-index parameters.
|
||
|
||
The optimizer has 2 param groups:
|
||
Group 0 (weight_decay): 168 = 28 layers × 6 (q,k,v,out,fc1,fc2)
|
||
Group 1 (no weight_decay): 58 = norms + embedding
|
||
|
||
We fuse q,k,v entries in group 0 (indices i*6+0,1,2 → one entry per layer).
|
||
Group 0 shrinks from 168 to 112 (28 layers × 4 params).
|
||
Group 1 stays at 58. Total: 170.
|
||
"""
|
||
old_state = opt_state["state"]
|
||
old_groups = opt_state["param_groups"]
|
||
|
||
group0_count = len(old_groups[0]["params"])
|
||
expected_g0 = N_LAYERS * OLD_PARAMS_PER_LAYER
|
||
if group0_count != expected_g0:
|
||
raise ValueError(
|
||
f"Group 0 has {group0_count} params, expected {expected_g0}. "
|
||
f"Cannot auto-detect QKV layout."
|
||
)
|
||
|
||
# Validate shapes for first layer
|
||
shapes = []
|
||
for j in range(OLD_PARAMS_PER_LAYER):
|
||
idx = old_groups[0]["params"][j]
|
||
shapes.append(list(old_state[idx]["exp_avg"].shape))
|
||
expected_shapes = [[3072, 3072], [1024, 3072], [1024, 3072],
|
||
[3072, 3072], [16384, 3072], [3072, 8192]]
|
||
if shapes != expected_shapes:
|
||
raise ValueError(
|
||
f"Layer 0 shapes {shapes} don't match expected {expected_shapes}. "
|
||
f"Cannot auto-detect QKV layout."
|
||
)
|
||
print(f" Shape validation passed for layer 0.")
|
||
|
||
new_state_entries = {}
|
||
new_idx = 0
|
||
|
||
# --- Group 0: fuse q/k/v per layer ---
|
||
for layer_i in range(N_LAYERS):
|
||
base = layer_i * OLD_PARAMS_PER_LAYER
|
||
q_opt_idx = old_groups[0]["params"][base + 0]
|
||
k_opt_idx = old_groups[0]["params"][base + 1]
|
||
v_opt_idx = old_groups[0]["params"][base + 2]
|
||
|
||
q_entry = old_state[q_opt_idx]
|
||
k_entry = old_state[k_opt_idx]
|
||
v_entry = old_state[v_opt_idx]
|
||
|
||
# Fuse QKV
|
||
fused_entry = {"step": q_entry["step"]}
|
||
for field in ["exp_avg", "exp_avg_sq"]:
|
||
if field in q_entry:
|
||
fused_entry[field] = torch.cat(
|
||
[q_entry[field], k_entry[field], v_entry[field]], dim=0
|
||
)
|
||
new_state_entries[new_idx] = fused_entry
|
||
if layer_i == 0:
|
||
print(f" Layer 0 QKV fused: exp_avg {list(fused_entry['exp_avg'].shape)}")
|
||
new_idx += 1
|
||
|
||
# Copy remaining params (out, fc1, fc2)
|
||
for offset in [3, 4, 5]:
|
||
opt_idx = old_groups[0]["params"][base + offset]
|
||
new_state_entries[new_idx] = old_state[opt_idx]
|
||
new_idx += 1
|
||
|
||
new_group0_count = new_idx # should be N_LAYERS * NEW_PARAMS_PER_LAYER = 112
|
||
print(f" Group 0: {group0_count} → {new_group0_count} params")
|
||
|
||
# --- Group 1: copy as-is (norms, embedding — no QKV) ---
|
||
group1_count = len(old_groups[1]["params"])
|
||
for j in range(group1_count):
|
||
opt_idx = old_groups[1]["params"][j]
|
||
if opt_idx in old_state:
|
||
new_state_entries[new_idx] = old_state[opt_idx]
|
||
new_idx += 1
|
||
print(f" Group 1: {group1_count} → {group1_count} params (unchanged)")
|
||
|
||
# Build new param_groups
|
||
new_groups = []
|
||
g0 = {k: v for k, v in old_groups[0].items() if k != "params"}
|
||
g0["params"] = list(range(0, new_group0_count))
|
||
new_groups.append(g0)
|
||
|
||
g1 = {k: v for k, v in old_groups[1].items() if k != "params"}
|
||
g1["params"] = list(range(new_group0_count, new_group0_count + group1_count))
|
||
new_groups.append(g1)
|
||
|
||
total = new_group0_count + group1_count
|
||
print(f" Total: {len(old_state)} → {total} optimizer params")
|
||
|
||
return {"state": new_state_entries, "param_groups": new_groups}
|
||
|
||
|
||
def migrate(ckpt_dir: Path) -> None:
|
||
model_path = ckpt_dir / "model.pt"
|
||
opt_path = ckpt_dir / "optimizer.pt"
|
||
|
||
if not model_path.exists():
|
||
raise FileNotFoundError(f"model.pt not found in {ckpt_dir}")
|
||
|
||
# --- Model migration ---
|
||
print(f"[1/2] Migrating model weights from {model_path} ...")
|
||
state = torch.load(model_path, map_location="cpu", weights_only=True)
|
||
|
||
has_old = any(".q_proj." in k for k in state)
|
||
has_new = any(".qkv_proj." in k for k in state)
|
||
|
||
if has_new and not has_old:
|
||
print(" Model already migrated. Skipping.")
|
||
elif has_old:
|
||
new_model_state = migrate_model(state)
|
||
torch.save(new_model_state, model_path)
|
||
print(f" Model saved.")
|
||
else:
|
||
raise RuntimeError("Model state has neither q_proj nor qkv_proj keys!")
|
||
|
||
# --- Optimizer migration ---
|
||
if opt_path.exists():
|
||
print(f"\n[2/2] Migrating optimizer states from {opt_path} ...")
|
||
opt = torch.load(opt_path, map_location="cpu", weights_only=True)
|
||
|
||
# Check if already migrated
|
||
total_params = sum(len(pg["params"]) for pg in opt["param_groups"])
|
||
expected_old = N_LAYERS * OLD_PARAMS_PER_LAYER + 58 # 168 + 58 = 226
|
||
expected_new = N_LAYERS * NEW_PARAMS_PER_LAYER + 58 # 112 + 58 = 170
|
||
|
||
if total_params == expected_old:
|
||
opt_backup = ckpt_dir / "optimizer.pt.backup_pre_qkv"
|
||
if not opt_backup.exists():
|
||
torch.save(opt, opt_backup)
|
||
print(f" Backup: {opt_backup}")
|
||
new_opt = migrate_optimizer(opt)
|
||
torch.save(new_opt, opt_path)
|
||
print(f" Optimizer saved.")
|
||
elif total_params == expected_new:
|
||
print(f" Optimizer already migrated ({total_params} params). Skipping.")
|
||
else:
|
||
print(f" [WARN] Unexpected param count {total_params} "
|
||
f"(expected old={expected_old} or new={expected_new}). "
|
||
f"Deleting optimizer.pt — optimizer will restart fresh.")
|
||
opt_path.unlink()
|
||
else:
|
||
print("\n[2/2] No optimizer.pt found. Optimizer will restart fresh.")
|
||
|
||
print("\nMigration complete!")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
if len(sys.argv) != 2:
|
||
print(__doc__)
|
||
sys.exit(1)
|
||
migrate(Path(sys.argv[1]))
|