File size: 4,673 Bytes
e8ca8bd d9d24eb e8ca8bd 88c5042 e8ca8bd | 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 | """Convert the BlinkDL rwkv7-g1d-0.1b .pth checkpoint into a HuggingFace
`trust_remote_code` model directory, re-sizing (and re-initializing) the
embedding + lm-head to the OLMo tokenizer vocabulary.
"""
"""Convert the BlinkDL rwkv7-g1d-0.1b .pth checkpoint into a HuggingFace
`trust_remote_code` model directory, re-sizing (and re-initializing) the
embedding + lm-head to the OLMo tokenizer vocabulary.
"""
"""
from huggingface_hub import snapshot_download
snapshot_download(
repo_id="allenai/Olmo-3-1025-7B",
local_dir="./olmo",
ignore_patterns=[
"*.bin",
"*.safetensors",
"*.pt",
"*.ckpt",
],
)
"""
import os
import sys
import math
import shutil
import torch
import torch.nn as nn
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from configuration_rwkv7 import RWKV7Config
from modeling_rwkv7 import RWKV7ForCausalLM
PTH = "/workspace/rwkv7-g1d-0.1b-20260129-ctx8192.pth"
OLMO = "/workspace/olmo"
OUT = "/workspace/rwkv7-g1d-olmo"
def main():
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(OLMO, trust_remote_code=True)
new_vocab = len(tok)
print(f"OLMo tokenizer vocab = {new_vocab}")
sd = torch.load(PTH, map_location="cpu")
old_vocab, n_embd = sd["emb.weight"].shape
print(f"checkpoint: vocab={old_vocab}, n_embd={n_embd}")
config = RWKV7Config(
vocab_size=new_vocab,
hidden_size=n_embd,
num_hidden_layers=12,#You need to modify it to fit your model.
head_size=64,
intermediate_size=sd["blocks.0.ffn.key.weight"].shape[0],
decay_lora=sd["blocks.0.att.w1"].shape[1],
aaa_lora=sd["blocks.0.att.a1"].shape[1],
mv_lora=sd["blocks.0.att.v1"].shape[1],
gate_lora=sd["blocks.0.att.g1"].shape[1],
eos_token_id=tok.eos_token_id,
pad_token_id=tok.pad_token_id,
torch_dtype="bfloat16",
architectures=["RWKV7ForCausalLM"],
auto_map={
"AutoConfig": "configuration_rwkv7.RWKV7Config",
"AutoModelForCausalLM": "modeling_rwkv7.RWKV7ForCausalLM",
},
)
# remap RWKV native keys -> HF module layout (rwkv.* + head)
remap = {}
for k, v in sd.items():
if k == "emb.weight" or k == "head.weight":
continue # re-initialized below
if k.startswith("blocks.") or k.startswith("ln_out."):
remap["rwkv." + k] = v
else:
remap["rwkv." + k] = v
model = RWKV7ForCausalLM(config).to(dtype=torch.bfloat16)
missing, unexpected = model.load_state_dict(remap, strict=False)
# emb/head are expected to be "missing" from `remap` (we drop them on purpose)
missing = [m for m in missing if not (m == "rwkv.emb.weight" or m == "head.weight")]
print("unexpected keys:", unexpected)
print("missing (besides emb/head):", missing)
assert not unexpected, f"unexpected keys present: {unexpected}"
assert not missing, f"unmatched keys: {missing}"
# Re-initialize embedding + head for the new (OLMo) vocabulary, using the
# RWKV reference init scheme.
with torch.no_grad():
emb = torch.empty(new_vocab, n_embd)
nn.init.uniform_(emb, a=-1e-4, b=1e-4)
model.rwkv.emb.weight.copy_(emb.to(torch.bfloat16))
head = torch.empty(new_vocab, n_embd)
scale = 0.5 * math.sqrt(new_vocab / n_embd) if new_vocab > n_embd else 0.5
nn.init.orthogonal_(head, gain=scale)
model.head.weight.copy_(head.to(torch.bfloat16))
print(f"re-initialized emb {tuple(model.rwkv.emb.weight.shape)} and head "
f"{tuple(model.head.weight.shape)} (head gain {scale:.4f})")
os.makedirs(OUT, exist_ok=True)
model.save_pretrained(OUT, safe_serialization=True)
config.save_pretrained(OUT)
# tokenizer + generation config
tok.save_pretrained(OUT)
# make sure the remote-code files sit alongside the weights
for fn in ["configuration_rwkv7.py", "modeling_rwkv7.py"]:
src = os.path.join(os.path.dirname(os.path.abspath(__file__)), fn)
dst = os.path.join(OUT, fn)
if os.path.abspath(src) != os.path.abspath(dst):
shutil.copy(src, dst)
print("saved to", OUT)
print("files:", sorted(os.listdir(OUT)))
if __name__ == "__main__":
main()
|