Spaces:
Sleeping
Sleeping
| """ | |
| quant_units.py -- load the int8-quantized datapath units shipped in | |
| `models_int8/`, and (re)verify any unit variant bit-exact over its full domain. | |
| The int8 files are produced by riscv/export_int8.py: the dominant 256x256 hidden | |
| weight is per-channel symmetric int8, everything else is fp16. Dequantizing | |
| reconstructs a normal BitMLP state_dict, so the rest of the stack (build_all / | |
| NeuralUnits / the rv32 core) is untouched -- this is purely an alternate weight | |
| source. All 13 units stay bit-exact, at ~28% of the fp32 size. | |
| """ | |
| import os | |
| import torch | |
| from x86_units import (BitMLP, u_ADC8, u_SBB8, u_logic, u_SHL1, u_SHR1, | |
| u_MASK8, u_NOT8, u_FLAGS8, u_prefix, u_modrm, u_sib) | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| INT8_DIR = os.path.join(HERE, "models_int8") | |
| # name -> truth-table generator, for exhaustive re-verification of either variant | |
| DOMAINS = { | |
| "ADC8": u_ADC8, "SBB8": u_SBB8, | |
| "AND8": lambda: u_logic(lambda a, b: a & b), | |
| "OR8": lambda: u_logic(lambda a, b: a | b), | |
| "XOR8": lambda: u_logic(lambda a, b: a ^ b), | |
| "SHL1": u_SHL1, "SHR1": u_SHR1, "MASK8": u_MASK8, "NOT8": u_NOT8, | |
| "FLAGS8": u_FLAGS8, "PREFIX": u_prefix, "MODRM": u_modrm, "SIB": u_sib, | |
| } | |
| def _dequant(blob): | |
| """int8 blob -> fp32 BitMLP state_dict.""" | |
| sd = {} | |
| for k, v in blob.items(): | |
| if k.endswith(".q"): | |
| sd[k[:-2]] = v.float() * blob[k[:-2] + ".s"].float() | |
| elif k.endswith(".s"): | |
| continue | |
| else: | |
| sd[k] = v.float() | |
| return sd | |
| def build_int8(int8_dir=INT8_DIR): | |
| """Like build_all(), but loads the int8 units (dequantized) -- no training.""" | |
| nets = {} | |
| for name in DOMAINS: | |
| blob = torch.load(os.path.join(int8_dir, f"{name}.pt"), weights_only=True) | |
| sd = _dequant(blob) | |
| nin = sd["net.0.weight"].shape[1] | |
| nout = sd["net.4.weight"].shape[0] | |
| net = BitMLP(nin, nout) | |
| net.load_state_dict(sd) | |
| net.eval() | |
| nets[name] = net | |
| return nets | |
| def verify(nets): | |
| """Exhaustively re-check every unit over its FULL domain. Returns | |
| (n_exact, n_total, list_of_failures).""" | |
| ok = 0; fails = [] | |
| for name, gen in DOMAINS.items(): | |
| X, Y = gen() | |
| Xt, Yt = torch.tensor(X), torch.tensor(Y) | |
| good = True | |
| with torch.no_grad(): | |
| for i in range(0, Xt.shape[0], 65536): | |
| if not bool(((nets[name](Xt[i:i+65536]) > 0).float() | |
| == Yt[i:i+65536]).all()): | |
| good = False; break | |
| ok += int(good) | |
| if not good: | |
| fails.append(name) | |
| return ok, len(DOMAINS), fails | |
| def dir_size(model_dir): | |
| """Total bytes of the 13 unit files in a model directory.""" | |
| tot = 0 | |
| for name in DOMAINS: | |
| p = os.path.join(model_dir, f"{name}.pt") | |
| if os.path.exists(p): | |
| tot += os.path.getsize(p) | |
| return tot | |