File size: 3,482 Bytes
948a05a | 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 90 91 92 93 94 95 | """Test SplitBit quantization: ternary, 4-bit, 8-bit, round-trip accuracy."""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import numpy as np
from splitbit_llm.model.quantization import SplitBitQuantizer, bits_per_weight, compression_ratio
def test_ternary_quantization():
"""Test ternary quantization round-trip."""
quantizer = SplitBitQuantizer(format="ternary")
weights = np.random.randn(64, 32).astype(np.float32) * 0.1
packed = quantizer.quantize(weights)
dequant = quantizer.dequantize(packed)
# Ternary loses precision but preserves sign and approximate magnitude
correlation = np.corrcoef(weights.flatten(), dequant.flatten())[0, 1]
assert correlation > 0.5, f"Ternary correlation too low: {correlation}"
print(f" Ternary correlation: {correlation:.3f}")
print(f" BPW: {quantizer.bpw:.3f}, Compression: {compression_ratio('ternary'):.1f}x")
def test_4bit_quantization():
"""Test 4-bit quantization round-trip."""
quantizer = SplitBitQuantizer(format="q4_k_m")
weights = np.random.randn(128, 64).astype(np.float32) * 0.1
packed = quantizer.quantize(weights)
dequant = quantizer.dequantize(packed)
# 4-bit should be closer to original
max_error = np.max(np.abs(weights - dequant))
rel_error = max_error / np.max(np.abs(weights))
assert rel_error < 0.2, f"4-bit relative error too high: {rel_error}"
print(f" 4-bit max relative error: {rel_error:.4f}")
print(f" BPW: {quantizer.bpw:.1f}, Compression: {compression_ratio('q4_k_m'):.1f}x")
def test_8bit_quantization():
"""Test 8-bit quantization round-trip."""
quantizer = SplitBitQuantizer(format="q8_0")
weights = np.random.randn(256, 128).astype(np.float32) * 0.1
packed = quantizer.quantize(weights)
dequant = quantizer.dequantize(packed)
# 8-bit should be very close
max_error = np.max(np.abs(weights - dequant))
rel_error = max_error / np.max(np.abs(weights))
assert rel_error < 0.02, f"8-bit relative error too high: {rel_error}"
print(f" 8-bit max relative error: {rel_error:.5f}")
def test_fp16_passthrough():
"""Test fp16 passthrough (no quantization)."""
quantizer = SplitBitQuantizer(format="fp16")
weights = np.random.randn(64, 32).astype(np.float32)
packed = quantizer.quantize(weights)
dequant = quantizer.dequantize(packed)
# fp16 should be nearly identical
max_error = np.max(np.abs(weights - dequant))
assert max_error < 0.01, f"fp16 error too high: {max_error}"
print(f" fp16 max error: {max_error:.6f}")
def test_bpw_table():
"""Test bits per weight table."""
assert bits_per_weight("ternary") > 1.5
assert bits_per_weight("q4_k_m") == 4.0
assert bits_per_weight("q8_0") == 8.0
assert bits_per_weight("fp16") == 16.0
assert compression_ratio("ternary") > 10.0
print(f" Ternary BPW: {bits_per_weight('ternary'):.3f}")
print(f" Ternary compression: {compression_ratio('ternary'):.1f}x")
if __name__ == "__main__":
print("Running quantization tests...")
test_ternary_quantization()
print(" ✓ test_ternary_quantization")
test_4bit_quantization()
print(" ✓ test_4bit_quantization")
test_8bit_quantization()
print(" ✓ test_8bit_quantization")
test_fp16_passthrough()
print(" ✓ test_fp16_passthrough")
test_bpw_table()
print(" ✓ test_bpw_table")
print("\nAll quantization tests passed!")
|