|
|
|
|
|
"""Packed unary converter: uint8 magnitudes + bitpacked signs + per-row scales.""" |
|
|
import os, json, sys, time |
|
|
import numpy as np |
|
|
from pathlib import Path |
|
|
|
|
|
def load_safetensors(model_dir): |
|
|
from safetensors.torch import load_file |
|
|
tensors = {} |
|
|
for f in sorted(Path(model_dir).glob("*.safetensors")): |
|
|
print(f" Loading {f.name}...") |
|
|
for k, v in load_file(str(f)).items(): |
|
|
tensors[k] = v.float().numpy() |
|
|
return tensors |
|
|
|
|
|
def quantize_packed(w, n_levels=7): |
|
|
out_dim, in_dim = w.shape |
|
|
chunks = (in_dim + 63) // 64 |
|
|
padded = chunks * 64 |
|
|
row_max = np.max(np.abs(w), axis=1, keepdims=True) |
|
|
row_max = np.where(row_max == 0, 1.0, row_max) |
|
|
scales = (row_max.flatten() / n_levels).astype(np.float32) |
|
|
mags = np.clip(np.round(np.abs(w / scales[:, None])), 0, n_levels).astype(np.uint8) |
|
|
signs = (w < 0) |
|
|
rmm = np.max(mags, axis=1).astype(np.uint8) |
|
|
if in_dim < padded: |
|
|
sp = np.zeros((out_dim, padded), dtype=bool) |
|
|
sp[:, :in_dim] = signs |
|
|
else: |
|
|
sp = signs |
|
|
bit_pos = np.uint64(1) << np.arange(64, dtype=np.uint64) |
|
|
sign_bits = np.bitwise_or.reduce(sp.reshape(out_dim, chunks, 64).astype(np.uint64) * bit_pos, axis=2) |
|
|
return mags, sign_bits, scales, rmm, np.mean(mags), np.mean(mags == 0) |
|
|
|
|
|
def convert(tensors, output_dir, n_levels=7): |
|
|
os.makedirs(output_dir, exist_ok=True) |
|
|
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, |
|
|
"n_levels":n_levels,"quant_type":"packed_unary"} |
|
|
linear_keys = [k for k in tensors if any(p in k for p in |
|
|
['q_proj.weight','k_proj.weight','v_proj.weight','o_proj.weight', |
|
|
'gate_proj.weight','up_proj.weight','down_proj.weight'])] |
|
|
other_keys = [k for k in tensors if k not in linear_keys] |
|
|
with open(os.path.join(output_dir, "config.json"), "w") as f: |
|
|
json.dump(config, f, indent=2) |
|
|
total_packed = total_orig = 0 |
|
|
all_avg = [] |
|
|
for key in linear_keys: |
|
|
w = tensors[key]; total_orig += w.nbytes |
|
|
t0 = time.time() |
|
|
mags, sb, sc, rmm, am, sp = quantize_packed(w, n_levels) |
|
|
dt = time.time() - t0 |
|
|
pfx = os.path.join(output_dir, key.replace(".", "_")) |
|
|
mags.tofile(pfx+".mags"); sb.tofile(pfx+".signs") |
|
|
sc.tofile(pfx+".scales"); rmm.tofile(pfx+".rmm") |
|
|
ub = mags.nbytes + sb.nbytes + sc.nbytes + rmm.nbytes |
|
|
total_packed += ub; all_avg.append(am) |
|
|
print(f" {key}: {w.shape} -> {ub/1024:.0f}KB (avg_mag={am:.2f}, {dt:.1f}s)") |
|
|
total_fp16 = 0 |
|
|
for key in other_keys: |
|
|
w = tensors[key].astype(np.float16) |
|
|
pfx = os.path.join(output_dir, key.replace(".", "_")) |
|
|
w.tofile(pfx+".fp16"); total_fp16 += w.nbytes |
|
|
manifest = {"packed":{k:list(tensors[k].shape) for k in linear_keys}, |
|
|
"fp16":{k:list(tensors[k].shape) for k in other_keys}} |
|
|
with open(os.path.join(output_dir, "manifest.json"), "w") as f: |
|
|
json.dump(manifest, f, indent=2) |
|
|
print(f"\n=== PACKED UNARY ===") |
|
|
print(f"Packed linear: {total_packed/1e6:.1f} MB | FP16 other: {total_fp16/1e6:.1f} MB") |
|
|
print(f"Total: {(total_packed+total_fp16)/1e6:.1f} MB | Avg mag: {np.mean(all_avg):.3f}") |
|
|
print(f"Expected speedup vs 7-plane: {7/np.mean(all_avg):.1f}x") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
model_dir = sys.argv[1] if len(sys.argv) > 1 else "deepseek-r1-1.5b-hf" |
|
|
output_dir = sys.argv[2] if len(sys.argv) > 2 else "deepseek-r1-1.5b-packed" |
|
|
tensors = load_safetensors(model_dir) |
|
|
convert(tensors, output_dir) |
|
|
print("Done!") |
|
|
|
