9c3d741a08
Move 17 library crates into crates/, CLI binary into bin/fxt, consolidate 10 test crates into testing/, split config crate from deployment config files. Root directory reduced from 38+ to ~17 directories. All Cargo.toml paths and build.rs proto refs updated. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
377 lines
12 KiB
Rust
377 lines
12 KiB
Rust
//! Per-Channel Quantization Tests
|
|
//!
|
|
//! Validates that per-channel quantization reduces quantization error from 2.5% to 1.5%
|
|
//! on attention weights and linear layer weights.
|
|
|
|
use candle_core::{DType, Device, Tensor};
|
|
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
|
|
use ml::MLError;
|
|
|
|
/// Test 1: Per-channel quantization reduces error on attention weights
|
|
#[test]
|
|
fn test_per_channel_attention_weight_accuracy() -> Result<(), MLError> {
|
|
let device = Device::Cpu;
|
|
|
|
// Create attention weight matrix [256, 256] (out_channels, in_channels)
|
|
let weight_data: Vec<f32> = (0..256 * 256)
|
|
.map(|i| ((i as f32) * 0.01).sin() * 0.5)
|
|
.collect();
|
|
let weight = Tensor::from_slice(&weight_data, (256, 256), &device)?;
|
|
|
|
// Test per-tensor quantization (per_channel = false)
|
|
let config_per_tensor = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: false,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer_per_tensor = Quantizer::new(config_per_tensor, device.clone());
|
|
let quantized_per_tensor = quantizer_per_tensor.quantize_tensor(&weight, "q_weight")?;
|
|
let dequantized_per_tensor = quantizer_per_tensor.dequantize_tensor(&quantized_per_tensor)?;
|
|
|
|
// Calculate per-tensor error
|
|
let error_per_tensor = calculate_relative_error(&weight, &dequantized_per_tensor)?;
|
|
|
|
// Test per-channel quantization (per_channel = true)
|
|
let config_per_channel = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer_per_channel = Quantizer::new(config_per_channel, device.clone());
|
|
let quantized_per_channel = quantizer_per_channel.quantize_tensor(&weight, "q_weight")?;
|
|
|
|
// Verify per-channel params were stored
|
|
assert!(
|
|
quantizer_per_channel.has_per_channel_params("q_weight"),
|
|
"Per-channel params should be stored"
|
|
);
|
|
|
|
// Dequantize using per-channel method
|
|
let dequantized_per_channel =
|
|
quantizer_per_channel.dequantize_tensor_per_channel(&quantized_per_channel, "q_weight")?;
|
|
|
|
// Calculate per-channel error
|
|
let error_per_channel = calculate_relative_error(&weight, &dequantized_per_channel)?;
|
|
|
|
println!(
|
|
"Per-tensor error: {:.4}%, Per-channel error: {:.4}%",
|
|
error_per_tensor * 100.0,
|
|
error_per_channel * 100.0
|
|
);
|
|
|
|
// Validate error reduction: per-channel should be lower than per-tensor
|
|
assert!(
|
|
error_per_channel < error_per_tensor,
|
|
"Per-channel error {:.4}% should be lower than per-tensor error {:.4}%",
|
|
error_per_channel * 100.0,
|
|
error_per_tensor * 100.0
|
|
);
|
|
|
|
// Target validation: per-channel error should be < 1.5%
|
|
assert!(
|
|
error_per_channel < 0.015,
|
|
"Per-channel error {:.4}% should be < 1.5%",
|
|
error_per_channel * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test 2: Per-channel quantization on linear layer weights
|
|
#[test]
|
|
fn test_per_channel_linear_layer_accuracy() -> Result<(), MLError> {
|
|
let device = Device::Cpu;
|
|
|
|
// Create linear layer weight [128, 256] (out_features, in_features)
|
|
let weight_data: Vec<f32> = (0..128 * 256)
|
|
.map(|i| ((i as f32) * 0.02).cos() * 0.3)
|
|
.collect();
|
|
let weight = Tensor::from_slice(&weight_data, (128, 256), &device)?;
|
|
|
|
// Per-channel quantization
|
|
let config = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer = Quantizer::new(config, device.clone());
|
|
|
|
// Quantize
|
|
let quantized = quantizer.quantize_tensor(&weight, "linear_weight")?;
|
|
|
|
// Verify quantized dtype is U8
|
|
assert_eq!(
|
|
quantized.data.dtype(),
|
|
DType::U8,
|
|
"Quantized data should be U8"
|
|
);
|
|
|
|
// Verify shape is preserved
|
|
assert_eq!(
|
|
quantized.data.dims(),
|
|
&[128, 256],
|
|
"Shape should be preserved"
|
|
);
|
|
|
|
// Dequantize
|
|
let dequantized = quantizer.dequantize_tensor_per_channel(&quantized, "linear_weight")?;
|
|
|
|
// Calculate error
|
|
let error = calculate_relative_error(&weight, &dequantized)?;
|
|
|
|
println!("Linear layer per-channel error: {:.4}%", error * 100.0);
|
|
|
|
// Validate error < 1.5%
|
|
assert!(
|
|
error < 0.015,
|
|
"Per-channel error {:.4}% should be < 1.5%",
|
|
error * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test 3: Per-channel parameters storage and retrieval
|
|
#[test]
|
|
fn test_per_channel_params_storage() -> Result<(), MLError> {
|
|
let device = Device::Cpu;
|
|
|
|
// Create weight [64, 128]
|
|
let weight_data: Vec<f32> = (0..64 * 128).map(|i| (i as f32) * 0.01).collect();
|
|
let weight = Tensor::from_slice(&weight_data, (64, 128), &device)?;
|
|
|
|
let config = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer = Quantizer::new(config, device.clone());
|
|
|
|
// Quantize
|
|
let _quantized = quantizer.quantize_tensor(&weight, "test_weight")?;
|
|
|
|
// Check params were stored
|
|
assert!(
|
|
quantizer.has_per_channel_params("test_weight"),
|
|
"Per-channel params should be stored"
|
|
);
|
|
|
|
// Retrieve params
|
|
let params = quantizer
|
|
.get_per_channel_params("test_weight")
|
|
.expect("Params should exist");
|
|
|
|
// Verify params shape
|
|
assert_eq!(
|
|
params.scales.len(),
|
|
64,
|
|
"Should have 64 scales (one per output channel)"
|
|
);
|
|
assert_eq!(params.zero_points.len(), 64, "Should have 64 zero points");
|
|
assert_eq!(params.min_vals.len(), 64, "Should have 64 min values");
|
|
assert_eq!(params.max_vals.len(), 64, "Should have 64 max values");
|
|
|
|
// Verify scales are positive
|
|
for (idx, scale) in params.scales.iter().enumerate() {
|
|
assert!(
|
|
*scale > 0.0,
|
|
"Scale {} should be positive, got {}",
|
|
idx,
|
|
scale
|
|
);
|
|
}
|
|
|
|
// Verify zero points are within valid range for symmetric quantization
|
|
for (idx, zp) in params.zero_points.iter().enumerate() {
|
|
assert_eq!(
|
|
*zp, 127,
|
|
"Zero point {} should be 127 (symmetric), got {}",
|
|
idx, zp
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test 4: Per-channel quantization comparison with per-tensor
|
|
#[test]
|
|
fn test_per_channel_vs_per_tensor_comparison() -> Result<(), MLError> {
|
|
let device = Device::Cpu;
|
|
|
|
// Create weight with varying scales across channels
|
|
// Each channel has different magnitude to amplify per-channel benefit
|
|
let mut weight_data = Vec::with_capacity(32 * 64);
|
|
for channel in 0..32 {
|
|
let scale = (channel + 1) as f32 * 0.1;
|
|
for i in 0..64 {
|
|
weight_data.push((i as f32 * 0.01).sin() * scale);
|
|
}
|
|
}
|
|
let weight = Tensor::from_slice(&weight_data, (32, 64), &device)?;
|
|
|
|
// Per-tensor quantization
|
|
let config_per_tensor = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: false,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer_per_tensor = Quantizer::new(config_per_tensor, device.clone());
|
|
let quantized_pt = quantizer_per_tensor.quantize_tensor(&weight, "w1")?;
|
|
let dequantized_pt = quantizer_per_tensor.dequantize_tensor(&quantized_pt)?;
|
|
let error_pt = calculate_relative_error(&weight, &dequantized_pt)?;
|
|
|
|
// Per-channel quantization
|
|
let config_per_channel = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer_per_channel = Quantizer::new(config_per_channel, device.clone());
|
|
let quantized_pc = quantizer_per_channel.quantize_tensor(&weight, "w1")?;
|
|
let dequantized_pc =
|
|
quantizer_per_channel.dequantize_tensor_per_channel(&quantized_pc, "w1")?;
|
|
let error_pc = calculate_relative_error(&weight, &dequantized_pc)?;
|
|
|
|
println!(
|
|
"Variable-scale weights - Per-tensor: {:.4}%, Per-channel: {:.4}%",
|
|
error_pt * 100.0,
|
|
error_pc * 100.0
|
|
);
|
|
|
|
// Per-channel should be significantly better for variable-scale weights
|
|
let improvement_ratio = error_pt / error_pc;
|
|
assert!(
|
|
improvement_ratio > 1.0,
|
|
"Per-channel should be better than per-tensor (improvement ratio: {:.2}x)",
|
|
improvement_ratio
|
|
);
|
|
|
|
// Per-channel should achieve target <1.5% error
|
|
assert!(
|
|
error_pc < 0.015,
|
|
"Per-channel error {:.4}% should be < 1.5%",
|
|
error_pc * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test 5: Matmul with per-channel dequantized weights
|
|
#[test]
|
|
fn test_per_channel_matmul_integration() -> Result<(), MLError> {
|
|
let device = Device::Cpu;
|
|
|
|
// Create weight [64, 128]
|
|
let weight_data: Vec<f32> = (0..64 * 128)
|
|
.map(|i| ((i as f32) * 0.01).sin() * 0.2)
|
|
.collect();
|
|
let weight = Tensor::from_slice(&weight_data, (64, 128), &device)?;
|
|
|
|
// Create input [2, 128] (batch_size=2, in_features=128)
|
|
let input_data: Vec<f32> = (0..2 * 128).map(|i| (i as f32) * 0.1).collect();
|
|
let input = Tensor::from_slice(&input_data, (2, 128), &device)?;
|
|
|
|
// F32 matmul (ground truth)
|
|
let output_f32 = input.matmul(&weight.t()?)?;
|
|
|
|
// Per-channel quantization
|
|
let config = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer = Quantizer::new(config, device.clone());
|
|
|
|
// Quantize weight
|
|
let quantized_weight = quantizer.quantize_tensor(&weight, "weight")?;
|
|
|
|
// Dequantize for inference
|
|
let dequantized_weight =
|
|
quantizer.dequantize_tensor_per_channel(&quantized_weight, "weight")?;
|
|
|
|
// INT8 matmul (with per-channel dequantization)
|
|
let output_int8 = input.matmul(&dequantized_weight.t()?)?;
|
|
|
|
// Calculate output error
|
|
let error = calculate_relative_error(&output_f32, &output_int8)?;
|
|
|
|
println!("Matmul output error (per-channel): {:.4}%", error * 100.0);
|
|
|
|
// Output error should be low
|
|
assert!(
|
|
error < 0.02,
|
|
"Matmul output error {:.4}% should be < 2.0%",
|
|
error * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test 6: Per-channel quantization rejects non-2D tensors
|
|
#[test]
|
|
fn test_per_channel_rejects_non_2d() -> Result<(), MLError> {
|
|
let device = Device::Cpu;
|
|
|
|
let config = QuantizationConfig {
|
|
quant_type: QuantizationType::Int8,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
calibration_samples: None,
|
|
};
|
|
let mut quantizer = Quantizer::new(config, device.clone());
|
|
|
|
// Test 1D tensor
|
|
let tensor_1d = Tensor::zeros((128,), DType::F32, &device)?;
|
|
let result_1d = quantizer.quantize_tensor_per_channel(&tensor_1d, "test");
|
|
assert!(
|
|
result_1d.is_err(),
|
|
"1D tensor should be rejected for per-channel quantization"
|
|
);
|
|
|
|
// Test 3D tensor
|
|
let tensor_3d = Tensor::zeros((2, 64, 128), DType::F32, &device)?;
|
|
let result_3d = quantizer.quantize_tensor_per_channel(&tensor_3d, "test");
|
|
assert!(
|
|
result_3d.is_err(),
|
|
"3D tensor should be rejected for per-channel quantization"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Helper: Calculate relative error between two tensors
|
|
fn calculate_relative_error(original: &Tensor, reconstructed: &Tensor) -> Result<f32, MLError> {
|
|
// Flatten both tensors
|
|
let orig_vec = original.flatten_all()?.to_vec1::<f32>()?;
|
|
let recon_vec = reconstructed.flatten_all()?.to_vec1::<f32>()?;
|
|
|
|
assert_eq!(orig_vec.len(), recon_vec.len());
|
|
|
|
// Calculate mean absolute error
|
|
let mae: f32 = orig_vec
|
|
.iter()
|
|
.zip(recon_vec.iter())
|
|
.map(|(o, r)| (o - r).abs())
|
|
.sum::<f32>()
|
|
/ orig_vec.len() as f32;
|
|
|
|
// Calculate mean of original values
|
|
let orig_mean = orig_vec.iter().sum::<f32>().abs() / orig_vec.len() as f32;
|
|
|
|
// Relative error
|
|
let relative_error = if orig_mean > 1e-8 {
|
|
mae / orig_mean
|
|
} else {
|
|
mae // If original is near zero, use absolute error
|
|
};
|
|
|
|
Ok(relative_error)
|
|
}
|