mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-07 04:53:47 +08:00
Daniel Koepping
118 lines
7.0 KiB
Python
118 lines
7.0 KiB
Python
import gc, itertools, unittest, weakref
|
|
import numpy as np
|
|
from tinygrad import Tensor, UOp, dtypes, function, Device
|
|
from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, QUANT_SIZES, HALFWORD_QUANTS, iq4_half_lut, _iq_grid
|
|
from tinygrad.llm.gguf import ggml_data_to_tensor
|
|
from test.helpers import not_support_multi_device
|
|
|
|
class QuantLinearMixin:
|
|
def _test_quant_linear(self, ggml_type, block_bytes, in_features=2048, out_features=64, token_counts=(1, 3, 32, 64, 128),
|
|
bias=False, custom=True, symbolic=False):
|
|
if custom and not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
|
|
rng = np.random.default_rng(42)
|
|
packed = rng.integers(0, 256, (out_features*in_features//256, block_bytes), dtype=np.uint8)
|
|
if ggml_type in (11, 14): packed[:, -2:] = np.array([0.001], dtype=np.float16).view(np.uint8)
|
|
elif ggml_type == 10:
|
|
packed[:, 80:82] = np.array([0.001], dtype=np.float16).view(np.uint8)
|
|
packed[:, 82:84] = np.array([0.0002], dtype=np.float16).view(np.uint8)
|
|
elif ggml_type == 20: packed.reshape(-1, 18)[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
|
|
else: packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
|
|
if ggml_type in (12, 13): packed[:, 2:4] = np.array([0.0002], dtype=np.float16).view(np.uint8)
|
|
raw = Tensor(np.pad(packed.flatten(), (4, 0))).contiguous().realize()[4:]
|
|
decoded = ggml_data_to_tensor(raw, out_features*in_features, ggml_type).reshape(out_features, in_features)
|
|
weight = decoded.numpy()
|
|
linear = Linear(in_features, out_features, bias=False)
|
|
linear.weight = decoded
|
|
bias_value = rng.normal(size=out_features).astype(np.float32) if bias else 0
|
|
if bias: linear.bias = Tensor(bias_value)
|
|
@function(allow_implicit=True)
|
|
def run(x:Tensor): return linear(x)
|
|
for tokens in token_counts:
|
|
with self.subTest(tokens=tokens):
|
|
x = rng.normal(size=(tokens, in_features)).astype(np.float32 if tokens == 3 else np.float16)
|
|
reference_x = x.astype(np.float32)
|
|
wmma = custom and (32 if symbolic else tokens) % 16 == 0 and out_features % 16 == 0
|
|
if wmma: reference_x = x.astype(np.float16).astype(np.float32)
|
|
if custom and not wmma:
|
|
grouped = reference_x.reshape(tokens, -1, 32)
|
|
scale = np.maximum(np.abs(grouped).max(-1, keepdims=True) / 127, 1e-8)
|
|
reference_x = (np.clip(np.rint(grouped/scale), -127, 127)*scale).reshape(tokens, in_features)
|
|
reference_w = weight.astype(np.float16).astype(np.float32) if wmma else weight
|
|
inp = Tensor(x) if not symbolic else Tensor(np.pad(x, ((0, 32-tokens), (0, 0)))).contiguous()[:
|
|
UOp.variable("wmma_tokens", 1, 32).bind(tokens)]
|
|
actual = (run if tokens == 1 or symbolic else linear)(inp)[:tokens].numpy()
|
|
self.assertEqual(linear.ggml_type, ggml_type if custom else None)
|
|
np.testing.assert_allclose(actual, reference_x @ reference_w.T + bias_value, rtol=3e-3, atol=2e-2)
|
|
if not symbolic and tokens == 3 and ggml_type not in (12, 13, 14, 23):
|
|
sym = Tensor(np.pad(x, ((0, 1), (0, 0)))).contiguous()[:UOp.variable("tokens", 1, 4).bind(3)]
|
|
np.testing.assert_allclose(linear(sym)[:3].numpy(), reference_x @ reference_w.T + bias_value, rtol=3e-3, atol=2e-2)
|
|
self.assertEqual(linear.ggml_type, ggml_type if custom else None)
|
|
|
|
class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
|
|
def test_quant_tables_not_retained(self):
|
|
# one _iq_grid table and the iq4 lut cover both table creation paths
|
|
for typ in (18, 23):
|
|
table = (iq4_half_lut(Device.DEFAULT) if typ == 23 else _iq_grid(Device.DEFAULT, typ)).realize()
|
|
ref = weakref.ref(table)
|
|
del table
|
|
gc.collect()
|
|
self.assertIsNone(ref())
|
|
|
|
@unittest.skipIf(not_support_multi_device(), "no multi")
|
|
def test_quant_tables_multidevice(self):
|
|
devices = (Device.DEFAULT, f"{Device.DEFAULT}:1")
|
|
for typ in (18, 23):
|
|
with self.subTest(ggml_type=typ):
|
|
single = iq4_half_lut(Device.DEFAULT) if typ == 23 else _iq_grid(Device.DEFAULT, typ)
|
|
table = (iq4_half_lut(devices) if typ == 23 else _iq_grid(devices, typ)).realize()
|
|
self.assertEqual(table.device, devices)
|
|
for i in range(len(devices)):
|
|
np.testing.assert_array_equal(Tensor(table.uop.mselect(i)).numpy(), single.numpy())
|
|
|
|
@unittest.skipIf(not_support_multi_device(), "no multi")
|
|
def test_quant_linear_multidevice(self):
|
|
if not amd_custom_kernels_supported(Device.DEFAULT): self.skipTest("RDNA3 required")
|
|
devices = (Device.DEFAULT, f"{Device.DEFAULT}:1")
|
|
rng = np.random.default_rng(42)
|
|
# Exercise both the grid-backed decode path and the IQ4 WMMA lookup table without a TP model.
|
|
for typ, tokens in ((18, 1), (23, 16)):
|
|
with self.subTest(ggml_type=typ):
|
|
raw = Tensor(rng.integers(0, 0x3c, 16*QUANT_SIZES[typ]+4, dtype=np.uint8))
|
|
x = Tensor(rng.normal(size=(tokens, 256)).astype(np.float32))
|
|
def run(device):
|
|
linear = Linear(256, 16, bias=False)
|
|
linear.weight = ggml_data_to_tensor(raw.to(device).realize()[4:], 16*256, typ).reshape(16, 256)
|
|
out = linear(x.to(device)).realize()
|
|
self.assertEqual(linear.ggml_type, typ)
|
|
return out
|
|
single, parallel = run(Device.DEFAULT).numpy(), run(devices)
|
|
for i in range(len(devices)):
|
|
np.testing.assert_allclose(Tensor(parallel.uop.mselect(i)).numpy(), single, atol=1e-5, rtol=1e-5)
|
|
|
|
def test_quant_weights_share_storage(self):
|
|
for (ggml_type, type_size), offset in itertools.product(QUANT_SIZES.items(), (0, 4)):
|
|
with self.subTest(ggml_type=ggml_type, offset=offset):
|
|
packed = np.arange(type_size + offset, dtype=np.uint8)
|
|
raw = Tensor(packed).realize()[offset:]
|
|
if raw.uop.contiguous_view() is None: self.skipTest("requires buffer views")
|
|
decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
|
|
linear = Linear(256, 1, bias=False)
|
|
linear.set_quantized(decoded)
|
|
linear.weight.realize()
|
|
self.assertEqual(linear.ggml_type, ggml_type)
|
|
self.assertEqual(linear.weight.dtype, dtypes.uint16 if ggml_type in HALFWORD_QUANTS else dtypes.uint32)
|
|
self.assertEqual(linear.weight.nbytes(), type_size)
|
|
np.testing.assert_array_equal(linear.weight.bitcast(dtypes.uint8).numpy(), packed[offset:])
|
|
raw.assign(raw.full_like(1)).realize()
|
|
np.testing.assert_array_equal(linear.weight.bitcast(dtypes.uint8).numpy(), np.ones(type_size, dtype=np.uint8))
|
|
|
|
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "slow on WEBGPU")
|
|
def test_quant_linear_fallback(self):
|
|
if amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("run with DISABLE_AMD_KERNELS=1")
|
|
# per-type dequant math on the generic path is covered by test_gguf, spot check a representative set here
|
|
for typ in (12, 14, 17, 23):
|
|
with self.subTest(ggml_type=typ):
|
|
self._test_quant_linear(typ, QUANT_SIZES[typ], in_features=256, out_features=16, token_counts=(1, 3), bias=True, custom=False)
|
|
|
|
if __name__ == "__main__": unittest.main()
|