mlx.core.quantize

quantize(w: array, /, group_size: int = 64, bits: int = 4, *, stream: None | Stream | Device = None) → tuple[array, array, array]

使用 bits 位对矩阵 w 中的每个元素进行量化。

注意，矩阵 w 中每一行的每 group_size 个元素会被一起量化。因此，w 的列数应能被 group_size 整除。特别是，w 的行被划分为大小为 group_size 的组，这些组被一起量化。

警告

quantize 当前仅支持尺寸为 32 的倍数的二维输入

形式上，对于矩阵 w 中一行连续的 \(g\) 个元素 \(w_1\) 到 \(w_g\)，我们按如下方式计算每个元素的量化表示 \(\hat{w_i}\)

\[\begin{split}\begin{aligned} \alpha &= \max_i w_i \\ \beta &= \min_i w_i \\ s &= \frac{\alpha - \beta}{2^b - 1} \\ \hat{w_i} &= \textrm{round}\left( \frac{w_i - \beta}{s}\right). \end{aligned}\end{split}\]

经过上述计算，\(\hat{w_i}\) 将占用 \(b\) 位，并从低位到高位打包到一个无符号 32 位整数中。例如，对于 4 位量化，我们将 8 个元素放入一个无符号 32 位整数中，其中第 1 个元素占用最低的 4 位，第 2 个元素占用第 4-7 位，依此类推。

为了能够对 w 的元素进行反量化，我们还需要保存 \(s\) 和 \(\beta\)，它们分别作为 scales 和 biases 返回。

参数:

• w (array) – 要量化的矩阵

• group_size (int, 可选) – w 中共享同一比例和偏差的组的大小。默认值: 64。

• bits (int, 可选) – 返回的量化矩阵中 w 的每个元素占用的位数。默认值: 4。

返回值:

一个包含以下内容的元组：

• w_q (array): w 的量化版本

• scales (array): 用于乘以每个元素的比例，即 \(s\)

• biases (array): 要添加到每个元素的偏差，即 \(\beta\)

返回类型:

tuple