c88e53c3b3
Replace the no-op apply_layer_norm() that returned input unchanged with real layer normalization using layer_norm_with_fallback. Initialize LayerNormParams with proper weight (ones) and bias (zeros) tensors in from_grn constructor instead of None placeholders.
336 lines
12 KiB
Rust
336 lines
12 KiB
Rust
//! Quantized Gated Residual Network for TFT
|
||
//!
|
||
//! INT8 quantization of GRN layers with careful handling of residual connections.
|
||
//! Target: 500MB → 125MB (75% reduction) with <5% accuracy loss.
|
||
|
||
use candle_core::{Device, Tensor};
|
||
use tracing::debug;
|
||
|
||
use crate::cuda_compat::manual_sigmoid;
|
||
use crate::memory_optimization::quantization::{QuantizationType, QuantizedTensor, Quantizer};
|
||
use crate::tft::gated_residual::GatedResidualNetwork;
|
||
use crate::MLError;
|
||
|
||
/// Quantized Gated Residual Network
|
||
///
|
||
/// Quantizes linear layers to INT8 while keeping skip connections in F32
|
||
/// for numerical precision. This provides 70-80% memory reduction with
|
||
/// minimal accuracy loss.
|
||
#[derive(Debug, Clone)]
|
||
pub struct QuantizedGatedResidualNetwork {
|
||
pub input_dim: usize,
|
||
pub output_dim: usize,
|
||
|
||
// Quantized linear layers
|
||
pub quantized_linear1: Option<QuantizedTensor>,
|
||
pub quantized_linear2: Option<QuantizedTensor>,
|
||
|
||
// Quantized GLU weights (linear, gate)
|
||
pub quantized_glu_weights: (Option<QuantizedTensor>, Option<QuantizedTensor>),
|
||
|
||
// Optional skip projection (kept in F32 for precision)
|
||
pub quantized_skip_proj: Option<QuantizedTensor>,
|
||
|
||
// Layer normalization (kept in F32)
|
||
layer_norm: Option<LayerNormParams>,
|
||
|
||
// Context projection (quantized)
|
||
quantized_context_proj: Option<QuantizedTensor>,
|
||
|
||
// Quantizer for dequantization
|
||
quantizer: Quantizer,
|
||
|
||
device: Device,
|
||
}
|
||
|
||
/// Layer normalization parameters stored for quantized model
|
||
#[derive(Debug, Clone)]
|
||
struct LayerNormParams {
|
||
normalized_shape: Vec<usize>,
|
||
weight: Option<Tensor>,
|
||
bias: Option<Tensor>,
|
||
eps: f64,
|
||
}
|
||
|
||
impl QuantizedGatedResidualNetwork {
|
||
/// Create quantized GRN from original GRN
|
||
pub fn from_grn(grn: &GatedResidualNetwork, mut quantizer: Quantizer) -> Result<Self, MLError> {
|
||
debug!(
|
||
"Quantizing GRN: input_dim={}, output_dim={}",
|
||
grn.input_dim, grn.output_dim
|
||
);
|
||
|
||
let device = quantizer.device().clone();
|
||
|
||
// Extract and quantize linear1 weights
|
||
let linear1_weight = Self::extract_linear_weight(grn, "linear1")?;
|
||
let quantized_linear1 = Some(quantizer.quantize_tensor(&linear1_weight, "linear1")?);
|
||
|
||
// Extract and quantize linear2 weights
|
||
let linear2_weight = Self::extract_linear_weight(grn, "linear2")?;
|
||
let quantized_linear2 = Some(quantizer.quantize_tensor(&linear2_weight, "linear2")?);
|
||
|
||
// Extract and quantize GLU weights
|
||
let glu_linear_weight = Self::extract_linear_weight(grn, "glu.linear")?;
|
||
let glu_gate_weight = Self::extract_linear_weight(grn, "glu.gate")?;
|
||
let quantized_glu_linear =
|
||
Some(quantizer.quantize_tensor(&glu_linear_weight, "glu.linear")?);
|
||
let quantized_glu_gate = Some(quantizer.quantize_tensor(&glu_gate_weight, "glu.gate")?);
|
||
|
||
// Extract skip projection if present (quantize)
|
||
let quantized_skip_proj = if grn.input_dim != grn.output_dim {
|
||
let skip_weight = Self::extract_linear_weight(grn, "skip_projection")?;
|
||
Some(quantizer.quantize_tensor(&skip_weight, "skip_projection")?)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
// Extract context projection if present
|
||
let quantized_context_proj = {
|
||
let ctx_weight = Self::extract_linear_weight(grn, "context_projection")?;
|
||
Some(quantizer.quantize_tensor(&ctx_weight, "context_projection")?)
|
||
};
|
||
|
||
// Store layer norm parameters (keep in F32)
|
||
// Initialize weight=ones and bias=zeros for proper layer normalization
|
||
let ln_weight = Tensor::ones(grn.output_dim, candle_core::DType::F32, &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create ln_weight: {}", e)))?;
|
||
let ln_bias = Tensor::zeros(grn.output_dim, candle_core::DType::F32, &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create ln_bias: {}", e)))?;
|
||
let layer_norm = Some(LayerNormParams {
|
||
normalized_shape: vec![grn.output_dim],
|
||
weight: Some(ln_weight),
|
||
bias: Some(ln_bias),
|
||
eps: 1e-5,
|
||
});
|
||
|
||
Ok(Self {
|
||
input_dim: grn.input_dim,
|
||
output_dim: grn.output_dim,
|
||
quantized_linear1,
|
||
quantized_linear2,
|
||
quantized_glu_weights: (quantized_glu_linear, quantized_glu_gate),
|
||
quantized_skip_proj,
|
||
layer_norm,
|
||
quantized_context_proj,
|
||
quantizer,
|
||
device,
|
||
})
|
||
}
|
||
|
||
/// Extract linear layer weight from GRN (helper for quantization)
|
||
fn extract_linear_weight(
|
||
_grn: &GatedResidualNetwork,
|
||
layer_name: &str,
|
||
) -> Result<Tensor, MLError> {
|
||
// In production, would extract actual weights from GRN layers
|
||
// For TDD, create placeholder weights
|
||
let device = Device::Cpu;
|
||
|
||
let (in_dim, out_dim) = match layer_name {
|
||
"linear1" => (128, 128), // Placeholder dimensions
|
||
"linear2" => (128, 128),
|
||
"glu.linear" => (128, 128),
|
||
"glu.gate" => (128, 128),
|
||
"skip_projection" => (64, 128),
|
||
"context_projection" => (128, 128),
|
||
_ => (128, 128),
|
||
};
|
||
|
||
// Create random weight tensor
|
||
let weight_data: Vec<f32> = (0..in_dim * out_dim)
|
||
.map(|i| (i as f32 * 0.01).sin())
|
||
.collect();
|
||
|
||
Tensor::from_slice(&weight_data, (out_dim, in_dim), &device)
|
||
.map_err(|e| MLError::ModelError(format!("Failed to create weight tensor: {}", e)))
|
||
}
|
||
|
||
/// Forward pass through quantized GRN
|
||
pub fn forward(
|
||
&self,
|
||
x: &Tensor,
|
||
context: Option<&Tensor>,
|
||
_quantizer: &Quantizer,
|
||
) -> Result<Tensor, MLError> {
|
||
// Dequantize and apply linear1
|
||
let linear1_weight = self.quantizer.dequantize_tensor(
|
||
self.quantized_linear1
|
||
.as_ref()
|
||
.ok_or_else(|| MLError::ModelError("Missing linear1".to_string()))?,
|
||
)?;
|
||
let mut hidden = self.apply_linear(x, &linear1_weight)?;
|
||
hidden = hidden.elu(1.0)?;
|
||
|
||
// Apply context if provided
|
||
if let (Some(ctx), Some(ctx_proj)) = (context, &self.quantized_context_proj) {
|
||
let ctx_weight = self.quantizer.dequantize_tensor(ctx_proj)?;
|
||
let ctx_out = self.apply_linear(ctx, &ctx_weight)?;
|
||
hidden = (&hidden + &ctx_out)?;
|
||
}
|
||
|
||
// Dequantize and apply linear2
|
||
let linear2_weight = self.quantizer.dequantize_tensor(
|
||
self.quantized_linear2
|
||
.as_ref()
|
||
.ok_or_else(|| MLError::ModelError("Missing linear2".to_string()))?,
|
||
)?;
|
||
hidden = self.apply_linear(&hidden, &linear2_weight)?;
|
||
|
||
// Apply GLU gating
|
||
let gated = self.apply_glu(&hidden)?;
|
||
|
||
// Skip connection (kept in F32 for precision)
|
||
let output = if let Some(skip_proj) = &self.quantized_skip_proj {
|
||
let skip_weight = self.quantizer.dequantize_tensor(skip_proj)?;
|
||
let skip = self.apply_linear(x, &skip_weight)?;
|
||
(&gated + &skip)?
|
||
} else {
|
||
(&gated + x)?
|
||
};
|
||
|
||
// Layer normalization (F32)
|
||
let normalized = self.apply_layer_norm(&output)?;
|
||
|
||
Ok(normalized)
|
||
}
|
||
|
||
/// Apply linear transformation: y = xW^T
|
||
fn apply_linear(&self, x: &Tensor, weight: &Tensor) -> Result<Tensor, MLError> {
|
||
// Matrix multiplication: [batch, in_dim] × [out_dim, in_dim]^T = [batch, out_dim]
|
||
let result = x.matmul(&weight.t()?)?;
|
||
Ok(result)
|
||
}
|
||
|
||
/// Apply Gated Linear Unit
|
||
fn apply_glu(&self, hidden: &Tensor) -> Result<Tensor, MLError> {
|
||
let (glu_linear_quant, glu_gate_quant) = &self.quantized_glu_weights;
|
||
|
||
let linear_weight = self.quantizer.dequantize_tensor(
|
||
glu_linear_quant
|
||
.as_ref()
|
||
.ok_or_else(|| MLError::ModelError("Missing GLU linear".to_string()))?,
|
||
)?;
|
||
let gate_weight = self.quantizer.dequantize_tensor(
|
||
glu_gate_quant
|
||
.as_ref()
|
||
.ok_or_else(|| MLError::ModelError("Missing GLU gate".to_string()))?,
|
||
)?;
|
||
|
||
let linear_out = self.apply_linear(hidden, &linear_weight)?;
|
||
let gate_out = self.apply_linear(hidden, &gate_weight)?;
|
||
let gate_activated = manual_sigmoid(&gate_out)?;
|
||
|
||
Ok((&linear_out * &gate_activated)?)
|
||
}
|
||
|
||
/// Apply layer normalization (kept in F32)
|
||
fn apply_layer_norm(&self, x: &Tensor) -> Result<Tensor, MLError> {
|
||
if let Some(ln) = &self.layer_norm {
|
||
crate::cuda_compat::layer_norm_with_fallback(
|
||
x,
|
||
&ln.normalized_shape,
|
||
ln.weight.as_ref(),
|
||
ln.bias.as_ref(),
|
||
ln.eps,
|
||
)
|
||
} else {
|
||
Ok(x.clone())
|
||
}
|
||
}
|
||
|
||
/// Get quantization type
|
||
pub fn quant_type(&self) -> QuantizationType {
|
||
self.quantized_linear1
|
||
.as_ref()
|
||
.map(|q| q.quant_type)
|
||
.unwrap_or(QuantizationType::None)
|
||
}
|
||
|
||
/// Calculate memory footprint in MB
|
||
pub fn memory_footprint_mb(&self) -> f64 {
|
||
let mut total_bytes = 0;
|
||
|
||
// Add quantized layer sizes
|
||
if let Some(q) = &self.quantized_linear1 {
|
||
total_bytes += q.memory_bytes();
|
||
}
|
||
if let Some(q) = &self.quantized_linear2 {
|
||
total_bytes += q.memory_bytes();
|
||
}
|
||
if let Some(q) = &self.quantized_glu_weights.0 {
|
||
total_bytes += q.memory_bytes();
|
||
}
|
||
if let Some(q) = &self.quantized_glu_weights.1 {
|
||
total_bytes += q.memory_bytes();
|
||
}
|
||
if let Some(q) = &self.quantized_skip_proj {
|
||
total_bytes += q.memory_bytes();
|
||
}
|
||
if let Some(q) = &self.quantized_context_proj {
|
||
total_bytes += q.memory_bytes();
|
||
}
|
||
|
||
// Add layer norm parameters (F32)
|
||
total_bytes += self.output_dim * 4 * 2; // weight + bias
|
||
|
||
total_bytes as f64 / (1024.0 * 1024.0)
|
||
}
|
||
}
|
||
|
||
// Duplicate device() method removed - already defined in quantization.rs
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::memory_optimization::quantization::QuantizationConfig;
|
||
use candle_core::DType;
|
||
use candle_nn::{VarBuilder, VarMap};
|
||
use std::sync::Arc;
|
||
|
||
#[test]
|
||
fn test_quantized_grn_creation() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = Arc::new(VarMap::new());
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let grn = GatedResidualNetwork::new(128, 128, vs.pp("grn"))?;
|
||
|
||
let quant_config = QuantizationConfig::default();
|
||
let quantizer = Quantizer::new(quant_config, device);
|
||
|
||
let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;
|
||
|
||
assert_eq!(quantized_grn.input_dim, 128);
|
||
assert_eq!(quantized_grn.output_dim, 128);
|
||
assert!(quantized_grn.quantized_linear1.is_some());
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn test_quantized_grn_memory_footprint() -> Result<(), MLError> {
|
||
let device = Device::Cpu;
|
||
let varmap = Arc::new(VarMap::new());
|
||
let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);
|
||
|
||
let grn = GatedResidualNetwork::new(512, 512, vs.pp("grn"))?;
|
||
|
||
let quant_config = QuantizationConfig::default();
|
||
let quantizer = Quantizer::new(quant_config, device);
|
||
|
||
let quantized_grn = QuantizedGatedResidualNetwork::from_grn(&grn, quantizer)?;
|
||
|
||
let memory_mb = quantized_grn.memory_footprint_mb();
|
||
|
||
// Should be ~1MB for INT8 quantization
|
||
assert!(
|
||
memory_mb < 2.0,
|
||
"Memory footprint too high: {} MB",
|
||
memory_mb
|
||
);
|
||
|
||
Ok(())
|
||
}
|
||
}
|