File size: 2,496 Bytes
97df2ae | 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 | //! Quantização dinâmica int8 por canal de saída (equivalente Rust do
//! `quanta/torchao_quant.py` — Teo 28.1: camadas linearizáveis com cota
//! h/√12). EXCLUSÕES estruturais (T10(iii)): codebooks EMA, cabeças de
//! conceitos, embeddings de nível — mantidas em f32.
use candle_core::{Device, Result, Tensor};
pub struct PesoInt8 {
pub q: Tensor, // i32 com valores em [-127,127] (cota int8)
pub escala: Tensor, // f32 (d_saida, 1)
pub forma: Vec<usize>,
}
/// Quantiza W (d_saida, d_entrada) por LINHA (canal de saída).
pub fn quantizar_linha(w: &Tensor) -> Result<PesoInt8> {
let dev = w.device();
let dados = w.flatten_all()?.to_vec1::<f32>()?;
let (s, e) = (w.dim(0)?, w.dim(1)?);
let mut q = vec![0i32; dados.len()];
let mut escala = vec![0f32; s];
for i in 0..s {
let linha = &dados[i * e..(i + 1) * e];
let amax = linha.iter().fold(0f32, |m, x| m.max(x.abs()));
let esc = (amax / 127.0).max(1e-12);
escala[i] = esc;
for j in 0..e {
let v = (linha[j] / esc).round().clamp(-128.0, 127.0);
q[i * e + j] = v as i32;
}
}
Ok(PesoInt8 {
q: Tensor::from_vec(q, (s, e), dev)?,
escala: Tensor::from_vec(escala, (s, 1), dev)?,
forma: vec![s, e],
})
}
/// y = x·Wᵀ desquantizado on-the-fly (por canal).
pub fn linear_int8(x: &Tensor, p: &PesoInt8, dev: &Device) -> Result<Tensor> {
let qf = p.q.to_dtype(candle_core::DType::F32)?;
let w = qf.broadcast_mul(&p.escala)?; // (s, e)
x.matmul(&w.transpose(0, 1)?)
}
/// Paridade int8 vs f32 (V14-análogo do Teo 28.1): erro relativo médio.
pub fn paridade(w: &Tensor, x: &Tensor, dev: &Device) -> Result<f64> {
let p = quantizar_linha(w)?;
let y_q = linear_int8(x, &p, dev)?;
let y_f = x.matmul(&w.transpose(0, 1)?)?;
let dif = (&y_q - &y_f)?.abs()?.mean_all()?.to_scalar::<f32>()? as f64;
let ref_ = y_f.abs()?.mean_all()?.to_scalar::<f32>()? as f64;
Ok(dif / ref_.max(1e-9))
}
#[cfg(test)]
mod testes {
use super::*;
#[test]
fn paridade_int8_dentro_da_cota() -> Result<()> {
let dev = Device::Cpu;
let w = Tensor::randn(0f32, 0.05, (64, 32), &dev)?;
let x = Tensor::randn(0f32, 1f32, (4, 32), &dev)?;
let erro = paridade(&w, &x, &dev)?;
// cota empírica do Teo 28.1 (h/√12 por canal, |w| ≤ 0.05·3σ)
assert!(erro < 0.05, "erro relativo int8 = {erro}");
Ok(())
}
}
|