初始化项目，由ModelHub XC社区提供模型
Model: pathcosmos/frankenstallm Source: Original Platform
This commit is contained in:
ModelHub XC
230
source/scripts/migrate_qkv_checkpoint.py
Normal file
230
source/scripts/migrate_qkv_checkpoint.py
Normal file
@@ -0,0 +1,230 @@
|
||||
#!/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]))
|
||||
Reference in New Issue
Block a user