File size: 6,480 Bytes
26d5b81 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | #!/usr/bin/env python3
"""
LMH2 → LMH1 转换器 / LMH1 生成器
模式1 (有lm_head.nfv1): 从 C++ LMH2 提取 w_embed/w_proj，新建 bridge
模式2 (无lm_head.nfv1): 从 model.nfv1 读取维度，全部 xavier 初始化
LMH1 张量: bridge.weight, bridge.bias, w_proj.weight, w_proj.bias, w_embed
用法:
# 模式1: 从LMH2转换
python3 scripts/lmh2_to_lmh1.py \
--nf-model output/model.nfv1 \
--lmh2 output/lm_head.nfv1 \
--output lm_head_lmh1.nfv1
# 模式2: 从model.nfv1直接生成 (无需lm_head)
python3 scripts/lmh2_to_lmh1.py \
--nf-model output/model.nfv1 \
--vocab-size 128000 \
--d-model 512 \
--output lm_head_lmh1.nfv1
"""
import argparse, struct, sys, os
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from infer_full import load_nfv1, load_lmh1
def load_lmh2(path):
weights = {}
with open(path, 'rb') as f:
magic = f.read(4)
if magic not in (b'LMH2', b'LMH1'):
raise ValueError(f"Bad magic: {magic} (expected LMH2 or LMH1)")
while True:
nl = struct.unpack('<I', f.read(4))[0]
if nl == 0:
break
name = f.read(nl).decode('utf-8', errors='replace')
ndim = struct.unpack('<I', f.read(4))[0]
shape = tuple(struct.unpack('<I', f.read(4))[0] for _ in range(ndim))
dsize = struct.unpack('<I', f.read(4))[0]
arr = np.frombuffer(f.read(dsize), dtype=np.float32).reshape(shape).copy()
weights[name] = arr
return weights, magic.decode('utf-8')
def save_lmh1(path, weights):
with open(path, 'wb') as f:
f.write(b'LMH1')
for name, arr in weights.items():
name_bytes = name.encode('utf-8')
f.write(struct.pack('<I', len(name_bytes)))
f.write(name_bytes)
f.write(struct.pack('<I', len(arr.shape)))
for d in arr.shape:
f.write(struct.pack('<I', d))
data = arr.astype(np.float32).tobytes()
f.write(struct.pack('<I', len(data)))
f.write(data)
f.write(struct.pack('<I', 0))
def xavier_uniform(shape, rng):
fan_in = shape[1] if len(shape) == 2 else shape[0]
fan_out = shape[0] if len(shape) == 2 else 1
limit = np.sqrt(6.0 / (fan_in + fan_out))
return rng.uniform(-limit, limit, size=shape).astype(np.float32)
def main():
parser = argparse.ArgumentParser(description='LMH2 → LMH1 converter / LMH1 generator')
parser.add_argument('--nf-model', required=True, help='NF 模型文件 (用于获取 hidden_dim)')
parser.add_argument('--lmh2', default='', help='C++ LMH2 格式 LM head 文件 (可选)')
parser.add_argument('--output', required=True, help='输出 LMH1 文件路径')
parser.add_argument('--vocab-size', type=int, default=128000, help='词表大小 (模式2需要)')
parser.add_argument('--d-model', type=int, default=0, help='LM d_model (0=从NF模型推断)')
parser.add_argument('--seed', type=int, default=42, help='随机种子')
args = parser.parse_args()
nf_w = load_nfv1(args.nf_model)
hidden_dim = nf_w['input_proj.weight'].shape[0]
print(f"NF hidden_dim: {hidden_dim}")
rng = np.random.RandomState(args.seed)
lmh1_w = {}
if args.lmh2 and os.path.exists(args.lmh2):
lmh2_w, fmt = load_lmh2(args.lmh2)
print(f"模式1: 从 {args.lmh2} 转换 (格式: {fmt}, {len(lmh2_w)} 个张量)")
if fmt == 'LMH1':
print("已经是 LMH1 格式，直接复制")
import shutil
shutil.copy2(args.lmh2, args.output)
print(f"输出: {args.output}")
return
vocab_size, d_model = lmh2_w['w_embed'].shape
print(f"vocab_size: {vocab_size}, d_model: {d_model}")
lmh1_w['w_embed'] = lmh2_w['w_embed']
print(f" w_embed: {lmh1_w['w_embed'].shape} (from LMH2)")
if 'w_proj.weight' in lmh2_w:
lmh1_w['w_proj.weight'] = lmh2_w['w_proj.weight']
lmh1_w['w_proj.bias'] = lmh2_w['w_proj.bias']
print(f" w_proj.weight: {lmh1_w['w_proj.weight'].shape} (from LMH2)")
print(f" w_proj.bias: {lmh1_w['w_proj.bias'].shape} (from LMH2)")
else:
lmh1_w['w_proj.weight'] = xavier_uniform((d_model, d_model), rng)
lmh1_w['w_proj.bias'] = np.zeros(d_model, dtype=np.float32)
print(f" w_proj.weight: {lmh1_w['w_proj.weight'].shape} (xavier init)")
print(f" w_proj.bias: {lmh1_w['w_proj.bias'].shape}")
else:
if args.lmh2:
print(f"警告: --lmh2 文件不存在: {args.lmh2}")
print(f"模式2: 从 model.nfv1 生成全新 LMH1")
d_model = args.d_model
if d_model == 0:
d_model = int(nf_w['output_fusion.up.weight'].shape[0])
print(f" d_model 从NF模型推断: {d_model}")
vocab_size = args.vocab_size
print(f"vocab_size: {vocab_size}, d_model: {d_model}")
lmh1_w['w_embed'] = xavier_uniform((vocab_size, d_model), rng)
print(f" w_embed: {lmh1_w['w_embed'].shape} (xavier init)")
lmh1_w['w_proj.weight'] = xavier_uniform((d_model, d_model), rng)
lmh1_w['w_proj.bias'] = np.zeros(d_model, dtype=np.float32)
print(f" w_proj.weight: {lmh1_w['w_proj.weight'].shape} (xavier init)")
print(f" w_proj.bias: {lmh1_w['w_proj.bias'].shape}")
lmh1_w['bridge.weight'] = xavier_uniform((d_model, hidden_dim), rng)
lmh1_w['bridge.bias'] = np.zeros(d_model, dtype=np.float32)
print(f" bridge.weight: {lmh1_w['bridge.weight'].shape} (xavier init)")
print(f" bridge.bias: {lmh1_w['bridge.bias'].shape}")
out_dir = os.path.dirname(os.path.abspath(args.output))
if out_dir:
os.makedirs(out_dir, exist_ok=True)
save_lmh1(args.output, lmh1_w)
print(f"\n输出: {args.output} (LMH1, {len(lmh1_w)} 个张量)")
verify_w = load_lmh1(args.output)
print(f"验证: {len(verify_w)} 个张量加载成功")
for name in lmh1_w:
if name in verify_w:
match = np.allclose(lmh1_w[name], verify_w[name])
print(f" {name}: {'OK' if match else 'MISMATCH'}")
else:
print(f" {name}: MISSING")
if __name__ == '__main__':
main()
|