File size: 5,466 Bytes
c335050 | 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 |
# scripts for converting pretrained hf model weights to fla style
import argparse
import os
import re
import torch
from transformers import AutoModelForCausalLM
import fla # noqa
from fla.models.rwkv7 import RWKV7Config
def convert(
rwkv7: str,
output: str,
precision: str = 'float32',
):
weights = torch.load(rwkv7, weights_only=True, map_location='cpu')
config = RWKV7Config()
config.vocab_size = weights['emb.weight'].shape[0] # 50304
config.hidden_size = weights['blocks.0.ffn.key.weight'].shape[1] # 768
config.hidden_ratio = weights['blocks.0.ffn.key.weight'].shape[0] / weights['blocks.0.ffn.key.weight'].shape[1] # 4.0
config.intermediate_size = weights['blocks.0.ffn.key.weight'].shape[0]
config.num_hidden_layers = 0
while f'blocks.{config.num_hidden_layers}.ffn.key.weight' in weights:
config.num_hidden_layers += 1
# 12
config.value_dim = [config.hidden_size] * config.num_hidden_layers
config.decay_low_rank_dim = weights['blocks.0.att.w1'].shape[1] # 64
config.gate_low_rank_dim = weights['blocks.0.att.g1'].shape[1] # 128
config.a_low_rank_dim = weights['blocks.0.att.a1'].shape[1] # 64
try:
config.v_low_rank_dim = weights['blocks.1.att.v1'].shape[1] # 32
except KeyError:
config.v_low_rank_dim = 32
if precision in ['bf16', 'bfloat16']:
precision = 'bfloat16'
dtype = torch.bfloat16
if precision in ['fp16', 'float16']:
precision = 'float16'
dtype = torch.float16
if precision in ['fp64', 'double', 'float64']:
precision = 'float64'
dtype = torch.float64
config.torch_dtype = precision
print(f"Creating model with config:\n{config}")
model = AutoModelForCausalLM.from_config(config).to(dtype=dtype)
print(model)
model_dict = model.state_dict()
model_names = [n for n in model_dict]
# these parameters may be present in pth file but are never used:
unused_names = ['blocks.0.att.v0', 'blocks.0.att.v1', 'blocks.0.att.v2']
# these parameters may or may not be present in pth file:
possible_absent_weights = [
'model.layers.0.pre_norm.weight', 'model.layers.0.pre_norm.bias',
]
# other parameters may raise a KeyError
def translate_into_fla(name):
transposed = False
emb_head = {
'emb.weight': 'model.embeddings.weight',
'ln_out.weight': 'model.norm.weight',
'ln_out.bias': 'model.norm.bias',
'head.weight': 'lm_head.weight',
}
proj = {
'receptance': 'r_proj',
'key': 'k_proj',
'value': 'v_proj',
'ln_x': 'g_norm',
'output': 'o_proj',
}
if name in unused_names:
return '', False
if name in emb_head:
return emb_head[name], False
name_compo = name.split('.')
assert name_compo[0] == 'blocks'
name_compo[0] = 'model.layers'
assert int(name_compo[1]) in range(config.num_hidden_layers)
name_compo[2] = {
'att': 'attn',
'ffn': 'ffn',
'ln0': 'pre_norm',
'ln1': 'attn_norm',
'ln2': 'ffn_norm',
}[name_compo[2]]
if re.match("[wvag][012]", name_compo[3]):
typ, num = name_compo[3]
name_compo[3] = f'{typ}_lora.lora.' + {
'0': '2.bias',
'1': '0.weight',
'2': '2.weight',
}[num]
transposed |= (num in ['1', '2'])
elif name_compo[2] == 'attn' and name_compo[3] in proj:
name_compo[3] = proj[name_compo[3]]
return '.'.join(name_compo), transposed
for name in weights:
fla_name, transposed = translate_into_fla(name)
print(f'{name:32} -> {fla_name:42}, {transposed}')
if not fla_name:
print('redundant parameters in source weight: ', name, '\n')
continue
weight = weights[name]
# print shape information
shape1 = list(weight.shape)
shape2 = list(model_dict[fla_name].shape)
print(f'{str(shape1):32} {str(shape2)}\n')
if transposed:
weight.t_()
if shape1 == [1, 1, config.hidden_size]:
weight.squeeze_()
if "attn.x_" in fla_name:
assert model_dict[fla_name].shape[2:] == weight.shape, \
f"Shape mismatch for {fla_name}: model_dict={model_dict[fla_name].shape}, weight={weight.shape}"
else:
assert model_dict[fla_name].shape == weight.shape, \
f"Shape mismatch for {fla_name}: model_dict={model_dict[fla_name].shape}, weight={weight.shape}"
model_dict[fla_name].data.copy_(weight)
model_names.remove(fla_name)
print("uninitialized parameters: ", model_names)
for n in model_names:
if n not in possible_absent_weights:
raise KeyError(n)
os.makedirs(output, exist_ok=True)
model.save_pretrained(output, max_shard_size="1000GB")
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Convert RWKV7')
parser.add_argument('--rwkv7', type=str, help='Path to the input model')
parser.add_argument('--output', type=str, help='Directory to save model')
parser.add_argument('--precision', type=str, default='float32')
args = parser.parse_args()
convert(args.rwkv7, args.output, precision=args.precision)
|