File size: 2,535 Bytes
19ed98b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import os, json, numpy as np, time, sys
from pathlib import Path
from safetensors import safe_open
import torch
sys.path.insert(0, "/root/ternary_engine")
from convert import quantize_weight_matrix
model_dir = "/root/ternary_engine/deepseek-r1-1.5b-hf"
output_dir = "/root/ternary_engine/deepseek-r1-1.5b-ternary"
alpha = 0.7
os.makedirs(output_dir, exist_ok=True)
tensors = {}
for f in sorted(Path(model_dir).glob("*.safetensors")):
print("Loading " + f.name)
with safe_open(str(f), framework="pt") as st:
for key in st.keys():
tensors[key] = st.get_tensor(key).float().numpy()
print("Loaded " + str(len(tensors)) + " tensors")
config = {
"hidden_size": 1536, "intermediate_size": 8960,
"num_attention_heads": 12, "num_key_value_heads": 2,
"num_hidden_layers": 28, "vocab_size": 151936,
"head_dim": 128, "rope_theta": 1000000.0,
"rms_norm_eps": 1e-6, "alpha": alpha,
}
ternary_manifest = {}
fp16_manifest = {}
linear_suffixes = ['q_proj.weight', 'k_proj.weight', 'v_proj.weight',
'o_proj.weight', 'gate_proj.weight', 'up_proj.weight',
'down_proj.weight']
total_tb = 0
total_ob = 0
for key, w in tensors.items():
prefix = os.path.join(output_dir, key.replace(".", "_"))
is_linear = any(key.endswith(s) for s in linear_suffixes)
if is_linear and len(w.shape) == 2:
out_dim, in_dim = w.shape
total_ob += w.nbytes
t0 = time.time()
pos, neg, scales, sparsity = quantize_weight_matrix(w, alpha)
dt = time.time() - t0
pos.tofile(prefix + ".pos")
neg.tofile(prefix + ".neg")
scales.tofile(prefix + ".scales")
tb = pos.nbytes + neg.nbytes + scales.nbytes
total_tb += tb
ratio = w.nbytes / tb
ternary_manifest[key] = list(w.shape)
print(" T %s: %s -> %dKB (%.1fx, %.0f%% sparse, %.1fs)" % (
key, str(w.shape), tb // 1024, ratio, sparsity * 100, dt))
else:
w16 = w.astype(np.float16)
w16.tofile(prefix + ".fp16")
fp16_manifest[key] = list(w.shape)
print(" F %s: %s -> %dKB" % (key, str(w.shape), w16.nbytes // 1024))
with open(os.path.join(output_dir, "config.json"), "w") as f:
json.dump(config, f, indent=2)
with open(os.path.join(output_dir, "manifest.json"), "w") as f:
json.dump({"ternary": ternary_manifest, "fp16": fp16_manifest}, f, indent=2)
print("")
print("Ternary: %.1fMB (from %.1fMB FP32)" % (total_tb / 1024 / 1024, total_ob / 1024 / 1024))
print("DONE")
|