File size: 14,923 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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 |
# scripts for converting pretrained hf model weights to fla style
# calling the code to make conversions for RWKV/rwkv-6-world-7b would achieve the following results:
# | Tasks |Version|Filter|n-shot| Metric | Value | |Stderr|
# |--------------|------:|------|-----:|---------------|------:|---|------|
# |arc_challenge | 1|none | 0|acc | 0.4130|± |0.0144|
# | | |none | 0|acc_norm | 0.4403|± |0.0145|
# |arc_easy | 1|none | 0|acc | 0.7382|± |0.0090|
# | | |none | 0|acc_norm | 0.7079|± |0.0093|
# |boolq | 2|none | 0|acc | 0.6823|± |0.0081|
# |copa | 1|none | 0|acc | 0.8700|± |0.0338|
# |hellaswag | 1|none | 0|acc | 0.5508|± |0.0050|
# | | |none | 0|acc_norm | 0.7171|± |0.0045|
# |lambada_openai| 1|none | 0|perplexity | 3.2989|± |0.0634|
# | | |none | 0|acc | 0.7493|± |0.0060|
# |openbookqa | 1|none | 0|acc | 0.3200|± |0.0209|
# | | |none | 0|acc_norm | 0.4440|± |0.0222|
# |piqa | 1|none | 0|acc | 0.7753|± |0.0097|
# | | |none | 0|acc_norm | 0.7894|± |0.0095|
# |sciq | 1|none | 0|acc | 0.9370|± |0.0077|
# | | |none | 0|acc_norm | 0.8860|± |0.0101|
# |winogrande | 1|none | 0|acc | 0.6867|± |0.0130|
import argparse
import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import fla # noqa
from fla.utils import device
def sizeof_fmt(num, suffix='B'):
for unit in ('', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi'):
if abs(num) < 1024.0:
return f'{num:.2f}{unit}{suffix}'
num /= 1024.0
return f'{num:.2f}Yi{suffix}'
def convert(
rwkv6: str,
config: str,
output: str,
):
torch.manual_seed(1)
AutoTokenizer.from_pretrained(rwkv6, trust_remote_code=True).save_pretrained(output)
rwkv6 = AutoModelForCausalLM.from_pretrained(rwkv6, trust_remote_code=True).to(device)
print(f"Loading rwkv6 ...\n{rwkv6}")
config = AutoConfig.from_pretrained(config)
model = AutoModelForCausalLM.from_config(config).to(device)
num_parameters = model.num_parameters()
print(f"Initializing the model from the config:\n{config}\n{model}")
print(f"Number of parameters in total: {num_parameters} ({sizeof_fmt(num_parameters)})")
print("Copying the weights from rwkv6 to the model ...")
print("rwkv6.rwkv.embeddings -> model.model.embeddings")
model.model.embeddings.weight.data.copy_(rwkv6.rwkv.embeddings.weight)
torch.testing.assert_close(model.model.embeddings.weight, rwkv6.rwkv.embeddings.weight)
for i in range(config.num_hidden_layers):
if hasattr(model.model.layers[i], 'pre_norm'):
if model.model.layers[i].pre_norm.weight is not None:
print(f"rwkv6.rwkv.blocks{i}.pre_ln.weight -> model.model.layers{i}.pre_norm.weight")
model.model.layers[i].pre_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.weight)
torch.testing.assert_close(model.model.layers[i].pre_norm.weight, rwkv6.rwkv.blocks[i].pre_ln.weight)
if model.model.layers[i].pre_norm.bias is not None:
print(f"rwkv6.rwkv.blocks{i}.pre_ln.bias -> model.model.layers{i}.pre_norm.bias")
model.model.layers[i].pre_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].pre_ln.bias)
torch.testing.assert_close(model.model.layers[i].pre_norm.bias, rwkv6.rwkv.blocks[i].pre_ln.bias)
model.model.layers[i].pre_norm.eps = rwkv6.rwkv.blocks[i].pre_ln.eps
if model.model.layers[i].attn_norm.weight is not None:
print(f"rwkv6.rwkv.blocks{i}.ln1.weight -> model.model.layers{i}.attn_norm.weight")
model.model.layers[i].attn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln1.weight)
torch.testing.assert_close(model.model.layers[i].attn_norm.weight, rwkv6.rwkv.blocks[i].ln1.weight)
if model.model.layers[i].attn_norm.bias is not None:
print(f"rwkv6.rwkv.blocks{i}.ln1.bias -> model.model.layers{i}.attn_norm.bias")
model.model.layers[i].attn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln1.bias)
torch.testing.assert_close(model.model.layers[i].attn_norm.bias, rwkv6.rwkv.blocks[i].ln1.bias)
model.model.layers[i].attn_norm.eps = rwkv6.rwkv.blocks[i].ln1.eps
print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_x -> model.model.layers.{i}.attn.x_proj0.mu")
model.model.layers[i].attn.x_proj[0].mu.data.copy_(rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1))
torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].mu,
rwkv6.rwkv.blocks[i].attention.time_maa_x.view(-1))
print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w1.weight -> model.model.layers{i}.attn.x_proj0.linear.weight")
ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w1.view(config.hidden_size, 5, -1).unbind(-2)
w = torch.cat((wr, ww, wk, wv, wg), -1).t()
model.model.layers[i].attn.x_proj[0].linear.weight.data.copy_(w)
torch.testing.assert_close(model.model.layers[i].attn.x_proj[0].linear.weight, w)
print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_w2.weight -> model.model.layers{i}.attn.x_proj2.weight")
ww, wk, wv, wr, wg = rwkv6.rwkv.blocks[i].attention.time_maa_w2.unbind(0)
w = torch.cat((wr, ww, wk, wv, wg), 0).t()
model.model.layers[i].attn.x_proj[2].weight.data.copy_(w)
torch.testing.assert_close(model.model.layers[i].attn.x_proj[2].weight, w)
print(f"rwkv6.rwkv.blocks{i}.attention.time_maa_wkvrg -> model.model.layers{i}.attn.x_bias")
bias = torch.stack((rwkv6.rwkv.blocks[i].attention.time_maa_r.view(-1),
rwkv6.rwkv.blocks[i].attention.time_maa_w.view(-1),
rwkv6.rwkv.blocks[i].attention.time_maa_k.view(-1),
rwkv6.rwkv.blocks[i].attention.time_maa_v.view(-1),
rwkv6.rwkv.blocks[i].attention.time_maa_g.view(-1)))
model.model.layers[i].attn.x_bias.data.copy_(bias)
torch.testing.assert_close(model.model.layers[i].attn.x_bias, bias)
print(f"rwkv6.rwkv.blocks{i}.attention.receptance.weight -> model.model.layers{i}.attn.r_proj.linear.weight")
model.model.layers[i].attn.r_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.receptance.weight)
torch.testing.assert_close(model.model.layers[i].attn.r_proj.linear.weight,
rwkv6.rwkv.blocks[i].attention.receptance.weight)
print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w1 -> model.model.layers{i}.attn.w_proj.linear.lora0.weight")
model.model.layers[i].attn.w_proj.linear.lora[0].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w1.t())
torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[0].weight,
rwkv6.rwkv.blocks[i].attention.time_decay_w1.t())
print(f"rwkv6.rwkv.blocks{i}.attention.time_decay_w2 -> model.model.layers{i}.attn.w_proj.linear.lora2.weight")
model.model.layers[i].attn.w_proj.linear.lora[2].weight.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay_w2.t())
torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].weight,
rwkv6.rwkv.blocks[i].attention.time_decay_w2.t())
print(f"rwkv6.rwkv.blocks{i}.attention.time_decay -> model.model.layers{i}.attn.w_proj.linear.lora2.bias")
model.model.layers[i].attn.w_proj.linear.lora[2].bias.data.copy_(rwkv6.rwkv.blocks[i].attention.time_decay.view(-1))
torch.testing.assert_close(model.model.layers[i].attn.w_proj.linear.lora[2].bias,
rwkv6.rwkv.blocks[i].attention.time_decay.view(-1))
print(f"rwkv6.rwkv.blocks{i}.attention.key.weight -> model.model.layers.{i}.attn.k_proj.linear.weight")
model.model.layers[i].attn.k_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.key.weight)
torch.testing.assert_close(model.model.layers[i].attn.k_proj.linear.weight,
rwkv6.rwkv.blocks[i].attention.key.weight)
print(f"rwkv6.rwkv.blocks{i}.attention.value.weight -> model.model.layers.{i}.attn.v_proj.linear.weight")
model.model.layers[i].attn.v_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.value.weight)
torch.testing.assert_close(model.model.layers[i].attn.v_proj.linear.weight,
rwkv6.rwkv.blocks[i].attention.value.weight)
print(f"rwkv6.rwkv.blocks{i}.attention.gate.weight -> model.model.layers.{i}.attn.g_proj.linear.weight")
model.model.layers[i].attn.g_proj.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.gate.weight)
torch.testing.assert_close(model.model.layers[i].attn.g_proj.linear.weight,
rwkv6.rwkv.blocks[i].attention.gate.weight)
print(f"rwkv6.rwkv.blocks{i}.attention.time_faaaa -> model.model.layers.{i}.attn.bonus")
bonus = rwkv6.rwkv.blocks[i].attention.time_faaaa.view(config.num_heads, -1)
model.model.layers[i].attn.bonus.data.copy_(bonus)
torch.testing.assert_close(model.model.layers[i].attn.bonus, bonus)
if model.model.layers[i].attn.g_norm.weight is not None:
print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.weight -> model.model.layers[i].attn.g_norm.weight")
model.model.layers[i].attn.g_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.weight)
torch.testing.assert_close(model.model.layers[i].attn.g_norm.weight, rwkv6.rwkv.blocks[i].attention.ln_x.weight)
if model.model.layers[i].attn.g_norm.bias is not None:
print(f"rwkv6.rwkv.blocks{i}.attention.ln_x.bias -> model.model.layers[i].attn.g_norm.bias")
model.model.layers[i].attn.g_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].attention.ln_x.bias)
torch.testing.assert_close(model.model.layers[i].attn.g_norm.bias, rwkv6.rwkv.blocks[i].attention.ln_x.bias)
model.model.layers[i].attn.g_norm.eps = rwkv6.rwkv.blocks[i].attention.ln_x.eps
print(f"rwkv6.rwkv.blocks{i}.attention.output.weight -> model.model.layers.{i}.attn.o_proj.weight")
model.model.layers[i].attn.o_proj.weight.data.copy_(rwkv6.rwkv.blocks[i].attention.output.weight)
torch.testing.assert_close(model.model.layers[i].attn.o_proj.weight, rwkv6.rwkv.blocks[i].attention.output.weight)
if model.model.layers[i].ffn_norm.weight is not None:
print(f"rwkv6.rwkv.blocks{i}.ln2.weight -> model.model.layers{i}.ffn_norm.weight")
model.model.layers[i].ffn_norm.weight.data.copy_(rwkv6.rwkv.blocks[i].ln2.weight)
torch.testing.assert_close(model.model.layers[i].ffn_norm.weight, rwkv6.rwkv.blocks[i].ln2.weight)
if model.model.layers[i].ffn_norm.bias is not None:
print(f"rwkv6.rwkv.blocks{i}.ln2.bias -> model.model.layers{i}.ffn_norm.bias")
model.model.layers[i].ffn_norm.bias.data.copy_(rwkv6.rwkv.blocks[i].ln2.bias)
torch.testing.assert_close(model.model.layers[i].ffn_norm.bias, rwkv6.rwkv.blocks[i].ln2.bias)
model.model.layers[i].ffn_norm.eps = rwkv6.rwkv.blocks[i].ln2.eps
print(f"rwkv6.rwkv.blocks{i}.feed_forward.key.weight -> model.model.layers.{i}.ffn.key.linear.weight")
model.model.layers[i].ffn.key.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.key.weight)
torch.testing.assert_close(model.model.layers[i].ffn.key.linear.weight,
rwkv6.rwkv.blocks[i].feed_forward.key.weight)
print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_k -> model.model.layers.{i}.ffn.key.mu")
model.model.layers[i].ffn.key.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1))
torch.testing.assert_close(model.model.layers[i].ffn.key.mu,
rwkv6.rwkv.blocks[i].feed_forward.time_maa_k.view(-1))
print(f"rwkv6.rwkv.blocks{i}.feed_forward.value.weight -> model.model.layers.{i}.ffn.value.weight")
model.model.layers[i].ffn.value.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.value.weight)
torch.testing.assert_close(model.model.layers[i].ffn.value.weight,
rwkv6.rwkv.blocks[i].feed_forward.value.weight)
print(f"rwkv6.rwkv.blocks{i}.feed_forward.receptance.weight -> model.model.layers.{i}.ffn.receptance.linear.weight")
model.model.layers[i].ffn.receptance.linear.weight.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.receptance.weight)
torch.testing.assert_close(model.model.layers[i].ffn.receptance.linear.weight,
rwkv6.rwkv.blocks[i].feed_forward.receptance.weight)
print(f"rwkv6.rwkv.blocks{i}.feed_forward.time_maa_r -> model.model.layers.{i}.ffn.receptance.mu")
model.model.layers[i].ffn.receptance.mu.data.copy_(rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1))
torch.testing.assert_close(model.model.layers[i].ffn.receptance.mu,
rwkv6.rwkv.blocks[i].feed_forward.time_maa_r.view(-1))
if model.model.norm.weight is not None:
print("rwkv6.rwkv.ln_out.weight -> model.model.norm.weight")
model.model.norm.weight.data.copy_(rwkv6.rwkv.ln_out.weight)
torch.testing.assert_close(model.model.norm.weight, rwkv6.rwkv.ln_out.weight)
if model.model.norm.bias is not None:
print("rwkv6.rwkv.ln_out.bias -> model.model.norm.bias")
model.model.norm.bias.data.copy_(rwkv6.rwkv.ln_out.bias)
torch.testing.assert_close(model.model.norm.bias, rwkv6.rwkv.ln_out.bias)
model.model.norm.eps = rwkv6.rwkv.ln_out.eps
if not model.config.tie_word_embeddings:
print("rwkv6.rwkv.head.weight -> model.lm_head.weight")
model.lm_head.weight.data.copy_(rwkv6.head.weight)
torch.testing.assert_close(model.lm_head.weight, rwkv6.head.weight)
print(f"Saving converted model \n{model}\n to {output} ...")
model.save_pretrained(output)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", default='RWKV/rwkv-6-world-7b')
parser.add_argument("--config", default='configs/rwkv6_7B.json')
parser.add_argument("--output", default='converted/rwkv6-7B')
args = parser.parse_args()
convert(args.model, args.config, args.output)
|