File size: 9,288 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 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 | #!/usr/bin/env python3
"""
PROPER UNARY CONVERTER — Global quantum, torch-based, BF16 support
Clips at P99.9 of |weights| instead of absmax to avoid wasting
quantization range on rare outliers. Values above clip point
saturate at K (still represented, just capped).
(c) 2026 OpenTransformers Ltd / Scott Bisset
"""
import torch, json, os, sys, gc, shutil
from safetensors import safe_open
import numpy as np
def scan_all_linears(model_dir):
"""Scan linear layers, return global stats."""
index_path = os.path.join(model_dir, "model.safetensors.index.json")
if os.path.exists(index_path):
index = json.load(open(index_path))
shards = sorted(set(index["weight_map"].values()))
else:
shards = ["model.safetensors"]
all_abs_samples = []
linear_names = []
global_max = 0.0
for shard in shards:
path = os.path.join(model_dir, shard)
print(f" Scanning {shard}...")
with safe_open(path, framework="pt") as f:
for name in f.keys():
t = f.get_tensor(name).float()
if t.dim() == 2 and "norm" not in name and "embed" not in name:
linear_names.append(name)
am = t.abs().max().item()
if am > global_max:
global_max = am
# Sample 2000 values for distribution
idx = torch.randint(0, t.numel(), (2000,))
all_abs_samples.append(t.flatten()[idx].abs())
all_abs = torch.cat(all_abs_samples)
return global_max, all_abs, linear_names, shards
def encode_to_proper_unary_torch(weight_f32, quantum, K):
"""
Encode [rows, cols] float32 tensor to proper unary.
Returns sign_packed [rows, chunks] uint64, slots_packed [K, rows, chunks] uint64
"""
rows, cols = weight_f32.shape
chunks = (cols + 63) // 64
inv_q = 1.0 / quantum
magnitudes = (weight_f32.abs() * inv_q).round().long().clamp(0, K)
signs = weight_f32 < 0
clip_count = int((weight_f32.abs() * inv_q > K).sum().item())
# Pack to uint64 bitplanes using numpy (torch lacks bit manipulation)
sign_packed = np.zeros((rows, chunks), dtype=np.uint64)
slots_packed = np.zeros((K, rows, chunks), dtype=np.uint64)
mags_np = magnitudes.numpy()
signs_np = signs.numpy()
for j in range(cols):
c = j // 64
bit = np.uint64(1) << np.uint64(j % 64)
# Sign
mask = signs_np[:, j]
sign_packed[mask, c] |= bit
# Unary slots: for each element, set slots 0..mag-1
col_mags = mags_np[:, j]
for p in range(K):
active = col_mags > p
slots_packed[p, active, c] |= bit
if (j + 1) % 256 == 0:
print(f" col {j+1}/{cols}", end="\r", flush=True)
print(f" {cols}/{cols} done, {clip_count} clipped")
return sign_packed, slots_packed, clip_count
def convert(model_dir, output_dir, K=32, clip_pct=99.9):
os.makedirs(output_dir, exist_ok=True)
config = json.load(open(os.path.join(model_dir, "config.json")))
print(f"Model: {config.get('_name_or_path', config.get('model_type', '?'))}")
print(f" Layers={config['num_hidden_layers']} Hidden={config['hidden_size']} Inter={config['intermediate_size']}")
# Scan
print("\nScanning weights...")
global_max, all_abs, linear_names, shards = scan_all_linears(model_dir)
# Pick quantum from clip percentile
clip_val = torch.quantile(all_abs, clip_pct / 100.0).item()
quantum = clip_val / K
print(f"\n Global absmax: {global_max:.6f}")
print(f" P{clip_pct} clip: {clip_val:.6f}")
print(f" K = {K}")
print(f" Quantum = {quantum:.8f}")
print(f" Values > clip ({clip_pct}%): saturate at K={K}")
# Distribution with chosen quantum
mags = (all_abs / quantum).round().clamp(0, K)
print(f"\n Mean magnitude: {mags.mean():.1f} slots")
print(f" Median: {mags.median():.1f} slots")
print(f" Zero fraction: {100*(mags==0).float().mean():.1f}%")
print(f" At K (clipped): {100*(mags==K).float().mean():.1f}%")
print(f" Unique levels: {len(mags.unique())} / {K+1}")
# Memory estimate
# Per linear: sign=rows*chunks*8 bytes, slots=K*rows*chunks*8 bytes
# Approx: (K+1) bits per element vs 16 bits BF16
bits_per_elem = K + 1 # K slot bits + 1 sign bit (stored in uint64 chunks)
ratio = bits_per_elem / 16.0
print(f"\n Bits per weight: {bits_per_elem}")
print(f" vs BF16 (16 bit): {ratio:.1f}x")
print(f" Original: ~7.6 GB → Estimated: ~{7.6 * ratio:.1f} GB")
# Build weight map
index_path = os.path.join(model_dir, "model.safetensors.index.json")
if os.path.exists(index_path):
weight_map = json.load(open(index_path))["weight_map"]
else:
weight_map = None
manifest = {
"format": "proper_unary",
"quantum": float(quantum),
"K": K,
"clip_pct": clip_pct,
"clip_val": float(clip_val),
"global_absmax": float(global_max),
"unary": {},
"fp16": [],
}
# Group linears by shard
shard_linears = {}
for name in linear_names:
shard = weight_map[name] if weight_map else "model.safetensors"
shard_linears.setdefault(shard, []).append(name)
total_unary_bytes = 0
total_fp16_bytes = 0
total_clipped = 0
done = 0
for shard in shards:
path = os.path.join(model_dir, shard)
shard_lins = shard_linears.get(shard, [])
print(f"\nProcessing {shard} ({len(shard_lins)} linear layers)...")
with safe_open(path, framework="pt") as f:
all_keys = list(f.keys())
# Non-linear weights → FP16
for name in all_keys:
if name in linear_names:
continue
fname = name.replace(".", "_") + ".fp16"
out_path = os.path.join(output_dir, fname)
if not os.path.exists(out_path):
t = f.get_tensor(name).half()
t.numpy().view(np.uint16).tofile(out_path)
sz = os.path.getsize(out_path)
total_fp16_bytes += sz
manifest["fp16"].append(name)
print(f" FP16: {name} {list(t.shape)} ({sz//1024}KB)")
# Linear weights → proper unary
for name in shard_lins:
fname = name.replace(".", "_")
sign_path = os.path.join(output_dir, f"{fname}.usign")
slots_path = os.path.join(output_dir, f"{fname}.uslots")
if os.path.exists(sign_path) and os.path.exists(slots_path):
t = f.get_tensor(name)
manifest["unary"][name] = list(t.shape)
total_unary_bytes += os.path.getsize(sign_path) + os.path.getsize(slots_path)
done += 1
print(f" Skip: {name}")
continue
t = f.get_tensor(name).float()
rows, cols = t.shape
print(f" Converting: {name} [{rows}x{cols}]...", flush=True)
sign_p, slots_p, clip_c = encode_to_proper_unary_torch(t, quantum, K)
total_clipped += clip_c
sign_p.tofile(sign_path)
slots_p.tofile(slots_path)
s_sz = os.path.getsize(sign_path)
sl_sz = os.path.getsize(slots_path)
total_unary_bytes += s_sz + sl_sz
manifest["unary"][name] = [rows, cols]
done += 1
print(f" sign={s_sz//1024}KB slots={sl_sz//1024}KB total={( s_sz+sl_sz)//1024//1024}MB")
del t, sign_p, slots_p
gc.collect()
# Copy config and tokenizer
for fname in os.listdir(model_dir):
if fname.endswith(('.json', '.txt', '.model')) and not fname.startswith('model.safetensors'):
src = os.path.join(model_dir, fname)
dst = os.path.join(output_dir, fname)
if not os.path.exists(dst):
shutil.copy2(src, dst)
manifest_path = os.path.join(output_dir, "manifest.json")
json.dump(manifest, open(manifest_path, "w"), indent=2)
total = total_unary_bytes + total_fp16_bytes
print(f"\n{'='*60}")
print(f"PROPER UNARY CONVERSION COMPLETE")
print(f"{'='*60}")
print(f" Quantum: {quantum:.8f}")
print(f" K: {K}")
print(f" Clip at P{clip_pct}: {clip_val:.6f}")
print(f" Linear layers: {done}")
print(f" Clipped vals: {total_clipped}")
print(f" Unary: {total_unary_bytes/1e9:.2f} GB")
print(f" FP16 (norms): {total_fp16_bytes/1e6:.1f} MB")
print(f" Total: {total/1e9:.2f} GB")
print(f" Original BF16: ~7.6 GB")
print(f" Ratio: {total/7.6e9:.1f}x")
print(f" Output dir: {output_dir}")
if __name__ == "__main__":
model_dir = sys.argv[1] if len(sys.argv) > 1 else "qwen3-4b-thinking-hf"
output_dir = sys.argv[2] if len(sys.argv) > 2 else "qwen3-4b-proper-unary"
K = int(sys.argv[3]) if len(sys.argv) > 3 else 32
clip = float(sys.argv[4]) if len(sys.argv) > 4 else 99.9
convert(model_dir, output_dir, K=K, clip_pct=clip)
|