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

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

@@ -1411,14 +1411,14 @@ dependencies = [
 
 [[package]]
 name = "quantize"
-version = "0.3.1"
+version = "0.3.2"
 dependencies = [
  "half",
 ]
 
 [[package]]
 name = "quantize-py"
-version = "0.3.1"
+version = "0.3.2"
 dependencies = [
  "half",
  "numpy",

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

@@ -23,7 +23,7 @@ resolver = "3"
 
 [workspace.package]
 # Shared with the Python package. Between releases it ends in -dev.
-version = "0.3.1"
+version = "0.3.2"
 
 [workspace.dependencies]
 half = "2"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quantize-py
-Version: 0.3.1
+Version: 0.3.2
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
@@ -27,7 +27,7 @@ python bindings for [quantize](https://github.com/aksheyd/quantize), a simple, f
 pip install quantize-py
 ```
 
-requires python 3.12 or newer. numpy is installed with it. where no wheel fits, like free-threaded python or alpine linux, pip builds it from source, which needs rust 1.88 or newer.
+requires python 3.12 or newer, or 3.14 or newer for free-threaded python. numpy is installed with it. where no wheel fits, like alpine linux, pip builds it from source, which needs rust 1.88 or newer.
 
 ```python
 from quantize import Scale, quantize
@@ -39,7 +39,7 @@ back = q.dequantize()  # [0.421, -0.098, 0.700, -0.498]
 dot = q.dot(weights)  # 0.926
 ```
 
-`bits` is the width of each code, from 2 to 16. `block` is how many values share one scale, and `scale` is how that scale is stored: `Scale.F32` (the default), `Scale.F16`, or `Scale.BF16`. values can be a list, a numpy array, or anything else `np.asarray` reads, like a pytorch tensor. a 2-d array keeps its shape, so `q.dequantize()` gives back a matrix and `q.matmul(x)` computes `x @ W.T`, like a linear layer. `q.shape` gives its rows and columns, and `len(q)` the number of values, as in rust.
+`bits` is the width of each code, from 2 to 16. `block` is how many values share one scale, and `scale` is how that scale is stored: `Scale.F32` (the default), `Scale.F16`, or `Scale.BF16`. values can be a list, a numpy array, or anything else `np.asarray` reads, like a pytorch tensor. a 2-d array keeps its shape, so `q.dequantize()` gives back a matrix and `q.matmul(x)` computes `x @ W.T`, like a linear layer. both can write into a float32 numpy array or pytorch tensor you pass, like `q.matmul(x, out=y)`, so a loop can reuse it. `q.shape` gives its rows and columns, and `len(q)` the number of values, as in rust.
 
 the scales count toward the size: 4-bit codes with one f16 scale per 32 values cost 4.5 bits per value, or 5 with the default f32 scale. `q.bits_per_element` reports it. above about 10 bits, use f32 scales with `asymmetric.quantize`, since f16 and bf16 zero-points cap its accuracy.
 
@@ -47,9 +47,9 @@ the other schemes return the same `Quantized` type:
 
 - `asymmetric.quantize(weights, bits=8, block=32)` adds a zero-point per block, for values that aren't centered on zero
 - `adaptive.quantize(weights, tolerance=0.1 * weights.std())` gives each block the fewest bits, from 2 to 8, that round every weight within `tolerance`, in the weights' own units. a tenth of their standard deviation gives about 5 bits a block. for a list, use `np.std(weights)`
-- `learned.refine(q, weights)` refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes `q` in place, so call `q.copy()` first to keep the original
-- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
-- `Scheme.Q4_32.quantize(weights)` picks a scheme at run time
+- `learned.refine(q, weights)` refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes `q` in place, so call `q.copy()` first to keep the original. it lets other threads run while it works, so threads can refit several layers at once
+- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place and lets other threads run. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
+- `Scheme.Q4_32.quantize(weights)` picks a scheme at run time, and `Scheme("symmetric(bits=4)")` reads one from text, like a config value, which `str(scheme)` writes
 
 a block with outliers can need more than 8 bits, which raises `ToleranceTooTightError`. retrying with its `smallest_tolerance` works, but loosens every block, not just that one:
 
@@ -77,6 +77,24 @@ q = Quantized.from_parts(**np.load("layer.npz"))
 
 to keep quantized values in a `torch.save` checkpoint, store `torch.frombuffer(bytearray(q.to_bytes()), dtype=torch.uint8)`, and load each back with `Quantized(t)`. the checkpoint is then as small as the bytes, and `torch.load` reads it without `add_safe_globals`. pickling `q` itself makes the checkpoint about 1.5 times larger, since `torch.save` stores bytes as text, and needs `torch.serialization.add_safe_globals([Quantized])` before `torch.load`.
 
+`q.matmul` runs on one core, but it lets other threads run while it multiplies, so threads can share out a batch. on an 8-core intel xeon, this multiplies a batch of 512 by a 4-bit 1536 × 576 matrix in 6 ms instead of 33 ms, with the same result, bit for bit:
+
+```python
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+pool = ThreadPoolExecutor()  # make it once, and reuse it for every layer
+
+def linear(q, x):  # x has shape (batch, columns)
+    # a piece per core, or pieces of 64 rows for a big batch, which stay in a core's cache
+    pieces = np.array_split(x, max(os.cpu_count(), len(x) // 64))
+    return np.concatenate(list(pool.map(q.matmul, pieces)))
+
+out = linear(q, x)
+```
+
+with a gil, while another thread keeps running python, each piece can wait up to `sys.getswitchinterval()`, 5 ms by default, to get the gil back, so splitting pays off only while your other threads are idle or in native code, or on free-threaded python. `dequantize`, `dot`, and `matmul` on 65,536 values or fewer, like one vector times a 256 × 256 matrix, keep the gil, so they don't wait. quantizing and refitting keep it only up to 4,096 values, so threads can quantize a model's layers in parallel.
+
 each value decodes as `code * scale`, or `(code - zero_point) * scale` with zero-points, using the scale and zero-point of its block. codes are signed and `bits` wide, and `q.codes` packs them low bits first. scales can be negative, since a symmetric block puts its value farthest from zero on the most negative code. `help(Quantized)` has the details.
 
 to build and test from a clone of the repo, with rust 1.88 or newer and [just](https://github.com/casey/just):

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

@@ -2,7 +2,7 @@
 
 a simple, fast quantization library usable as a [rust crate](https://crates.io/crates/quantize) or [python package](https://pypi.org/project/quantize-py/).
 
-quantization stores numbers in fewer bits, trading a little accuracy for a lot less memory.
+quantization stores numbers in fewer bits, trading some accuracy for a lot less memory.
 
 ## learn
 
@@ -74,7 +74,7 @@ quantize two random matrices, reconstruct them, then matmul. mse is the mean squ
 
 <!-- comparison:end -->
 
-on WikiText-2 with SmolLM-135M, 4-bit perplexity is 26.44 against candle `Q4_0`'s 26.46, and fp32 is 18.93. perplexity is roughly how many tokens the model is choosing between when it guesses the next one, so lower is better. reproducing these numbers downloads the model and dataset from hugging face, then scores the whole test set, which takes hours:
+on WikiText-2, a set of wikipedia articles, the small language model SmolLM-135M has a 4-bit perplexity of 26.44 against candle `Q4_0`'s 26.46, and 18.93 in fp32. perplexity is roughly how many tokens the model is choosing between when it guesses the next one, so lower is better. reproducing these numbers downloads the model and dataset from hugging face, then scores the whole test set, which takes hours:
 
 ```
 cargo run --release -p benchmarks --example wikitext --features workload
@@ -88,7 +88,7 @@ add `-- --max-tokens 512` to score only the first 512 tokens. that takes about a
 cargo run --release -p benchmarks --example throughput
 ```
 
-quantize and dequantize with f16 scales. both libraries allocate their output on every call. aarch64 is an apple M5 Max and x86_64 an intel xeon. the hand-written simd only targets aarch64, so x86_64 runs plain loops and is slower. the run also times this crate's `dot`, `matmul` on 16 inputs, and row-by-row decoding of an adaptive matrix, which aren't compared with candle.
+quantize and dequantize with f16 scales. both libraries allocate their output on every call. aarch64 is an apple M5 Max and x86_64 an intel xeon. the hand-written simd only targets aarch64, so x86_64 runs plain loops and is slower than aarch64. the run also times this crate's `dot`, `matmul` on 16 inputs and on one, and row-by-row decoding of an adaptive matrix, which aren't compared with candle.
 
 <!-- speed:start -->
 
@@ -96,8 +96,8 @@ quantize and dequantize with f16 scales. both libraries allocate their output on
 | --- | ---: | ---: | ---: | ---: |
 | quantize, aarch64 | 0.30 | 0.30 | 0.08 | 0.07 |
 | candle, aarch64 | 0.41 | 0.43 | 0.20 | 0.25 |
-| quantize, x86_64 | 3.90 | 3.74 | 1.06 | 0.36 |
-| candle, x86_64 | 2.22 | 4.29 | 0.45 | 0.45 |
+| quantize, x86_64 | 3.93 | 3.77 | 0.37 | 0.26 |
+| candle, x86_64 | 2.25 | 4.28 | 0.45 | 0.44 |
 
 <!-- speed:end -->
 

{quantize_py-0.3.1 → quantize_py-0.3.2}/python/README.md

@@ -6,7 +6,7 @@ python bindings for [quantize](https://github.com/aksheyd/quantize), a simple, f
 pip install quantize-py
 ```
 
-requires python 3.12 or newer. numpy is installed with it. where no wheel fits, like free-threaded python or alpine linux, pip builds it from source, which needs rust 1.88 or newer.
+requires python 3.12 or newer, or 3.14 or newer for free-threaded python. numpy is installed with it. where no wheel fits, like alpine linux, pip builds it from source, which needs rust 1.88 or newer.
 
 ```python
 from quantize import Scale, quantize
@@ -18,7 +18,7 @@ back = q.dequantize()  # [0.421, -0.098, 0.700, -0.498]
 dot = q.dot(weights)  # 0.926
 ```
 
-`bits` is the width of each code, from 2 to 16. `block` is how many values share one scale, and `scale` is how that scale is stored: `Scale.F32` (the default), `Scale.F16`, or `Scale.BF16`. values can be a list, a numpy array, or anything else `np.asarray` reads, like a pytorch tensor. a 2-d array keeps its shape, so `q.dequantize()` gives back a matrix and `q.matmul(x)` computes `x @ W.T`, like a linear layer. `q.shape` gives its rows and columns, and `len(q)` the number of values, as in rust.
+`bits` is the width of each code, from 2 to 16. `block` is how many values share one scale, and `scale` is how that scale is stored: `Scale.F32` (the default), `Scale.F16`, or `Scale.BF16`. values can be a list, a numpy array, or anything else `np.asarray` reads, like a pytorch tensor. a 2-d array keeps its shape, so `q.dequantize()` gives back a matrix and `q.matmul(x)` computes `x @ W.T`, like a linear layer. both can write into a float32 numpy array or pytorch tensor you pass, like `q.matmul(x, out=y)`, so a loop can reuse it. `q.shape` gives its rows and columns, and `len(q)` the number of values, as in rust.
 
 the scales count toward the size: 4-bit codes with one f16 scale per 32 values cost 4.5 bits per value, or 5 with the default f32 scale. `q.bits_per_element` reports it. above about 10 bits, use f32 scales with `asymmetric.quantize`, since f16 and bf16 zero-points cap its accuracy.
 
@@ -26,9 +26,9 @@ the other schemes return the same `Quantized` type:
 
 - `asymmetric.quantize(weights, bits=8, block=32)` adds a zero-point per block, for values that aren't centered on zero
 - `adaptive.quantize(weights, tolerance=0.1 * weights.std())` gives each block the fewest bits, from 2 to 8, that round every weight within `tolerance`, in the weights' own units. a tenth of their standard deviation gives about 5 bits a block. for a list, use `np.std(weights)`
-- `learned.refine(q, weights)` refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes `q` in place, so call `q.copy()` first to keep the original
-- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
-- `Scheme.Q4_32.quantize(weights)` picks a scheme at run time
+- `learned.refine(q, weights)` refits each block's scale, and its zero-point if it has one, to lower the mean squared error. it changes `q` in place, so call `q.copy()` first to keep the original. it lets other threads run while it works, so threads can refit several layers at once
+- `learned.alternate(q, weights)` refits too, then rounds each value to the nearest code on its block's new line, and repeats until no code moves. it also changes `q` in place and lets other threads run. both can raise the worst error past an adaptive tensor's tolerance, and lowering the error of the weights doesn't always lower the error of a model's outputs, so check those too
+- `Scheme.Q4_32.quantize(weights)` picks a scheme at run time, and `Scheme("symmetric(bits=4)")` reads one from text, like a config value, which `str(scheme)` writes
 
 a block with outliers can need more than 8 bits, which raises `ToleranceTooTightError`. retrying with its `smallest_tolerance` works, but loosens every block, not just that one:
 
@@ -56,6 +56,24 @@ q = Quantized.from_parts(**np.load("layer.npz"))
 
 to keep quantized values in a `torch.save` checkpoint, store `torch.frombuffer(bytearray(q.to_bytes()), dtype=torch.uint8)`, and load each back with `Quantized(t)`. the checkpoint is then as small as the bytes, and `torch.load` reads it without `add_safe_globals`. pickling `q` itself makes the checkpoint about 1.5 times larger, since `torch.save` stores bytes as text, and needs `torch.serialization.add_safe_globals([Quantized])` before `torch.load`.
 
+`q.matmul` runs on one core, but it lets other threads run while it multiplies, so threads can share out a batch. on an 8-core intel xeon, this multiplies a batch of 512 by a 4-bit 1536 × 576 matrix in 6 ms instead of 33 ms, with the same result, bit for bit:
+
+```python
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+pool = ThreadPoolExecutor()  # make it once, and reuse it for every layer
+
+def linear(q, x):  # x has shape (batch, columns)
+    # a piece per core, or pieces of 64 rows for a big batch, which stay in a core's cache
+    pieces = np.array_split(x, max(os.cpu_count(), len(x) // 64))
+    return np.concatenate(list(pool.map(q.matmul, pieces)))
+
+out = linear(q, x)
+```
+
+with a gil, while another thread keeps running python, each piece can wait up to `sys.getswitchinterval()`, 5 ms by default, to get the gil back, so splitting pays off only while your other threads are idle or in native code, or on free-threaded python. `dequantize`, `dot`, and `matmul` on 65,536 values or fewer, like one vector times a 256 × 256 matrix, keep the gil, so they don't wait. quantizing and refitting keep it only up to 4,096 values, so threads can quantize a model's layers in parallel.
+
 each value decodes as `code * scale`, or `(code - zero_point) * scale` with zero-points, using the scale and zero-point of its block. codes are signed and `bits` wide, and `q.codes` packs them low bits first. scales can be negative, since a symmetric block puts its value farthest from zero on the most negative code. `help(Quantized)` has the details.
 
 to build and test from a clone of the repo, with rust 1.88 or newer and [just](https://github.com/casey/just):

quantize_py-0.3.2/python/src/gil.rs

@@ -0,0 +1,49 @@
+//! When calls let other threads run.
+
+use pyo3::marker::Ungil;
+use pyo3::prelude::*;
+
+/// Calls that read a tensor, like `dot` and `matmul`, keep the GIL when they
+/// go through at most this many values, like one vector times a 256 × 256
+/// matrix. They take a tenth of a millisecond or less, so holding it barely
+/// delays other threads. Giving it up would make them wait to get it back
+/// while another thread keeps running Python, for up to the switch interval,
+/// `sys.getswitchinterval()`, which is 5 ms by default.
+const MOST_VALUES_READ_KEEPING_THE_GIL: usize = 65_536;
+
+/// Quantizing or refitting a value takes 10 to 20 times as long as reading
+/// it, so those calls keep the GIL only up to this many values, like a
+/// 64 × 64 matrix, which take about as long as the largest reads that keep
+/// it. Larger ones give it up, so that threads quantizing a model's layers
+/// run in parallel instead of taking turns.
+const MOST_VALUES_QUANTIZED_KEEPING_THE_GIL: usize = 4_096;
+
+/// Run `work`, which reads a tensor, letting other threads run meanwhile
+/// unless it's small. `values` counts each value it goes through, once for
+/// each vector it's multiplied by.
+pub fn detach_if_large<T, F>(py: Python<'_>, values: usize, work: F) -> T
+where
+    F: Ungil + FnOnce() -> T,
+    T: Ungil,
+{
+    if values <= MOST_VALUES_READ_KEEPING_THE_GIL {
+        work()
+    } else {
+        py.detach(work)
+    }
+}
+
+/// Run `work`, which quantizes or refits a tensor, letting other threads run
+/// meanwhile unless it's tiny. `values` counts each value it goes through,
+/// once for each pass over it.
+pub fn detach_unless_tiny<T, F>(py: Python<'_>, values: usize, work: F) -> T
+where
+    F: Ungil + FnOnce() -> T,
+    T: Ungil,
+{
+    if values <= MOST_VALUES_QUANTIZED_KEEPING_THE_GIL {
+        work()
+    } else {
+        py.detach(work)
+    }
+}

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

@@ -13,8 +13,11 @@ use crate::error::{InvalidBitsError, InvalidBlockError, length_mismatch};
 const CODES_TYPE: &str = "codes must be a 1-D signed integer array or a sequence of int; packed Quantized.codes is uint8 and must not be passed here — use unpacked_codes";
 const PACKED_CODES_TYPE: &str = "codes must be a 1-D uint8 array, like Quantized.codes";
 const BYTES_TYPE: &str = "data must be bytes, like to_bytes returns, or a 1-D uint8 array";
-const OUT_TYPE: &str = "out must be a writable C-contiguous native-endian float32 array";
+const OUT_TYPE: &str = "out must be a float32 numpy array or pytorch tensor";
 const OUT_CONTIG: &str = "out must be writable and C-contiguous";
+const OUT_OVERLAPS_INPUTS: &str = "out can't share memory with inputs";
+const OUT_IN_USE: &str =
+    "out is in use by another call, like one on another thread; give each call its own out";
 
 fn is_native_dtype(arr: &Bound<'_, PyUntypedArray>) -> PyResult<bool> {
     arr.dtype().getattr("isnative")?.extract()
@@ -75,7 +78,13 @@ fn read_f32<'py>(
     let float32 = numpy.getattr("float32")?;
     let contiguous = numpy.call_method1("ascontiguousarray", (array, float32))?;
     let typed = contiguous.cast::<PyArrayDyn<f32>>()?;
-    Ok((typed.try_readonly()?, array.shape().to_vec()))
+    // Reading fails only while another call writes into the array.
+    let readonly = typed.try_readonly().map_err(|_| {
+        PyValueError::new_err(format!(
+            "{argument} is being written by another call, like a matmul or dequantize with out= on another thread; read it once that call returns"
+        ))
+    })?;
+    Ok((readonly, array.shape().to_vec()))
 }
 
 /// `obj`, or if it's a PyTorch tensor, a tensor that `numpy.asarray` reads:
@@ -205,16 +214,56 @@ fn read_uint8(obj: &Bound<'_, PyAny>, message: &'static str) -> PyResult<Vec<u8>
     Ok(bytes.try_readonly()?.as_array().to_vec())
 }
 
-/// Borrow `out` for writing, after checking it has exactly `shape`.
+/// `out` as a NumPy array to write into: itself, or for a PyTorch tensor,
+/// the array that `out.numpy()` gives, which shares the tensor's memory. A
+/// tensor must be float32 and on the CPU, and can't require grad, since
+/// writing into it would go around autograd.
+fn out_array<'py>(out: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
+    if !is_torch_tensor(out)? {
+        return Ok(out.clone());
+    }
+    let device = out.getattr("device")?.str()?;
+    if device.to_cow()? != "cpu" {
+        return Err(PyTypeError::new_err(format!(
+            "out must be a tensor on the cpu, got one on {device}"
+        )));
+    }
+    if out.getattr("requires_grad")?.is_truthy()? {
+        return Err(PyTypeError::new_err(
+            "out can't be a tensor that requires grad",
+        ));
+    }
+    if out.getattr("dtype")?.str()?.to_cow()? != "torch.float32" {
+        return Err(out_type_error(out)?);
+    }
+    out.call_method0("numpy")
+}
+
+/// A `TypeError` that says what `out` must be, and what it is: its type,
+/// and its dtype if it has one, like `numpy.ndarray with dtype float64`.
+fn out_type_error(out: &Bound<'_, PyAny>) -> PyResult<PyErr> {
+    let mut description = out.get_type().fully_qualified_name()?.to_string();
+    if let Ok(dtype) = out.getattr("dtype") {
+        description += &format!(" with dtype {}", dtype.str()?);
+    }
+    let message = format!("{OUT_TYPE}, got {description}");
+    Ok(PyTypeError::new_err(message))
+}
+
+/// Borrow `out`, a NumPy array or PyTorch tensor, for writing, after
+/// checking it has exactly `shape` and shares no memory with `inputs`, the
+/// values the call reads while it writes `out`.
 pub fn as_writable_f32_out<'py>(
     obj: &Bound<'py, PyAny>,
     shape: &[usize],
+    inputs: &[f32],
 ) -> PyResult<numpy::PyReadwriteArrayDyn<'py, f32>> {
-    let arr = obj
-        .cast::<PyArrayDyn<f32>>()
-        .map_err(|_| PyTypeError::new_err(OUT_TYPE))?;
+    let array = out_array(obj)?;
+    let Ok(arr) = array.cast::<PyArrayDyn<f32>>() else {
+        return Err(out_type_error(obj)?);
+    };
     if !is_native_dtype(arr.as_untyped())? {
-        return Err(PyTypeError::new_err(OUT_TYPE));
+        return Err(out_type_error(obj)?);
     }
     let flags = arr.getattr("flags")?;
     let c_contiguous: bool = flags.getattr("c_contiguous")?.extract()?;
@@ -223,8 +272,30 @@ pub fn as_writable_f32_out<'py>(
         return Err(PyValueError::new_err(OUT_CONTIG));
     }
     check_shape(obj.py(), "out", shape, arr.shape())?;
+    // This comes before the borrow, which fails on an `out` that `inputs`
+    // was read from in place, as if another call held it.
+    if shares_memory(arr, inputs) {
+        return Err(PyValueError::new_err(OUT_OVERLAPS_INPUTS));
+    }
+    // Its flags are checked above, so this fails only while another call
+    // reads or writes it.
     arr.try_readwrite()
-        .map_err(|_| PyValueError::new_err(OUT_CONTIG))
+        .map_err(|_| PyValueError::new_err(OUT_IN_USE))
+}
+
+/// Whether `out` and `inputs`, both C-contiguous float32, share memory.
+/// Each fills one unbroken run of bytes, so they share memory exactly when
+/// the runs overlap. The memory is compared, not the arrays, since two
+/// arrays over one buffer, like `t.numpy()` and the array read from a
+/// tensor `t`, are separate objects. Comparing addresses keeps the GIL,
+/// which `numpy.may_share_memory` gives up.
+fn shares_memory(out: &Bound<'_, PyArrayDyn<f32>>, inputs: &[f32]) -> bool {
+    let out_start = out.data().addr();
+    let out_end = out_start + out.len() * size_of::<f32>();
+    let inputs_start = inputs.as_ptr().addr();
+    let inputs_end = inputs_start + size_of_val(inputs);
+    // The bytes both fill run from the later start to the earlier end.
+    out_start.max(inputs_start) < out_end.min(inputs_end)
 }
 
 /// Check that `argument`, which came in with shape `got`, has exactly the

{quantize_py-0.3.1 → quantize_py-0.3.2}/python/src/learned.rs

@@ -3,9 +3,15 @@
 use pyo3::prelude::*;
 
 use crate::error::{from_quantize, length_mismatch};
+use crate::gil::detach_unless_tiny;
 use crate::input::{as_f32_array, as_f32_values, as_i32_codes, check_shape};
 use crate::quantized::PyQuantized;
 
+// `refine` and `alternate` refit a snapshot of the tensor, with the GIL
+// released unless it's small, then store it. `Arc::make_mut` copies the
+// shared tensor before the refit changes it, so until it's stored, other
+// threads read the tensor as it was.
+
 /// Refit each block's scale, and its zero-point if it has one, to lower the
 /// mean squared error against `values`, the numbers `quantized` was
 /// quantized from, in the same shape. The codes don't move, so the tensor
@@ -16,18 +22,23 @@ use crate::quantized::PyQuantized;
 ///
 /// This changes `quantized` in place and returns it. Call
 /// `quantized.copy()` first to keep the original.
+///
+/// Other threads keep running while it refits, unless the tensor is small,
+/// and see `quantized` as it was until it returns. If two threads refit the
+/// same tensor at once, both start from the tensor as it was, and it keeps
+/// the result of whichever finishes last.
 #[pyfunction]
 pub fn refine<'py>(
     quantized: Bound<'py, PyQuantized>,
     values: Bound<'py, PyAny>,
 ) -> PyResult<Bound<'py, PyQuantized>> {
     let (array, values_shape) = as_f32_array(&values)?;
-    let tensor_shape = quantized.borrow().inner.shape();
-    check_shape(values.py(), "values", &tensor_shape, &values_shape)?;
-    quantized
-        .borrow_mut()
-        .refine(array.as_slice()?)
+    let mut refined = quantized.get().snapshot();
+    check_shape(values.py(), "values", &refined.shape(), &values_shape)?;
+    let values = array.as_slice()?;
+    detach_unless_tiny(quantized.py(), values.len(), || refined.refine(values))
         .map_err(from_quantize)?;
+    quantized.get().store(refined);
     Ok(quantized)
 }
 
@@ -40,18 +51,23 @@ pub fn refine<'py>(
 /// land past the tolerance it was quantized with.
 ///
 /// Like `refine`, this changes `quantized` in place, so call
-/// `quantized.copy()` first to keep the original. Unlike `refine`, it returns
-/// whether the codes settled: `True` once no code moves, or `False` if it
-/// stopped after 100 passes. Call it again while it returns `False`.
+/// `quantized.copy()` first to keep the original, and other threads keep
+/// running while it refits, as `refine` describes. Unlike `refine`, it
+/// returns whether the codes settled: `True` once no code moves, or `False`
+/// if it stopped after 100 passes. Call it again while it returns `False`.
 #[pyfunction]
 pub fn alternate(quantized: Bound<'_, PyQuantized>, values: Bound<'_, PyAny>) -> PyResult<bool> {
     let (array, values_shape) = as_f32_array(&values)?;
-    let tensor_shape = quantized.borrow().inner.shape();
-    check_shape(values.py(), "values", &tensor_shape, &values_shape)?;
-    quantized
-        .borrow_mut()
-        .alternate(array.as_slice()?)
-        .map_err(from_quantize)
+    let mut alternated = quantized.get().snapshot();
+    check_shape(values.py(), "values", &alternated.shape(), &values_shape)?;
+    let values = array.as_slice()?;
+    let values_in_100_passes = values.len().saturating_mul(100);
+    let settled = detach_unless_tiny(quantized.py(), values_in_100_passes, || {
+        alternated.alternate(values)
+    })
+    .map_err(from_quantize)?;
+    quantized.get().store(alternated);
+    Ok(settled)
 }
 
 /// The `(scale, zero_point)` that best fit

{quantize_py-0.3.1 → quantize_py-0.3.2}/python/src/lib.rs

@@ -1,6 +1,7 @@
 //! Native Python bindings.
 
 mod error;
+mod gil;
 mod input;
 mod learned;
 mod quantized;
@@ -15,6 +16,7 @@ use crate::error::{
     NotAMatrixError, QuantizeError, ScaleOutOfRangeError, ShapeMismatchError,
     ToleranceTooTightError,
 };
+use crate::gil::detach_unless_tiny;
 use crate::input::as_f32_array;
 use crate::learned::{alternate, fit_scale_and_zero_point, refine};
 use crate::quantized::PyQuantized;
@@ -30,7 +32,9 @@ fn quantize_values(
     let (array, shape) = as_f32_array(&values)?;
     let values = array.as_slice()?;
     let scheme = scheme(values.len());
-    py.detach(|| PyQuantized::from_scheme(scheme, values, &shape, scale))
+    detach_unless_tiny(py, values.len(), || {
+        PyQuantized::from_scheme(scheme, values, &shape, scale)
+    })
 }
 
 // Each `text_signature` repeats its `signature` so that `help()` shows the
@@ -147,11 +151,10 @@ fn adaptive_quantize(
     quantize_values(py, values, scale, |_| Scheme::Adaptive { block, tolerance })
 }
 
-// `refine` and `alternate` borrow a tensor mutably while they change it, and
-// only the GIL keeps other threads from borrowing it at the same time, which
-// would panic. So free-threaded Python turns the GIL back on when it imports
-// this module.
-#[pymodule(gil_used = true)]
+// Threads share each tensor through a lock instead of relying on the GIL, so
+// free-threaded Python leaves the GIL off when it imports this module, and
+// its threads call into it in parallel.
+#[pymodule(gil_used = false)]
 #[pyo3(name = "_native")]
 fn native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("__version__", env!("CARGO_PKG_VERSION"))?;

{quantize_py-0.3.1 → quantize_py-0.3.2}/python/src/quantized/inner.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::sync::{Arc, Mutex, PoisonError};
 
 use half::{bf16, f16};
 use pyo3::prelude::*;
@@ -11,8 +11,8 @@ use crate::scale::PyScale;
 /// The tensor, with one of three scale types, behind an `Arc` so that a
 /// clone shares its codes and scales instead of copying them. `refine` and
 /// `alternate` change it through `Arc::make_mut`, which copies a shared
-/// tensor first, so a `copy()`, or a `dot` or `matmul` still running, keeps
-/// the values it had.
+/// tensor first, so a `copy()`, or a `dot`, `matmul`, or `dequantize` still
+/// running, keeps the values it had.
 #[derive(Clone, PartialEq)]
 pub(crate) enum QuantizedInner {
     F32(Arc<Quantized<f32>>),
@@ -160,10 +160,15 @@ impl QuantizedInner {
 /// `Quantized(data)` loads a tensor that `to_bytes` saved, like `from_bytes`.
 /// Pickles load through it, so `torch.load` accepts them once
 /// `torch.serialization.add_safe_globals([Quantized])` allows the class.
-#[pyclass(name = "Quantized", module = "quantize", eq)]
-#[derive(PartialEq)]
+#[pyclass(name = "Quantized", module = "quantize", eq, frozen)]
 pub struct PyQuantized {
-    pub(crate) inner: QuantizedInner,
+    // `frozen` drops PyO3's borrow flag, which makes a refit and another
+    // thread's read fail when they overlap without the GIL. Threads share the
+    // tensor through this lock instead, held only to clone the `Arc` or to put
+    // a refit in its place, never while computing or calling Python. Each
+    // store replaces the tensor in one step, so even a lock that a panic
+    // poisoned holds a whole tensor.
+    inner: Mutex<QuantizedInner>,
 }
 
 impl PyQuantized {
@@ -173,24 +178,38 @@ impl PyQuantized {
         shape: &[usize],
         scale: PyScale,
     ) -> PyResult<Self> {
-        Ok(Self {
-            inner: QuantizedInner::from_scheme(scheme, values, shape, scale)?,
-        })
+        QuantizedInner::from_scheme(scheme, values, shape, scale).map(Self::from)
     }
 
     pub fn len(&self) -> usize {
-        self.inner.len()
+        self.snapshot().len()
+    }
+
+    /// The tensor as it is now, sharing its codes and scales. A refit stored
+    /// later puts a new tensor in its place, so this one keeps its values.
+    pub(crate) fn snapshot(&self) -> QuantizedInner {
+        self.inner
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .clone()
     }
 
-    pub fn refine(&mut self, values: &[f32]) -> quantize::Result<()> {
-        self.inner.refine(values)
+    /// Put `inner`, such as a refit of a snapshot, in place of the tensor.
+    pub(crate) fn store(&self, inner: QuantizedInner) {
+        *self.inner.lock().unwrap_or_else(PoisonError::into_inner) = inner;
     }
+}
 
-    pub fn alternate(&mut self, values: &[f32]) -> quantize::Result<bool> {
-        self.inner.alternate(values)
+impl From<QuantizedInner> for PyQuantized {
+    fn from(inner: QuantizedInner) -> Self {
+        Self {
+            inner: Mutex::new(inner),
+        }
     }
+}
 
-    pub(crate) fn dequantize_into(&self, out: &mut [f32]) -> quantize::Result<()> {
-        with_inner!(&self.inner, |quantized| quantized.dequantize_into(out))
+impl PartialEq for PyQuantized {
+    fn eq(&self, other: &Self) -> bool {
+        self.snapshot() == other.snapshot()
     }
 }