File size: 4,310 Bytes
3fd1a35 | 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 | """Calibration preserves the runtime INT4 contract and invalidates stale assets."""
import json
from pathlib import Path
import sys
import tempfile
import unittest
import numpy as np
import torch
from safetensors.torch import save_file
sys.path.insert(0,str(Path(__file__).resolve().parents[1]/'tools'))
from calibrate_w4 import choose_scales,safe_rows
from quantize_model import AssetWriter, source_linear_names
class CalibrationTests(unittest.TestCase):
def test_weighted_error_and_zero_rows(self):
torch.manual_seed(42)
w = torch.randn(65,128)
w[:,0] *= 7
w[-1].zero_()
h = torch.linspace(.1,2,128)
s, old, new = choose_scales(w,h,rows=32)
baseline_q, baseline_s = AssetWriter._quantize_rows(w)
ref = ((w-baseline_q*baseline_s[:,None]).square()*h).sum(1)
got = ((w-torch.round(w/s[:,None]).clamp(-8,7)*s[:,None]).square()*h).sum(1)
self.assertTrue(torch.all(got <= ref+1e-5))
self.assertLess(new,old)
self.assertTrue(torch.isfinite(s).all() and (s>0).all())
self.assertEqual(float(got[-1]),0.)
def test_mapping(self):
self.assertEqual(source_linear_names('model.layers.3.attention.o_proj'),
['model.layers.3.attention.dense.weight'])
self.assertEqual(source_linear_names('model.layers.2.attention.o_proj'),
['model.layers.2.attention.o_proj.weight'])
self.assertEqual(source_linear_names('model.layers.1.mlp.experts.7.gate_up'),
['model.layers.1.mlp.experts.7.gate_proj.weight',
'model.layers.1.mlp.experts.7.up_proj.weight'])
def test_accumulator_constraint(self):
q=torch.full((1,2048),7.)
self.assertFalse(bool(safe_rows(q,1)[0]))
self.assertTrue(bool(safe_rows(q,4)[0]))
with self.assertRaises(RuntimeError):
AssetWriter._safe_splits(q.to(torch.int8),required_parts=1)
self.assertEqual(AssetWriter._safe_splits(q.to(torch.int8),required_parts=4),4)
torch.manual_seed(14)
w=torch.randn(8,2048)
s,_,_=choose_scales(w,torch.ones(2048),k_splits=1)
self.assertTrue(safe_rows(torch.round(w/s[:,None]).clamp(-8,7),1).all())
def test_packing_padding_correction_and_cache(self):
with tempfile.TemporaryDirectory(prefix='ling3-calibration-') as directory:
root=Path(directory); path=root/'scales.safetensors'
name='model.layers.0.attention.qkvfgb'
names=source_linear_names(name)
torch.manual_seed(5)
tensors=[torch.randn(16,32) for _ in names]
scales={n:torch.full((16,),.2+i*.01) for i,n in enumerate(names)}
save_file(scales,str(path))
writer=AssetWriter(root/'assets-output',path)
writer.linear(name,tensors,pad_n=True)
folder=writer.assets/'linear'/name
data=(folder/'weight-int4-kn.bin').read_bytes()
packed=np.frombuffer(data,dtype=np.uint8)
raw=np.stack((packed&15,packed>>4),axis=1).reshape(32,128).T.astype(np.int8)
raw[raw>=8]-=16
expected=torch.round(torch.cat(tensors)/torch.cat(list(scales.values()))[:,None]).clamp(-8,7)
np.testing.assert_array_equal(raw[:96],expected.numpy())
self.assertFalse(raw[96:].any())
corr=np.fromfile(folder/'correction-i32.bin',dtype=np.int32)
np.testing.assert_array_equal(corr,8*raw.astype(np.int32).sum(1))
old_sha=json.loads((folder/'spec.json').read_text())['calibration_sha256']
save_file({n:s*1.3 for n,s in scales.items()},str(path))
writer=AssetWriter(root/'assets-output',path)
writer.linear(name,tensors,pad_n=True)
self.assertNotEqual(old_sha,json.loads((folder/'spec.json').read_text())['calibration_sha256'])
self.assertNotEqual(data,(folder/'weight-int4-kn.bin').read_bytes())
# A non-calibrated export must not silently reuse calibrated bytes either.
writer=AssetWriter(root/'assets-output')
writer.linear(name,tensors,pad_n=True)
self.assertNotIn('calibration_sha256',json.loads((folder/'spec.json').read_text()))
if __name__=='__main__':
unittest.main()
|