372 lines
15 KiB
Python
372 lines
15 KiB
Python
# ==============================================================================
|
|
# JiRack -> GGUF converter, 1.5B edition (stage 1: .pt -> HuggingFace folder)
|
|
# COPYRIGHT (c) 2026 Konstantin Vladimirovich Grabko.
|
|
#
|
|
# Verified against JiRackTernaryUltra_1b.py [DS1.5-1]:
|
|
# vocab_size 151936, hidden 1536, n_layers 28, n_heads 12, n_kv_heads 2,
|
|
# head_dim 128 (12*128=1536 -- the converter's hardcoded 128 is correct),
|
|
# rope_theta 10000.0 (same as 7B), rms_eps 1e-6,
|
|
# tie_word_embeddings = FALSE [DS1.5-2] -- lm_head ships separately, and
|
|
# this converter auto-detects that from the presence of lm_head.weight.
|
|
#
|
|
# Pipeline is two stages:
|
|
#
|
|
# Stage 1 (THIS SCRIPT, run in venv_ji):
|
|
# model.pt -> HF folder (model.safetensors + config.json + tokenizer)
|
|
#
|
|
# Stage 2 (llama.cpp, run once per model):
|
|
# python convert_hf_to_gguf.py <hf_folder> \
|
|
# --outfile jirack_1p5b.gguf --outtype bf16
|
|
# ./build/bin/llama-quantize jirack_1p5b.gguf \
|
|
# jirack_1p5b.Q4_K_M.gguf Q4_K_M
|
|
#
|
|
# Key points handled here:
|
|
# * config.json is derived from the ACTUAL tensor shapes in the checkpoint,
|
|
# so vocab (151936 vs 7B's 152064) and any Net2Net-expanded FFN width are
|
|
# picked up automatically -- no stock-config copying.
|
|
# * lambda_ buffers (ternary fake-quant training machinery) are dropped --
|
|
# at inference you run set_lambda(0.0) anyway, so the stored weights ARE
|
|
# the full-precision weights; the exported model is a plain Qwen2 dense.
|
|
# * Keys: HF naming passes through; JiRack native naming
|
|
# (token_emb / blocks.N.* / ffn_w1-w3-w2) is remapped automatically.
|
|
#
|
|
# EDIT THE THREE PATHS BELOW.
|
|
# ==============================================================================
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import sys
|
|
|
|
import torch
|
|
|
|
# ========================= EDIT THESE =========================
|
|
CKPT_PATH = "model.pt"
|
|
TOKENIZER_DIR = "."
|
|
OUTPUT_DIR = "."
|
|
# rope_theta cannot be inferred from tensor shapes -- set per base model:
|
|
# DeepSeek-R1-Distill-Qwen-1.5B -> 10000.0 (same as 7B)
|
|
# DeepSeek-R1-Distill-Qwen-14B -> 1000000.0
|
|
# DeepSeek-R1-Distill-Qwen-32B -> 1000000.0
|
|
ROPE_THETA = 10000.0
|
|
MAX_POSITION = 131072
|
|
RMS_NORM_EPS = 1e-6
|
|
|
|
# Q2_0 = 2-bit ternary {-1, 0, +1} quantization, one fp16 scale per group of
|
|
# weights -- the real encoding for BitNet-style ternary weights, once inference
|
|
# actually runs true ternary rather than bf16 dense. For now Q4_K_M remains
|
|
# the practical choice; the Q2_0 command is just printed ready for later.
|
|
EMIT_Q2_0_CMD = True
|
|
Q2_0_GROUP = 64 # 64 = mainline llama.cpp, no fork needed.
|
|
# ================================================================
|
|
|
|
# HF Qwen2 key patterns we expect to find (N = layer index)
|
|
HF_LAYER_KEYS = [
|
|
"model.layers.{n}.self_attn.q_proj.weight",
|
|
"model.layers.{n}.self_attn.q_proj.bias",
|
|
"model.layers.{n}.self_attn.k_proj.weight",
|
|
"model.layers.{n}.self_attn.k_proj.bias",
|
|
"model.layers.{n}.self_attn.v_proj.weight",
|
|
"model.layers.{n}.self_attn.v_proj.bias",
|
|
"model.layers.{n}.self_attn.o_proj.weight",
|
|
"model.layers.{n}.mlp.gate_proj.weight",
|
|
"model.layers.{n}.mlp.up_proj.weight",
|
|
"model.layers.{n}.mlp.down_proj.weight",
|
|
"model.layers.{n}.input_layernorm.weight",
|
|
"model.layers.{n}.post_attention_layernorm.weight",
|
|
]
|
|
HF_TOP_KEYS = [
|
|
"model.embed_tokens.weight",
|
|
"model.norm.weight",
|
|
"lm_head.weight",
|
|
]
|
|
|
|
|
|
def load_state_dict(path):
|
|
print(f"📥 Loading checkpoint: {path}")
|
|
ckpt = torch.load(path, map_location="cpu", weights_only=False)
|
|
sd = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
|
|
if not isinstance(sd, dict):
|
|
sys.exit("❌ Checkpoint is not a state_dict and has no 'model' key.")
|
|
return sd
|
|
|
|
|
|
def drop_training_buffers(sd):
|
|
dropped = [k for k in sd if k.endswith("lambda_")]
|
|
for k in dropped:
|
|
del sd[k]
|
|
if dropped:
|
|
print(f"🧹 Dropped {len(dropped)} lambda_ buffers (ternary training machinery).")
|
|
return sd
|
|
|
|
|
|
def normalize_keys(sd):
|
|
"""Pass HF-style keys through; try trivial prefix fixes; else abort with a listing."""
|
|
keys = list(sd.keys())
|
|
|
|
# Case 1: already HF-style
|
|
if "model.embed_tokens.weight" in sd:
|
|
print("✅ Keys already use HF (Qwen2) naming -- no remap needed.")
|
|
return sd
|
|
|
|
# Case 2: same names but without the leading 'model.' (e.g. 'embed_tokens.weight')
|
|
if "embed_tokens.weight" in sd:
|
|
print("🔁 Keys look HF-like without the 'model.' prefix -- adding it.")
|
|
out = {}
|
|
for k, v in sd.items():
|
|
if k == "lm_head.weight":
|
|
out[k] = v
|
|
else:
|
|
out["model." + k] = v
|
|
if "model.embed_tokens.weight" in out:
|
|
return out
|
|
|
|
# Case 3: JiRack native naming (token_emb / blocks.N.* / ffn_w1-w3-w2)
|
|
if "token_emb.weight" in sd and any(k.startswith("blocks.") for k in sd):
|
|
print("🔁 JiRack native naming detected -- remapping to HF (Qwen2) keys.")
|
|
hidden = sd["token_emb.weight"].shape[1]
|
|
block_map = {
|
|
"norm1.weight": "input_layernorm.weight",
|
|
"norm2.weight": "post_attention_layernorm.weight",
|
|
"q_proj.weight": "self_attn.q_proj.weight",
|
|
"q_proj.bias": "self_attn.q_proj.bias",
|
|
"k_proj.weight": "self_attn.k_proj.weight",
|
|
"k_proj.bias": "self_attn.k_proj.bias",
|
|
"v_proj.weight": "self_attn.v_proj.weight",
|
|
"v_proj.bias": "self_attn.v_proj.bias",
|
|
"out_proj.weight": "self_attn.o_proj.weight",
|
|
"ffn_w1.weight": "mlp.gate_proj.weight", # SwiGLU gate
|
|
"ffn_w3.weight": "mlp.up_proj.weight", # SwiGLU up
|
|
"ffn_w2.weight": "mlp.down_proj.weight", # SwiGLU down
|
|
}
|
|
out = {"model.embed_tokens.weight": sd["token_emb.weight"]}
|
|
leftovers = {}
|
|
blk_pat = re.compile(r"^blocks\.(\d+)\.(.+)$")
|
|
for k, v in sd.items():
|
|
if k == "token_emb.weight":
|
|
continue
|
|
m = blk_pat.match(k)
|
|
if m:
|
|
idx, sub = m.group(1), m.group(2)
|
|
if sub == "out_proj.bias":
|
|
sys.exit("❌ out_proj has a bias -- Qwen2 arch has no o_proj "
|
|
"bias, this checkpoint isn't Qwen2-compatible as-is.")
|
|
if sub not in block_map:
|
|
sys.exit(f"❌ Unknown per-block key: {k} -- send this back.")
|
|
out[f"model.layers.{idx}.{block_map[sub]}"] = v
|
|
else:
|
|
leftovers[k] = v
|
|
# classify the remaining top-level keys by tensor shape
|
|
for k, v in leftovers.items():
|
|
shp = tuple(v.shape)
|
|
if len(shp) == 1 and shp[0] == hidden:
|
|
print(f" final norm : {k} -> model.norm.weight")
|
|
out["model.norm.weight"] = v
|
|
elif len(shp) == 2 and shp[1] == hidden:
|
|
print(f" lm head : {k} -> lm_head.weight")
|
|
out["lm_head.weight"] = v
|
|
else:
|
|
sys.exit(f"❌ Unexplained top-level key: {k} {shp} -- send back.")
|
|
if "model.norm.weight" not in out:
|
|
sys.exit("❌ No final-norm tensor found (1-D, size=hidden). Send the "
|
|
"full key list (the tail beyond the first 80).")
|
|
print(f"✅ Remapped {len(out)} tensors to HF naming.")
|
|
return out
|
|
|
|
# Case 4: unknown naming -- print everything and stop
|
|
print("❌ Unrecognized key naming scheme. Full key list (first 80):")
|
|
for k in keys[:80]:
|
|
print(" ", k, tuple(sd[k].shape) if hasattr(sd[k], "shape") else "")
|
|
print(f" ... total {len(keys)} keys")
|
|
sys.exit(
|
|
"\nSend this key list back and I'll add the exact JiRack->HF mapping "
|
|
"to normalize_keys()."
|
|
)
|
|
|
|
|
|
def infer_config(sd):
|
|
"""Derive Qwen2 config.json entirely from tensor shapes."""
|
|
embed = sd["model.embed_tokens.weight"]
|
|
vocab_size, hidden_size = embed.shape
|
|
|
|
layer_ids = set()
|
|
pat = re.compile(r"^model\.layers\.(\d+)\.")
|
|
for k in sd:
|
|
m = pat.match(k)
|
|
if m:
|
|
layer_ids.add(int(m.group(1)))
|
|
num_layers = max(layer_ids) + 1
|
|
|
|
q_w = sd["model.layers.0.self_attn.q_proj.weight"] # [n_heads*head_dim, hidden]
|
|
k_w = sd["model.layers.0.self_attn.k_proj.weight"] # [n_kv*head_dim, hidden]
|
|
gate = sd["model.layers.0.mlp.gate_proj.weight"] # [intermediate, hidden]
|
|
intermediate_size = gate.shape[0]
|
|
|
|
# Qwen2 1.5B/7B/14B/32B all use head_dim=128 (1.5B: 12*128=1536)
|
|
head_dim = 128
|
|
num_attention_heads = q_w.shape[0] // head_dim
|
|
num_key_value_heads = k_w.shape[0] // head_dim
|
|
|
|
# sanity: every layer's FFN must have the same (expanded) width
|
|
widths = {sd[f"model.layers.{i}.mlp.gate_proj.weight"].shape[0] for i in layer_ids}
|
|
if len(widths) != 1:
|
|
sys.exit(f"❌ Inconsistent FFN widths across layers: {sorted(widths)}")
|
|
|
|
tie = "lm_head.weight" not in sd
|
|
cfg = {
|
|
"architectures": ["Qwen2ForCausalLM"],
|
|
"model_type": "qwen2",
|
|
"vocab_size": vocab_size,
|
|
"hidden_size": hidden_size,
|
|
"intermediate_size": intermediate_size,
|
|
"num_hidden_layers": num_layers,
|
|
"num_attention_heads": num_attention_heads,
|
|
"num_key_value_heads": num_key_value_heads,
|
|
"hidden_act": "silu",
|
|
"max_position_embeddings": MAX_POSITION,
|
|
"rms_norm_eps": RMS_NORM_EPS,
|
|
"rope_theta": ROPE_THETA,
|
|
"tie_word_embeddings": tie,
|
|
"torch_dtype": "bfloat16",
|
|
"use_cache": True,
|
|
"bos_token_id": 151646,
|
|
"eos_token_id": 151643,
|
|
}
|
|
print("🧾 Inferred config from tensor shapes:")
|
|
for k in ("vocab_size", "hidden_size", "intermediate_size", "num_hidden_layers",
|
|
"num_attention_heads", "num_key_value_heads", "tie_word_embeddings"):
|
|
print(f" {k} = {cfg[k]}")
|
|
print(f" rope_theta = {ROPE_THETA} (from the EDIT block -- verify for this base model!)")
|
|
return cfg
|
|
|
|
|
|
def save_hf(sd, cfg):
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
print(f"🔄 Casting weights to bf16 on {device.upper()} ...")
|
|
for k in sd:
|
|
t = sd[k]
|
|
if torch.is_tensor(t) and t.is_floating_point():
|
|
sd[k] = t.to(device=device, dtype=torch.bfloat16).cpu().contiguous()
|
|
|
|
try:
|
|
from safetensors.torch import save_file
|
|
# single-file safetensors; llama.cpp's converter handles it fine
|
|
path = os.path.join(OUTPUT_DIR, "model.safetensors")
|
|
print(f"💾 Saving {path} ...")
|
|
save_file(sd, path, metadata={"format": "pt"})
|
|
except ImportError:
|
|
# fallback: pytorch_model.bin, also accepted by convert_hf_to_gguf.py
|
|
path = os.path.join(OUTPUT_DIR, "pytorch_model.bin")
|
|
print(f"⚠️ safetensors not installed -- saving {path} instead (also works).")
|
|
torch.save(sd, path)
|
|
|
|
with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
|
|
json.dump(cfg, f, indent=2)
|
|
with open(os.path.join(OUTPUT_DIR, "generation_config.json"), "w") as f:
|
|
json.dump({"bos_token_id": cfg["bos_token_id"],
|
|
"eos_token_id": cfg["eos_token_id"],
|
|
"do_sample": True, "temperature": 0.6, "top_p": 0.95}, f, indent=2)
|
|
|
|
print(f"📎 Copying tokenizer from {TOKENIZER_DIR} ...")
|
|
same_dir = os.path.abspath(TOKENIZER_DIR) == os.path.abspath(OUTPUT_DIR)
|
|
if same_dir:
|
|
print(" TOKENIZER_DIR == OUTPUT_DIR -- tokenizer files are already in "
|
|
"place, skipping copy.")
|
|
copied = sum(
|
|
1 for name in os.listdir(TOKENIZER_DIR)
|
|
if name.startswith(("tokenizer", "special_tokens", "added_tokens",
|
|
"vocab", "merges", "chat_template"))
|
|
)
|
|
else:
|
|
copied = 0
|
|
for name in os.listdir(TOKENIZER_DIR):
|
|
if name.startswith(("tokenizer", "special_tokens", "added_tokens", "vocab", "merges", "chat_template")):
|
|
shutil.copy2(os.path.join(TOKENIZER_DIR, name), os.path.join(OUTPUT_DIR, name))
|
|
copied += 1
|
|
if copied == 0:
|
|
sys.exit(f"❌ No tokenizer files found in {TOKENIZER_DIR}")
|
|
print(f" copied {copied} tokenizer files.")
|
|
|
|
|
|
def verify(cfg):
|
|
"""Cross-check tokenizer length vs embedding rows."""
|
|
try:
|
|
from transformers import AutoTokenizer
|
|
tok = AutoTokenizer.from_pretrained(OUTPUT_DIR)
|
|
n = len(tok)
|
|
rows = cfg["vocab_size"]
|
|
if n > rows:
|
|
sys.exit(f"❌ Tokenizer has {n} tokens but embedding matrix only {rows} rows -- "
|
|
f"resize the checkpoint before converting.")
|
|
print(f"✅ Tokenizer check: {n} tokens <= {rows} embedding rows "
|
|
f"({rows - n} spare rows).")
|
|
except Exception as e:
|
|
print(f"⚠️ Could not verify tokenizer ({e}) -- continuing anyway.")
|
|
|
|
|
|
def main():
|
|
if not os.path.exists(CKPT_PATH):
|
|
sys.exit(f"❌ {CKPT_PATH} not found")
|
|
sd = load_state_dict(CKPT_PATH)
|
|
sd = drop_training_buffers(sd)
|
|
sd = normalize_keys(sd)
|
|
cfg = infer_config(sd)
|
|
save_hf(sd, cfg)
|
|
verify(cfg)
|
|
|
|
print("\n" + "=" * 78)
|
|
print("✅ Stage 1 done. HF model at:", OUTPUT_DIR)
|
|
print("=" * 78)
|
|
|
|
out_norm = OUTPUT_DIR.rstrip("/")
|
|
gguf_base = "jirack_1p5b" if out_norm in ("", ".") else out_norm
|
|
|
|
q2_0_block = ""
|
|
if EMIT_Q2_0_CMD:
|
|
suffix = "Q2_0" if Q2_0_GROUP == 64 else f"Q2_0_g{Q2_0_GROUP}"
|
|
fork_note = (
|
|
"group-64 is in mainline llama.cpp -- no fork needed, CPU/Metal ready."
|
|
if Q2_0_GROUP == 64 else
|
|
"group-128 needs a CUDA fork -- not needed on CPU-only."
|
|
)
|
|
q2_0_block = """
|
|
Ternary quantization (Q2_0, 2 bits/weight, {{-1,0,+1}} + fp16 group scale --
|
|
this is the real encoding for BitNet-style ternary weights, once your model
|
|
actually runs true ternary at inference rather than bf16 dense):
|
|
{fork_note}
|
|
|
|
./build/bin/llama-quantize {gguf} {gguf_q2} {suffix}
|
|
""".format(
|
|
fork_note=fork_note,
|
|
gguf=gguf_base + ".gguf",
|
|
gguf_q2=gguf_base + f".{suffix}.gguf",
|
|
suffix=suffix,
|
|
)
|
|
|
|
print("""
|
|
Stage 2 -- make the GGUF (one-time llama.cpp setup, then per model):
|
|
|
|
git clone https://github.com/ggml-org/llama.cpp /mnt/nfs_share/llama.cpp
|
|
cd /mnt/nfs_share/llama.cpp
|
|
pip install -r requirements.txt
|
|
|
|
python convert_hf_to_gguf.py {out} \\
|
|
--outfile {gguf} --outtype bf16
|
|
|
|
Optional dense quantization (build llama.cpp first: cmake -B build && cmake --build build -j):
|
|
|
|
./build/bin/llama-quantize {gguf} {gguf_q} Q4_K_M
|
|
{q2_0_block}""".format(
|
|
out=OUTPUT_DIR,
|
|
gguf=gguf_base + ".gguf",
|
|
gguf_q=gguf_base + ".Q4_K_M.gguf",
|
|
q2_0_block=q2_0_block,
|
|
))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|