JiRackTernaryPro_1b / pack_v2.py
kgrabko's picture
Final version upload of JiRackTernaryPro 1B weights
5ba17ce verified
Raw
History Blame Contribute Delete
1.12 kB
import torch
import os
def pack_for_prod(input_path, output_path):
print(f"📦 Загружаем исходник: {input_path}")
sd = torch.load(input_path, map_location="cpu")
if "model_state_dict" in sd: sd = sd["model_state_dict"]
new_sd = {}
for name, weight in sd.items():
if any(x in name for x in ["q_proj", "k_proj", "v_proj", "out_proj", "ffn_w"]):
print(f" -> Пакуем слой: {name}")
# Твоя формула квантования
gamma = weight.abs().mean().clamp(min=1e-9)
w_quant = torch.round(weight / gamma).clamp(-1, 1).to(torch.int8)
new_sd[name] = w_quant
# СОХРАНЯЕМ ГАММУ С ТОЧНЫМ ИМЕНЕМ
gamma_key = name.replace(".weight", "_gamma")
new_sd[gamma_key] = gamma.half()
else:
new_sd[name] = weight.half()
torch.save(new_sd, output_path)
print(f"✅ Готово! Файл сохранен в {output_path}")
pack_for_prod("jiarck_pro_1b_model.pt", "jirack_pro_1b_prod.pt")