quantize-py 0.3.2__tar.gz → 0.4.0__tar.gz

{quantize_py-0.3.2 → quantize_py-0.4.0}/Cargo.lock

@@ -1411,14 +1411,24 @@ dependencies = [
 
 [[package]]
 name = "quantize"
-version = "0.3.2"
+version = "0.4.0"
+dependencies = [
+ "half",
+ "rayon",
+]
+
+[[package]]
+name = "quantize-files"
+version = "0.4.0"
 dependencies = [
  "half",
+ "quantize",
+ "serde_json",
 ]
 
 [[package]]
 name = "quantize-py"
-version = "0.3.2"
+version = "0.4.0"
 dependencies = [
  "half",
  "numpy",
@@ -1778,6 +1788,17 @@ version = "1.16.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9395f0f0eee849a9b707b2f06bb92a6a422090e2123bb2ef8e87a0e61892a8e"
 
+[[package]]
+name = "smollm"
+version = "0.0.0"
+dependencies = [
+ "hf-hub",
+ "quantize",
+ "quantize-files",
+ "serde_json",
+ "tokenizers",
+]
+
 [[package]]
 name = "socks"
 version = "0.3.4"

{quantize_py-0.3.2 → quantize_py-0.4.0}/Cargo.toml

@@ -23,7 +23,7 @@ resolver = "3"
 
 [workspace.package]
 # Shared with the Python package. Between releases it ends in -dev.
-version = "0.3.2"
+version = "0.4.0"
 
 [workspace.dependencies]
 half = "2"
@@ -34,13 +34,12 @@ numpy = "0.29"
 all = { level = "deny", priority = -1 }
 
 [features]
-default = ["std"]
-# Does nothing, since the crate always needs the standard library. Kept so
-# that a Cargo.toml that names it still builds.
-std = []
+# Splits large matmul calls across every core.
+rayon = ["dep:rayon"]
 
 [dependencies]
 half = { workspace = true }
+rayon = { version = "1", optional = true }
 
 [lints]
 workspace = true

{quantize_py-0.3.2 → quantize_py-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quantize-py
-Version: 0.3.2
+Version: 0.4.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13

{quantize_py-0.3.2 → quantize_py-0.4.0}/README.md

@@ -48,6 +48,8 @@ let dot = q.dot(&weights).unwrap(); // 0.926
 
 for a weight matrix, `q.set_shape(rows, columns)` records its shape, so `q.matmul(&inputs)` can multiply a batch of inputs by it, like a linear layer, and `q.dequantize_row_into(row, &mut out)` can decode one row, like an embedding lookup. `q.to_bytes()` saves a tensor, shape included, and `Quantized::from_bytes` loads it back.
 
+to run `matmul` on every core, turn on the `rayon` feature with `cargo add quantize --features rayon`. a big enough call splits across the cores, by the matrix's rows, by vectors, or both, and the results are the same, bit for bit. smaller calls stay on one core. a program that already keeps every core busy calling `matmul` loses up to about a quarter, most of all for one vector through a large matrix, and one with cores to spare gains. without it, the only dependency is `half`.
+
 codes can be 2 to 16 bits, and scales `f32`, `f16`, or `bf16`. the rest is in the [api docs](https://docs.rs/quantize). for python, see [python/](https://github.com/aksheyd/quantize/tree/main/python).
 
 ---

{quantize_py-0.3.2 → quantize_py-0.4.0}/python/src/input.rs

@@ -195,6 +195,27 @@ pub fn as_packed_codes(obj: &Bound<'_, PyAny>) -> PyResult<Vec<u8>> {
     read_uint8(obj, PACKED_CODES_TYPE)
 }
 
+/// Read block widths: a 1-D uint8 array, like `Quantized.block_bits`
+/// returns, or a sequence of ints. The array is copied at once, since reading
+/// it as a sequence would convert one NumPy scalar at a time.
+pub fn as_block_bits(obj: &Bound<'_, PyAny>) -> PyResult<Vec<u32>> {
+    match obj.cast::<PyArray1<u8>>() {
+        Ok(widths) => {
+            let widths = widths.try_readonly()?;
+            Ok(widths
+                .as_array()
+                .iter()
+                .map(|&bits| u32::from(bits))
+                .collect())
+        }
+        Err(_) => obj.extract().inspect_err(|error: &PyErr| {
+            // Errors like "Can't extract `str` to `Vec`" don't name the
+            // argument, so add the note PyO3 adds to the arguments it reads.
+            let _ = error.add_note(obj.py(), "while processing 'block_bits'");
+        }),
+    }
+}
+
 /// Read the bytes that `to_bytes` saved: `bytes` or another bytes-like
 /// object, or a 1-D uint8 array.
 pub fn as_bytes(obj: &Bound<'_, PyAny>) -> PyResult<Vec<u8>> {

{quantize_py-0.3.2 → quantize_py-0.4.0}/python/src/quantized/methods.rs

@@ -10,7 +10,7 @@ use super::parts::Parts;
 use crate::error::from_quantize;
 use crate::gil::detach_if_large;
 use crate::input::{
-    as_bytes, as_f32_array, as_f32_matmul_values, as_f32_values, as_packed_codes,
+    as_block_bits, as_bytes, as_f32_array, as_f32_matmul_values, as_f32_values, as_packed_codes,
     as_writable_f32_out, check_shape,
 };
 use crate::scale::PyScale;
@@ -235,14 +235,18 @@ impl PyQuantized {
             .map(|bits| bits.to_vec().into_pyarray(py)))
     }
 
-    /// Bytes held by the codes, scales, zero-points, and block widths.
+    /// Bytes held by the codes, scales, zero-points, and block widths. An
+    /// adaptive matrix also keeps 8 bytes a row in memory, to find where each
+    /// row starts, which this doesn't count.
     #[getter]
     fn nbytes(&self) -> usize {
         with_inner!(&self.snapshot(), |quantized| quantized.nbytes())
     }
 
     /// Bits per value, counting the scales: 4-bit codes with one f16 scale
-    /// per 32 values cost 4.5.
+    /// per 32 values cost 4.5. An adaptive matrix also keeps 8 bytes a row in
+    /// memory, to find where each row starts, which this doesn't count: at
+    /// 64 columns, that's 1 more bit per value.
     #[getter]
     fn bits_per_element(&self) -> f32 {
         with_inner!(&self.snapshot(), |quantized| quantized.bits_per_element())
@@ -270,8 +274,12 @@ impl PyQuantized {
         scale: PyScale,
         zero_points: Option<Bound<'_, PyAny>>,
         bits: Option<u32>,
-        block_bits: Option<Vec<u32>>,
+        block_bits: Option<Bound<'_, PyAny>>,
     ) -> PyResult<Self> {
+        let block_bits = match block_bits {
+            Some(block_bits) => Some(as_block_bits(&block_bits)?),
+            None => None,
+        };
         let zero_points = match zero_points {
             Some(zero_points) => as_f32_values(&zero_points)?.to_vec()?,
             None => Vec::new(),

{quantize_py-0.3.2 → quantize_py-0.4.0}/python/src/quantized/parts.rs

@@ -23,14 +23,19 @@ pub(crate) struct Parts {
 }
 
 impl Parts {
-    /// Build the variant that `kind` names, then [`Quantized::validate`] it,
-    /// so parts that don't fit together are an error instead of a tensor
-    /// that decodes out of bounds.
+    /// Build the variant that `kind` names as a flat vector, then
+    /// [`Quantized::validate`] it, so parts that don't fit together are an
+    /// error instead of a tensor that decodes out of bounds. A 2-D `shape`
+    /// goes on last, through `set_shape`, which finds where each adaptive row
+    /// starts.
     pub(crate) fn into_quantized<S: Scale>(self) -> PyResult<Quantized<S>> {
-        let (len, columns) = match self.shape[..] {
+        let (len, rows_and_columns) = match self.shape[..] {
             [len] => (len, None),
+            // Named here, since validating the empty tensor first would blame
+            // its scales instead.
+            [_, 0] => return Err(from_quantize(Error::ShapeMismatch { len: 0, columns: 0 })),
             [rows, columns] => match rows.checked_mul(columns) {
-                Some(len) => (len, Some(columns)),
+                Some(len) => (len, Some((rows, columns))),
                 None => return Err(PyErr::new::<QuantizeError, _>(SHAPE_TOO_LARGE)),
             },
             _ => return Err(PyErr::new::<QuantizeError, _>(SHAPE_DIMENSIONS)),
@@ -39,13 +44,13 @@ impl Parts {
         let scales: Vec<S> = self.scales.into_iter().map(S::from_f32).collect();
         let zero_points: Vec<S> = self.zero_points.into_iter().map(S::from_f32).collect();
 
-        let quantized = match (self.kind.as_str(), self.bits, self.block_bits) {
+        let mut quantized = match (self.kind.as_str(), self.bits, self.block_bits) {
             ("symmetric", Some(bits), None) if zero_points.is_empty() => Quantized::Symmetric {
                 scales,
                 codes: packed_codes(self.codes, bits, len)?,
                 block,
                 len,
-                columns,
+                columns: None,
             },
             ("asymmetric", Some(bits), None) => Quantized::Asymmetric {
                 scales,
@@ -53,7 +58,7 @@ impl Parts {
                 codes: packed_codes(self.codes, bits, len)?,
                 block,
                 len,
-                columns,
+                columns: None,
             },
             ("adaptive", None, Some(block_bits)) => Quantized::Adaptive {
                 scales,
@@ -62,11 +67,15 @@ impl Parts {
                 block_bits: block_widths(block_bits)?,
                 block,
                 len,
-                columns,
+                columns: None,
+                row_starts: Vec::new(),
             },
             _ => return Err(PyErr::new::<QuantizeError, _>(KIND_PARTS)),
         };
         quantized.validate().map_err(from_quantize)?;
+        if let Some((rows, columns)) = rows_and_columns {
+            quantized.set_shape(rows, columns).map_err(from_quantize)?;
+        }
         Ok(quantized)
     }
 }

{quantize_py-0.3.2 → quantize_py-0.4.0}/python/tests/test_quantize.py

@@ -759,6 +759,7 @@ def test_from_parts_rebuilds_parts_saved_with_numpy_as_the_readme_says():
             quantize(weights, bits=4, block=32, scale=scale),
             asymmetric.quantize(weights, bits=5, block=16, scale=scale),
             adaptive.quantize(weights.ravel(), block=32, tolerance=0.001, scale=scale),
+            adaptive.quantize(weights, block=32, tolerance=0.001, scale=scale),
             quantize([], scale=scale),
         ]:
             parts = {
@@ -826,3 +827,31 @@ def test_adaptive_block_widths_take_one_byte_each():
             block_bits=[300, 8, 8],
             scale="f32",
         )
+
+
+def test_from_parts_reads_block_bits_from_an_array_or_any_sequence_of_ints():
+    quantized = adaptive.quantize(weight_matrix(3, 30), block=8, tolerance=0.001)
+    parts = {
+        "kind": "adaptive",
+        "shape": quantized.shape,
+        "block": quantized.block,
+        "codes": quantized.codes,
+        "scales": quantized.scales,
+        "zero_points": quantized.zero_points,
+        "scale": "f32",
+    }
+    widths = quantized.block_bits
+    non_contiguous_widths = np.repeat(widths, 2)[::2]
+    for block_bits in [
+        widths,
+        non_contiguous_widths,
+        widths.tolist(),
+        tuple(widths),
+        widths.astype(np.int64),
+    ]:
+        assert Quantized.from_parts(**parts, block_bits=block_bits) == quantized
+    with pytest.raises(InvalidBitsError, match="got 17"):
+        Quantized.from_parts(**parts, block_bits=np.full_like(widths, 17))
+    with pytest.raises(TypeError) as raised:
+        Quantized.from_parts(**parts, block_bits="4")
+    assert raised.value.__notes__ == ["while processing 'block_bits'"]

{quantize_py-0.3.2 → quantize_py-0.4.0}/src/lib.rs

@@ -25,9 +25,10 @@
 //! 32 values adds 16 / 32 = 0.5 bits to each value: 4-bit codes cost 4.5 bits
 //! per value.
 //!
-//! `BITS` and `BLOCK` are const generics. To choose them at run time, call the
-//! scheme's `quantize_with`, like [`symmetric::quantize_with`], which takes
-//! them as ordinary arguments.
+//! `BITS` and `BLOCK` are const generics, so a width outside 2 to 16 or a
+//! block of 0 stops the build instead of returning an error. To choose them at
+//! run time, call the scheme's `quantize_with`, like
+//! [`symmetric::quantize_with`], which takes them as ordinary arguments.
 //!
 //! [`quantize`] is symmetric: each block gets one scale. Everything below
 //! uses the same [`Quantized`] type:
@@ -64,9 +65,25 @@
 //!
 //! ## Features
 //!
-//! The crate needs the standard library. Its `std` feature, on by default,
-//! does nothing, and is kept so that a `Cargo.toml` that names it still
-//! builds.
+//! The `rayon` feature, off by default, splits a large
+//! [`matmul`](Quantized::matmul) or [`matmul_into`](Quantized::matmul_into)
+//! across every core with the [rayon](https://docs.rs/rayon) crate, and the
+//! results are the same, bit for bit. A call splits the matrix's rows, its
+//! inputs, or both, as `matmul` explains. A call too small to split stays on
+//! the thread that made it, as every call does without the feature.
+//!
+//! A program that already keeps every core busy calling `matmul` has no core
+//! free for a share, so a split can't make its calls finish sooner. Its
+//! smaller calls run as they did, and a call big enough to split costs up to
+//! about a quarter more, the most for one input through a large matrix. A
+//! program with cores to spare gains instead, since the shares run on them.
+//! Rayon starts one thread per core: set `RAYON_NUM_THREADS` to use fewer, or
+//! to 1 to stop splitting, which holds for any other rayon code in the program
+//! too. Cargo turns a feature on for every user of a crate once anything in
+//! the build asks for it, so a dependency can turn this one on for you.
+//!
+//! Turn it on with `cargo add quantize --features rayon`. Without it, the
+//! crate's only dependency is `half`.
 
 #![warn(missing_docs)]
 
@@ -87,4 +104,6 @@ pub use shared::scheme::Scheme;
 pub use shared::tensor::Quantized;
 pub use symmetric::{quantize, quantize_tensor};
 
+#[cfg(feature = "rayon")]
+pub(crate) use shared::cores;
 pub(crate) use shared::{decode, error, packed, scale, tensor};

{quantize_py-0.3.2 → quantize_py-0.4.0}/src/methods/adaptive/mod.rs

@@ -24,21 +24,25 @@ use crate::tensor::Quantized;
 /// times past on blocks far from zero, as [`asymmetric`](crate::asymmetric)
 /// explains.
 ///
+/// `BLOCK` must be at least 1, or the build stops, as in
+/// [`symmetric::quantize`](crate::symmetric::quantize).
+///
 /// # Errors
 ///
 /// [`Error::InvalidTolerance`] if `tolerance` is not finite and `> 0`.
 /// [`Error::ToleranceTooTight`] if even 8 bits can't round a block within
 /// `tolerance`, with the smallest tolerance that every block meets.
-/// [`Error::InvalidBlock`] if `BLOCK == 0`.
 /// [`Error::ScaleOutOfRange`] if `S` can't hold a block's scale or zero-point.
 pub fn quantize<S: Scale, const BLOCK: usize>(
     values: &[f32],
     tolerance: f32,
 ) -> Result<Quantized<S>> {
+    const { assert!(BLOCK >= 1, "BLOCK must be at least 1") };
     quantize_with::<S>(values, BLOCK, tolerance)
 }
 
-/// Runtime-block variant of [`quantize`].
+/// [`quantize`] with the block size chosen at run time. It returns the same
+/// errors, plus [`Error::InvalidBlock`] if `block` is 0.
 pub fn quantize_with<S: Scale>(
     values: &[f32],
     block: usize,
@@ -55,6 +59,7 @@ pub fn quantize_with<S: Scale>(
             block,
             len: 0,
             columns: None,
+            row_starts: Vec::new(),
         });
     }
 
@@ -98,6 +103,7 @@ pub fn quantize_with<S: Scale>(
         block,
         len: values.len(),
         columns: None,
+        row_starts: Vec::new(),
     })
 }
 

{quantize_py-0.3.2 → quantize_py-0.4.0}/src/methods/asymmetric/mod.rs

@@ -15,18 +15,30 @@ use crate::tensor::Quantized;
 
 /// Quantize into blocks of `BLOCK` using one bit width for every block.
 ///
+/// `BITS` must be from 2 to 16, and `BLOCK` at least 1, or the build stops,
+/// as in [`symmetric::quantize`](crate::symmetric::quantize).
+///
 /// # Errors
 ///
-/// [`crate::Error::InvalidBits`] or [`crate::Error::InvalidBlock`], and
 /// [`crate::Error::ScaleOutOfRange`] if `S` can't hold a block's scale or
 /// zero-point.
 pub fn quantize<S: Scale, const BITS: u32, const BLOCK: usize>(
     values: &[f32],
 ) -> Result<Quantized<S>> {
+    const {
+        assert!(2 <= BITS && BITS <= 16, "BITS must be from 2 to 16");
+        assert!(BLOCK >= 1, "BLOCK must be at least 1");
+    }
     quantize_with::<S>(values, BITS, BLOCK)
 }
 
-/// Runtime-width variant of [`quantize`].
+/// [`quantize`] with the bit width and block size chosen at run time.
+///
+/// # Errors
+///
+/// [`crate::Error::InvalidBits`] or [`crate::Error::InvalidBlock`], and
+/// [`crate::Error::ScaleOutOfRange`] if `S` can't hold a block's scale or
+/// zero-point.
 pub fn quantize_with<S: Scale>(values: &[f32], bits: u32, block: usize) -> Result<Quantized<S>> {
     check_bits(bits)?;
     check_block(block)?;
@@ -58,8 +70,10 @@ pub fn quantize_with<S: Scale>(values: &[f32], bits: u32, block: usize) -> Resul
     })
 }
 
-/// Quantize the entire tensor with one scale and one zero-point.
+/// Quantize the entire tensor with one scale and one zero-point. `BITS` must
+/// be from 2 to 16, as in [`quantize`].
 pub fn quantize_tensor<S: Scale, const BITS: u32>(values: &[f32]) -> Result<Quantized<S>> {
+    const { assert!(2 <= BITS && BITS <= 16, "BITS must be from 2 to 16") };
     quantize_with::<S>(values, BITS, values.len().max(1))
 }
 

{quantize_py-0.3.2 → quantize_py-0.4.0}/src/methods/symmetric/mod.rs

@@ -8,17 +8,32 @@ use crate::tensor::Quantized;
 
 /// Quantize `values` into fixed-size blocks of `BLOCK` with `BITS`-wide codes.
 ///
+/// `BITS` must be from 2 to 16, and `BLOCK` at least 1. Anything else stops
+/// the build, though `cargo check` doesn't catch it:
+///
+/// ```compile_fail
+/// let q = quantize::quantize::<f32, 1, 32>(&[0.5]);
+/// ```
+///
 /// # Errors
 ///
-/// [`crate::Error::InvalidBits`] or [`crate::Error::InvalidBlock`], and
 /// [`crate::Error::ScaleOutOfRange`] if `S` can't hold a block's scale.
 pub fn quantize<S: Scale, const BITS: u32, const BLOCK: usize>(
     values: &[f32],
 ) -> Result<Quantized<S>> {
+    const {
+        assert!(2 <= BITS && BITS <= 16, "BITS must be from 2 to 16");
+        assert!(BLOCK >= 1, "BLOCK must be at least 1");
+    }
     quantize_with::<S>(values, BITS, BLOCK)
 }
 
-/// Runtime-width variant of [`quantize`].
+/// [`quantize`] with the bit width and block size chosen at run time.
+///
+/// # Errors
+///
+/// [`crate::Error::InvalidBits`] or [`crate::Error::InvalidBlock`], and
+/// [`crate::Error::ScaleOutOfRange`] if `S` can't hold a block's scale.
 pub fn quantize_with<S: Scale>(values: &[f32], bits: u32, block: usize) -> Result<Quantized<S>> {
     check_bits(bits)?;
     check_block(block)?;
@@ -41,8 +56,10 @@ pub fn quantize_with<S: Scale>(values: &[f32], bits: u32, block: usize) -> Resul
     })
 }
 
-/// Quantize the entire tensor with a single scale.
+/// Quantize the entire tensor with a single scale. `BITS` must be from 2 to
+/// 16, as in [`quantize`].
 pub fn quantize_tensor<S: Scale, const BITS: u32>(values: &[f32]) -> Result<Quantized<S>> {
+    const { assert!(2 <= BITS && BITS <= 16, "BITS must be from 2 to 16") };
     quantize_with::<S>(values, BITS, values.len().max(1))
 }
 
@@ -379,13 +396,12 @@ mod tests {
         let mut q = quantize::<f32, 8, 32>(&w).unwrap();
         q.set_shape(2, 32).unwrap();
         let inputs = [0.0_f32; 48];
-        assert!(matches!(
-            q.matmul(&inputs),
-            Err(crate::Error::ShapeMismatch {
-                len: 48,
-                columns: 32
-            })
-        ));
+        let mismatch = crate::Error::InputMismatch {
+            columns: 32,
+            got: 48,
+        };
+        assert_eq!(q.matmul(&inputs), Err(mismatch.clone()));
+        assert_eq!(q.matmul_into(&inputs, &mut [0.0; 2]), Err(mismatch));
     }
 
     #[test]

{quantize_py-0.3.2 → quantize_py-0.4.0}/src/shared/bytes.rs

@@ -113,7 +113,10 @@ impl<S: Scale> Quantized<S> {
 
         check_block(block)?;
         let blocks = len.div_ceil(block);
-        let quantized = match kind {
+        // Every kind loads as a flat vector first. The shape goes on last,
+        // through `set_shape`, since it finds where each adaptive row starts
+        // by reading widths that must be checked first.
+        let mut quantized = match kind {
             SYMMETRIC => {
                 let scales = reader.scales(blocks)?;
                 let codes = reader.codes(code_bits, len)?;
@@ -122,7 +125,7 @@ impl<S: Scale> Quantized<S> {
                     codes,
                     block,
                     len,
-                    columns,
+                    columns: None,
                 }
             }
             ASYMMETRIC => {
@@ -135,7 +138,7 @@ impl<S: Scale> Quantized<S> {
                     codes,
                     block,
                     len,
-                    columns,
+                    columns: None,
                 }
             }
             ADAPTIVE => {
@@ -150,7 +153,8 @@ impl<S: Scale> Quantized<S> {
                     block_bits,
                     block,
                     len,
-                    columns,
+                    columns: None,
+                    row_starts: Vec::new(),
                 }
             }
             _ => return Err(malformed("unknown kind")),
@@ -164,6 +168,12 @@ impl<S: Scale> Quantized<S> {
             Ordering::Equal => {}
         }
         quantized.validate()?;
+        if let Some(columns) = columns {
+            if !len.is_multiple_of(columns) {
+                return Err(Error::ShapeMismatch { len, columns });
+            }
+            quantized.set_shape(len / columns, columns)?;
+        }
         Ok(quantized)
     }
 }
@@ -233,11 +243,15 @@ mod tests {
         five_by_sixteen.set_shape(5, 16).unwrap();
         let mut two_by_forty = asymmetric::quantize_with(&values, 8, 32).unwrap();
         two_by_forty.set_shape(2, 40).unwrap();
-        let tensors: [Quantized<S>; 5] = [
+        // Its second row starts partway through a block.
+        let mut adaptive_two_by_forty = adaptive::quantize_with(&values, 32, 0.01).unwrap();
+        adaptive_two_by_forty.set_shape(2, 40).unwrap();
+        let tensors: [Quantized<S>; 6] = [
             symmetric::quantize_with(&values, 4, 32).unwrap(),
             five_by_sixteen,
             two_by_forty,
             adaptive::quantize_with(&values, 32, 0.01).unwrap(),
+            adaptive_two_by_forty,
             symmetric::quantize_with(&[], 8, 32).unwrap(),
         ];
         for quantized in tensors {