Improve activation quant function

huggingface · Dec 31, 2024 · c00f3c8 · c00f3c8
1 parent f2a3809
commit c00f3c8
Showing 1 changed file with 3 additions and 10 deletions.
diff --git a/candle-nn/src/bit_linear.rs b/candle-nn/src/bit_linear.rs
@@ -42,17 +42,10 @@ fn weight_quant(x: &Tensor) -> Result<Tensor> {
 }
 
 fn activation_quant(x: &Tensor) -> Result<Tensor> {
-    let scale = (127.0
-        / x.abs()?
-            .max(D::Minus1)?
-            .max(D::Minus1)?
-            .clamp(1e-5, f32::INFINITY)?)?
-    .to_dtype(x.dtype())?;
+    let scale = x.abs()?.max_keepdim(D::Minus1)?.clamp(1e-5, f32::INFINITY)?;
+    let scale = (127.0 / scale)?;
 
-    let y = x
-        .broadcast_mul(&scale.unsqueeze(D::Minus1)?.unsqueeze(D::Minus1)?)?
-        .clamp(-128.0, 127.0)?
-        .broadcast_div(&scale)?;
+    let y = (x.broadcast_mul(&scale))?.round()?.clamp(-128., 127.)?.broadcast_div(&scale)?;
 
     Ok(y)
 }