jirack_to_gguf_14b.py
14.8 KB · 374 lines · python Raw
1 # ==============================================================================
2 # JiRack -> GGUF converter, 14B edition (stage 1: .pt -> HuggingFace folder)
3 # COPYRIGHT (c) 2026 Konstantin Vladimirovich Grabko.
4 #
5 # HF still reports it as Qwen2ForCausalLM -- same as all other sizes here).
6 # Verified against JiRackTernaryUltra_14b.py [DS14-1..5]:
7 # vocab_size 152064, hidden 5120, intermediate 13824, n_layers 48,
8 # n_heads 40, n_kv_heads 8, head_dim 128 (40*128=5120 -- the converter's
9 # hardcoded 128 is correct), rope_theta 1000000.0 (NOT 7B's 10000.0!),
10 # rms_eps 1e-5 (NOT 7B's 1e-6!), tie_word_embeddings = FALSE (separate
11 # lm_head, same as 7B/32B) -- this converter auto-detects that from the
12 # presence of lm_head.weight. QKV bias=True (same as all sizes).
13 #
14 # Pipeline is two stages:
15 #
16 # Stage 1 (THIS SCRIPT, run in venv_ji):
17 # model.pt -> HF folder (model.safetensors + config.json + tokenizer)
18 #
19 # Stage 2 (llama.cpp, run once per model):
20 # python convert_hf_to_gguf.py <hf_folder> \
21 # --outfile jirack_14b.gguf --outtype bf16
22 # ./build/bin/llama-quantize jirack_14b.gguf \
23 # jirack_14b.Q4_K_M.gguf Q4_K_M
24 #
25 # Key points handled here:
26 # * config.json is derived from the ACTUAL tensor shapes in the checkpoint,
27 # so vocab (151936 vs 7B's 152064) and any Net2Net-expanded FFN width are
28 # picked up automatically -- no stock-config copying.
29 # * lambda_ buffers (ternary fake-quant training machinery) are dropped --
30 # at inference you run set_lambda(0.0) anyway, so the stored weights ARE
31 # the full-precision weights; the exported model is a plain Qwen2 dense.
32 # * Keys: HF naming passes through; JiRack native naming
33 # (token_emb / blocks.N.* / ffn_w1-w3-w2) is remapped automatically.
34 #
35 # EDIT THE THREE PATHS BELOW.
36 # ==============================================================================
37
38 import json
39 import os
40 import re
41 import shutil
42 import sys
43
44 import torch
45
46 # ========================= EDIT THESE =========================
47 CKPT_PATH = "model.pt"
48 TOKENIZER_DIR = "."
49 OUTPUT_DIR = "."
50 # rope_theta cannot be inferred from tensor shapes -- set per base model:
51 # DeepSeek-R1-Distill-Qwen-1.5B -> 10000.0 (same as 7B)
52 # DeepSeek-R1-Distill-Qwen-14B -> 1000000.0 <-- this file
53 # DeepSeek-R1-Distill-Qwen-32B -> 1000000.0
54 ROPE_THETA = 1000000.0 # [DS14-2]
55 MAX_POSITION = 131072
56 RMS_NORM_EPS = 1e-5 # [DS14-3]
57
58 # Q2_0 = 2-bit ternary {-1, 0, +1} quantization, one fp16 scale per group of
59 # weights -- the real encoding for BitNet-style ternary weights, once inference
60 # actually runs true ternary rather than bf16 dense. For now Q4_K_M remains
61 # the practical choice; the Q2_0 command is just printed ready for later.
62 EMIT_Q2_0_CMD = True
63 Q2_0_GROUP = 64 # 64 = mainline llama.cpp, no fork needed.
64 # ================================================================
65
66 # HF Qwen2 key patterns we expect to find (N = layer index)
67 HF_LAYER_KEYS = [
68 "model.layers.{n}.self_attn.q_proj.weight",
69 "model.layers.{n}.self_attn.q_proj.bias",
70 "model.layers.{n}.self_attn.k_proj.weight",
71 "model.layers.{n}.self_attn.k_proj.bias",
72 "model.layers.{n}.self_attn.v_proj.weight",
73 "model.layers.{n}.self_attn.v_proj.bias",
74 "model.layers.{n}.self_attn.o_proj.weight",
75 "model.layers.{n}.mlp.gate_proj.weight",
76 "model.layers.{n}.mlp.up_proj.weight",
77 "model.layers.{n}.mlp.down_proj.weight",
78 "model.layers.{n}.input_layernorm.weight",
79 "model.layers.{n}.post_attention_layernorm.weight",
80 ]
81 HF_TOP_KEYS = [
82 "model.embed_tokens.weight",
83 "model.norm.weight",
84 "lm_head.weight",
85 ]
86
87
88 def load_state_dict(path):
89 print(f"📥 Loading checkpoint: {path}")
90 ckpt = torch.load(path, map_location="cpu", weights_only=False)
91 sd = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
92 if not isinstance(sd, dict):
93 sys.exit("❌ Checkpoint is not a state_dict and has no 'model' key.")
94 return sd
95
96
97 def drop_training_buffers(sd):
98 dropped = [k for k in sd if k.endswith("lambda_")]
99 for k in dropped:
100 del sd[k]
101 if dropped:
102 print(f"🧹 Dropped {len(dropped)} lambda_ buffers (ternary training machinery).")
103 return sd
104
105
106 def normalize_keys(sd):
107 """Pass HF-style keys through; try trivial prefix fixes; else abort with a listing."""
108 keys = list(sd.keys())
109
110 # Case 1: already HF-style
111 if "model.embed_tokens.weight" in sd:
112 print("✅ Keys already use HF (Qwen2) naming -- no remap needed.")
113 return sd
114
115 # Case 2: same names but without the leading 'model.' (e.g. 'embed_tokens.weight')
116 if "embed_tokens.weight" in sd:
117 print("🔁 Keys look HF-like without the 'model.' prefix -- adding it.")
118 out = {}
119 for k, v in sd.items():
120 if k == "lm_head.weight":
121 out[k] = v
122 else:
123 out["model." + k] = v
124 if "model.embed_tokens.weight" in out:
125 return out
126
127 # Case 3: JiRack native naming (token_emb / blocks.N.* / ffn_w1-w3-w2)
128 if "token_emb.weight" in sd and any(k.startswith("blocks.") for k in sd):
129 print("🔁 JiRack native naming detected -- remapping to HF (Qwen2) keys.")
130 hidden = sd["token_emb.weight"].shape[1]
131 block_map = {
132 "norm1.weight": "input_layernorm.weight",
133 "norm2.weight": "post_attention_layernorm.weight",
134 "q_proj.weight": "self_attn.q_proj.weight",
135 "q_proj.bias": "self_attn.q_proj.bias",
136 "k_proj.weight": "self_attn.k_proj.weight",
137 "k_proj.bias": "self_attn.k_proj.bias",
138 "v_proj.weight": "self_attn.v_proj.weight",
139 "v_proj.bias": "self_attn.v_proj.bias",
140 "out_proj.weight": "self_attn.o_proj.weight",
141 "ffn_w1.weight": "mlp.gate_proj.weight", # SwiGLU gate
142 "ffn_w3.weight": "mlp.up_proj.weight", # SwiGLU up
143 "ffn_w2.weight": "mlp.down_proj.weight", # SwiGLU down
144 }
145 out = {"model.embed_tokens.weight": sd["token_emb.weight"]}
146 leftovers = {}
147 blk_pat = re.compile(r"^blocks\.(\d+)\.(.+)$")
148 for k, v in sd.items():
149 if k == "token_emb.weight":
150 continue
151 m = blk_pat.match(k)
152 if m:
153 idx, sub = m.group(1), m.group(2)
154 if sub == "out_proj.bias":
155 sys.exit("❌ out_proj has a bias -- Qwen2 arch has no o_proj "
156 "bias, this checkpoint isn't Qwen2-compatible as-is.")
157 if sub not in block_map:
158 sys.exit(f"❌ Unknown per-block key: {k} -- send this back.")
159 out[f"model.layers.{idx}.{block_map[sub]}"] = v
160 else:
161 leftovers[k] = v
162 # classify the remaining top-level keys by tensor shape
163 for k, v in leftovers.items():
164 shp = tuple(v.shape)
165 if len(shp) == 1 and shp[0] == hidden:
166 print(f" final norm : {k} -> model.norm.weight")
167 out["model.norm.weight"] = v
168 elif len(shp) == 2 and shp[1] == hidden:
169 print(f" lm head : {k} -> lm_head.weight")
170 out["lm_head.weight"] = v
171 else:
172 sys.exit(f"❌ Unexplained top-level key: {k} {shp} -- send back.")
173 if "model.norm.weight" not in out:
174 sys.exit("❌ No final-norm tensor found (1-D, size=hidden). Send the "
175 "full key list (the tail beyond the first 80).")
176 print(f"✅ Remapped {len(out)} tensors to HF naming.")
177 return out
178
179 # Case 4: unknown naming -- print everything and stop
180 print("❌ Unrecognized key naming scheme. Full key list (first 80):")
181 for k in keys[:80]:
182 print(" ", k, tuple(sd[k].shape) if hasattr(sd[k], "shape") else "")
183 print(f" ... total {len(keys)} keys")
184 sys.exit(
185 "\nSend this key list back and I'll add the exact JiRack->HF mapping "
186 "to normalize_keys()."
187 )
188
189
190 def infer_config(sd):
191 """Derive Qwen2 config.json entirely from tensor shapes."""
192 embed = sd["model.embed_tokens.weight"]
193 vocab_size, hidden_size = embed.shape
194
195 layer_ids = set()
196 pat = re.compile(r"^model\.layers\.(\d+)\.")
197 for k in sd:
198 m = pat.match(k)
199 if m:
200 layer_ids.add(int(m.group(1)))
201 num_layers = max(layer_ids) + 1
202
203 q_w = sd["model.layers.0.self_attn.q_proj.weight"] # [n_heads*head_dim, hidden]
204 k_w = sd["model.layers.0.self_attn.k_proj.weight"] # [n_kv*head_dim, hidden]
205 gate = sd["model.layers.0.mlp.gate_proj.weight"] # [intermediate, hidden]
206 intermediate_size = gate.shape[0]
207
208 # Qwen2 1.5B/7B/14B/32B all use head_dim=128 (1.5B: 12*128=1536)
209 head_dim = 128
210 num_attention_heads = q_w.shape[0] // head_dim
211 num_key_value_heads = k_w.shape[0] // head_dim
212
213 # sanity: every layer's FFN must have the same (expanded) width
214 widths = {sd[f"model.layers.{i}.mlp.gate_proj.weight"].shape[0] for i in layer_ids}
215 if len(widths) != 1:
216 sys.exit(f"❌ Inconsistent FFN widths across layers: {sorted(widths)}")
217
218 tie = "lm_head.weight" not in sd
219 cfg = {
220 "architectures": ["Qwen2ForCausalLM"],
221 "model_type": "qwen2",
222 "vocab_size": vocab_size,
223 "hidden_size": hidden_size,
224 "intermediate_size": intermediate_size,
225 "num_hidden_layers": num_layers,
226 "num_attention_heads": num_attention_heads,
227 "num_key_value_heads": num_key_value_heads,
228 "hidden_act": "silu",
229 "max_position_embeddings": MAX_POSITION,
230 "rms_norm_eps": RMS_NORM_EPS,
231 "rope_theta": ROPE_THETA,
232 "tie_word_embeddings": tie,
233 "torch_dtype": "bfloat16",
234 "use_cache": True,
235 "bos_token_id": 151646,
236 "eos_token_id": 151643,
237 }
238 print("🧾 Inferred config from tensor shapes:")
239 for k in ("vocab_size", "hidden_size", "intermediate_size", "num_hidden_layers",
240 "num_attention_heads", "num_key_value_heads", "tie_word_embeddings"):
241 print(f" {k} = {cfg[k]}")
242 print(f" rope_theta = {ROPE_THETA} (from the EDIT block -- verify for this base model!)")
243 return cfg
244
245
246 def save_hf(sd, cfg):
247 os.makedirs(OUTPUT_DIR, exist_ok=True)
248
249 device = "cuda" if torch.cuda.is_available() else "cpu"
250 print(f"🔄 Casting weights to bf16 on {device.upper()} ...")
251 for k in sd:
252 t = sd[k]
253 if torch.is_tensor(t) and t.is_floating_point():
254 sd[k] = t.to(device=device, dtype=torch.bfloat16).cpu().contiguous()
255
256 try:
257 from safetensors.torch import save_file
258 # single-file safetensors; llama.cpp's converter handles it fine
259 path = os.path.join(OUTPUT_DIR, "model.safetensors")
260 print(f"💾 Saving {path} ...")
261 save_file(sd, path, metadata={"format": "pt"})
262 except ImportError:
263 # fallback: pytorch_model.bin, also accepted by convert_hf_to_gguf.py
264 path = os.path.join(OUTPUT_DIR, "pytorch_model.bin")
265 print(f"⚠️ safetensors not installed -- saving {path} instead (also works).")
266 torch.save(sd, path)
267
268 with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
269 json.dump(cfg, f, indent=2)
270 with open(os.path.join(OUTPUT_DIR, "generation_config.json"), "w") as f:
271 json.dump({"bos_token_id": cfg["bos_token_id"],
272 "eos_token_id": cfg["eos_token_id"],
273 "do_sample": True, "temperature": 0.6, "top_p": 0.95}, f, indent=2)
274
275 print(f"📎 Copying tokenizer from {TOKENIZER_DIR} ...")
276 same_dir = os.path.abspath(TOKENIZER_DIR) == os.path.abspath(OUTPUT_DIR)
277 if same_dir:
278 print(" TOKENIZER_DIR == OUTPUT_DIR -- tokenizer files are already in "
279 "place, skipping copy.")
280 copied = sum(
281 1 for name in os.listdir(TOKENIZER_DIR)
282 if name.startswith(("tokenizer", "special_tokens", "added_tokens",
283 "vocab", "merges", "chat_template"))
284 )
285 else:
286 copied = 0
287 for name in os.listdir(TOKENIZER_DIR):
288 if name.startswith(("tokenizer", "special_tokens", "added_tokens", "vocab", "merges", "chat_template")):
289 shutil.copy2(os.path.join(TOKENIZER_DIR, name), os.path.join(OUTPUT_DIR, name))
290 copied += 1
291 if copied == 0:
292 sys.exit(f"❌ No tokenizer files found in {TOKENIZER_DIR}")
293 print(f" copied {copied} tokenizer files.")
294
295
296 def verify(cfg):
297 """Cross-check tokenizer length vs embedding rows."""
298 try:
299 from transformers import AutoTokenizer
300 tok = AutoTokenizer.from_pretrained(OUTPUT_DIR)
301 n = len(tok)
302 rows = cfg["vocab_size"]
303 if n > rows:
304 sys.exit(f"❌ Tokenizer has {n} tokens but embedding matrix only {rows} rows -- "
305 f"resize the checkpoint before converting.")
306 print(f"✅ Tokenizer check: {n} tokens <= {rows} embedding rows "
307 f"({rows - n} spare rows).")
308 except Exception as e:
309 print(f"⚠️ Could not verify tokenizer ({e}) -- continuing anyway.")
310
311
312 def main():
313 if not os.path.exists(CKPT_PATH):
314 sys.exit(f"❌ {CKPT_PATH} not found")
315 sd = load_state_dict(CKPT_PATH)
316 sd = drop_training_buffers(sd)
317 sd = normalize_keys(sd)
318 cfg = infer_config(sd)
319 save_hf(sd, cfg)
320 verify(cfg)
321
322 print("\n" + "=" * 78)
323 print("✅ Stage 1 done. HF model at:", OUTPUT_DIR)
324 print("=" * 78)
325
326 out_norm = OUTPUT_DIR.rstrip("/")
327 gguf_base = "jirack_14b" if out_norm in ("", ".") else out_norm
328
329 q2_0_block = ""
330 if EMIT_Q2_0_CMD:
331 suffix = "Q2_0" if Q2_0_GROUP == 64 else f"Q2_0_g{Q2_0_GROUP}"
332 fork_note = (
333 "group-64 is in mainline llama.cpp -- no fork needed, CPU/Metal ready."
334 if Q2_0_GROUP == 64 else
335 "group-128 needs a CUDA fork -- not needed on CPU-only."
336 )
337 q2_0_block = """
338 Ternary quantization (Q2_0, 2 bits/weight, {{-1,0,+1}} + fp16 group scale --
339 this is the real encoding for BitNet-style ternary weights, once your model
340 actually runs true ternary at inference rather than bf16 dense):
341 {fork_note}
342
343 ./build/bin/llama-quantize {gguf} {gguf_q2} {suffix}
344 """.format(
345 fork_note=fork_note,
346 gguf=gguf_base + ".gguf",
347 gguf_q2=gguf_base + f".{suffix}.gguf",
348 suffix=suffix,
349 )
350
351 print("""
352 Stage 2 -- make the GGUF (one-time llama.cpp setup, then per model):
353
354 git clone https://github.com/ggml-org/llama.cpp /mnt/nfs_share/llama.cpp
355 cd /mnt/nfs_share/llama.cpp
356 pip install -r requirements.txt
357
358 python convert_hf_to_gguf.py {out} \\
359 --outfile {gguf} --outtype bf16
360
361 Optional dense quantization (build llama.cpp first: cmake -B build && cmake --build build -j):
362
363 ./build/bin/llama-quantize {gguf} {gguf_q} Q4_K_M
364 {q2_0_block}""".format(
365 out=OUTPUT_DIR,
366 gguf=gguf_base + ".gguf",
367 gguf_q=gguf_base + ".Q4_K_M.gguf",
368 q2_0_block=q2_0_block,
369 ))
370
371
372 if __name__ == "__main__":
373 main()
374