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