| |
| |
| import json, shutil, os, torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
| from transformers import Qwen2Config, Qwen2ForCausalLM |
| SRC="model"; DST="qwen25_3b_la" |
| os.makedirs(DST, exist_ok=True) |
| idx=json.load(open(f"{SRC}/model.safetensors.index.json"))["weight_map"] |
| llm={}; extra={} |
| files={} |
| for k,f in idx.items(): files.setdefault(f,[]).append(k) |
| for f,keys in files.items(): |
| with safe_open(f"{SRC}/{f}","pt") as sf: |
| for k in keys: |
| if k.startswith("language_model."): llm[k[len("language_model."):]]=sf.get_tensor(k) |
| elif k.startswith("mlp1."): extra[k]=sf.get_tensor(k) |
| tied = torch.equal(llm["lm_head.weight"], llm["model.embed_tokens.weight"]) |
| print("lm_head == embed_tokens:", tied, "| vocab rows:", llm["model.embed_tokens.weight"].shape) |
| cfg=json.load(open(f"{SRC}/config.json"))["text_config"] |
| qc=Qwen2Config(vocab_size=cfg["vocab_size"],hidden_size=cfg["hidden_size"],intermediate_size=cfg["intermediate_size"], |
| num_hidden_layers=cfg["num_hidden_layers"],num_attention_heads=cfg["num_attention_heads"],num_key_value_heads=cfg["num_key_value_heads"], |
| max_position_embeddings=cfg["max_position_embeddings"],rms_norm_eps=cfg["rms_norm_eps"],rope_theta=cfg["rope_theta"], |
| tie_word_embeddings=tied,use_sliding_window=False,sliding_window=None,bos_token_id=cfg["bos_token_id"],eos_token_id=cfg["eos_token_id"],torch_dtype="bfloat16") |
| qc.architectures=["Qwen2ForCausalLM"] |
| qc.save_pretrained(DST) |
| if tied: llm.pop("lm_head.weight") |
| save_file({k:v.contiguous() for k,v in llm.items()}, f"{DST}/model.safetensors", metadata={"format":"pt"}) |
| save_file({k:v.contiguous() for k,v in extra.items()}, f"{DST}/mlp1.safetensors") |
| for f in ["vocab.json","merges.txt","tokenizer_config.json","added_tokens.json","special_tokens_map.json","generation_config.json"]: |
| shutil.copy(f"{SRC}/{f}", f"{DST}/{f}") |
| |
| m=Qwen2ForCausalLM.from_pretrained(DST, torch_dtype=torch.float32).eval() |
| ids=torch.randint(0,150000,(1,32)) |
| with torch.no_grad(): lg=m(ids).logits |
| print("reloaded ok, logits", tuple(lg.shape), "finite:", torch.isfinite(lg).all().item()) |
| torch.save({"ids":ids,"logits":lg}, "ref_logits_plain.pt") |
|
|
