842 lines
25 KiB
Rust
842 lines
25 KiB
Rust
//! Comprehensive tests for PT-BitNet Phase 0 ternary quantization
|
|
//!
|
|
//! Test coverage based on ADR-017 (AD-1, AD-18):
|
|
//! - Ternary packing/unpacking roundtrips
|
|
//! - Absmean quantization correctness
|
|
//! - Dequantization accuracy
|
|
//! - Full tensor quantization
|
|
//! - Edge cases and error conditions
|
|
|
|
use super::{
|
|
dequantize_bitnet_t158, pack_ternary, quantize_tensor, unpack_ternary, PtBitnetConfig,
|
|
TernaryTensor,
|
|
};
|
|
|
|
// ============================================================================
|
|
// Test Constants
|
|
// ============================================================================
|
|
|
|
const EPSILON: f32 = 1e-6;
|
|
const BLOCK_SIZE: usize = 256;
|
|
|
|
// ============================================================================
|
|
// 1. Ternary Packing Roundtrip Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_pack_unpack_simple_roundtrip() {
|
|
// Simple 4-element ternary array
|
|
let ternary = vec![1i8, 0, -1, 1];
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 4);
|
|
|
|
assert_eq!(ternary, unpacked, "Packing roundtrip failed for [1, 0, -1, 1]");
|
|
}
|
|
|
|
#[test]
|
|
fn test_pack_all_zeros() {
|
|
let ternary = vec![0i8; 256];
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 256);
|
|
|
|
assert_eq!(ternary, unpacked);
|
|
assert!(unpacked.iter().all(|&x| x == 0), "All zeros should remain all zeros");
|
|
}
|
|
|
|
#[test]
|
|
fn test_pack_all_ones() {
|
|
let ternary = vec![1i8; 256];
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 256);
|
|
|
|
assert_eq!(ternary, unpacked);
|
|
assert!(unpacked.iter().all(|&x| x == 1), "All +1 should remain all +1");
|
|
}
|
|
|
|
#[test]
|
|
fn test_pack_all_neg_ones() {
|
|
let ternary = vec![-1i8; 256];
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 256);
|
|
|
|
assert_eq!(ternary, unpacked);
|
|
assert!(unpacked.iter().all(|&x| x == -1), "All -1 should remain all -1");
|
|
}
|
|
|
|
#[test]
|
|
fn test_pack_one_block_256_elements() {
|
|
// One full block (256 elements) with alternating pattern
|
|
let mut ternary = Vec::with_capacity(256);
|
|
for i in 0..256 {
|
|
ternary.push(match i % 3 {
|
|
0 => 1,
|
|
1 => 0,
|
|
2 => -1,
|
|
_ => unreachable!(),
|
|
});
|
|
}
|
|
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 256);
|
|
|
|
assert_eq!(ternary, unpacked, "256-element block roundtrip failed");
|
|
|
|
// Verify storage size: 256 elements * 2 bits = 64 bytes
|
|
assert_eq!(packed.len(), 64, "Packed size should be 64 bytes for 256 elements");
|
|
}
|
|
|
|
#[test]
|
|
fn test_pack_non_aligned_size() {
|
|
// 100 elements (not divisible by 128, the typical packing boundary)
|
|
let mut ternary = Vec::with_capacity(100);
|
|
for i in 0..100 {
|
|
ternary.push(if i % 2 == 0 { 1 } else { -1 });
|
|
}
|
|
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 100);
|
|
|
|
assert_eq!(
|
|
ternary.len(),
|
|
unpacked.len(),
|
|
"Unpacked length should match original"
|
|
);
|
|
assert_eq!(ternary, unpacked, "Non-aligned size roundtrip failed");
|
|
}
|
|
|
|
#[test]
|
|
fn test_pack_large_tensor() {
|
|
// Multiple blocks (1024 elements = 4 blocks)
|
|
let ternary: Vec<i8> = (0..1024)
|
|
.map(|i| match i % 5 {
|
|
0 | 1 => 1,
|
|
2 | 3 => -1,
|
|
4 => 0,
|
|
_ => unreachable!(),
|
|
})
|
|
.collect();
|
|
|
|
let packed = pack_ternary(&ternary);
|
|
let unpacked = unpack_ternary(&packed, 1024);
|
|
|
|
assert_eq!(ternary, unpacked, "Large tensor roundtrip failed");
|
|
}
|
|
|
|
// ============================================================================
|
|
// 2. Absmean Quantization Correctness Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_quantize_uniform_random() {
|
|
// Uniform random weights in [-1, 1] should produce all ternary values
|
|
let weights = vec![0.5, -0.3, 0.1, -0.7, 0.9, -0.1, 0.0, 0.4];
|
|
let ternary = quantize_absmean(&weights);
|
|
|
|
// All outputs must be in {-1, 0, +1}
|
|
for &t in &ternary {
|
|
assert!(
|
|
t == -1 || t == 0 || t == 1,
|
|
"Quantized value {} not in ternary set",
|
|
t
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_all_zeros() {
|
|
let weights = vec![0.0; 256];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// All ternary values should be zero
|
|
assert!(
|
|
ternary.iter().all(|&x| x == 0),
|
|
"All-zero input should produce all-zero ternary"
|
|
);
|
|
|
|
// Scale should be near epsilon (avoiding division by zero)
|
|
assert!(
|
|
scale < 1e-5,
|
|
"Scale for all-zero weights should be near epsilon, got {}",
|
|
scale
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_large_positive() {
|
|
// Large positive weights should quantize to all +1
|
|
let weights = vec![10.0; 256];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// All should be +1
|
|
assert!(
|
|
ternary.iter().all(|&x| x == 1),
|
|
"Large positive weights should quantize to +1"
|
|
);
|
|
|
|
// Scale should be approximately 10.0 (mean absolute value)
|
|
assert!(
|
|
(scale - 10.0).abs() < 0.1,
|
|
"Scale should be ~10.0, got {}",
|
|
scale
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_large_negative() {
|
|
// Large negative weights should quantize to all -1
|
|
let weights = vec![-10.0; 256];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// All should be -1
|
|
assert!(
|
|
ternary.iter().all(|&x| x == -1),
|
|
"Large negative weights should quantize to -1"
|
|
);
|
|
|
|
// Scale should be approximately 10.0 (mean absolute value)
|
|
assert!(
|
|
(scale - 10.0).abs() < 0.1,
|
|
"Scale should be ~10.0, got {}",
|
|
scale
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_known_example() {
|
|
// From ADR: W_ternary = RoundClip(W / (mean(|W|) + epsilon), -1, 1)
|
|
// Example: weights = [0.5, -0.3, 0.1, -0.7]
|
|
// gamma = mean(|W|) = (0.5 + 0.3 + 0.1 + 0.7) / 4 = 0.4
|
|
// normalized = [1.25, -0.75, 0.25, -1.75]
|
|
// ternary = [1, -1, 0, -1] (after clamp and round)
|
|
|
|
let weights = vec![0.5, -0.3, 0.1, -0.7];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// Verify scale is approximately 0.4
|
|
assert!(
|
|
(scale - 0.4).abs() < 0.01,
|
|
"Expected scale ~0.4, got {}",
|
|
scale
|
|
);
|
|
|
|
// Verify ternary values
|
|
// 1.25 -> 1, -0.75 -> -1, 0.25 -> 0, -1.75 -> -1
|
|
assert_eq!(ternary[0], 1, "0.5/0.4 = 1.25 should round to 1");
|
|
assert_eq!(ternary[1], -1, "-0.3/0.4 = -0.75 should round to -1");
|
|
assert_eq!(ternary[2], 0, "0.1/0.4 = 0.25 should round to 0");
|
|
assert_eq!(ternary[3], -1, "-0.7/0.4 = -1.75 should clamp to -1");
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_scale_calculation() {
|
|
// Verify scale = mean(|weights|)
|
|
let weights = vec![1.0, -2.0, 3.0, -4.0];
|
|
let (_, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
let expected_scale = (1.0 + 2.0 + 3.0 + 4.0) / 4.0; // = 2.5
|
|
assert!(
|
|
(scale - expected_scale).abs() < EPSILON,
|
|
"Scale should be mean of absolute values: expected {}, got {}",
|
|
expected_scale,
|
|
scale
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// 3. Dequantization Correctness Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_dequantize_simple() {
|
|
let ternary = vec![1i8, 0, -1];
|
|
let scale = 2.0;
|
|
|
|
let dequantized = dequantize_ternary(&ternary, scale);
|
|
|
|
assert_eq!(dequantized.len(), 3);
|
|
assert!((dequantized[0] - 2.0).abs() < EPSILON, "1 * 2.0 = 2.0");
|
|
assert!((dequantized[1] - 0.0).abs() < EPSILON, "0 * 2.0 = 0.0");
|
|
assert!((dequantized[2] - (-2.0)).abs() < EPSILON, "-1 * 2.0 = -2.0");
|
|
}
|
|
|
|
#[test]
|
|
fn test_dequantize_packed_data() {
|
|
// Pack known ternary data, then dequantize
|
|
let ternary = vec![1i8, 0, -1, 1];
|
|
let packed = pack_ternary(&ternary);
|
|
let scale = 3.5;
|
|
|
|
let unpacked = unpack_ternary(&packed, 4);
|
|
let dequantized = dequantize_ternary(&unpacked, scale);
|
|
|
|
assert_eq!(dequantized.len(), 4);
|
|
assert!((dequantized[0] - 3.5).abs() < EPSILON);
|
|
assert!((dequantized[1] - 0.0).abs() < EPSILON);
|
|
assert!((dequantized[2] - (-3.5)).abs() < EPSILON);
|
|
assert!((dequantized[3] - 3.5).abs() < EPSILON);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_dequantize_roundtrip_mse() {
|
|
// Quantize -> Dequantize should have bounded MSE
|
|
let weights = vec![0.5, -0.3, 0.1, -0.7, 0.9, -0.1, 0.4, -0.5];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
let dequantized = dequantize_ternary(&ternary, scale);
|
|
|
|
// Compute MSE
|
|
let mse: f32 = weights
|
|
.iter()
|
|
.zip(dequantized.iter())
|
|
.map(|(&w, &d)| (w - d).powi(2))
|
|
.sum::<f32>()
|
|
/ weights.len() as f32;
|
|
|
|
// MSE should be reasonable (ternary quantization is lossy)
|
|
// For absmean, expect MSE < 0.5 for normalized weights
|
|
assert!(
|
|
mse < 0.5,
|
|
"MSE too high: {} (weights may not reconstruct well)",
|
|
mse
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dequantize_full_block() {
|
|
// Dequantize a full 256-element block
|
|
let ternary: Vec<i8> = (0..256).map(|i| if i % 2 == 0 { 1 } else { -1 }).collect();
|
|
let scale = 1.5;
|
|
|
|
let dequantized = dequantize_ternary(&ternary, scale);
|
|
|
|
assert_eq!(dequantized.len(), 256);
|
|
for (i, &val) in dequantized.iter().enumerate() {
|
|
let expected = if i % 2 == 0 { 1.5 } else { -1.5 };
|
|
assert!(
|
|
(val - expected).abs() < EPSILON,
|
|
"Element {} incorrect: expected {}, got {}",
|
|
i,
|
|
expected,
|
|
val
|
|
);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// 4. Full Tensor Quantization Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_tensor_quantize_256x256() {
|
|
// 256x256 random tensor (65536 elements)
|
|
let mut weights = Vec::with_capacity(65536);
|
|
for i in 0..65536 {
|
|
let val = ((i as f32) * 0.001).sin(); // Pseudo-random in [-1, 1]
|
|
weights.push(val);
|
|
}
|
|
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
// Verify shape preserved
|
|
assert_eq!(
|
|
tensor.num_elements(),
|
|
65536,
|
|
"Tensor should preserve element count"
|
|
);
|
|
|
|
// Verify sparsity is in valid range
|
|
let sparsity = tensor.sparsity();
|
|
assert!(
|
|
sparsity >= 0.0 && sparsity <= 1.0,
|
|
"Sparsity {} out of range [0, 1]",
|
|
sparsity
|
|
);
|
|
|
|
// For uniform random, expect ~1/3 zeros (rough heuristic)
|
|
assert!(
|
|
sparsity > 0.15 && sparsity < 0.5,
|
|
"Sparsity {} seems unrealistic for uniform random input",
|
|
sparsity
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_memory_bytes() {
|
|
let weights = vec![0.5; 256];
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
// Expected memory:
|
|
// - Packed data: 256 elements * 2 bits / 8 = 64 bytes
|
|
// - Scales: 1 block * 4 bytes (f32) = 4 bytes
|
|
// Total: 68 bytes
|
|
let expected_bytes = 64 + 4;
|
|
|
|
assert_eq!(
|
|
tensor.memory_bytes(),
|
|
expected_bytes,
|
|
"Memory calculation incorrect"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_sparsity_calculation() {
|
|
// Known sparsity: 50% zeros
|
|
let weights: Vec<f32> = (0..256)
|
|
.map(|i| if i % 2 == 0 { 0.0 } else { 1.0 })
|
|
.collect();
|
|
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
let sparsity = tensor.sparsity();
|
|
|
|
// Should be close to 0.5 (half zeros)
|
|
assert!(
|
|
(sparsity - 0.5).abs() < 0.1,
|
|
"Expected sparsity ~0.5, got {}",
|
|
sparsity
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_block_alignment() {
|
|
// 512 elements = 2 blocks of 256
|
|
let weights = vec![1.0; 512];
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
// Should have 2 scale factors (one per block)
|
|
assert_eq!(
|
|
tensor.num_blocks(),
|
|
2,
|
|
"Expected 2 blocks for 512 elements"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_non_aligned_padding() {
|
|
// 300 elements (256 + 44) should create 2 blocks with padding
|
|
let weights = vec![0.5; 300];
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
// Should pad to 2 full blocks (512 elements)
|
|
let num_blocks = (300 + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
|
assert_eq!(
|
|
tensor.num_blocks(),
|
|
num_blocks,
|
|
"Non-aligned tensor should pad to full blocks"
|
|
);
|
|
|
|
// Original element count should be preserved
|
|
assert_eq!(tensor.num_elements(), 300);
|
|
}
|
|
|
|
// ============================================================================
|
|
// 5. TernaryTensor Properties Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_ternary_tensor_properties() {
|
|
let weights: Vec<f32> = (0..512).map(|i| (i as f32) * 0.01).collect();
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
// Memory bytes should match calculation
|
|
let num_blocks = (512 + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
|
let packed_bytes = num_blocks * BLOCK_SIZE * 2 / 8; // 2 bits per element
|
|
let scale_bytes = num_blocks * 4; // f32 scales
|
|
let expected = packed_bytes + scale_bytes;
|
|
|
|
assert_eq!(tensor.memory_bytes(), expected);
|
|
|
|
// Sparsity should be in valid range
|
|
assert!(tensor.sparsity() >= 0.0 && tensor.sparsity() <= 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ternary_tensor_uniform_random_sparsity() {
|
|
// Uniform random should have ~1/3 sparsity
|
|
let mut weights = Vec::with_capacity(2048);
|
|
for i in 0..2048 {
|
|
weights.push(((i as f32) * 1.234).sin());
|
|
}
|
|
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
let sparsity = tensor.sparsity();
|
|
|
|
// Rough heuristic: 20-45% zeros for uniform random
|
|
assert!(
|
|
sparsity > 0.2 && sparsity < 0.45,
|
|
"Uniform random sparsity {} outside expected range [0.2, 0.45]",
|
|
sparsity
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// 6. Config Validation Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_config_default_values() {
|
|
let config = PtBitnetConfig::default();
|
|
|
|
assert_eq!(config.block_size, 256, "Default block size should be 256");
|
|
assert!(
|
|
config.calibration_samples > 0,
|
|
"Calibration samples must be > 0"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "block_size must be > 0")]
|
|
fn test_config_invalid_block_size() {
|
|
let _config = PtBitnetConfig {
|
|
block_size: 0,
|
|
..Default::default()
|
|
};
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "calibration_samples must be > 0")]
|
|
fn test_config_invalid_calibration_samples() {
|
|
let _config = PtBitnetConfig {
|
|
calibration_samples: 0,
|
|
..Default::default()
|
|
};
|
|
}
|
|
|
|
// ============================================================================
|
|
// 7. Edge Case Tests
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_empty_input() {
|
|
let weights: Vec<f32> = vec![];
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
assert_eq!(tensor.num_elements(), 0);
|
|
assert_eq!(tensor.num_blocks(), 0);
|
|
assert_eq!(tensor.sparsity(), 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_single_element() {
|
|
let weights = vec![0.5];
|
|
let tensor = TernaryTensor::quantize(&weights, BLOCK_SIZE);
|
|
|
|
assert_eq!(tensor.num_elements(), 1);
|
|
// Should create 1 block (padded)
|
|
assert_eq!(tensor.num_blocks(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_very_large_values() {
|
|
let weights = vec![f32::MAX, f32::MAX, f32::MAX, f32::MAX];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// Should all quantize to +1
|
|
assert!(ternary.iter().all(|&x| x == 1), "f32::MAX should quantize to +1");
|
|
|
|
// Scale should be approximately f32::MAX
|
|
assert!(scale > 1e30, "Scale should be very large");
|
|
|
|
// Dequantization should not produce NaN
|
|
let dequantized = dequantize_ternary(&ternary, scale);
|
|
assert!(
|
|
dequantized.iter().all(|&x| !x.is_nan()),
|
|
"Dequantization should not produce NaN"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_subnormal_floats() {
|
|
// Very small positive values (subnormal range)
|
|
let weights = vec![1e-40, -1e-40, 1e-39, -1e-39];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// Should quantize reasonably (may be all zeros or small values)
|
|
assert!(ternary.iter().all(|&x| x >= -1 && x <= 1));
|
|
|
|
// Scale should be tiny but not zero
|
|
assert!(scale > 0.0, "Scale should be > 0 even for subnormal inputs");
|
|
}
|
|
|
|
#[test]
|
|
fn test_nan_handling() {
|
|
// NaN should not crash, but behavior is implementation-defined
|
|
let weights = vec![f32::NAN, 1.0, -1.0, 0.0];
|
|
let result = std::panic::catch_unwind(|| {
|
|
quantize_absmean_with_scale(&weights)
|
|
});
|
|
|
|
// Should either panic or handle gracefully
|
|
// At minimum, should not produce infinite loop or segfault
|
|
if let Ok((ternary, scale)) = result {
|
|
// If it succeeds, output should not contain NaN
|
|
assert!(
|
|
!scale.is_nan() || scale == 0.0,
|
|
"Scale should not be NaN unless handled explicitly"
|
|
);
|
|
assert!(
|
|
ternary.iter().all(|&x| x >= -1 && x <= 1),
|
|
"Ternary values must be in valid range"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_infinity_handling() {
|
|
let weights = vec![f32::INFINITY, f32::NEG_INFINITY, 1.0, -1.0];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// Infinities should quantize to ±1
|
|
assert_eq!(ternary[0], 1, "INFINITY should quantize to +1");
|
|
assert_eq!(ternary[1], -1, "NEG_INFINITY should quantize to -1");
|
|
|
|
// Scale should be finite (or handled gracefully)
|
|
// Implementation may cap scale to avoid overflow
|
|
assert!(
|
|
scale.is_finite() || scale > 1e30,
|
|
"Scale should be finite or very large"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_mixed_magnitudes() {
|
|
// Mix of very large and very small values
|
|
let weights = vec![1000.0, 0.001, -1000.0, -0.001, 0.0];
|
|
let (ternary, scale) = quantize_absmean_with_scale(&weights);
|
|
|
|
// Should produce valid ternary values
|
|
assert!(ternary.iter().all(|&x| x >= -1 && x <= 1));
|
|
|
|
// Scale should be dominated by large values
|
|
assert!(scale > 100.0, "Scale should reflect large values");
|
|
|
|
// Small values should quantize to 0
|
|
assert_eq!(
|
|
ternary[1], 0,
|
|
"0.001 compared to scale ~500 should be 0"
|
|
);
|
|
assert_eq!(ternary[3], 0, "-0.001 should be 0");
|
|
}
|
|
|
|
// ============================================================================
|
|
// 8. Layer Filter Tests (per ADR-017 AD-2)
|
|
// ============================================================================
|
|
|
|
#[test]
|
|
fn test_should_quantize_expert_layers() {
|
|
// MoE expert FFN layers (gate_proj, up_proj, down_proj) should be quantized
|
|
use super::LayerMask;
|
|
|
|
let layer_mask = LayerMask::ExpertsOnly;
|
|
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.mlp.gate_proj.weight", &layer_mask),
|
|
"gate_proj should be quantized"
|
|
);
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.mlp.up_proj.weight", &layer_mask),
|
|
"up_proj should be quantized"
|
|
);
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.mlp.down_proj.weight", &layer_mask),
|
|
"down_proj should be quantized"
|
|
);
|
|
assert!(
|
|
should_quantize_layer("model.layers.15.block_sparse_moe.experts.7.w3.weight", &layer_mask),
|
|
"Expert w3 (up_proj) should be quantized"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_not_quantize_router() {
|
|
// Router and gate layers must remain in FP16 per ADR-017 (AD-2)
|
|
use super::LayerMask;
|
|
|
|
let layer_mask = LayerMask::ExpertsOnly;
|
|
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.mlp.router.weight", &layer_mask),
|
|
"Router should NOT be quantized"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.block_sparse_moe.gate.weight", &layer_mask),
|
|
"MoE gate should NOT be quantized"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_not_quantize_embed() {
|
|
// Embeddings and LM head must remain in FP16 per ADR-017 (AD-2)
|
|
use super::LayerMask;
|
|
|
|
let layer_mask = LayerMask::ExpertsOnly;
|
|
|
|
assert!(
|
|
!should_quantize_layer("model.embed_tokens.weight", &layer_mask),
|
|
"Embed tokens should NOT be quantized"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("lm_head.weight", &layer_mask),
|
|
"LM head should NOT be quantized"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.embeddings.word_embeddings", &layer_mask),
|
|
"Word embeddings should NOT be quantized"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_not_quantize_norm() {
|
|
// Normalization layers must remain in FP16 per ADR-017 (AD-2)
|
|
use super::LayerMask;
|
|
|
|
let layer_mask = LayerMask::ExpertsOnly;
|
|
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.input_layernorm.weight", &layer_mask),
|
|
"Input layernorm should NOT be quantized"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.post_attention_layernorm.weight", &layer_mask),
|
|
"Post-attention layernorm should NOT be quantized"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.norm.weight", &layer_mask),
|
|
"Final norm should NOT be quantized"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.self_attn.layer_norm", &layer_mask),
|
|
"Self-attention layer_norm should NOT be quantized"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_mask_all() {
|
|
// LayerMask::All should quantize all linear layers except protected ones
|
|
use super::LayerMask;
|
|
|
|
let layer_mask = LayerMask::All;
|
|
|
|
// Should quantize attention projections
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.self_attn.q_proj.weight", &layer_mask),
|
|
"Query projection should be quantized with LayerMask::All"
|
|
);
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.self_attn.k_proj.weight", &layer_mask),
|
|
"Key projection should be quantized with LayerMask::All"
|
|
);
|
|
|
|
// Should still protect router/embed/norm
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.mlp.router.weight", &layer_mask),
|
|
"Router should be protected even with LayerMask::All"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.embed_tokens.weight", &layer_mask),
|
|
"Embeddings should be protected even with LayerMask::All"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_layer_mask_custom() {
|
|
// LayerMask::Custom should match specified patterns only
|
|
use super::LayerMask;
|
|
|
|
let layer_mask = LayerMask::Custom(vec!["w1".to_string(), "w3".to_string()]);
|
|
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.mlp.experts.0.w1.weight", &layer_mask),
|
|
"w1 should match custom pattern"
|
|
);
|
|
assert!(
|
|
should_quantize_layer("model.layers.0.mlp.experts.0.w3.weight", &layer_mask),
|
|
"w3 should match custom pattern"
|
|
);
|
|
assert!(
|
|
!should_quantize_layer("model.layers.0.mlp.experts.0.w2.weight", &layer_mask),
|
|
"w2 should NOT match custom pattern"
|
|
);
|
|
}
|
|
|
|
/// Helper function for layer filtering logic (matches ADR-017 AD-2 specification)
|
|
fn should_quantize_layer(layer_name: &str, mask: &super::LayerMask) -> bool {
|
|
use super::LayerMask;
|
|
|
|
match mask {
|
|
LayerMask::ExpertsOnly => {
|
|
// Quantize MoE expert FFN layers only (gate_proj, up_proj, down_proj, w1, w2, w3)
|
|
// Exclude: router, gate, embed, norm, lm_head
|
|
let is_expert_ffn = layer_name.contains("gate_proj")
|
|
|| layer_name.contains("up_proj")
|
|
|| layer_name.contains("down_proj")
|
|
|| (layer_name.contains("experts")
|
|
&& (layer_name.contains(".w1.") || layer_name.contains(".w2.") || layer_name.contains(".w3.")));
|
|
|
|
let is_protected = layer_name.contains("router")
|
|
|| layer_name.contains(".gate.") // MoE gate (not gate_proj)
|
|
|| layer_name.contains("embed")
|
|
|| layer_name.contains("lm_head")
|
|
|| layer_name.contains("norm");
|
|
|
|
is_expert_ffn && !is_protected
|
|
}
|
|
LayerMask::All => {
|
|
// Quantize all linear layers except protected ones
|
|
let is_protected = layer_name.contains("router")
|
|
|| layer_name.contains("embed")
|
|
|| layer_name.contains("lm_head")
|
|
|| layer_name.contains("norm");
|
|
|
|
!is_protected
|
|
}
|
|
LayerMask::Custom(patterns) => {
|
|
// Match any custom pattern
|
|
patterns.iter().any(|p| layer_name.contains(p))
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Helper Functions
|
|
// ============================================================================
|
|
|
|
/// Helper to quantize weights using absmean method
|
|
/// Returns both ternary values and scale factor
|
|
fn quantize_absmean_with_scale(weights: &[f32]) -> (Vec<i8>, f32) {
|
|
if weights.is_empty() {
|
|
return (vec![], 0.0);
|
|
}
|
|
|
|
// Compute absmean scale: gamma = mean(|W|) + epsilon
|
|
let absmean: f32 = weights.iter().map(|&w| w.abs()).sum::<f32>() / weights.len() as f32;
|
|
let scale = absmean + EPSILON;
|
|
|
|
// Quantize: W_ternary = RoundClip(W / scale, -1, 1)
|
|
let ternary: Vec<i8> = weights
|
|
.iter()
|
|
.map(|&w| {
|
|
let normalized = w / scale;
|
|
// Round and clip to {-1, 0, +1}
|
|
if normalized >= 0.5 {
|
|
1
|
|
} else if normalized <= -0.5 {
|
|
-1
|
|
} else {
|
|
0
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
(ternary, scale)
|
|
}
|
|
|
|
/// Helper to quantize weights (scale not needed)
|
|
fn quantize_absmean(weights: &[f32]) -> Vec<i8> {
|
|
let (ternary, _scale) = quantize_absmean_with_scale(weights);
|
|
ternary
|
|
}
|
|
|
|
/// Helper to dequantize ternary values
|
|
fn dequantize_ternary(ternary: &[i8], scale: f32) -> Vec<f32> {
|
|
ternary.iter().map(|&t| (t as f32) * scale).collect()
|
|
}
|