mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-20 07:46:08 -05:00
Vijay Janapa Reddi
a71e0eded5
Moved memoization (KV-cache) after compression to align with optimization tier milestones. Changes: - Module 15: Quantization (was 16) - Module 16: Compression (was 17) - Module 17: Memoization (was 15) Pedagogical Rationale: This creates clear alignment with the optimization milestone structure: - M06 (Profiling): Module 14 - M07 (Compression): Modules 15-16 (Quantization + Compression) - M08 (Acceleration): Modules 17-18 (Memoization/KV-cache + Acceleration) Before: Students learned KV-cache before understanding why models are slow After: Students profile → compress → then optimize with KV-cache Updated milestone reference in profile_kv_cache.py: Module 15 → Module 17
128 KiB
128 KiB
In [ ]:
#| default_exp optimization.quantizationIn [ ]:
#| export
import numpy as np
import time
from typing import Tuple, Dict, List, Optional
import warnings
# Import dependencies from other modules
from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Linear
from tinytorch.core.activations import ReLU
print("✅ Quantization module imports complete")In [ ]:
def quantize_int8(tensor: Tensor) -> Tuple[Tensor, float, int]:
"""
Quantize FP32 tensor to INT8 using symmetric quantization.
TODO: Implement INT8 quantization with scale and zero_point calculation
APPROACH:
1. Find min/max values in tensor data
2. Calculate scale: (max_val - min_val) / 255 (INT8 range: -128 to 127)
3. Calculate zero_point: offset to map FP32 zero to INT8 zero
4. Apply quantization formula: round((value - zero_point) / scale)
5. Clamp to INT8 range [-128, 127]
EXAMPLE:
>>> tensor = Tensor([[-1.0, 0.0, 2.0], [0.5, 1.5, -0.5]])
>>> q_tensor, scale, zero_point = quantize_int8(tensor)
>>> print(f"Scale: {scale:.4f}, Zero point: {zero_point}")
Scale: 0.0118, Zero point: 42
HINTS:
- Use np.round() for quantization
- Clamp with np.clip(values, -128, 127)
- Handle edge case where min_val == max_val (set scale=1.0)
"""
### BEGIN SOLUTION
data = tensor.data
# Step 1: Find dynamic range
min_val = float(np.min(data))
max_val = float(np.max(data))
# Step 2: Handle edge case (constant tensor)
if abs(max_val - min_val) < 1e-8:
scale = 1.0
zero_point = 0
quantized_data = np.zeros_like(data, dtype=np.int8)
return Tensor(quantized_data), scale, zero_point
# Step 3: Calculate scale and zero_point for standard quantization
# Map [min_val, max_val] to [-128, 127] (INT8 range)
scale = (max_val - min_val) / 255.0
zero_point = int(np.round(-128 - min_val / scale))
# Clamp zero_point to valid INT8 range
zero_point = int(np.clip(zero_point, -128, 127))
# Step 4: Apply quantization formula: q = (x / scale) + zero_point
quantized_data = np.round(data / scale + zero_point)
# Step 5: Clamp to INT8 range and convert to int8
quantized_data = np.clip(quantized_data, -128, 127).astype(np.int8)
return Tensor(quantized_data), scale, zero_point
### END SOLUTION
def test_unit_quantize_int8():
"""🔬 Test INT8 quantization implementation."""
print("🔬 Unit Test: INT8 Quantization...")
# Test basic quantization
tensor = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
q_tensor, scale, zero_point = quantize_int8(tensor)
# Verify quantized values are in INT8 range
assert np.all(q_tensor.data >= -128)
assert np.all(q_tensor.data <= 127)
assert isinstance(scale, float)
assert isinstance(zero_point, int)
# Test dequantization preserves approximate values
dequantized = scale * (q_tensor.data - zero_point)
error = np.mean(np.abs(tensor.data - dequantized))
assert error < 0.2, f"Quantization error too high: {error}"
# Test edge case: constant tensor
constant_tensor = Tensor([[2.0, 2.0], [2.0, 2.0]])
q_const, scale_const, zp_const = quantize_int8(constant_tensor)
assert scale_const == 1.0
print("✅ INT8 quantization works correctly!")
test_unit_quantize_int8()In [ ]:
def dequantize_int8(q_tensor: Tensor, scale: float, zero_point: int) -> Tensor:
"""
Dequantize INT8 tensor back to FP32.
TODO: Implement dequantization using the inverse formula
APPROACH:
1. Apply inverse quantization: scale * quantized_value + zero_point * scale
2. Return as new FP32 Tensor
EXAMPLE:
>>> q_tensor = Tensor([[-42, 0, 85]]) # INT8 values
>>> scale, zero_point = 0.0314, 64
>>> fp32_tensor = dequantize_int8(q_tensor, scale, zero_point)
>>> print(fp32_tensor.data)
[[-1.31, 2.01, 2.67]] # Approximate original values
HINT:
- Formula: dequantized = scale * quantized + zero_point * scale
"""
### BEGIN SOLUTION
# Apply inverse quantization formula
dequantized_data = scale * q_tensor.data + zero_point * scale
return Tensor(dequantized_data.astype(np.float32))
### END SOLUTION
def test_unit_dequantize_int8():
"""🔬 Test INT8 dequantization implementation."""
print("🔬 Unit Test: INT8 Dequantization...")
# Test round-trip: quantize → dequantize
original = Tensor([[-1.5, 0.0, 3.2], [1.1, -0.8, 2.7]])
q_tensor, scale, zero_point = quantize_int8(original)
restored = dequantize_int8(q_tensor, scale, zero_point)
# Verify round-trip error is small
error = np.mean(np.abs(original.data - restored.data))
assert error < 2.0, f"Round-trip error too high: {error}"
# Verify output is float32
assert restored.data.dtype == np.float32
print("✅ INT8 dequantization works correctly!")
test_unit_dequantize_int8()In [ ]:
def analyze_quantization_error():
"""📊 Analyze quantization error across different distributions."""
print("📊 Analyzing Quantization Error Across Distributions...")
distributions = {
'uniform': np.random.uniform(-1, 1, (1000,)),
'normal': np.random.normal(0, 0.5, (1000,)),
'outliers': np.concatenate([np.random.normal(0, 0.1, (900,)),
np.random.uniform(-2, 2, (100,))]),
'sparse': np.random.choice([0, 0, 0, 1], size=(1000,)) * np.random.normal(0, 1, (1000,))
}
results = {}
for name, data in distributions.items():
# Quantize and measure error
original = Tensor(data)
q_tensor, scale, zero_point = quantize_int8(original)
restored = dequantize_int8(q_tensor, scale, zero_point)
# Calculate metrics
mse = np.mean((original.data - restored.data) ** 2)
max_error = np.max(np.abs(original.data - restored.data))
results[name] = {
'mse': mse,
'max_error': max_error,
'scale': scale,
'range_ratio': (np.max(data) - np.min(data)) / scale if scale > 0 else 0
}
print(f"{name:8}: MSE={mse:.6f}, Max Error={max_error:.4f}, Scale={scale:.4f}")
print("\n💡 Insights:")
print("- Uniform: Low error, good scale utilization")
print("- Normal: Higher error at distribution tails")
print("- Outliers: Poor quantization due to extreme values")
print("- Sparse: Wasted quantization levels on zeros")
return results
# Analyze quantization quality
error_analysis = analyze_quantization_error()In [ ]:
class QuantizedLinear:
"""Quantized version of Linear layer using INT8 arithmetic."""
def __init__(self, linear_layer: Linear):
"""
Create quantized version of existing linear layer.
TODO: Quantize weights and bias, store quantization parameters
APPROACH:
1. Quantize weights using quantize_int8
2. Quantize bias if it exists
3. Store original layer reference for forward pass
4. Store quantization parameters for dequantization
IMPLEMENTATION STRATEGY:
- Store quantized weights, scales, and zero points
- Implement forward pass using dequantized computation (educational approach)
- Production: Would use INT8 matrix multiplication libraries
"""
### BEGIN SOLUTION
self.original_layer = linear_layer
# Quantize weights
self.q_weight, self.weight_scale, self.weight_zero_point = quantize_int8(linear_layer.weight)
# Quantize bias if it exists
if linear_layer.bias is not None:
self.q_bias, self.bias_scale, self.bias_zero_point = quantize_int8(linear_layer.bias)
else:
self.q_bias = None
self.bias_scale = None
self.bias_zero_point = None
# Store input quantization parameters (set during calibration)
self.input_scale = None
self.input_zero_point = None
### END SOLUTION
def calibrate(self, sample_inputs: List[Tensor]):
"""
Calibrate input quantization parameters using sample data.
TODO: Calculate optimal input quantization parameters
APPROACH:
1. Collect statistics from sample inputs
2. Calculate optimal scale and zero_point for inputs
3. Store for use in forward pass
"""
### BEGIN SOLUTION
# Collect all input values
all_values = []
for inp in sample_inputs:
all_values.extend(inp.data.flatten())
all_values = np.array(all_values)
# Calculate input quantization parameters
min_val = float(np.min(all_values))
max_val = float(np.max(all_values))
if abs(max_val - min_val) < 1e-8:
self.input_scale = 1.0
self.input_zero_point = 0
else:
self.input_scale = (max_val - min_val) / 255.0
self.input_zero_point = int(np.round(-128 - min_val / self.input_scale))
self.input_zero_point = np.clip(self.input_zero_point, -128, 127)
### END SOLUTION
def forward(self, x: Tensor) -> Tensor:
"""
Forward pass with quantized computation.
TODO: Implement quantized forward pass
APPROACH:
1. Quantize input (if calibrated)
2. Dequantize weights and input for computation (educational approach)
3. Perform matrix multiplication
4. Return FP32 result
NOTE: Production quantization uses INT8 GEMM libraries for speed
"""
### BEGIN SOLUTION
# For educational purposes, we dequantize and compute in FP32
# Production systems use specialized INT8 GEMM operations
# Dequantize weights
weight_fp32 = dequantize_int8(self.q_weight, self.weight_scale, self.weight_zero_point)
# Perform computation (same as original layer)
result = x.matmul(weight_fp32)
# Add bias if it exists
if self.q_bias is not None:
bias_fp32 = dequantize_int8(self.q_bias, self.bias_scale, self.bias_zero_point)
result = Tensor(result.data + bias_fp32.data)
return result
### END SOLUTION
def __call__(self, x: Tensor) -> Tensor:
"""Allows the quantized linear layer to be called like a function."""
return self.forward(x)
def parameters(self) -> List[Tensor]:
"""Return quantized parameters."""
params = [self.q_weight]
if self.q_bias is not None:
params.append(self.q_bias)
return params
def memory_usage(self) -> Dict[str, float]:
"""Calculate memory usage in bytes."""
### BEGIN SOLUTION
# Original FP32 usage
original_weight_bytes = self.original_layer.weight.data.size * 4 # 4 bytes per FP32
original_bias_bytes = 0
if self.original_layer.bias is not None:
original_bias_bytes = self.original_layer.bias.data.size * 4
# Quantized INT8 usage
quantized_weight_bytes = self.q_weight.data.size * 1 # 1 byte per INT8
quantized_bias_bytes = 0
if self.q_bias is not None:
quantized_bias_bytes = self.q_bias.data.size * 1
# Add overhead for scales and zero points (small)
overhead_bytes = 8 * 2 # 2 floats + 2 ints for weight/bias quantization params
return {
'original_bytes': original_weight_bytes + original_bias_bytes,
'quantized_bytes': quantized_weight_bytes + quantized_bias_bytes + overhead_bytes,
'compression_ratio': (original_weight_bytes + original_bias_bytes) /
(quantized_weight_bytes + quantized_bias_bytes + overhead_bytes)
}
### END SOLUTION
def test_unit_quantized_linear():
"""🔬 Test QuantizedLinear implementation."""
print("🔬 Unit Test: QuantizedLinear...")
# Create original linear layer
original = Linear(4, 3)
original.weight = Tensor(np.random.randn(4, 3) * 0.5) # Smaller range for testing
original.bias = Tensor(np.random.randn(3) * 0.1)
# Create quantized version
quantized = QuantizedLinear(original)
# Test forward pass
x = Tensor(np.random.randn(2, 4) * 0.5)
# Original forward pass
original_output = original.forward(x)
# Quantized forward pass
quantized_output = quantized.forward(x)
# Compare outputs (should be close but not identical due to quantization)
error = np.mean(np.abs(original_output.data - quantized_output.data))
assert error < 1.0, f"Quantization error too high: {error}"
# Test memory usage
memory_info = quantized.memory_usage()
assert memory_info['compression_ratio'] > 3.0, "Should achieve ~4× compression"
print(f" Memory reduction: {memory_info['compression_ratio']:.1f}×")
print("✅ QuantizedLinear works correctly!")
test_unit_quantized_linear()In [ ]:
def quantize_model(model, calibration_data: Optional[List[Tensor]] = None) -> None:
"""
Quantize all Linear layers in a model in-place.
TODO: Replace all Linear layers with QuantizedLinear versions
APPROACH:
1. Find all Linear layers in the model
2. Replace each with QuantizedLinear version
3. If calibration data provided, calibrate input quantization
4. Handle Sequential containers properly
EXAMPLE:
>>> model = Sequential(Linear(10, 5), ReLU(), Linear(5, 2))
>>> quantize_model(model)
>>> # Now model uses quantized layers
HINT:
- Handle Sequential.layers list for layer replacement
- Use isinstance(layer, Linear) to identify layers to quantize
"""
### BEGIN SOLUTION
if hasattr(model, 'layers'): # Sequential model
for i, layer in enumerate(model.layers):
if isinstance(layer, Linear):
# Replace with quantized version
quantized_layer = QuantizedLinear(layer)
# Calibrate if data provided
if calibration_data is not None:
# Run forward passes to get intermediate activations
sample_inputs = []
for data in calibration_data[:10]: # Use first 10 samples for efficiency
# Forward through layers up to this point
x = data
for j in range(i):
if hasattr(model.layers[j], 'forward'):
x = model.layers[j].forward(x)
sample_inputs.append(x)
quantized_layer.calibrate(sample_inputs)
model.layers[i] = quantized_layer
elif isinstance(model, Linear): # Single Linear layer
# Can't replace in-place for single layer, user should handle
raise ValueError("Cannot quantize single Linear layer in-place. Use QuantizedLinear directly.")
else:
raise ValueError(f"Unsupported model type: {type(model)}")
### END SOLUTION
def test_unit_quantize_model():
"""🔬 Test model quantization implementation."""
print("🔬 Unit Test: Model Quantization...")
# Create test model
model = Sequential(
Linear(4, 8),
ReLU(),
Linear(8, 3)
)
# Initialize weights
model.layers[0].weight = Tensor(np.random.randn(4, 8) * 0.5)
model.layers[0].bias = Tensor(np.random.randn(8) * 0.1)
model.layers[2].weight = Tensor(np.random.randn(8, 3) * 0.5)
model.layers[2].bias = Tensor(np.random.randn(3) * 0.1)
# Test original model
x = Tensor(np.random.randn(2, 4))
original_output = model.forward(x)
# Create calibration data
calibration_data = [Tensor(np.random.randn(1, 4)) for _ in range(5)]
# Quantize model
quantize_model(model, calibration_data)
# Verify layers were replaced
assert isinstance(model.layers[0], QuantizedLinear)
assert isinstance(model.layers[1], ReLU) # Should remain unchanged
assert isinstance(model.layers[2], QuantizedLinear)
# Test quantized model
quantized_output = model.forward(x)
# Compare outputs
error = np.mean(np.abs(original_output.data - quantized_output.data))
print(f" Model quantization error: {error:.4f}")
assert error < 2.0, f"Model quantization error too high: {error}"
print("✅ Model quantization works correctly!")
test_unit_quantize_model()In [ ]:
def compare_model_sizes(original_model, quantized_model) -> Dict[str, float]:
"""
Compare memory usage between original and quantized models.
TODO: Calculate comprehensive memory comparison
APPROACH:
1. Count parameters in both models
2. Calculate bytes used (FP32 vs INT8)
3. Include quantization overhead
4. Return comparison metrics
"""
### BEGIN SOLUTION
# Count original model parameters
original_params = 0
original_bytes = 0
if hasattr(original_model, 'layers'):
for layer in original_model.layers:
if hasattr(layer, 'parameters'):
params = layer.parameters()
for param in params:
original_params += param.data.size
original_bytes += param.data.size * 4 # 4 bytes per FP32
# Count quantized model parameters
quantized_params = 0
quantized_bytes = 0
if hasattr(quantized_model, 'layers'):
for layer in quantized_model.layers:
if isinstance(layer, QuantizedLinear):
memory_info = layer.memory_usage()
quantized_bytes += memory_info['quantized_bytes']
params = layer.parameters()
for param in params:
quantized_params += param.data.size
elif hasattr(layer, 'parameters'):
# Non-quantized layers
params = layer.parameters()
for param in params:
quantized_params += param.data.size
quantized_bytes += param.data.size * 4
compression_ratio = original_bytes / quantized_bytes if quantized_bytes > 0 else 1.0
memory_saved = original_bytes - quantized_bytes
return {
'original_params': original_params,
'quantized_params': quantized_params,
'original_bytes': original_bytes,
'quantized_bytes': quantized_bytes,
'compression_ratio': compression_ratio,
'memory_saved_mb': memory_saved / (1024 * 1024),
'memory_saved_percent': (memory_saved / original_bytes) * 100 if original_bytes > 0 else 0
}
### END SOLUTION
def test_unit_compare_model_sizes():
"""🔬 Test model size comparison."""
print("🔬 Unit Test: Model Size Comparison...")
# Create and quantize a model for testing
original_model = Sequential(Linear(100, 50), ReLU(), Linear(50, 10))
original_model.layers[0].weight = Tensor(np.random.randn(100, 50))
original_model.layers[0].bias = Tensor(np.random.randn(50))
original_model.layers[2].weight = Tensor(np.random.randn(50, 10))
original_model.layers[2].bias = Tensor(np.random.randn(10))
# Create quantized copy
quantized_model = Sequential(Linear(100, 50), ReLU(), Linear(50, 10))
quantized_model.layers[0].weight = Tensor(np.random.randn(100, 50))
quantized_model.layers[0].bias = Tensor(np.random.randn(50))
quantized_model.layers[2].weight = Tensor(np.random.randn(50, 10))
quantized_model.layers[2].bias = Tensor(np.random.randn(10))
quantize_model(quantized_model)
# Compare sizes
comparison = compare_model_sizes(original_model, quantized_model)
# Verify compression achieved
assert comparison['compression_ratio'] > 2.0, "Should achieve significant compression"
assert comparison['memory_saved_percent'] > 50, "Should save >50% memory"
print(f" Compression ratio: {comparison['compression_ratio']:.1f}×")
print(f" Memory saved: {comparison['memory_saved_percent']:.1f}%")
print("✅ Model size comparison works correctly!")
test_unit_compare_model_sizes()In [ ]:
def analyze_quantization_performance():
"""📊 Comprehensive analysis of quantization benefits and trade-offs."""
print("📊 Analyzing Quantization Performance Across Model Sizes...")
# Test different model configurations
configs = [
{'name': 'Small MLP', 'layers': [64, 32, 10], 'batch_size': 32},
{'name': 'Medium MLP', 'layers': [512, 256, 128, 10], 'batch_size': 64},
{'name': 'Large MLP', 'layers': [2048, 1024, 512, 10], 'batch_size': 128},
]
results = []
for config in configs:
print(f"\n🔍 Testing {config['name']}...")
# Create original model
layers = []
for i in range(len(config['layers']) - 1):
layers.append(Linear(config['layers'][i], config['layers'][i+1]))
if i < len(config['layers']) - 2: # Add ReLU except for last layer
layers.append(ReLU())
original_model = Sequential(*layers)
# Initialize weights
for layer in original_model.layers:
if isinstance(layer, Linear):
layer.weight = Tensor(np.random.randn(*layer.weight.shape) * 0.1)
layer.bias = Tensor(np.random.randn(*layer.bias.shape) * 0.01)
# Create quantized copy
quantized_model = Sequential(*layers)
for i, layer in enumerate(original_model.layers):
if isinstance(layer, Linear):
quantized_model.layers[i].weight = Tensor(layer.weight.data.copy())
quantized_model.layers[i].bias = Tensor(layer.bias.data.copy())
# Generate calibration data
input_size = config['layers'][0]
calibration_data = [Tensor(np.random.randn(1, input_size)) for _ in range(10)]
# Quantize model
quantize_model(quantized_model, calibration_data)
# Measure performance
test_input = Tensor(np.random.randn(config['batch_size'], input_size))
# Time original model
start_time = time.time()
for _ in range(10):
original_output = original_model.forward(test_input)
original_time = (time.time() - start_time) / 10
# Time quantized model
start_time = time.time()
for _ in range(10):
quantized_output = quantized_model.forward(test_input)
quantized_time = (time.time() - start_time) / 10
# Calculate accuracy preservation (using MSE as proxy)
mse = np.mean((original_output.data - quantized_output.data) ** 2)
relative_error = np.sqrt(mse) / (np.std(original_output.data) + 1e-8)
# Memory comparison
memory_comparison = compare_model_sizes(original_model, quantized_model)
result = {
'name': config['name'],
'original_time': original_time * 1000, # Convert to ms
'quantized_time': quantized_time * 1000,
'speedup': original_time / quantized_time if quantized_time > 0 else 1.0,
'compression_ratio': memory_comparison['compression_ratio'],
'relative_error': relative_error,
'memory_saved_mb': memory_comparison['memory_saved_mb']
}
results.append(result)
print(f" Speedup: {result['speedup']:.1f}×")
print(f" Compression: {result['compression_ratio']:.1f}×")
print(f" Error: {result['relative_error']:.1%}")
print(f" Memory saved: {result['memory_saved_mb']:.1f}MB")
# Summary analysis
print(f"\n📈 QUANTIZATION PERFORMANCE SUMMARY")
print("=" * 50)
avg_speedup = np.mean([r['speedup'] for r in results])
avg_compression = np.mean([r['compression_ratio'] for r in results])
avg_error = np.mean([r['relative_error'] for r in results])
total_memory_saved = sum([r['memory_saved_mb'] for r in results])
print(f"Average speedup: {avg_speedup:.1f}×")
print(f"Average compression: {avg_compression:.1f}×")
print(f"Average relative error: {avg_error:.1%}")
print(f"Total memory saved: {total_memory_saved:.1f}MB")
print(f"\n💡 Key Insights:")
print(f"- Quantization achieves ~{avg_compression:.0f}× memory reduction")
print(f"- Typical speedup: {avg_speedup:.1f}× (varies by hardware)")
print(f"- Accuracy loss: <{avg_error:.1%} for well-calibrated models")
print(f"- Best for: Memory-constrained deployment")
return results
# Run comprehensive performance analysis
performance_results = analyze_quantization_performance()In [ ]:
def visualize_quantization_effects():
"""📊 Visualize the effects of quantization on weight distributions."""
print("📊 Visualizing Quantization Effects on Weight Distributions...")
# Create sample weight tensors with different characteristics
weight_types = {
'Normal': np.random.normal(0, 0.1, (1000,)),
'Uniform': np.random.uniform(-0.2, 0.2, (1000,)),
'Sparse': np.random.choice([0, 0, 0, 1], (1000,)) * np.random.normal(0, 0.15, (1000,)),
'Heavy-tailed': np.concatenate([
np.random.normal(0, 0.05, (800,)),
np.random.uniform(-0.5, 0.5, (200,))
])
}
fig, axes = plt.subplots(2, 2, figsize=(12, 8))
axes = axes.flatten()
for idx, (name, weights) in enumerate(weight_types.items()):
# Original weights
original_tensor = Tensor(weights)
# Quantize and dequantize
q_tensor, scale, zero_point = quantize_int8(original_tensor)
restored_tensor = dequantize_int8(q_tensor, scale, zero_point)
# Plot histograms
ax = axes[idx]
ax.hist(weights, bins=50, alpha=0.6, label='Original', density=True)
ax.hist(restored_tensor.data, bins=50, alpha=0.6, label='Quantized', density=True)
ax.set_title(f'{name} Weights\nScale: {scale:.4f}')
ax.set_xlabel('Weight Value')
ax.set_ylabel('Density')
ax.legend()
ax.grid(True, alpha=0.3)
# Calculate and display error metrics
mse = np.mean((weights - restored_tensor.data) ** 2)
ax.text(0.02, 0.98, f'MSE: {mse:.6f}', transform=ax.transAxes,
verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
plt.tight_layout()
plt.savefig('/tmp/claude/quantization_effects.png', dpi=100, bbox_inches='tight')
plt.show()
print("💡 Observations:")
print("- Normal: Smooth quantization, good preservation")
print("- Uniform: Excellent quantization, full range utilized")
print("- Sparse: Many wasted quantization levels on zeros")
print("- Heavy-tailed: Outliers dominate scale, poor precision for small weights")
# Visualize quantization effects
visualize_quantization_effects()In [ ]:
def analyze_quantization_strategies():
"""📊 Compare different quantization strategies and their trade-offs."""
print("📊 Analyzing Advanced Quantization Strategies...")
# Create test model and data
model = Sequential(Linear(128, 64), ReLU(), Linear(64, 10))
model.layers[0].weight = Tensor(np.random.randn(128, 64) * 0.1)
model.layers[0].bias = Tensor(np.random.randn(64) * 0.01)
model.layers[2].weight = Tensor(np.random.randn(64, 10) * 0.1)
model.layers[2].bias = Tensor(np.random.randn(10) * 0.01)
test_input = Tensor(np.random.randn(32, 128))
original_output = model.forward(test_input)
strategies = {}
# Strategy 1: Per-tensor quantization (what we implemented)
print("\n🔍 Strategy 1: Per-Tensor Quantization")
model_copy = Sequential(Linear(128, 64), ReLU(), Linear(64, 10))
for i, layer in enumerate(model.layers):
if isinstance(layer, Linear):
model_copy.layers[i].weight = Tensor(layer.weight.data.copy())
model_copy.layers[i].bias = Tensor(layer.bias.data.copy())
quantize_model(model_copy)
output1 = model_copy.forward(test_input)
error1 = np.mean((original_output.data - output1.data) ** 2)
strategies['per_tensor'] = {'mse': error1, 'description': 'Single scale per tensor'}
print(f" MSE: {error1:.6f}")
# Strategy 2: Per-channel quantization simulation
print("\n🔍 Strategy 2: Per-Channel Quantization (simulated)")
# Simulate by quantizing each output channel separately
def per_channel_quantize(tensor):
"""Simulate per-channel quantization for 2D weight matrices."""
if len(tensor.shape) < 2:
return quantize_int8(tensor)
quantized_data = np.zeros_like(tensor.data, dtype=np.int8)
scales = []
zero_points = []
for i in range(tensor.shape[1]): # Per output channel
channel_tensor = Tensor(tensor.data[:, i:i+1])
q_channel, scale, zp = quantize_int8(channel_tensor)
quantized_data[:, i] = q_channel.data.flatten()
scales.append(scale)
zero_points.append(zp)
return Tensor(quantized_data), scales, zero_points
# Apply per-channel quantization to weights
total_error = 0
for layer in model.layers:
if isinstance(layer, Linear):
q_weight, scales, zps = per_channel_quantize(layer.weight)
# Simulate dequantization and error
for i in range(layer.weight.shape[1]):
original_channel = layer.weight.data[:, i]
restored_channel = scales[i] * q_weight.data[:, i] + zps[i] * scales[i]
total_error += np.mean((original_channel - restored_channel) ** 2)
strategies['per_channel'] = {'mse': total_error, 'description': 'Scale per output channel'}
print(f" MSE: {total_error:.6f}")
# Strategy 3: Mixed precision simulation
print("\n🔍 Strategy 3: Mixed Precision")
# Keep sensitive layers in FP32, quantize others
sensitive_layers = [0] # First layer often most sensitive
mixed_error = 0
for i, layer in enumerate(model.layers):
if isinstance(layer, Linear):
if i in sensitive_layers:
# Keep in FP32 (no quantization error)
pass
else:
# Quantize layer
q_weight, scale, zp = quantize_int8(layer.weight)
restored = dequantize_int8(q_weight, scale, zp)
mixed_error += np.mean((layer.weight.data - restored.data) ** 2)
strategies['mixed_precision'] = {'mse': mixed_error, 'description': 'FP32 sensitive + INT8 others'}
print(f" MSE: {mixed_error:.6f}")
# Compare strategies
print(f"\n📊 QUANTIZATION STRATEGY COMPARISON")
print("=" * 60)
for name, info in strategies.items():
print(f"{name:15}: MSE={info['mse']:.6f} | {info['description']}")
# Find best strategy
best_strategy = min(strategies.items(), key=lambda x: x[1]['mse'])
print(f"\n🏆 Best Strategy: {best_strategy[0]} (MSE: {best_strategy[1]['mse']:.6f})")
print(f"\n💡 Production Insights:")
print("- Per-channel: Better accuracy, more complex implementation")
print("- Mixed precision: Optimal accuracy/efficiency trade-off")
print("- Per-tensor: Simplest, good for most applications")
print("- Hardware support varies: INT8 GEMM, per-channel scales")
return strategies
# Analyze quantization strategies
strategy_analysis = analyze_quantization_strategies()In [ ]:
def test_module():
"""
Comprehensive test of entire quantization module functionality.
This final test runs before module summary to ensure:
- All quantization functions work correctly
- Model quantization preserves functionality
- Memory savings are achieved
- Module is ready for integration with TinyTorch
"""
print("🧪 RUNNING MODULE INTEGRATION TEST")
print("=" * 50)
# Run all unit tests
print("Running unit tests...")
test_unit_quantize_int8()
test_unit_dequantize_int8()
test_unit_quantized_linear()
test_unit_quantize_model()
test_unit_compare_model_sizes()
print("\nRunning integration scenarios...")
# Test realistic usage scenario
print("🔬 Integration Test: End-to-end quantization workflow...")
# Create a realistic model
model = Sequential(
Linear(784, 128), # MNIST-like input
ReLU(),
Linear(128, 64),
ReLU(),
Linear(64, 10) # 10-class output
)
# Initialize with realistic weights
for layer in model.layers:
if isinstance(layer, Linear):
# Xavier initialization
fan_in, fan_out = layer.weight.shape
std = np.sqrt(2.0 / (fan_in + fan_out))
layer.weight = Tensor(np.random.randn(fan_in, fan_out) * std)
layer.bias = Tensor(np.zeros(fan_out))
# Generate realistic calibration data
calibration_data = [Tensor(np.random.randn(1, 784) * 0.1) for _ in range(20)]
# Test original model
test_input = Tensor(np.random.randn(8, 784) * 0.1)
original_output = model.forward(test_input)
# Quantize the model
quantize_model(model, calibration_data)
# Test quantized model
quantized_output = model.forward(test_input)
# Verify functionality is preserved
assert quantized_output.shape == original_output.shape, "Output shape mismatch"
# Verify reasonable accuracy preservation
mse = np.mean((original_output.data - quantized_output.data) ** 2)
relative_error = np.sqrt(mse) / (np.std(original_output.data) + 1e-8)
assert relative_error < 0.1, f"Accuracy degradation too high: {relative_error:.3f}"
# Verify memory savings
# Create equivalent original model for comparison
original_model = Sequential(
Linear(784, 128),
ReLU(),
Linear(128, 64),
ReLU(),
Linear(64, 10)
)
for i, layer in enumerate(model.layers):
if isinstance(layer, QuantizedLinear):
# Restore original weights for comparison
original_model.layers[i].weight = dequantize_int8(
layer.q_weight, layer.weight_scale, layer.weight_zero_point
)
if layer.q_bias is not None:
original_model.layers[i].bias = dequantize_int8(
layer.q_bias, layer.bias_scale, layer.bias_zero_point
)
memory_comparison = compare_model_sizes(original_model, model)
assert memory_comparison['compression_ratio'] > 2.0, "Insufficient compression achieved"
print(f"✅ Compression achieved: {memory_comparison['compression_ratio']:.1f}×")
print(f"✅ Accuracy preserved: {relative_error:.1%} relative error")
print(f"✅ Memory saved: {memory_comparison['memory_saved_mb']:.1f}MB")
# Test edge cases
print("🔬 Testing edge cases...")
# Test constant tensor quantization
constant_tensor = Tensor([[1.0, 1.0], [1.0, 1.0]])
q_const, scale_const, zp_const = quantize_int8(constant_tensor)
assert scale_const == 1.0, "Constant tensor quantization failed"
# Test zero tensor
zero_tensor = Tensor([[0.0, 0.0], [0.0, 0.0]])
q_zero, scale_zero, zp_zero = quantize_int8(zero_tensor)
restored_zero = dequantize_int8(q_zero, scale_zero, zp_zero)
assert np.allclose(restored_zero.data, 0.0, atol=1e-6), "Zero tensor restoration failed"
print("✅ Edge cases handled correctly!")
print("\n" + "=" * 50)
print("🎉 ALL TESTS PASSED! Module ready for export.")
print("📈 Quantization system provides:")
print(f" • {memory_comparison['compression_ratio']:.1f}× memory reduction")
print(f" • <{relative_error:.1%} accuracy loss")
print(f" • Production-ready INT8 quantization")
print("Run: tito module complete 17")
# Call the comprehensive test
test_module()In [ ]:
if __name__ == "__main__":
print("🚀 Running Quantization module...")
test_module()
print("✅ Module validation complete!")In [ ]:
#| export
class QuantizationComplete:
"""
Complete quantization system for milestone use.
Provides INT8 quantization with calibration for 4× memory reduction.
"""
@staticmethod
def quantize_tensor(tensor: Tensor) -> Tuple[Tensor, float, int]:
"""Quantize FP32 tensor to INT8."""
data = tensor.data
min_val = float(np.min(data))
max_val = float(np.max(data))
if abs(max_val - min_val) < 1e-8:
return Tensor(np.zeros_like(data, dtype=np.int8)), 1.0, 0
scale = (max_val - min_val) / 255.0
zero_point = int(np.round(-128 - min_val / scale))
zero_point = int(np.clip(zero_point, -128, 127))
quantized_data = np.round(data / scale + zero_point)
quantized_data = np.clip(quantized_data, -128, 127).astype(np.int8)
return Tensor(quantized_data), scale, zero_point
@staticmethod
def dequantize_tensor(q_tensor: Tensor, scale: float, zero_point: int) -> Tensor:
"""Dequantize INT8 tensor back to FP32."""
dequantized_data = (q_tensor.data.astype(np.float32) - zero_point) * scale
return Tensor(dequantized_data)
@staticmethod
def quantize_model(model, calibration_data: Optional[List[Tensor]] = None) -> Dict[str, any]:
"""
Quantize all Linear layers in a model.
Returns dictionary with quantization info and memory savings.
"""
quantized_layers = {}
original_size = 0
quantized_size = 0
# Iterate through model parameters
if hasattr(model, 'parameters'):
for i, param in enumerate(model.parameters()):
param_size = param.data.nbytes
original_size += param_size
# Quantize parameter
q_param, scale, zp = QuantizationComplete.quantize_tensor(param)
quantized_size += q_param.data.nbytes
quantized_layers[f'param_{i}'] = {
'quantized': q_param,
'scale': scale,
'zero_point': zp,
'original_shape': param.data.shape
}
return {
'quantized_layers': quantized_layers,
'original_size_mb': original_size / (1024 * 1024),
'quantized_size_mb': quantized_size / (1024 * 1024),
'compression_ratio': original_size / quantized_size if quantized_size > 0 else 1.0
}
@staticmethod
def compare_models(original_model, quantized_info: Dict) -> Dict[str, float]:
"""Compare memory usage between original and quantized models."""
return {
'original_mb': quantized_info['original_size_mb'],
'quantized_mb': quantized_info['quantized_size_mb'],
'compression_ratio': quantized_info['compression_ratio'],
'memory_saved_mb': quantized_info['original_size_mb'] - quantized_info['quantized_size_mb']
}
# Convenience functions for backward compatibility
def quantize_int8(tensor: Tensor) -> Tuple[Tensor, float, int]:
"""Quantize FP32 tensor to INT8."""
return QuantizationComplete.quantize_tensor(tensor)
def dequantize_int8(q_tensor: Tensor, scale: float, zero_point: int) -> Tensor:
"""Dequantize INT8 tensor back to FP32."""
return QuantizationComplete.dequantize_tensor(q_tensor, scale, zero_point)
def quantize_model(model, calibration_data: Optional[List[Tensor]] = None) -> Dict[str, any]:
"""Quantize entire model to INT8."""
return QuantizationComplete.quantize_model(model, calibration_data)