50 lines
2.1 KiB
Python
50 lines
2.1 KiB
Python
import unittest
|
|
|
|
try:
|
|
import torch
|
|
except ImportError: # the planning host has no torch
|
|
torch = None
|
|
|
|
if torch is not None:
|
|
from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4,
|
|
to_gpu_int4)
|
|
|
|
|
|
@unittest.skipIf(torch is None, "needs torch")
|
|
class Quant(unittest.TestCase):
|
|
def setUp(self):
|
|
torch.manual_seed(0)
|
|
self.w = torch.randn(64, 256) * 0.02
|
|
|
|
def test_layouts_match_what_quantize_returns(self):
|
|
for quant in ("int8", "int4"):
|
|
got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)]
|
|
self.assertEqual(got, layout(quant, 64, 256))
|
|
|
|
def test_round_trip_error(self):
|
|
for quant, bound in (("int8", 0.01), ("int4", 0.11)):
|
|
err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm()
|
|
self.assertLess(err / self.w.norm(), bound)
|
|
|
|
def test_int4_clipping_beats_plain_round_to_nearest(self):
|
|
g = self.w.unflatten(-1, (-1, GROUP))
|
|
scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float()
|
|
plain = ((g / scale).round() * scale - g).norm()
|
|
clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm()
|
|
self.assertLess(clipped, plain)
|
|
|
|
def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self):
|
|
q, _ = quantize(self.w, "int4")
|
|
nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32)
|
|
self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1)))
|
|
self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q))
|
|
|
|
def test_cpu_matmul_matches_dequantized_weights(self):
|
|
x = torch.randn(3, 256, dtype=torch.bfloat16)
|
|
for quant in ("int8", "int4"):
|
|
q, scale = quantize(self.w, quant)
|
|
ref = x.float() @ dequantize(q, scale, quant, torch.float32).t()
|
|
stored = to_cpu_int4(q) if quant == "int4" else q
|
|
got = cpu_matmul(quant, stored, scale)(x).float()
|
|
self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant)
|