File size: 6,995 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | #!/usr/bin/env python3
"""
Convert model weights to UNARY (base-1) thermometer encoding.
True unary: magnitude N = N consecutive 1-bits across N bitplanes.
Each bitplane contributes equally (value=1), NOT binary powers.
Weight 0.3 with scale -> magnitude 5 -> planes 0,1,2,3,4 have bit set
Weight -0.1 with scale -> magnitude 2, sign=neg -> planes 0,1 set + sign bit
More precision than ternary (N+1 levels vs 3), still no multiplication.
(c) 2026 OpenTransformers Ltd / Scott Bisset
"""
import os
import json
import numpy as np
from pathlib import Path
import time
def load_safetensors(model_dir):
"""Load all tensors from safetensors files."""
import torch
from safetensors.torch import load_file
tensors = {}
for f in sorted(Path(model_dir).glob("*.safetensors")):
print(f"Loading {f.name}...")
state = load_file(str(f))
for key, val in state.items():
tensors[key] = val.float().numpy()
return tensors
def quantize_matrix_unary(weight, n_planes=7):
"""Quantize weight matrix to unary thermometer encoding.
n_planes determines max magnitude (and precision levels = n_planes + 1).
n_planes=7 gives 8 levels: {0,1,2,3,4,5,6,7} * sign = 15 distinct values.
Returns: sign_bits, mag_planes, scales, sparsity
"""
w = weight.astype(np.float32)
out_dim, in_dim = w.shape
chunks = ((in_dim + 63) // 64)
padded = chunks * 64
# Per-row quantization
row_max = np.max(np.abs(w), axis=1, keepdims=True)
row_max = np.where(row_max == 0, 1.0, row_max)
# Scale to [0, n_planes] range per row
scales = (row_max.flatten() / n_planes).astype(np.float32)
# Quantize to integer magnitudes
w_scaled = w / scales[:, None] # Now in [-n_planes, +n_planes]
magnitudes = np.round(np.abs(w_scaled)).astype(np.int32)
magnitudes = np.clip(magnitudes, 0, n_planes)
signs = (w < 0) # True = negative
# Sparsity (magnitude 0)
sparsity = np.mean(magnitudes == 0)
# Pad to multiple of 64
if in_dim < padded:
magnitudes = np.concatenate([magnitudes, np.zeros((out_dim, padded - in_dim), dtype=np.int32)], axis=1)
signs = np.concatenate([signs, np.zeros((out_dim, padded - in_dim), dtype=bool)], axis=1)
# Pack sign bits - vectorized
bit_positions = (np.uint64(1) << np.arange(64, dtype=np.uint64))
signs_r = signs.reshape(out_dim, chunks, 64).astype(np.uint64)
sign_bits = np.bitwise_or.reduce(signs_r * bit_positions, axis=2) # [out_dim, chunks]
# Pack magnitude planes - thermometer encoding
# Plane p has bit set where magnitude > p (i.e., magnitude >= p+1)
mag_planes = np.zeros((n_planes, out_dim, chunks), dtype=np.uint64)
for p in range(n_planes):
active = (magnitudes >= (p + 1)) # [out_dim, padded]
active_r = active.reshape(out_dim, chunks, 64).astype(np.uint64)
mag_planes[p] = np.bitwise_or.reduce(active_r * bit_positions, axis=2)
return sign_bits, mag_planes, scales, sparsity
def save_unary_model(tensors, output_dir, n_planes=7):
"""Convert and save full model to unary format."""
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_planes": n_planes,
"quant_type": "unary",
}
ternary_keys = []
keep_keys = []
for key in tensors:
if any(p in key for p in ['q_proj.weight', 'k_proj.weight', 'v_proj.weight',
'o_proj.weight', 'gate_proj.weight', 'up_proj.weight',
'down_proj.weight']):
ternary_keys.append(key)
else:
keep_keys.append(key)
print(f"\nUnary layers: {len(ternary_keys)} (n_planes={n_planes}, levels={n_planes+1})")
print(f"FP16 layers: {len(keep_keys)}")
with open(os.path.join(output_dir, "config.json"), "w") as f:
json.dump(config, f, indent=2)
total_unary_bytes = 0
total_original_bytes = 0
for key in ternary_keys:
w = tensors[key]
out_dim, in_dim = w.shape
total_original_bytes += w.nbytes
t0 = time.time()
sign_bits, mag_planes, scales, sparsity = quantize_matrix_unary(w, n_planes)
dt = time.time() - t0
prefix = os.path.join(output_dir, key.replace(".", "_"))
sign_bits.tofile(prefix + ".sign")
mag_planes.tofile(prefix + ".planes")
scales.tofile(prefix + ".scales")
unary_bytes = sign_bits.nbytes + mag_planes.nbytes + scales.nbytes
total_unary_bytes += unary_bytes
ratio = w.nbytes / unary_bytes
# Calculate effective bits per weight
bpw = (unary_bytes * 8) / (out_dim * in_dim)
print(f" {key}: {w.shape} -> unary ({unary_bytes/1024:.0f}KB, "
f"{ratio:.1f}x compress, {bpw:.2f} bpw, {sparsity:.1%} sparse, {dt:.1f}s)")
total_fp16_bytes = 0
for key in keep_keys:
w = tensors[key].astype(np.float16)
prefix = os.path.join(output_dir, key.replace(".", "_"))
w.tofile(prefix + ".fp16")
total_fp16_bytes += w.nbytes
print(f" {key}: {w.shape} -> fp16 ({w.nbytes/1024:.0f}KB)")
manifest = {
"unary": {k: list(tensors[k].shape) for k in ternary_keys},
"fp16": {k: list(tensors[k].shape) for k in keep_keys},
}
with open(os.path.join(output_dir, "manifest.json"), "w") as f:
json.dump(manifest, f, indent=2)
total_bytes = total_unary_bytes + total_fp16_bytes
avg_bpw = (total_unary_bytes * 8) / sum(np.prod(tensors[k].shape) for k in ternary_keys)
print(f"\n=== Summary ===")
print(f"Original FP32 linear weights: {total_original_bytes/1024/1024:.1f} MB")
print(f"Unary linear weights: {total_unary_bytes/1024/1024:.1f} MB")
print(f"FP16 other weights: {total_fp16_bytes/1024/1024:.1f} MB")
print(f"Total model size: {total_bytes/1024/1024:.1f} MB")
print(f"Average bits per weight (linear): {avg_bpw:.2f}")
print(f"Compression vs FP32: {(total_original_bytes + total_fp16_bytes)/total_bytes:.1f}x")
print(f"Precision levels: {n_planes + 1} (vs ternary=3, INT4=16)")
if __name__ == "__main__":
import sys
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-unary"
n_planes = int(sys.argv[3]) if len(sys.argv) > 3 else 7
print(f"Loading model from {model_dir}...")
tensors = load_safetensors(model_dir)
print(f"Converting to unary (n_planes={n_planes})...")
save_unary_model(tensors, output_dir, n_planes)
print("Done!")
|