| """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()) |
| |
| 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() |
|
|
